From 3dcdbcc34dbf64a296a78a8252c0b42d7137cc3c Mon Sep 17 00:00:00 2001 From: Conrad Ludgate Date: Mon, 4 Nov 2024 13:29:13 +0000 Subject: [PATCH 01/52] remove aws-lc-rs dep and fix storage_broker tls (#9613) It seems the ecosystem is not so keen on moving to aws-lc-rs as it's build setup is more complicated than ring (requiring cmake). Eventually I expect the ecosystem should pivot to https://github.com/ctz/graviola/tree/main/rustls-graviola as it stabilises (it has a very simply build step and license), but for now let's try not have a headache of juggling two crypto libs. I also noticed that tonic will just fail with tls without a default provider, so I added some defensive code for that. --- Cargo.lock | 150 +++--------------- Cargo.toml | 4 +- libs/postgres_backend/tests/simple_select.rs | 6 +- proxy/src/bin/pg_sni_router.rs | 17 +- proxy/src/compute.rs | 18 +-- proxy/src/config.rs | 6 +- proxy/src/proxy/tests/mod.rs | 28 ++-- storage_broker/Cargo.toml | 1 + storage_broker/src/lib.rs | 6 + .../src/scan_safekeeper_metadata.rs | 6 +- workspace_hack/Cargo.toml | 9 +- 11 files changed, 72 insertions(+), 179 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c5af247e8b..44ef6d960c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -310,33 +310,6 @@ dependencies = [ "zeroize", ] -[[package]] -name = "aws-lc-rs" -version = "1.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f95446d919226d587817a7d21379e6eb099b97b45110a7f272a444ca5c54070" -dependencies = [ - "aws-lc-sys", - "mirai-annotations", - "paste", - "zeroize", -] - -[[package]] -name = "aws-lc-sys" -version = "0.21.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b3ddc4a5b231dd6958b140ff3151b6412b3f4321fab354f399eec8f14b06df62" -dependencies = [ - "bindgen 0.69.5", - "cc", - "cmake", - "dunce", - "fs_extra", - "libc", - "paste", -] - [[package]] name = "aws-runtime" version = "1.4.3" @@ -942,29 +915,6 @@ dependencies = [ "serde", ] -[[package]] -name = "bindgen" -version = "0.69.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "271383c67ccabffb7381723dea0672a673f292304fcb45c01cc648c7a8d58088" -dependencies = [ - "bitflags 2.4.1", - "cexpr", - "clang-sys", - "itertools 0.10.5", - "lazy_static", - "lazycell", - "log", - "prettyplease", - "proc-macro2", - "quote", - "regex", - "rustc-hash", - "shlex", - "syn 2.0.52", - "which", -] - [[package]] name = "bindgen" version = "0.70.1" @@ -974,7 +924,7 @@ dependencies = [ "bitflags 2.4.1", "cexpr", "clang-sys", - "itertools 0.10.5", + "itertools 0.12.1", "log", "prettyplease", "proc-macro2", @@ -1220,15 +1170,6 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2da6da31387c7e4ef160ffab6d5e7f00c42626fe39aea70a7b0f1773f7dd6c1b" -[[package]] -name = "cmake" -version = "0.1.51" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fb1e43aa7fd152b1f968787f7dbcdeb306d1867ff373c69955211876c053f91a" -dependencies = [ - "cc", -] - [[package]] name = "colorchoice" version = "1.0.0" @@ -1329,9 +1270,9 @@ dependencies = [ [[package]] name = "const-oid" -version = "0.9.5" +version = "0.9.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "28c122c3980598d243d63d9a704629a2d748d101f278052ff068be5a4423ab6f" +checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8" [[package]] name = "const-random" @@ -1815,12 +1756,6 @@ dependencies = [ "syn 2.0.52", ] -[[package]] -name = "dunce" -version = "1.0.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" - [[package]] name = "dyn-clone" version = "1.0.14" @@ -2125,12 +2060,6 @@ dependencies = [ "tokio-util", ] -[[package]] -name = "fs_extra" -version = "1.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" - [[package]] name = "fsevent-sys" version = "4.1.0" @@ -2484,15 +2413,6 @@ dependencies = [ "digest", ] -[[package]] -name = "home" -version = "0.5.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3d1354bf6b7235cb4a0576c2619fd4ed18183f689b12b006a0ee7329eeff9a5" -dependencies = [ - "windows-sys 0.52.0", -] - [[package]] name = "hostname" version = "0.4.0" @@ -2988,12 +2908,6 @@ dependencies = [ "spin", ] -[[package]] -name = "lazycell" -version = "1.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" - [[package]] name = "libc" version = "0.2.150" @@ -3224,12 +3138,6 @@ dependencies = [ "windows-sys 0.48.0", ] -[[package]] -name = "mirai-annotations" -version = "1.12.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c9be0862c1b3f26a88803c4a49de6889c10e608b3ee9344e6ef5b45fb37ad3d1" - [[package]] name = "multimap" version = "0.8.3" @@ -4147,7 +4055,7 @@ dependencies = [ "bytes", "once_cell", "pq_proto", - "rustls 0.23.7", + "rustls 0.23.16", "rustls-pemfile 2.1.1", "serde", "thiserror", @@ -4176,7 +4084,7 @@ name = "postgres_ffi" version = "0.1.0" dependencies = [ "anyhow", - "bindgen 0.70.1", + "bindgen", "bytes", "crc32c", "env_logger", @@ -4314,7 +4222,7 @@ checksum = "0c1318b19085f08681016926435853bbf7858f9c082d0999b80550ff5d9abe15" dependencies = [ "bytes", "heck 0.5.0", - "itertools 0.10.5", + "itertools 0.12.1", "log", "multimap", "once_cell", @@ -4334,7 +4242,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e9552f850d5f0964a4e4d0bf306459ac29323ddfbae05e35a7c0d35cb0803cc5" dependencies = [ "anyhow", - "itertools 0.10.5", + "itertools 0.12.1", "proc-macro2", "quote", "syn 2.0.52", @@ -4422,7 +4330,7 @@ dependencies = [ "rsa", "rstest", "rustc-hash", - "rustls 0.23.7", + "rustls 0.23.16", "rustls-native-certs 0.8.0", "rustls-pemfile 2.1.1", "scopeguard", @@ -5106,23 +5014,22 @@ dependencies = [ "log", "ring", "rustls-pki-types", - "rustls-webpki 0.102.2", + "rustls-webpki 0.102.8", "subtle", "zeroize", ] [[package]] name = "rustls" -version = "0.23.7" +version = "0.23.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ebbbdb961df0ad3f2652da8f3fdc4b36122f568f968f45ad3316f26c025c677b" +checksum = "eee87ff5d9b36712a58574e12e9f0ea80f915a5b0ac518d322b24a465617925e" dependencies = [ - "aws-lc-rs", "log", "once_cell", "ring", "rustls-pki-types", - "rustls-webpki 0.102.2", + "rustls-webpki 0.102.8", "subtle", "zeroize", ] @@ -5202,11 +5109,10 @@ dependencies = [ [[package]] name = "rustls-webpki" -version = "0.102.2" +version = "0.102.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "faaa0a62740bedb9b2ef5afa303da42764c012f743917351dc9a237ea1663610" +checksum = "64ca1bc8749bd4cf37b5ce386cc146580777b4e8572c7b97baf22c83f444bee9" dependencies = [ - "aws-lc-rs", "ring", "rustls-pki-types", "untrusted", @@ -5823,6 +5729,7 @@ dependencies = [ "once_cell", "parking_lot 0.12.1", "prost", + "rustls 0.23.16", "tokio", "tonic", "tonic-build", @@ -5905,7 +5812,7 @@ dependencies = [ "postgres_ffi", "remote_storage", "reqwest 0.12.4", - "rustls 0.23.7", + "rustls 0.23.16", "rustls-native-certs 0.8.0", "serde", "serde_json", @@ -6338,7 +6245,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04fb792ccd6bbcd4bba408eb8a292f70fc4a3589e5d793626f45190e6454b6ab" dependencies = [ "ring", - "rustls 0.23.7", + "rustls 0.23.16", "tokio", "tokio-postgres", "tokio-rustls 0.26.0", @@ -6372,7 +6279,7 @@ version = "0.26.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c7bc40d0e5a97695bb96e27995cd3a08538541b0a846f65bba7a359f36700d4" dependencies = [ - "rustls 0.23.7", + "rustls 0.23.16", "rustls-pki-types", "tokio", ] @@ -6781,7 +6688,7 @@ dependencies = [ "base64 0.22.1", "log", "once_cell", - "rustls 0.23.7", + "rustls 0.23.16", "rustls-pki-types", "url", "webpki-roots 0.26.1", @@ -6985,7 +6892,7 @@ name = "walproposer" version = "0.1.0" dependencies = [ "anyhow", - "bindgen 0.70.1", + "bindgen", "postgres_ffi", "utils", ] @@ -7160,18 +7067,6 @@ dependencies = [ "rustls-pki-types", ] -[[package]] -name = "which" -version = "4.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" -dependencies = [ - "either", - "home", - "once_cell", - "rustix", -] - [[package]] name = "whoami" version = "1.5.1" @@ -7431,7 +7326,7 @@ dependencies = [ "hyper-util", "indexmap 1.9.3", "indexmap 2.0.1", - "itertools 0.10.5", + "itertools 0.12.1", "lazy_static", "libc", "log", @@ -7452,8 +7347,7 @@ dependencies = [ "regex-automata 0.4.3", "regex-syntax 0.8.2", "reqwest 0.12.4", - "rustls 0.23.7", - "rustls-webpki 0.102.2", + "rustls 0.23.16", "scopeguard", "serde", "serde_json", diff --git a/Cargo.toml b/Cargo.toml index 7f9a766ff9..e5f7719e7f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -143,7 +143,7 @@ reqwest-retry = "0.5" routerify = "3" rpds = "0.13" rustc-hash = "1.1.0" -rustls = "0.23" +rustls = { version = "0.23.16", default-features = false } rustls-pemfile = "2" scopeguard = "1.1" sysinfo = "0.29.2" @@ -174,7 +174,7 @@ tokio = { version = "1.17", features = ["macros"] } tokio-epoll-uring = { git = "https://github.com/neondatabase/tokio-epoll-uring.git" , branch = "main" } tokio-io-timeout = "1.2.0" tokio-postgres-rustls = "0.12.0" -tokio-rustls = "0.26" +tokio-rustls = { version = "0.26.0", default-features = false, features = ["tls12", "ring"]} tokio-stream = "0.1" tokio-tar = "0.3" tokio-util = { version = "0.7.10", features = ["io", "rt"] } diff --git a/libs/postgres_backend/tests/simple_select.rs b/libs/postgres_backend/tests/simple_select.rs index 9d3031d699..3fcfbf4a03 100644 --- a/libs/postgres_backend/tests/simple_select.rs +++ b/libs/postgres_backend/tests/simple_select.rs @@ -2,7 +2,7 @@ use once_cell::sync::Lazy; use postgres_backend::{AuthType, Handler, PostgresBackend, QueryError}; use pq_proto::{BeMessage, RowDescriptor}; -use rustls::crypto::aws_lc_rs; +use rustls::crypto::ring; use std::io::Cursor; use std::sync::Arc; use tokio::io::{AsyncRead, AsyncWrite}; @@ -94,7 +94,7 @@ async fn simple_select_ssl() { let (client_sock, server_sock) = make_tcp_pair().await; let server_cfg = - rustls::ServerConfig::builder_with_provider(Arc::new(aws_lc_rs::default_provider())) + rustls::ServerConfig::builder_with_provider(Arc::new(ring::default_provider())) .with_safe_default_protocol_versions() .expect("aws_lc_rs should support the default protocol versions") .with_no_client_auth() @@ -110,7 +110,7 @@ async fn simple_select_ssl() { }); let client_cfg = - rustls::ClientConfig::builder_with_provider(Arc::new(aws_lc_rs::default_provider())) + rustls::ClientConfig::builder_with_provider(Arc::new(ring::default_provider())) .with_safe_default_protocol_versions() .expect("aws_lc_rs should support the default protocol versions") .with_root_certificates({ diff --git a/proxy/src/bin/pg_sni_router.rs b/proxy/src/bin/pg_sni_router.rs index 025053d3cb..71783ee452 100644 --- a/proxy/src/bin/pg_sni_router.rs +++ b/proxy/src/bin/pg_sni_router.rs @@ -15,7 +15,7 @@ use proxy::context::RequestMonitoring; use proxy::metrics::{Metrics, ThreadPoolMetrics}; use proxy::proxy::{copy_bidirectional_client_compute, run_until_cancelled, ErrorSource}; use proxy::stream::{PqStream, Stream}; -use rustls::crypto::aws_lc_rs; +use rustls::crypto::ring; use rustls::pki_types::PrivateKeyDer; use tokio::io::{AsyncRead, AsyncWrite}; use tokio::net::TcpListener; @@ -105,14 +105,13 @@ async fn main() -> anyhow::Result<()> { let first_cert = cert_chain.first().context("missing certificate")?; let tls_server_end_point = TlsServerEndPoint::new(first_cert)?; - let tls_config = rustls::ServerConfig::builder_with_provider(Arc::new( - aws_lc_rs::default_provider(), - )) - .with_protocol_versions(&[&rustls::version::TLS13, &rustls::version::TLS12]) - .context("aws_lc_rs should support TLS1.2 and TLS1.3")? - .with_no_client_auth() - .with_single_cert(cert_chain, key)? - .into(); + let tls_config = + rustls::ServerConfig::builder_with_provider(Arc::new(ring::default_provider())) + .with_protocol_versions(&[&rustls::version::TLS13, &rustls::version::TLS12]) + .context("ring should support TLS1.2 and TLS1.3")? + .with_no_client_auth() + .with_single_cert(cert_chain, key)? + .into(); (tls_config, tls_server_end_point) } diff --git a/proxy/src/compute.rs b/proxy/src/compute.rs index b97942ee5d..65b6dd215b 100644 --- a/proxy/src/compute.rs +++ b/proxy/src/compute.rs @@ -8,7 +8,7 @@ use itertools::Itertools; use once_cell::sync::OnceCell; use pq_proto::StartupMessageParams; use rustls::client::danger::ServerCertVerifier; -use rustls::crypto::aws_lc_rs; +use rustls::crypto::ring; use rustls::pki_types::InvalidDnsNameError; use thiserror::Error; use tokio::net::TcpStream; @@ -266,12 +266,12 @@ impl ConnCfg { } } +type RustlsStream = >::Stream; + pub(crate) struct PostgresConnection { /// Socket connected to a compute node. - pub(crate) stream: tokio_postgres::maybe_tls_stream::MaybeTlsStream< - tokio::net::TcpStream, - tokio_postgres_rustls::RustlsStream, - >, + pub(crate) stream: + tokio_postgres::maybe_tls_stream::MaybeTlsStream, /// PostgreSQL connection parameters. pub(crate) params: std::collections::HashMap, /// Query cancellation token. @@ -298,9 +298,9 @@ impl ConnCfg { let client_config = if allow_self_signed_compute { // Allow all certificates for creating the connection let verifier = Arc::new(AcceptEverythingVerifier); - rustls::ClientConfig::builder_with_provider(Arc::new(aws_lc_rs::default_provider())) + rustls::ClientConfig::builder_with_provider(Arc::new(ring::default_provider())) .with_safe_default_protocol_versions() - .expect("aws_lc_rs should support the default protocol versions") + .expect("ring should support the default protocol versions") .dangerous() .with_custom_certificate_verifier(verifier) } else { @@ -308,9 +308,9 @@ impl ConnCfg { .get_or_try_init(load_certs) .map_err(ConnectionError::TlsCertificateError)? .clone(); - rustls::ClientConfig::builder_with_provider(Arc::new(aws_lc_rs::default_provider())) + rustls::ClientConfig::builder_with_provider(Arc::new(ring::default_provider())) .with_safe_default_protocol_versions() - .expect("aws_lc_rs should support the default protocol versions") + .expect("ring should support the default protocol versions") .with_root_certificates(root_store) }; let client_config = client_config.with_no_client_auth(); diff --git a/proxy/src/config.rs b/proxy/src/config.rs index 5183f22fa3..2870e100b7 100644 --- a/proxy/src/config.rs +++ b/proxy/src/config.rs @@ -7,7 +7,7 @@ use anyhow::{bail, ensure, Context, Ok}; use clap::ValueEnum; use itertools::Itertools; use remote_storage::RemoteStorageConfig; -use rustls::crypto::aws_lc_rs::{self, sign}; +use rustls::crypto::ring::{self, sign}; use rustls::pki_types::{CertificateDer, PrivateKeyDer}; use sha2::{Digest, Sha256}; use tracing::{error, info}; @@ -127,9 +127,9 @@ pub fn configure_tls( // allow TLS 1.2 to be compatible with older client libraries let mut config = - rustls::ServerConfig::builder_with_provider(Arc::new(aws_lc_rs::default_provider())) + rustls::ServerConfig::builder_with_provider(Arc::new(ring::default_provider())) .with_protocol_versions(&[&rustls::version::TLS13, &rustls::version::TLS12]) - .context("aws_lc_rs should support TLS1.2 and TLS1.3")? + .context("ring should support TLS1.2 and TLS1.3")? .with_no_client_auth() .with_cert_resolver(cert_resolver.clone()); diff --git a/proxy/src/proxy/tests/mod.rs b/proxy/src/proxy/tests/mod.rs index fe62fee204..abb0599d08 100644 --- a/proxy/src/proxy/tests/mod.rs +++ b/proxy/src/proxy/tests/mod.rs @@ -9,11 +9,12 @@ use async_trait::async_trait; use http::StatusCode; use retry::{retry_after, ShouldRetryWakeCompute}; use rstest::rstest; -use rustls::crypto::aws_lc_rs; +use rustls::crypto::ring; use rustls::pki_types; +use tokio::io::DuplexStream; use tokio_postgres::config::SslMode; use tokio_postgres::tls::{MakeTlsConnect, NoTls}; -use tokio_postgres_rustls::{MakeRustlsConnect, RustlsStream}; +use tokio_postgres_rustls::MakeRustlsConnect; use super::connect_compute::ConnectMechanism; use super::retry::CouldRetry; @@ -69,19 +70,12 @@ struct ClientConfig<'a> { hostname: &'a str, } +type TlsConnect = >::TlsConnect; + impl ClientConfig<'_> { - fn make_tls_connect( - self, - ) -> anyhow::Result< - impl tokio_postgres::tls::TlsConnect< - S, - Error = impl std::fmt::Debug + use, - Future = impl Send + use, - Stream = RustlsStream, - > + use, - > { + fn make_tls_connect(self) -> anyhow::Result> { let mut mk = MakeRustlsConnect::new(self.config); - let tls = MakeTlsConnect::::make_tls_connect(&mut mk, self.hostname)?; + let tls = MakeTlsConnect::::make_tls_connect(&mut mk, self.hostname)?; Ok(tls) } } @@ -95,9 +89,9 @@ fn generate_tls_config<'a>( let tls_config = { let config = - rustls::ServerConfig::builder_with_provider(Arc::new(aws_lc_rs::default_provider())) + rustls::ServerConfig::builder_with_provider(Arc::new(ring::default_provider())) .with_safe_default_protocol_versions() - .context("aws_lc_rs should support the default protocol versions")? + .context("ring should support the default protocol versions")? .with_no_client_auth() .with_single_cert(vec![cert.clone()], key.clone_key())? .into(); @@ -116,9 +110,9 @@ fn generate_tls_config<'a>( let client_config = { let config = - rustls::ClientConfig::builder_with_provider(Arc::new(aws_lc_rs::default_provider())) + rustls::ClientConfig::builder_with_provider(Arc::new(ring::default_provider())) .with_safe_default_protocol_versions() - .context("aws_lc_rs should support the default protocol versions")? + .context("ring should support the default protocol versions")? .with_root_certificates({ let mut store = rustls::RootCertStore::empty(); store.add(ca)?; diff --git a/storage_broker/Cargo.toml b/storage_broker/Cargo.toml index 2d19472c36..17d4aed63b 100644 --- a/storage_broker/Cargo.toml +++ b/storage_broker/Cargo.toml @@ -28,6 +28,7 @@ tokio = { workspace = true, features = ["rt-multi-thread"] } tracing.workspace = true metrics.workspace = true utils.workspace = true +rustls.workspace = true workspace_hack.workspace = true diff --git a/storage_broker/src/lib.rs b/storage_broker/src/lib.rs index bc632a39f7..3ac40f6e14 100644 --- a/storage_broker/src/lib.rs +++ b/storage_broker/src/lib.rs @@ -52,6 +52,12 @@ where // If schema starts with https, start encrypted connection; do plain text // otherwise. if let Some("https") = tonic_endpoint.uri().scheme_str() { + // if there's no default provider and both ring+aws-lc-rs are enabled + // this the tls settings on tonic will not work. + // erroring is ok. + rustls::crypto::ring::default_provider() + .install_default() + .ok(); let tls = ClientTlsConfig::new(); tonic_endpoint = tonic_endpoint.tls_config(tls)?; } diff --git a/storage_scrubber/src/scan_safekeeper_metadata.rs b/storage_scrubber/src/scan_safekeeper_metadata.rs index 6c312d0036..403b4590a8 100644 --- a/storage_scrubber/src/scan_safekeeper_metadata.rs +++ b/storage_scrubber/src/scan_safekeeper_metadata.rs @@ -6,7 +6,7 @@ use once_cell::sync::OnceCell; use pageserver_api::shard::TenantShardId; use postgres_ffi::{XLogFileName, PG_TLI}; use remote_storage::GenericRemoteStorage; -use rustls::crypto::aws_lc_rs; +use rustls::crypto::ring; use serde::Serialize; use tokio_postgres::types::PgLsn; use tracing::{debug, error, info}; @@ -256,9 +256,9 @@ async fn load_timelines_from_db( // Use rustls (Neon requires TLS) let root_store = TLS_ROOTS.get_or_try_init(load_certs)?.clone(); let client_config = - rustls::ClientConfig::builder_with_provider(Arc::new(aws_lc_rs::default_provider())) + rustls::ClientConfig::builder_with_provider(Arc::new(ring::default_provider())) .with_safe_default_protocol_versions() - .context("aws_lc_rs should support the default protocol versions")? + .context("ring should support the default protocol versions")? .with_root_certificates(root_store) .with_no_client_auth(); let tls_connector = tokio_postgres_rustls::MakeRustlsConnect::new(client_config); diff --git a/workspace_hack/Cargo.toml b/workspace_hack/Cargo.toml index 28c51b8ac1..8d83d9d9e2 100644 --- a/workspace_hack/Cargo.toml +++ b/workspace_hack/Cargo.toml @@ -47,7 +47,7 @@ hyper-dff4ba8e3ae991db = { package = "hyper", version = "1", features = ["full"] hyper-util = { version = "0.1", features = ["client-legacy", "server-auto", "service"] } indexmap-dff4ba8e3ae991db = { package = "indexmap", version = "1", default-features = false, features = ["std"] } indexmap-f595c2ba2a3f28df = { package = "indexmap", version = "2", features = ["serde"] } -itertools = { version = "0.10" } +itertools = { version = "0.12" } lazy_static = { version = "1", default-features = false, features = ["spin_no_std"] } libc = { version = "0.2", features = ["extra_traits", "use_std"] } log = { version = "0.4", default-features = false, features = ["std"] } @@ -65,8 +65,7 @@ regex = { version = "1" } regex-automata = { version = "0.4", default-features = false, features = ["dfa-onepass", "hybrid", "meta", "nfa-backtrack", "perf-inline", "perf-literal", "unicode"] } regex-syntax = { version = "0.8" } reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "stream"] } -rustls = { version = "0.23", features = ["ring"] } -rustls-webpki = { version = "0.102", default-features = false, features = ["aws_lc_rs", "ring", "std"] } +rustls = { version = "0.23", default-features = false, features = ["logging", "ring", "std", "tls12"] } scopeguard = { version = "1" } serde = { version = "1", features = ["alloc", "derive"] } serde_json = { version = "1", features = ["alloc", "raw_value"] } @@ -80,7 +79,7 @@ tikv-jemalloc-sys = { version = "0.5" } time = { version = "0.3", features = ["macros", "serde-well-known"] } tokio = { version = "1", features = ["fs", "io-std", "io-util", "macros", "net", "process", "rt-multi-thread", "signal", "test-util"] } tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev = "20031d7a9ee1addeae6e0968e3899ae6bf01cee2", features = ["with-serde_json-1"] } -tokio-rustls = { version = "0.26", features = ["ring"] } +tokio-rustls = { version = "0.26", default-features = false, features = ["logging", "ring", "tls12"] } tokio-stream = { version = "0.1", features = ["net"] } tokio-util = { version = "0.7", features = ["codec", "compat", "io", "rt"] } toml_edit = { version = "0.22", features = ["serde"] } @@ -106,7 +105,7 @@ half = { version = "2", default-features = false, features = ["num-traits"] } hashbrown = { version = "0.14", features = ["raw"] } indexmap-dff4ba8e3ae991db = { package = "indexmap", version = "1", default-features = false, features = ["std"] } indexmap-f595c2ba2a3f28df = { package = "indexmap", version = "2", features = ["serde"] } -itertools = { version = "0.10" } +itertools = { version = "0.12" } libc = { version = "0.2", features = ["extra_traits", "use_std"] } log = { version = "0.4", default-features = false, features = ["std"] } memchr = { version = "2" } From 8ad1dbce7252d1ed91543ddda6cd4c7c8ade414d Mon Sep 17 00:00:00 2001 From: Conrad Ludgate Date: Mon, 4 Nov 2024 14:04:56 +0000 Subject: [PATCH 02/52] [proxy]: parse proxy protocol TLVs with aws/azure support (#9610) AWS/azure private link shares extra information in the "TLV" values of the proxy protocol v2 header. This code doesn't action on it, but it parses it as appropriate. --- Cargo.lock | 2 + proxy/Cargo.toml | 2 + proxy/src/bin/pg_sni_router.rs | 6 +- proxy/src/console_redirect_proxy.rs | 6 +- proxy/src/context/mod.rs | 26 ++-- proxy/src/context/parquet.rs | 2 +- proxy/src/protocol2.rs | 186 ++++++++++++++++++++++++---- proxy/src/proxy/mod.rs | 10 +- proxy/src/serverless/mod.rs | 56 ++++++--- 9 files changed, 236 insertions(+), 60 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 44ef6d960c..484769bd16 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4341,6 +4341,8 @@ dependencies = [ "smallvec", "smol_str", "socket2", + "strum", + "strum_macros", "subtle", "thiserror", "tikv-jemalloc-ctl", diff --git a/proxy/Cargo.toml b/proxy/Cargo.toml index e25d2fcbab..2580b1cf8a 100644 --- a/proxy/Cargo.toml +++ b/proxy/Cargo.toml @@ -74,6 +74,8 @@ sha2 = { workspace = true, features = ["asm", "oid"] } smol_str.workspace = true smallvec.workspace = true socket2.workspace = true +strum.workspace = true +strum_macros.workspace = true subtle.workspace = true thiserror.workspace = true tikv-jemallocator.workspace = true diff --git a/proxy/src/bin/pg_sni_router.rs b/proxy/src/bin/pg_sni_router.rs index 71783ee452..ef5b5e8509 100644 --- a/proxy/src/bin/pg_sni_router.rs +++ b/proxy/src/bin/pg_sni_router.rs @@ -13,6 +13,7 @@ use itertools::Itertools; use proxy::config::TlsServerEndPoint; use proxy::context::RequestMonitoring; use proxy::metrics::{Metrics, ThreadPoolMetrics}; +use proxy::protocol2::ConnectionInfo; use proxy::proxy::{copy_bidirectional_client_compute, run_until_cancelled, ErrorSource}; use proxy::stream::{PqStream, Stream}; use rustls::crypto::ring; @@ -178,7 +179,10 @@ async fn task_main( info!(%peer_addr, "serving"); let ctx = RequestMonitoring::new( session_id, - peer_addr.ip(), + ConnectionInfo { + addr: peer_addr, + extra: None, + }, proxy::metrics::Protocol::SniRouter, "sni", ); diff --git a/proxy/src/console_redirect_proxy.rs b/proxy/src/console_redirect_proxy.rs index 81d1d70958..243ef07854 100644 --- a/proxy/src/console_redirect_proxy.rs +++ b/proxy/src/console_redirect_proxy.rs @@ -11,7 +11,7 @@ use crate::config::{ProxyConfig, ProxyProtocolV2}; use crate::context::RequestMonitoring; use crate::error::ReportableError; use crate::metrics::{Metrics, NumClientConnectionsGuard}; -use crate::protocol2::read_proxy_protocol; +use crate::protocol2::{read_proxy_protocol, ConnectionInfo}; use crate::proxy::connect_compute::{connect_to_compute, TcpMechanism}; use crate::proxy::handshake::{handshake, HandshakeData}; use crate::proxy::passthrough::ProxyPassthrough; @@ -65,8 +65,8 @@ pub async fn task_main( error!("proxy protocol header not supported"); return; } - Ok((socket, Some(addr))) => (socket, addr.ip()), - Ok((socket, None)) => (socket, peer_addr.ip()), + Ok((socket, Some(info))) => (socket, info), + Ok((socket, None)) => (socket, ConnectionInfo{ addr: peer_addr, extra: None }), }; match socket.inner.set_nodelay(true) { diff --git a/proxy/src/context/mod.rs b/proxy/src/context/mod.rs index ca3b808a1b..2a6c9c5969 100644 --- a/proxy/src/context/mod.rs +++ b/proxy/src/context/mod.rs @@ -19,6 +19,7 @@ use crate::intern::{BranchIdInt, ProjectIdInt}; use crate::metrics::{ ConnectOutcome, InvalidEndpointsGroup, LatencyTimer, Metrics, Protocol, Waiting, }; +use crate::protocol2::ConnectionInfo; use crate::types::{DbName, EndpointId, RoleName}; pub mod parquet; @@ -40,7 +41,7 @@ pub struct RequestMonitoring( ); struct RequestMonitoringInner { - pub(crate) peer_addr: IpAddr, + pub(crate) conn_info: ConnectionInfo, pub(crate) session_id: Uuid, pub(crate) protocol: Protocol, first_packet: chrono::DateTime, @@ -84,7 +85,7 @@ impl Clone for RequestMonitoring { fn clone(&self) -> Self { let inner = self.0.try_lock().expect("should not deadlock"); let new = RequestMonitoringInner { - peer_addr: inner.peer_addr, + conn_info: inner.conn_info.clone(), session_id: inner.session_id, protocol: inner.protocol, first_packet: inner.first_packet, @@ -117,7 +118,7 @@ impl Clone for RequestMonitoring { impl RequestMonitoring { pub fn new( session_id: Uuid, - peer_addr: IpAddr, + conn_info: ConnectionInfo, protocol: Protocol, region: &'static str, ) -> Self { @@ -125,13 +126,13 @@ impl RequestMonitoring { "connect_request", %protocol, ?session_id, - %peer_addr, + %conn_info, ep = tracing::field::Empty, role = tracing::field::Empty, ); let inner = RequestMonitoringInner { - peer_addr, + conn_info, session_id, protocol, first_packet: Utc::now(), @@ -162,7 +163,11 @@ impl RequestMonitoring { #[cfg(test)] pub(crate) fn test() -> Self { - RequestMonitoring::new(Uuid::now_v7(), [127, 0, 0, 1].into(), Protocol::Tcp, "test") + use std::net::SocketAddr; + let ip = IpAddr::from([127, 0, 0, 1]); + let addr = SocketAddr::new(ip, 5432); + let conn_info = ConnectionInfo { addr, extra: None }; + RequestMonitoring::new(Uuid::now_v7(), conn_info, Protocol::Tcp, "test") } pub(crate) fn console_application_name(&self) -> String { @@ -286,7 +291,12 @@ impl RequestMonitoring { } pub(crate) fn peer_addr(&self) -> IpAddr { - self.0.try_lock().expect("should not deadlock").peer_addr + self.0 + .try_lock() + .expect("should not deadlock") + .conn_info + .addr + .ip() } pub(crate) fn cold_start_info(&self) -> ColdStartInfo { @@ -362,7 +372,7 @@ impl RequestMonitoringInner { } fn has_private_peer_addr(&self) -> bool { - match self.peer_addr { + match self.conn_info.addr.ip() { IpAddr::V4(ip) => ip.is_private(), IpAddr::V6(_) => false, } diff --git a/proxy/src/context/parquet.rs b/proxy/src/context/parquet.rs index 3432ac5ff6..adbb74c8e5 100644 --- a/proxy/src/context/parquet.rs +++ b/proxy/src/context/parquet.rs @@ -121,7 +121,7 @@ impl From<&RequestMonitoringInner> for RequestData { fn from(value: &RequestMonitoringInner) -> Self { Self { session_id: value.session_id, - peer_addr: value.peer_addr.to_string(), + peer_addr: value.conn_info.addr.ip().to_string(), timestamp: value.first_packet.naive_utc(), username: value.user.as_deref().map(String::from), application_name: value.application.as_deref().map(String::from), diff --git a/proxy/src/protocol2.rs b/proxy/src/protocol2.rs index ef2391cdd8..d1084ca2ff 100644 --- a/proxy/src/protocol2.rs +++ b/proxy/src/protocol2.rs @@ -1,12 +1,15 @@ //! Proxy Protocol V2 implementation +//! Compatible with +use core::fmt; use std::io; -use std::net::SocketAddr; +use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr}; use std::pin::Pin; use std::task::{Context, Poll}; -use bytes::BytesMut; +use bytes::{Buf, Bytes, BytesMut}; use pin_project_lite::pin_project; +use strum_macros::FromRepr; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, ReadBuf}; pin_project! { @@ -58,9 +61,35 @@ const HEADER: [u8; 12] = [ 0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A, ]; +#[derive(PartialEq, Eq, Clone, Debug)] +pub struct ConnectionInfo { + pub addr: SocketAddr, + pub extra: Option, +} + +impl fmt::Display for ConnectionInfo { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match &self.extra { + None => self.addr.ip().fmt(f), + Some(ConnectionInfoExtra::Aws { vpce_id }) => { + write!(f, "vpce_id[{vpce_id:?}]:addr[{}]", self.addr.ip()) + } + Some(ConnectionInfoExtra::Azure { link_id }) => { + write!(f, "link_id[{link_id}]:addr[{}]", self.addr.ip()) + } + } + } +} + +#[derive(PartialEq, Eq, Clone, Debug)] +pub enum ConnectionInfoExtra { + Aws { vpce_id: Bytes }, + Azure { link_id: u32 }, +} + pub(crate) async fn read_proxy_protocol( mut read: T, -) -> std::io::Result<(ChainRW, Option)> { +) -> std::io::Result<(ChainRW, Option)> { let mut buf = BytesMut::with_capacity(128); while buf.len() < 16 { let bytes_read = read.read_buf(&mut buf).await?; @@ -164,22 +193,107 @@ pub(crate) async fn read_proxy_protocol( // - destination layer 3 address in network byte order // - source layer 4 address if any, in network byte order (port) // - destination layer 4 address if any, in network byte order (port) - let addresses = buf.split_to(remaining_length as usize); - let socket = match address_length { + let mut header = buf.split_to(usize::from(remaining_length)); + let mut addr = header.split_to(usize::from(address_length)); + let socket = match addr.len() { 12 => { - let src_addr: [u8; 4] = addresses[0..4].try_into().unwrap(); - let src_port = u16::from_be_bytes(addresses[8..10].try_into().unwrap()); + let src_addr = Ipv4Addr::from_bits(addr.get_u32()); + let _dst_addr = Ipv4Addr::from_bits(addr.get_u32()); + let src_port = addr.get_u16(); + let _dst_port = addr.get_u16(); Some(SocketAddr::from((src_addr, src_port))) } 36 => { - let src_addr: [u8; 16] = addresses[0..16].try_into().unwrap(); - let src_port = u16::from_be_bytes(addresses[32..34].try_into().unwrap()); + let src_addr = Ipv6Addr::from_bits(addr.get_u128()); + let _dst_addr = Ipv6Addr::from_bits(addr.get_u128()); + let src_port = addr.get_u16(); + let _dst_port = addr.get_u16(); Some(SocketAddr::from((src_addr, src_port))) } _ => None, }; - Ok((ChainRW { inner: read, buf }, socket)) + let mut extra = None; + + while let Some(mut tlv) = read_tlv(&mut header) { + match Pp2Kind::from_repr(tlv.kind) { + Some(Pp2Kind::Aws) => { + if tlv.value.is_empty() { + tracing::warn!("invalid aws tlv: no subtype"); + } + let subtype = tlv.value.get_u8(); + match Pp2AwsType::from_repr(subtype) { + Some(Pp2AwsType::VpceId) => { + extra = Some(ConnectionInfoExtra::Aws { vpce_id: tlv.value }); + } + None => { + tracing::warn!("unknown aws tlv: subtype={subtype}"); + } + } + } + Some(Pp2Kind::Azure) => { + if tlv.value.is_empty() { + tracing::warn!("invalid azure tlv: no subtype"); + } + let subtype = tlv.value.get_u8(); + match Pp2AzureType::from_repr(subtype) { + Some(Pp2AzureType::PrivateEndpointLinkId) => { + if tlv.value.len() != 4 { + tracing::warn!("invalid azure link_id: {:?}", tlv.value); + } + extra = Some(ConnectionInfoExtra::Azure { + link_id: tlv.value.get_u32_le(), + }); + } + None => { + tracing::warn!("unknown azure tlv: subtype={subtype}"); + } + } + } + Some(kind) => { + tracing::debug!("unused tlv[{kind:?}]: {:?}", tlv.value); + } + None => { + tracing::debug!("unknown tlv: {tlv:?}"); + } + } + } + + let conn_info = socket.map(|addr| ConnectionInfo { addr, extra }); + + Ok((ChainRW { inner: read, buf }, conn_info)) +} + +#[derive(FromRepr, Debug, Copy, Clone)] +#[repr(u8)] +enum Pp2Kind { + // The following are defined by https://www.haproxy.org/download/3.1/doc/proxy-protocol.txt + // we don't use these but it would be interesting to know what's available + Alpn = 0x01, + Authority = 0x02, + Crc32C = 0x03, + Noop = 0x04, + UniqueId = 0x05, + Ssl = 0x20, + NetNs = 0x30, + + /// + Aws = 0xEA, + + /// + Azure = 0xEE, +} + +#[derive(FromRepr, Debug, Copy, Clone)] +#[repr(u8)] +enum Pp2AwsType { + VpceId = 0x01, +} + +#[derive(FromRepr, Debug, Copy, Clone)] +#[repr(u8)] +enum Pp2AzureType { + PrivateEndpointLinkId = 0x01, } impl AsyncRead for ChainRW { @@ -216,6 +330,25 @@ impl ChainRW { } } +#[derive(Debug)] +struct Tlv { + kind: u8, + value: Bytes, +} + +fn read_tlv(b: &mut BytesMut) -> Option { + if b.len() < 3 { + return None; + } + let kind = b.get_u8(); + let len = usize::from(b.get_u16()); + if b.len() < len { + return None; + } + let value = b.split_to(len).freeze(); + Some(Tlv { kind, value }) +} + #[cfg(test)] mod tests { use tokio::io::AsyncReadExt; @@ -242,7 +375,7 @@ mod tests { let extra_data = [0x55; 256]; - let (mut read, addr) = read_proxy_protocol(header.chain(extra_data.as_slice())) + let (mut read, info) = read_proxy_protocol(header.chain(extra_data.as_slice())) .await .unwrap(); @@ -250,7 +383,9 @@ mod tests { read.read_to_end(&mut bytes).await.unwrap(); assert_eq!(bytes, extra_data); - assert_eq!(addr, Some(([127, 0, 0, 1], 65535).into())); + + let info = info.unwrap(); + assert_eq!(info.addr, ([127, 0, 0, 1], 65535).into()); } #[tokio::test] @@ -273,7 +408,7 @@ mod tests { let extra_data = [0x55; 256]; - let (mut read, addr) = read_proxy_protocol(header.chain(extra_data.as_slice())) + let (mut read, info) = read_proxy_protocol(header.chain(extra_data.as_slice())) .await .unwrap(); @@ -281,9 +416,11 @@ mod tests { read.read_to_end(&mut bytes).await.unwrap(); assert_eq!(bytes, extra_data); + + let info = info.unwrap(); assert_eq!( - addr, - Some(([15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0], 257).into()) + info.addr, + ([15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0], 257).into() ); } @@ -291,30 +428,31 @@ mod tests { async fn test_invalid() { let data = [0x55; 256]; - let (mut read, addr) = read_proxy_protocol(data.as_slice()).await.unwrap(); + let (mut read, info) = read_proxy_protocol(data.as_slice()).await.unwrap(); let mut bytes = vec![]; read.read_to_end(&mut bytes).await.unwrap(); assert_eq!(bytes, data); - assert_eq!(addr, None); + assert_eq!(info, None); } #[tokio::test] async fn test_short() { let data = [0x55; 10]; - let (mut read, addr) = read_proxy_protocol(data.as_slice()).await.unwrap(); + let (mut read, info) = read_proxy_protocol(data.as_slice()).await.unwrap(); let mut bytes = vec![]; read.read_to_end(&mut bytes).await.unwrap(); assert_eq!(bytes, data); - assert_eq!(addr, None); + assert_eq!(info, None); } #[tokio::test] async fn test_large_tlv() { let tlv = vec![0x55; 32768]; - let len = (12 + tlv.len() as u16).to_be_bytes(); + let tlv_len = (tlv.len() as u16).to_be_bytes(); + let len = (12 + 3 + tlv.len() as u16).to_be_bytes(); let header = super::HEADER // Proxy command, Inet << 4 | Stream @@ -330,11 +468,13 @@ mod tests { // dst port .chain([1, 1].as_slice()) // TLV + .chain([255].as_slice()) + .chain(tlv_len.as_slice()) .chain(tlv.as_slice()); let extra_data = [0xaa; 256]; - let (mut read, addr) = read_proxy_protocol(header.chain(extra_data.as_slice())) + let (mut read, info) = read_proxy_protocol(header.chain(extra_data.as_slice())) .await .unwrap(); @@ -342,6 +482,8 @@ mod tests { read.read_to_end(&mut bytes).await.unwrap(); assert_eq!(bytes, extra_data); - assert_eq!(addr, Some(([55, 56, 57, 58], 65535).into())); + + let info = info.unwrap(); + assert_eq!(info.addr, ([55, 56, 57, 58], 65535).into()); } } diff --git a/proxy/src/proxy/mod.rs b/proxy/src/proxy/mod.rs index 2970d93393..922646d889 100644 --- a/proxy/src/proxy/mod.rs +++ b/proxy/src/proxy/mod.rs @@ -28,7 +28,7 @@ use crate::config::{ProxyConfig, ProxyProtocolV2, TlsConfig}; use crate::context::RequestMonitoring; use crate::error::ReportableError; use crate::metrics::{Metrics, NumClientConnectionsGuard}; -use crate::protocol2::read_proxy_protocol; +use crate::protocol2::{read_proxy_protocol, ConnectionInfo}; use crate::proxy::handshake::{handshake, HandshakeData}; use crate::rate_limiter::EndpointRateLimiter; use crate::stream::{PqStream, Stream}; @@ -87,7 +87,7 @@ pub async fn task_main( let endpoint_rate_limiter2 = endpoint_rate_limiter.clone(); connections.spawn(async move { - let (socket, peer_addr) = match read_proxy_protocol(socket).await { + let (socket, conn_info) = match read_proxy_protocol(socket).await { Err(e) => { warn!("per-client task finished with an error: {e:#}"); return; @@ -100,8 +100,8 @@ pub async fn task_main( warn!("proxy protocol header not supported"); return; } - Ok((socket, Some(addr))) => (socket, addr.ip()), - Ok((socket, None)) => (socket, peer_addr.ip()), + Ok((socket, Some(info))) => (socket, info), + Ok((socket, None)) => (socket, ConnectionInfo { addr: peer_addr, extra: None }), }; match socket.inner.set_nodelay(true) { @@ -114,7 +114,7 @@ pub async fn task_main( let ctx = RequestMonitoring::new( session_id, - peer_addr, + conn_info, crate::metrics::Protocol::Tcp, &config.region, ); diff --git a/proxy/src/serverless/mod.rs b/proxy/src/serverless/mod.rs index edbb0347d3..4b60ddf60f 100644 --- a/proxy/src/serverless/mod.rs +++ b/proxy/src/serverless/mod.rs @@ -44,10 +44,10 @@ use tracing::{info, warn, Instrument}; use utils::http::error::ApiError; use crate::cancellation::CancellationHandlerMain; -use crate::config::ProxyConfig; +use crate::config::{ProxyConfig, ProxyProtocolV2}; use crate::context::RequestMonitoring; use crate::metrics::Metrics; -use crate::protocol2::{read_proxy_protocol, ChainRW}; +use crate::protocol2::{read_proxy_protocol, ChainRW, ConnectionInfo}; use crate::proxy::run_until_cancelled; use crate::rate_limiter::EndpointRateLimiter; use crate::serverless::backend::PoolingBackend; @@ -180,7 +180,7 @@ pub async fn task_main( peer_addr, )) .await; - let Some((conn, peer_addr)) = startup_result else { + let Some((conn, conn_info)) = startup_result else { return; }; @@ -192,7 +192,7 @@ pub async fn task_main( endpoint_rate_limiter, conn_token, conn, - peer_addr, + conn_info, session_id, )) .await; @@ -240,7 +240,7 @@ async fn connection_startup( session_id: uuid::Uuid, conn: TcpStream, peer_addr: SocketAddr, -) -> Option<(AsyncRW, IpAddr)> { +) -> Option<(AsyncRW, ConnectionInfo)> { // handle PROXY protocol let (conn, peer) = match read_proxy_protocol(conn).await { Ok(c) => c, @@ -250,17 +250,32 @@ async fn connection_startup( } }; - let peer_addr = peer.unwrap_or(peer_addr).ip(); - let has_private_peer_addr = match peer_addr { + let conn_info = match peer { + None if config.proxy_protocol_v2 == ProxyProtocolV2::Required => { + tracing::warn!("missing required proxy protocol header"); + return None; + } + Some(_) if config.proxy_protocol_v2 == ProxyProtocolV2::Rejected => { + tracing::warn!("proxy protocol header not supported"); + return None; + } + Some(info) => info, + None => ConnectionInfo { + addr: peer_addr, + extra: None, + }, + }; + + let has_private_peer_addr = match conn_info.addr.ip() { IpAddr::V4(ip) => ip.is_private(), IpAddr::V6(_) => false, }; - info!(?session_id, %peer_addr, "accepted new TCP connection"); + info!(?session_id, %conn_info, "accepted new TCP connection"); // try upgrade to TLS, but with a timeout. let conn = match timeout(config.handshake_timeout, tls_acceptor.accept(conn)).await { Ok(Ok(conn)) => { - info!(?session_id, %peer_addr, "accepted new TLS connection"); + info!(?session_id, %conn_info, "accepted new TLS connection"); conn } // The handshake failed @@ -268,7 +283,7 @@ async fn connection_startup( if !has_private_peer_addr { Metrics::get().proxy.tls_handshake_failures.inc(); } - warn!(?session_id, %peer_addr, "failed to accept TLS connection: {e:?}"); + warn!(?session_id, %conn_info, "failed to accept TLS connection: {e:?}"); return None; } // The handshake timed out @@ -276,12 +291,12 @@ async fn connection_startup( if !has_private_peer_addr { Metrics::get().proxy.tls_handshake_failures.inc(); } - warn!(?session_id, %peer_addr, "failed to accept TLS connection: {e:?}"); + warn!(?session_id, %conn_info, "failed to accept TLS connection: {e:?}"); return None; } }; - Some((conn, peer_addr)) + Some((conn, conn_info)) } /// Handles HTTP connection @@ -297,7 +312,7 @@ async fn connection_handler( endpoint_rate_limiter: Arc, cancellation_token: CancellationToken, conn: AsyncRW, - peer_addr: IpAddr, + conn_info: ConnectionInfo, session_id: uuid::Uuid, ) { let session_id = AtomicTake::new(session_id); @@ -306,6 +321,7 @@ async fn connection_handler( let http_cancellation_token = CancellationToken::new(); let _cancel_connection = http_cancellation_token.clone().drop_guard(); + let conn_info2 = conn_info.clone(); let server = Builder::new(TokioExecutor::new()); let conn = server.serve_connection_with_upgrades( hyper_util::rt::TokioIo::new(conn), @@ -340,7 +356,7 @@ async fn connection_handler( connections.clone(), cancellation_handler.clone(), session_id, - peer_addr, + conn_info2.clone(), http_request_token, endpoint_rate_limiter.clone(), ) @@ -365,7 +381,7 @@ async fn connection_handler( // On cancellation, trigger the HTTP connection handler to shut down. let res = match select(pin!(cancellation_token.cancelled()), pin!(conn)).await { Either::Left((_cancelled, mut conn)) => { - tracing::debug!(%peer_addr, "cancelling connection"); + tracing::debug!(%conn_info, "cancelling connection"); conn.as_mut().graceful_shutdown(); conn.await } @@ -373,8 +389,8 @@ async fn connection_handler( }; match res { - Ok(()) => tracing::info!(%peer_addr, "HTTP connection closed"), - Err(e) => tracing::warn!(%peer_addr, "HTTP connection error {e}"), + Ok(()) => tracing::info!(%conn_info, "HTTP connection closed"), + Err(e) => tracing::warn!(%conn_info, "HTTP connection error {e}"), } } @@ -386,7 +402,7 @@ async fn request_handler( ws_connections: TaskTracker, cancellation_handler: Arc, session_id: uuid::Uuid, - peer_addr: IpAddr, + conn_info: ConnectionInfo, // used to cancel in-flight HTTP requests. not used to cancel websockets http_cancellation_token: CancellationToken, endpoint_rate_limiter: Arc, @@ -404,7 +420,7 @@ async fn request_handler( { let ctx = RequestMonitoring::new( session_id, - peer_addr, + conn_info, crate::metrics::Protocol::Ws, &config.region, ); @@ -439,7 +455,7 @@ async fn request_handler( } else if request.uri().path() == "/sql" && *request.method() == Method::POST { let ctx = RequestMonitoring::new( session_id, - peer_addr, + conn_info, crate::metrics::Protocol::Http, &config.region, ); From 0d5a512825705cce3f2c30707e27b3e437a9bc91 Mon Sep 17 00:00:00 2001 From: Erik Grinaker Date: Mon, 4 Nov 2024 16:22:46 +0100 Subject: [PATCH 03/52] safekeeper: add walreceiver metrics (#9450) ## Problem We don't have any observability for Safekeeper WAL receiver queues. ## Summary of changes Adds a few WAL receiver metrics: * `safekeeper_wal_receivers`: gauge of currently connected WAL receivers. * `safekeeper_wal_receiver_queue_depth`: histogram of queue depths per receiver, sampled every 5 seconds. * `safekeeper_wal_receiver_queue_depth_total`: gauge of total queued messages across all receivers. * `safekeeper_wal_receiver_queue_size_total`: gauge of total queued message sizes across all receivers. There are already metrics for ingested WAL volume: `written_wal_bytes` counter per timeline, and `safekeeper_write_wal_bytes` per-request histogram. --- libs/metrics/src/lib.rs | 81 +++++++++++++++++++++++++++++++++++ safekeeper/src/metrics.rs | 52 +++++++++++++++++++--- safekeeper/src/receive_wal.rs | 46 +++++++++++++++++++- safekeeper/src/safekeeper.rs | 64 +++++++++++++++++++++++++++ 4 files changed, 234 insertions(+), 9 deletions(-) diff --git a/libs/metrics/src/lib.rs b/libs/metrics/src/lib.rs index 64e56cb691..0f6c2a0937 100644 --- a/libs/metrics/src/lib.rs +++ b/libs/metrics/src/lib.rs @@ -110,6 +110,23 @@ static MAXRSS_KB: Lazy = Lazy::new(|| { pub const DISK_FSYNC_SECONDS_BUCKETS: &[f64] = &[0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1.0, 5.0, 10.0, 30.0]; +/// Constructs histogram buckets that are powers of two starting at 1 (i.e. 2^0), covering the end +/// points. For example, passing start=5,end=20 yields 4,8,16,32 as does start=4,end=32. +pub fn pow2_buckets(start: usize, end: usize) -> Vec { + assert_ne!(start, 0); + assert!(start <= end); + let start = match start.checked_next_power_of_two() { + Some(n) if n == start => n, // start already power of two + Some(n) => n >> 1, // power of two below start + None => panic!("start too large"), + }; + let end = end.checked_next_power_of_two().expect("end too large"); + std::iter::successors(Some(start), |n| n.checked_mul(2)) + .take_while(|n| n <= &end) + .map(|n| n as f64) + .collect() +} + pub struct BuildInfo { pub revision: &'static str, pub build_tag: &'static str, @@ -595,3 +612,67 @@ where self.dec.collect_into(metadata, labels, name, &mut enc.0) } } + +#[cfg(test)] +mod tests { + use super::*; + + const POW2_BUCKETS_MAX: usize = 1 << (usize::BITS - 1); + + #[test] + fn pow2_buckets_cases() { + assert_eq!(pow2_buckets(1, 1), vec![1.0]); + assert_eq!(pow2_buckets(1, 2), vec![1.0, 2.0]); + assert_eq!(pow2_buckets(1, 3), vec![1.0, 2.0, 4.0]); + assert_eq!(pow2_buckets(1, 4), vec![1.0, 2.0, 4.0]); + assert_eq!(pow2_buckets(1, 5), vec![1.0, 2.0, 4.0, 8.0]); + assert_eq!(pow2_buckets(1, 6), vec![1.0, 2.0, 4.0, 8.0]); + assert_eq!(pow2_buckets(1, 7), vec![1.0, 2.0, 4.0, 8.0]); + assert_eq!(pow2_buckets(1, 8), vec![1.0, 2.0, 4.0, 8.0]); + assert_eq!( + pow2_buckets(1, 200), + vec![1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0, 128.0, 256.0] + ); + + assert_eq!(pow2_buckets(1, 8), vec![1.0, 2.0, 4.0, 8.0]); + assert_eq!(pow2_buckets(2, 8), vec![2.0, 4.0, 8.0]); + assert_eq!(pow2_buckets(3, 8), vec![2.0, 4.0, 8.0]); + assert_eq!(pow2_buckets(4, 8), vec![4.0, 8.0]); + assert_eq!(pow2_buckets(5, 8), vec![4.0, 8.0]); + assert_eq!(pow2_buckets(6, 8), vec![4.0, 8.0]); + assert_eq!(pow2_buckets(7, 8), vec![4.0, 8.0]); + assert_eq!(pow2_buckets(8, 8), vec![8.0]); + assert_eq!(pow2_buckets(20, 200), vec![16.0, 32.0, 64.0, 128.0, 256.0]); + + // Largest valid values. + assert_eq!( + pow2_buckets(1, POW2_BUCKETS_MAX).len(), + usize::BITS as usize + ); + assert_eq!(pow2_buckets(POW2_BUCKETS_MAX, POW2_BUCKETS_MAX).len(), 1); + } + + #[test] + #[should_panic] + fn pow2_buckets_zero_start() { + pow2_buckets(0, 1); + } + + #[test] + #[should_panic] + fn pow2_buckets_end_lt_start() { + pow2_buckets(2, 1); + } + + #[test] + #[should_panic] + fn pow2_buckets_end_overflow_min() { + pow2_buckets(1, POW2_BUCKETS_MAX + 1); + } + + #[test] + #[should_panic] + fn pow2_buckets_end_overflow_max() { + pow2_buckets(1, usize::MAX); + } +} diff --git a/safekeeper/src/metrics.rs b/safekeeper/src/metrics.rs index e8fdddcdc1..bb56e923f8 100644 --- a/safekeeper/src/metrics.rs +++ b/safekeeper/src/metrics.rs @@ -5,23 +5,23 @@ use std::{ time::{Instant, SystemTime}, }; -use ::metrics::{register_histogram, GaugeVec, Histogram, IntGauge, DISK_FSYNC_SECONDS_BUCKETS}; use anyhow::Result; use futures::Future; use metrics::{ core::{AtomicU64, Collector, Desc, GenericCounter, GenericGaugeVec, Opts}, + pow2_buckets, proto::MetricFamily, - register_histogram_vec, register_int_counter, register_int_counter_pair, - register_int_counter_pair_vec, register_int_counter_vec, register_int_gauge, Gauge, - HistogramVec, IntCounter, IntCounterPair, IntCounterPairVec, IntCounterVec, IntGaugeVec, + register_histogram, register_histogram_vec, register_int_counter, register_int_counter_pair, + register_int_counter_pair_vec, register_int_counter_vec, register_int_gauge, Gauge, GaugeVec, + Histogram, HistogramVec, IntCounter, IntCounterPair, IntCounterPairVec, IntCounterVec, + IntGauge, IntGaugeVec, DISK_FSYNC_SECONDS_BUCKETS, }; use once_cell::sync::Lazy; - use postgres_ffi::XLogSegNo; -use utils::pageserver_feedback::PageserverFeedback; -use utils::{id::TenantTimelineId, lsn::Lsn}; +use utils::{id::TenantTimelineId, lsn::Lsn, pageserver_feedback::PageserverFeedback}; use crate::{ + receive_wal::MSG_QUEUE_SIZE, state::{TimelineMemState, TimelinePersistentState}, GlobalTimelines, }; @@ -204,6 +204,44 @@ pub static WAL_BACKUP_TASKS: Lazy = Lazy::new(|| { ) .expect("Failed to register safekeeper_wal_backup_tasks_finished_total counter") }); +pub static WAL_RECEIVERS: Lazy = Lazy::new(|| { + register_int_gauge!( + "safekeeper_wal_receivers", + "Number of currently connected WAL receivers (i.e. connected computes)" + ) + .expect("Failed to register safekeeper_wal_receivers") +}); +pub static WAL_RECEIVER_QUEUE_DEPTH: Lazy = Lazy::new(|| { + // Use powers of two buckets, but add a bucket at 0 and the max queue size to track empty and + // full queues respectively. + let mut buckets = pow2_buckets(1, MSG_QUEUE_SIZE); + buckets.insert(0, 0.0); + buckets.insert(buckets.len() - 1, (MSG_QUEUE_SIZE - 1) as f64); + assert!(buckets.len() <= 12, "too many histogram buckets"); + + register_histogram!( + "safekeeper_wal_receiver_queue_depth", + "Number of queued messages per WAL receiver (sampled every 5 seconds)", + buckets + ) + .expect("Failed to register safekeeper_wal_receiver_queue_depth histogram") +}); +pub static WAL_RECEIVER_QUEUE_DEPTH_TOTAL: Lazy = Lazy::new(|| { + register_int_gauge!( + "safekeeper_wal_receiver_queue_depth_total", + "Total number of queued messages across all WAL receivers", + ) + .expect("Failed to register safekeeper_wal_receiver_queue_depth_total gauge") +}); +// TODO: consider adding a per-receiver queue_size histogram. This will require wrapping the Tokio +// MPSC channel to update counters on send, receive, and drop, while forwarding all other methods. +pub static WAL_RECEIVER_QUEUE_SIZE_TOTAL: Lazy = Lazy::new(|| { + register_int_gauge!( + "safekeeper_wal_receiver_queue_size_total", + "Total memory byte size of queued messages across all WAL receivers", + ) + .expect("Failed to register safekeeper_wal_receiver_queue_size_total gauge") +}); // Metrics collected on operations on the storage repository. #[derive(strum_macros::EnumString, strum_macros::Display, strum_macros::IntoStaticStr)] diff --git a/safekeeper/src/receive_wal.rs b/safekeeper/src/receive_wal.rs index 2410e22f45..a0a96c6e99 100644 --- a/safekeeper/src/receive_wal.rs +++ b/safekeeper/src/receive_wal.rs @@ -3,6 +3,10 @@ //! sends replies back. use crate::handler::SafekeeperPostgresHandler; +use crate::metrics::{ + WAL_RECEIVERS, WAL_RECEIVER_QUEUE_DEPTH, WAL_RECEIVER_QUEUE_DEPTH_TOTAL, + WAL_RECEIVER_QUEUE_SIZE_TOTAL, +}; use crate::safekeeper::AcceptorProposerMessage; use crate::safekeeper::ProposerAcceptorMessage; use crate::safekeeper::ServerInfo; @@ -86,6 +90,7 @@ impl WalReceivers { }; self.update_num(&shared); + WAL_RECEIVERS.inc(); WalReceiverGuard { id: pos, @@ -144,6 +149,7 @@ impl WalReceivers { let mut shared = self.mutex.lock(); shared.slots[id] = None; self.update_num(&shared); + WAL_RECEIVERS.dec(); } /// Broadcast pageserver feedback to connected walproposers. @@ -390,6 +396,7 @@ async fn read_network_loop( loop { let started = Instant::now(); + let size = next_msg.size(); match msg_tx.send_timeout(next_msg, SLOW_THRESHOLD).await { Ok(()) => {} // Slow send, log a message and keep trying. Log context has timeline ID. @@ -409,6 +416,11 @@ async fn read_network_loop( // WalAcceptor terminated. Err(SendTimeoutError::Closed(_)) => return Ok(()), } + + // Update metrics. Will be decremented in WalAcceptor. + WAL_RECEIVER_QUEUE_DEPTH_TOTAL.inc(); + WAL_RECEIVER_QUEUE_SIZE_TOTAL.add(size as i64); + next_msg = read_message(pgb_reader).await?; } } @@ -466,6 +478,12 @@ async fn network_write( /// walproposer, even when it's writing a steady stream of messages. const FLUSH_INTERVAL: Duration = Duration::from_secs(1); +/// The metrics computation interval. +/// +/// The Prometheus poll interval is 60 seconds at the time of writing. We sample the queue depth +/// every 5 seconds, for 12 samples per poll. This will give a count of up to 12x active timelines. +const METRICS_INTERVAL: Duration = Duration::from_secs(5); + /// Encapsulates a task which takes messages from msg_rx, processes and pushes /// replies to reply_tx. /// @@ -512,12 +530,15 @@ impl WalAcceptor { async fn run(&mut self) -> anyhow::Result<()> { let walreceiver_guard = self.tli.get_walreceivers().register(self.conn_id); - // Periodically flush the WAL. + // Periodically flush the WAL and compute metrics. let mut flush_ticker = tokio::time::interval(FLUSH_INTERVAL); flush_ticker.set_missed_tick_behavior(MissedTickBehavior::Delay); flush_ticker.tick().await; // skip the initial, immediate tick - // Tracks unflushed appends. + let mut metrics_ticker = tokio::time::interval(METRICS_INTERVAL); + metrics_ticker.set_missed_tick_behavior(MissedTickBehavior::Skip); + + // Tracks whether we have unflushed appends. let mut dirty = false; loop { @@ -529,6 +550,10 @@ impl WalAcceptor { break; }; + // Update gauge metrics. + WAL_RECEIVER_QUEUE_DEPTH_TOTAL.dec(); + WAL_RECEIVER_QUEUE_SIZE_TOTAL.sub(msg.size() as i64); + // Update walreceiver state in shmem for reporting. if let ProposerAcceptorMessage::Elected(_) = &msg { walreceiver_guard.get().status = WalReceiverStatus::Streaming; @@ -565,6 +590,12 @@ impl WalAcceptor { .process_msg(&ProposerAcceptorMessage::FlushWAL) .await? } + + // Update histogram metrics periodically. + _ = metrics_ticker.tick() => { + WAL_RECEIVER_QUEUE_DEPTH.observe(self.msg_rx.len() as f64); + None // no reply + } }; // Send reply, if any. @@ -585,3 +616,14 @@ impl WalAcceptor { Ok(()) } } + +/// On drop, drain msg_rx and update metrics to avoid leaks. +impl Drop for WalAcceptor { + fn drop(&mut self) { + self.msg_rx.close(); // prevent further sends + while let Ok(msg) = self.msg_rx.try_recv() { + WAL_RECEIVER_QUEUE_DEPTH_TOTAL.dec(); + WAL_RECEIVER_QUEUE_SIZE_TOTAL.sub(msg.size() as i64); + } + } +} diff --git a/safekeeper/src/safekeeper.rs b/safekeeper/src/safekeeper.rs index b3e006ab05..cf41d7a0ab 100644 --- a/safekeeper/src/safekeeper.rs +++ b/safekeeper/src/safekeeper.rs @@ -422,6 +422,70 @@ impl ProposerAcceptorMessage { _ => bail!("unknown proposer-acceptor message tag: {}", tag), } } + + /// The memory size of the message, including byte slices. + pub fn size(&self) -> usize { + const BASE_SIZE: usize = std::mem::size_of::(); + + // For most types, the size is just the base enum size including the nested structs. Some + // types also contain byte slices; add them. + // + // We explicitly list all fields, to draw attention here when new fields are added. + let mut size = BASE_SIZE; + size += match self { + Self::Greeting(ProposerGreeting { + protocol_version: _, + pg_version: _, + proposer_id: _, + system_id: _, + timeline_id: _, + tenant_id: _, + tli: _, + wal_seg_size: _, + }) => 0, + + Self::VoteRequest(VoteRequest { term: _ }) => 0, + + Self::Elected(ProposerElected { + term: _, + start_streaming_at: _, + term_history: _, + timeline_start_lsn: _, + }) => 0, + + Self::AppendRequest(AppendRequest { + h: + AppendRequestHeader { + term: _, + term_start_lsn: _, + begin_lsn: _, + end_lsn: _, + commit_lsn: _, + truncate_lsn: _, + proposer_uuid: _, + }, + wal_data, + }) => wal_data.len(), + + Self::NoFlushAppendRequest(AppendRequest { + h: + AppendRequestHeader { + term: _, + term_start_lsn: _, + begin_lsn: _, + end_lsn: _, + commit_lsn: _, + truncate_lsn: _, + proposer_uuid: _, + }, + wal_data, + }) => wal_data.len(), + + Self::FlushWAL => 0, + }; + + size + } } /// Acceptor -> Proposer messages From 06113e94e6e0108882dd04820a37a83ac68bf59a Mon Sep 17 00:00:00 2001 From: Christian Schwarz Date: Mon, 4 Nov 2024 17:42:08 +0100 Subject: [PATCH 04/52] fix(test_regress): always use storcon virtual pageserver API to set tenant config (#9622) Problem ------- Tests that directly call the Pageserver Management API to set tenant config are flaky if the Pageserver is managed by Storcon because Storcon is the source of truth and may (theoretically) reconcile a tenant at any time. Solution -------- Switch all users of `set_tenant_config`/`patch_tenant_config_client_side` to use the `env.storage_controller.pageserver_api()` Future Work ----------- Prevent regressions from creeping in. And generally clean up up tenant configuration. Maybe we can avoid the Pageserver having a default tenant config at all and put the default into Storcon instead? * => https://github.com/neondatabase/neon/issues/9621 Refs ---- fixes https://github.com/neondatabase/neon/issues/9522 --- test_runner/fixtures/pageserver/http.py | 11 ++++++++++ .../regress/test_attach_tenant_config.py | 20 +++++++++++-------- .../regress/test_disk_usage_eviction.py | 15 ++++++++------ .../regress/test_ingestion_layer_size.py | 2 +- .../regress/test_layers_from_future.py | 2 +- .../test_pageserver_crash_consistency.py | 4 +++- .../test_pageserver_getpage_throttle.py | 6 +++--- test_runner/regress/test_s3_restore.py | 4 +++- .../regress/test_timeline_detach_ancestor.py | 2 +- 9 files changed, 44 insertions(+), 22 deletions(-) diff --git a/test_runner/fixtures/pageserver/http.py b/test_runner/fixtures/pageserver/http.py index 175a1870d4..57a5d6875e 100644 --- a/test_runner/fixtures/pageserver/http.py +++ b/test_runner/fixtures/pageserver/http.py @@ -404,6 +404,12 @@ class PageserverHttpClient(requests.Session, MetricsGetter): return res.json() def set_tenant_config(self, tenant_id: Union[TenantId, TenantShardId], config: dict[str, Any]): + """ + Only use this via storage_controller.pageserver_api(). + + Storcon is the authority on tenant config - changes you make directly + against pageserver may be reconciled away at any time. + """ assert "tenant_id" not in config.keys() res = self.put( f"http://localhost:{self.port}/v1/tenant/config", @@ -417,6 +423,11 @@ class PageserverHttpClient(requests.Session, MetricsGetter): inserts: Optional[dict[str, Any]] = None, removes: Optional[list[str]] = None, ): + """ + Only use this via storage_controller.pageserver_api(). + + See `set_tenant_config` for more information. + """ current = self.tenant_config(tenant_id).tenant_specific_overrides if inserts is not None: current.update(inserts) diff --git a/test_runner/regress/test_attach_tenant_config.py b/test_runner/regress/test_attach_tenant_config.py index 83d003a5cc..64de7626f4 100644 --- a/test_runner/regress/test_attach_tenant_config.py +++ b/test_runner/regress/test_attach_tenant_config.py @@ -176,17 +176,21 @@ def test_fully_custom_config(positive_env: NeonEnv): "lsn_lease_length_for_ts": "5s", } - ps_http = env.pageserver.http_client() + vps_http = env.storage_controller.pageserver_api() - initial_tenant_config = ps_http.tenant_config(env.initial_tenant) - assert initial_tenant_config.tenant_specific_overrides == {} + initial_tenant_config = vps_http.tenant_config(env.initial_tenant) + assert [ + (key, val) + for key, val in initial_tenant_config.tenant_specific_overrides.items() + if val is not None + ] == [] assert set(initial_tenant_config.effective_config.keys()) == set( fully_custom_config.keys() ), "ensure we cover all config options" (tenant_id, _) = env.create_tenant() - ps_http.set_tenant_config(tenant_id, fully_custom_config) - our_tenant_config = ps_http.tenant_config(tenant_id) + vps_http.set_tenant_config(tenant_id, fully_custom_config) + our_tenant_config = vps_http.tenant_config(tenant_id) assert our_tenant_config.tenant_specific_overrides == fully_custom_config assert set(our_tenant_config.effective_config.keys()) == set( fully_custom_config.keys() @@ -199,10 +203,10 @@ def test_fully_custom_config(positive_env: NeonEnv): == {k: True for k in fully_custom_config.keys()} ), "ensure our custom config has different values than the default config for all config options, so we know we overrode everything" - ps_http.tenant_detach(tenant_id) + env.pageserver.tenant_detach(tenant_id) env.pageserver.tenant_attach(tenant_id, config=fully_custom_config) - assert ps_http.tenant_config(tenant_id).tenant_specific_overrides == fully_custom_config - assert set(ps_http.tenant_config(tenant_id).effective_config.keys()) == set( + assert vps_http.tenant_config(tenant_id).tenant_specific_overrides == fully_custom_config + assert set(vps_http.tenant_config(tenant_id).effective_config.keys()) == set( fully_custom_config.keys() ), "ensure we cover all config options" diff --git a/test_runner/regress/test_disk_usage_eviction.py b/test_runner/regress/test_disk_usage_eviction.py index 72866766de..c8d3b2ff3e 100644 --- a/test_runner/regress/test_disk_usage_eviction.py +++ b/test_runner/regress/test_disk_usage_eviction.py @@ -38,21 +38,24 @@ def test_min_resident_size_override_handling( neon_env_builder: NeonEnvBuilder, config_level_override: int ): env = neon_env_builder.init_start() + vps_http = env.storage_controller.pageserver_api() ps_http = env.pageserver.http_client() def assert_config(tenant_id, expect_override, expect_effective): + # talk to actual pageserver to _get_ the config, workaround for + # https://github.com/neondatabase/neon/issues/9621 config = ps_http.tenant_config(tenant_id) assert config.tenant_specific_overrides.get("min_resident_size_override") == expect_override assert config.effective_config.get("min_resident_size_override") == expect_effective def assert_overrides(tenant_id, default_tenant_conf_value): - ps_http.set_tenant_config(tenant_id, {"min_resident_size_override": 200}) + vps_http.set_tenant_config(tenant_id, {"min_resident_size_override": 200}) assert_config(tenant_id, 200, 200) - ps_http.set_tenant_config(tenant_id, {"min_resident_size_override": 0}) + vps_http.set_tenant_config(tenant_id, {"min_resident_size_override": 0}) assert_config(tenant_id, 0, 0) - ps_http.set_tenant_config(tenant_id, {}) + vps_http.set_tenant_config(tenant_id, {}) assert_config(tenant_id, None, default_tenant_conf_value) if config_level_override is not None: @@ -72,7 +75,7 @@ def test_min_resident_size_override_handling( # Also ensure that specifying the paramter to create_tenant works, in addition to http-level recconfig. tenant_id, _ = env.create_tenant(conf={"min_resident_size_override": "100"}) assert_config(tenant_id, 100, 100) - ps_http.set_tenant_config(tenant_id, {}) + vps_http.set_tenant_config(tenant_id, {}) assert_config(tenant_id, None, config_level_override) @@ -457,10 +460,10 @@ def test_pageserver_respects_overridden_resident_size( assert ( du_by_timeline[large_tenant] > min_resident_size ), "ensure the larger tenant will get a haircut" - ps_http.patch_tenant_config_client_side( + env.neon_env.storage_controller.pageserver_api().patch_tenant_config_client_side( small_tenant[0], {"min_resident_size_override": min_resident_size} ) - ps_http.patch_tenant_config_client_side( + env.neon_env.storage_controller.pageserver_api().patch_tenant_config_client_side( large_tenant[0], {"min_resident_size_override": min_resident_size} ) diff --git a/test_runner/regress/test_ingestion_layer_size.py b/test_runner/regress/test_ingestion_layer_size.py index 2edbf4d6d3..646dac8e6e 100644 --- a/test_runner/regress/test_ingestion_layer_size.py +++ b/test_runner/regress/test_ingestion_layer_size.py @@ -81,7 +81,7 @@ def test_ingesting_large_batches_of_images(neon_env_builder: NeonEnvBuilder, bui print_layer_size_histogram(post_ingest) # since all we have are L0s, we should be getting nice L1s and images out of them now - ps_http.patch_tenant_config_client_side( + env.storage_controller.pageserver_api().patch_tenant_config_client_side( env.initial_tenant, { "compaction_threshold": 1, diff --git a/test_runner/regress/test_layers_from_future.py b/test_runner/regress/test_layers_from_future.py index 2536ec1b3c..309e0f3015 100644 --- a/test_runner/regress/test_layers_from_future.py +++ b/test_runner/regress/test_layers_from_future.py @@ -127,7 +127,7 @@ def test_issue_5878(neon_env_builder: NeonEnvBuilder): ), "sanity check for what above loop is supposed to do" # create the image layer from the future - ps_http.patch_tenant_config_client_side( + env.storage_controller.pageserver_api().patch_tenant_config_client_side( tenant_id, {"image_creation_threshold": image_creation_threshold}, None ) assert ps_http.tenant_config(tenant_id).effective_config["image_creation_threshold"] == 1 diff --git a/test_runner/regress/test_pageserver_crash_consistency.py b/test_runner/regress/test_pageserver_crash_consistency.py index ac46d3e62a..fcae7983f4 100644 --- a/test_runner/regress/test_pageserver_crash_consistency.py +++ b/test_runner/regress/test_pageserver_crash_consistency.py @@ -46,7 +46,9 @@ def test_local_only_layers_after_crash(neon_env_builder: NeonEnvBuilder, pg_bin: for sk in env.safekeepers: sk.stop() - pageserver_http.patch_tenant_config_client_side(tenant_id, {"compaction_threshold": 3}) + env.storage_controller.pageserver_api().patch_tenant_config_client_side( + tenant_id, {"compaction_threshold": 3} + ) # hit the exit failpoint with pytest.raises(ConnectionError, match="Remote end closed connection without response"): pageserver_http.timeline_checkpoint(tenant_id, timeline_id) diff --git a/test_runner/regress/test_pageserver_getpage_throttle.py b/test_runner/regress/test_pageserver_getpage_throttle.py index 6811d09cff..f1aad85fe9 100644 --- a/test_runner/regress/test_pageserver_getpage_throttle.py +++ b/test_runner/regress/test_pageserver_getpage_throttle.py @@ -146,13 +146,13 @@ def test_throttle_fair_config_is_settable_but_ignored_in_mgmt_api(neon_env_build To be removed after https://github.com/neondatabase/neon/pull/8539 is rolled out. """ env = neon_env_builder.init_start() - ps_http = env.pageserver.http_client() + vps_http = env.storage_controller.pageserver_api() # with_fair config should still be settable - ps_http.set_tenant_config( + vps_http.set_tenant_config( env.initial_tenant, {"timeline_get_throttle": throttle_config_with_field_fair_set}, ) - conf = ps_http.tenant_config(env.initial_tenant) + conf = vps_http.tenant_config(env.initial_tenant) assert_throttle_config_with_field_fair_set(conf.effective_config["timeline_get_throttle"]) assert_throttle_config_with_field_fair_set( conf.tenant_specific_overrides["timeline_get_throttle"] diff --git a/test_runner/regress/test_s3_restore.py b/test_runner/regress/test_s3_restore.py index bedc9b5865..7a9e6d62b2 100644 --- a/test_runner/regress/test_s3_restore.py +++ b/test_runner/regress/test_s3_restore.py @@ -52,7 +52,9 @@ def test_tenant_s3_restore( tenant_id = env.initial_tenant # now lets create the small layers - ps_http.set_tenant_config(tenant_id, many_small_layers_tenant_config()) + env.storage_controller.pageserver_api().set_tenant_config( + tenant_id, many_small_layers_tenant_config() + ) # Default tenant and the one we created assert ps_http.get_metric_value("pageserver_tenant_manager_slots", {"mode": "attached"}) == 1 diff --git a/test_runner/regress/test_timeline_detach_ancestor.py b/test_runner/regress/test_timeline_detach_ancestor.py index d467c59e62..ed47f9432b 100644 --- a/test_runner/regress/test_timeline_detach_ancestor.py +++ b/test_runner/regress/test_timeline_detach_ancestor.py @@ -511,7 +511,7 @@ def test_compaction_induced_by_detaches_in_history( assert len(delta_layers(branch_timeline_id)) == 5 - client.patch_tenant_config_client_side( + env.storage_controller.pageserver_api().patch_tenant_config_client_side( env.initial_tenant, {"compaction_threshold": 5}, None ) From 81d1bb19410a1c9a5c6bc673059ea5d27103a5be Mon Sep 17 00:00:00 2001 From: Conrad Ludgate Date: Mon, 4 Nov 2024 17:28:10 +0000 Subject: [PATCH 05/52] quieten aws_config logs (#9626) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit logs during aws authentication are soooo noisy in staging 🙃 --- proxy/src/logging.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/proxy/src/logging.rs b/proxy/src/logging.rs index 11921867e4..74d2b9a1d0 100644 --- a/proxy/src/logging.rs +++ b/proxy/src/logging.rs @@ -18,6 +18,7 @@ pub async fn init() -> anyhow::Result { let env_filter = EnvFilter::builder() .with_default_directive(LevelFilter::INFO.into()) .from_env_lossy() + .add_directive("aws_config=info".parse().unwrap()) .add_directive("azure_core::policies::transport=off".parse().unwrap()); let fmt_layer = tracing_subscriber::fmt::layer() From 59879985b4dea9eae57c99e1c07dbbab684be3f4 Mon Sep 17 00:00:00 2001 From: Folke Behrens Date: Mon, 4 Nov 2024 19:56:40 +0100 Subject: [PATCH 06/52] proxy: Wrap JWT errors in separate AuthError variant (#9625) * Also rename `AuthFailed` variant to `PasswordFailed`. * Before this all JWT errors end up in `AuthError::AuthFailed()`, expects a username and also causes cache invalidation. --- proxy/src/auth/backend/classic.rs | 2 +- proxy/src/auth/backend/hacks.rs | 2 +- proxy/src/auth/backend/mod.rs | 4 ++-- proxy/src/auth/mod.rs | 20 +++++++++++++------- proxy/src/serverless/backend.rs | 10 ++++------ 5 files changed, 21 insertions(+), 17 deletions(-) diff --git a/proxy/src/auth/backend/classic.rs b/proxy/src/auth/backend/classic.rs index de32a06e9e..6d26c99832 100644 --- a/proxy/src/auth/backend/classic.rs +++ b/proxy/src/auth/backend/classic.rs @@ -51,7 +51,7 @@ pub(super) async fn authenticate( sasl::Outcome::Success(key) => key, sasl::Outcome::Failure(reason) => { info!("auth backend failed with an error: {reason}"); - return Err(auth::AuthError::auth_failed(&*creds.user)); + return Err(auth::AuthError::password_failed(&*creds.user)); } }; diff --git a/proxy/src/auth/backend/hacks.rs b/proxy/src/auth/backend/hacks.rs index 28bdacd769..1411d908a5 100644 --- a/proxy/src/auth/backend/hacks.rs +++ b/proxy/src/auth/backend/hacks.rs @@ -46,7 +46,7 @@ pub(crate) async fn authenticate_cleartext( sasl::Outcome::Success(key) => key, sasl::Outcome::Failure(reason) => { info!("auth backend failed with an error: {reason}"); - return Err(auth::AuthError::auth_failed(&*info.user)); + return Err(auth::AuthError::password_failed(&*info.user)); } }; diff --git a/proxy/src/auth/backend/mod.rs b/proxy/src/auth/backend/mod.rs index 17334b9cbb..0eb68e6412 100644 --- a/proxy/src/auth/backend/mod.rs +++ b/proxy/src/auth/backend/mod.rs @@ -349,7 +349,7 @@ async fn auth_quirks( { Ok(keys) => Ok(keys), Err(e) => { - if e.is_auth_failed() { + if e.is_password_failed() { // The password could have been changed, so we invalidate the cache. cached_entry.invalidate(); } @@ -376,7 +376,7 @@ async fn authenticate_with_secret( crate::sasl::Outcome::Success(key) => key, crate::sasl::Outcome::Failure(reason) => { info!("auth backend failed with an error: {reason}"); - return Err(auth::AuthError::auth_failed(&*info.user)); + return Err(auth::AuthError::password_failed(&*info.user)); } }; diff --git a/proxy/src/auth/mod.rs b/proxy/src/auth/mod.rs index 7a373dd825..2bd7a2da3d 100644 --- a/proxy/src/auth/mod.rs +++ b/proxy/src/auth/mod.rs @@ -21,6 +21,7 @@ pub(crate) use flow::*; use thiserror::Error; use tokio::time::error::Elapsed; +use crate::auth::backend::jwt::JwtError; use crate::control_plane; use crate::error::{ReportableError, UserFacingError}; @@ -55,7 +56,7 @@ pub(crate) enum AuthError { MissingEndpointName, #[error("password authentication failed for user '{0}'")] - AuthFailed(Box), + PasswordFailed(Box), /// Errors produced by e.g. [`crate::stream::PqStream`]. #[error(transparent)] @@ -76,6 +77,9 @@ pub(crate) enum AuthError { #[error("Disconnected due to inactivity after {0}.")] ConfirmationTimeout(humantime::Duration), + + #[error(transparent)] + Jwt(#[from] JwtError), } impl AuthError { @@ -83,8 +87,8 @@ impl AuthError { AuthError::BadAuthMethod(name.into()) } - pub(crate) fn auth_failed(user: impl Into>) -> Self { - AuthError::AuthFailed(user.into()) + pub(crate) fn password_failed(user: impl Into>) -> Self { + AuthError::PasswordFailed(user.into()) } pub(crate) fn ip_address_not_allowed(ip: IpAddr) -> Self { @@ -95,8 +99,8 @@ impl AuthError { AuthError::TooManyConnections } - pub(crate) fn is_auth_failed(&self) -> bool { - matches!(self, AuthError::AuthFailed(_)) + pub(crate) fn is_password_failed(&self) -> bool { + matches!(self, AuthError::PasswordFailed(_)) } pub(crate) fn user_timeout(elapsed: Elapsed) -> Self { @@ -114,7 +118,7 @@ impl UserFacingError for AuthError { Self::Web(e) => e.to_string_client(), Self::GetAuthInfo(e) => e.to_string_client(), Self::Sasl(e) => e.to_string_client(), - Self::AuthFailed(_) => self.to_string(), + Self::PasswordFailed(_) => self.to_string(), Self::BadAuthMethod(_) => self.to_string(), Self::MalformedPassword(_) => self.to_string(), Self::MissingEndpointName => self.to_string(), @@ -123,6 +127,7 @@ impl UserFacingError for AuthError { Self::TooManyConnections => self.to_string(), Self::UserTimeout(_) => self.to_string(), Self::ConfirmationTimeout(_) => self.to_string(), + Self::Jwt(_) => self.to_string(), } } } @@ -133,7 +138,7 @@ impl ReportableError for AuthError { Self::Web(e) => e.get_error_kind(), Self::GetAuthInfo(e) => e.get_error_kind(), Self::Sasl(e) => e.get_error_kind(), - Self::AuthFailed(_) => crate::error::ErrorKind::User, + Self::PasswordFailed(_) => crate::error::ErrorKind::User, Self::BadAuthMethod(_) => crate::error::ErrorKind::User, Self::MalformedPassword(_) => crate::error::ErrorKind::User, Self::MissingEndpointName => crate::error::ErrorKind::User, @@ -142,6 +147,7 @@ impl ReportableError for AuthError { Self::TooManyConnections => crate::error::ErrorKind::RateLimit, Self::UserTimeout(_) => crate::error::ErrorKind::User, Self::ConfirmationTimeout(_) => crate::error::ErrorKind::User, + Self::Jwt(_) => crate::error::ErrorKind::User, } } } diff --git a/proxy/src/serverless/backend.rs b/proxy/src/serverless/backend.rs index 07e0e30148..c89e0f0232 100644 --- a/proxy/src/serverless/backend.rs +++ b/proxy/src/serverless/backend.rs @@ -81,7 +81,7 @@ impl PoolingBackend { None => { // If we don't have an authentication secret, for the http flow we can just return an error. info!("authentication info not found"); - return Err(AuthError::auth_failed(&*user_info.user)); + return Err(AuthError::password_failed(&*user_info.user)); } }; let ep = EndpointIdInt::from(&user_info.endpoint); @@ -99,7 +99,7 @@ impl PoolingBackend { } crate::sasl::Outcome::Failure(reason) => { info!("auth backend failed with an error: {reason}"); - Err(AuthError::auth_failed(&*user_info.user)) + Err(AuthError::password_failed(&*user_info.user)) } }; res.map(|key| ComputeCredentials { @@ -126,8 +126,7 @@ impl PoolingBackend { &**console, &jwt, ) - .await - .map_err(|e| AuthError::auth_failed(e.to_string()))?; + .await?; Ok(ComputeCredentials { info: user_info.clone(), @@ -146,8 +145,7 @@ impl PoolingBackend { &StaticAuthRules, &jwt, ) - .await - .map_err(|e| AuthError::auth_failed(e.to_string()))?; + .await?; Ok(ComputeCredentials { info: user_info.clone(), From 1085fe57d308b017ed0d079a4890f34c7e9cd52a Mon Sep 17 00:00:00 2001 From: Folke Behrens Date: Mon, 4 Nov 2024 20:19:26 +0100 Subject: [PATCH 07/52] proxy: Rewrite ControlPlaneEvent as enum (#9627) --- proxy/src/cache/endpoints.rs | 110 ++++++++++++++++++----------------- 1 file changed, 58 insertions(+), 52 deletions(-) diff --git a/proxy/src/cache/endpoints.rs b/proxy/src/cache/endpoints.rs index 12c33169bf..400c76291e 100644 --- a/proxy/src/cache/endpoints.rs +++ b/proxy/src/cache/endpoints.rs @@ -1,7 +1,7 @@ use std::convert::Infallible; +use std::future::pending; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; -use std::time::Duration; use dashmap::DashSet; use redis::streams::{StreamReadOptions, StreamReadReply}; @@ -19,25 +19,38 @@ use crate::rate_limiter::GlobalRateLimiter; use crate::redis::connection_with_credentials_provider::ConnectionWithCredentialsProvider; use crate::types::EndpointId; +#[allow(clippy::enum_variant_names)] #[derive(Deserialize, Debug, Clone)] -pub(crate) struct ControlPlaneEventKey { - endpoint_created: Option, - branch_created: Option, - project_created: Option, +#[serde(tag = "type", rename_all(deserialize = "snake_case"))] +enum ControlPlaneEvent { + EndpointCreated { endpoint_created: EndpointCreated }, + BranchCreated { branch_created: BranchCreated }, + ProjectCreated { project_created: ProjectCreated }, } + #[derive(Deserialize, Debug, Clone)] struct EndpointCreated { endpoint_id: String, } + #[derive(Deserialize, Debug, Clone)] struct BranchCreated { branch_id: String, } + #[derive(Deserialize, Debug, Clone)] struct ProjectCreated { project_id: String, } +impl TryFrom<&Value> for ControlPlaneEvent { + type Error = anyhow::Error; + fn try_from(value: &Value) -> Result { + let json = String::from_redis_value(value)?; + Ok(serde_json::from_str(&json)?) + } +} + pub struct EndpointsCache { config: EndpointCacheConfig, endpoints: DashSet, @@ -60,6 +73,7 @@ impl EndpointsCache { ready: AtomicBool::new(false), } } + pub(crate) async fn is_valid(&self, ctx: &RequestMonitoring, endpoint: &EndpointId) -> bool { if !self.ready.load(Ordering::Acquire) { return true; @@ -74,6 +88,7 @@ impl EndpointsCache { } !rejected } + fn should_reject(&self, endpoint: &EndpointId) -> bool { if endpoint.is_endpoint() { !self.endpoints.contains(&EndpointIdInt::from(endpoint)) @@ -87,33 +102,28 @@ impl EndpointsCache { .contains(&ProjectIdInt::from(&endpoint.as_project())) } } - fn insert_event(&self, key: ControlPlaneEventKey) { - // Do not do normalization here, we expect the events to be normalized. - if let Some(endpoint_created) = key.endpoint_created { - self.endpoints - .insert(EndpointIdInt::from(&endpoint_created.endpoint_id.into())); - Metrics::get() - .proxy - .redis_events_count - .inc(RedisEventsCount::EndpointCreated); - } - if let Some(branch_created) = key.branch_created { - self.branches - .insert(BranchIdInt::from(&branch_created.branch_id.into())); - Metrics::get() - .proxy - .redis_events_count - .inc(RedisEventsCount::BranchCreated); - } - if let Some(project_created) = key.project_created { - self.projects - .insert(ProjectIdInt::from(&project_created.project_id.into())); - Metrics::get() - .proxy - .redis_events_count - .inc(RedisEventsCount::ProjectCreated); - } + + fn insert_event(&self, event: ControlPlaneEvent) { + let counter = match event { + ControlPlaneEvent::EndpointCreated { endpoint_created } => { + self.endpoints + .insert(EndpointIdInt::from(&endpoint_created.endpoint_id.into())); + RedisEventsCount::EndpointCreated + } + ControlPlaneEvent::BranchCreated { branch_created } => { + self.branches + .insert(BranchIdInt::from(&branch_created.branch_id.into())); + RedisEventsCount::BranchCreated + } + ControlPlaneEvent::ProjectCreated { project_created } => { + self.projects + .insert(ProjectIdInt::from(&project_created.project_id.into())); + RedisEventsCount::ProjectCreated + } + }; + Metrics::get().proxy.redis_events_count.inc(counter); } + pub async fn do_read( &self, mut con: ConnectionWithCredentialsProvider, @@ -131,12 +141,13 @@ impl EndpointsCache { } if cancellation_token.is_cancelled() { info!("cancellation token is cancelled, exiting"); - tokio::time::sleep(Duration::from_secs(60 * 60 * 24 * 7)).await; - // 1 week. + // Maintenance tasks run forever. Sleep forever when canceled. + pending::<()>().await; } tokio::time::sleep(self.config.retry_interval).await; } } + async fn read_from_stream( &self, con: &mut ConnectionWithCredentialsProvider, @@ -162,10 +173,7 @@ impl EndpointsCache { ) .await } - fn parse_key_value(value: &Value) -> anyhow::Result { - let s: String = FromRedisValue::from_redis_value(value)?; - Ok(serde_json::from_str(&s)?) - } + async fn batch_read( &self, conn: &mut ConnectionWithCredentialsProvider, @@ -196,27 +204,25 @@ impl EndpointsCache { anyhow::bail!("Cannot read from redis stream {}", self.config.stream_name); } - let res = res.keys.pop().expect("Checked length above"); - let len = res.ids.len(); - for x in res.ids { + let key = res.keys.pop().expect("Checked length above"); + let len = key.ids.len(); + for stream_id in key.ids { total += 1; - for (_, v) in x.map { - let key = match Self::parse_key_value(&v) { - Ok(x) => x, - Err(e) => { + for value in stream_id.map.values() { + match value.try_into() { + Ok(event) => self.insert_event(event), + Err(err) => { Metrics::get().proxy.redis_errors_total.inc(RedisErrors { channel: &self.config.stream_name, }); - tracing::error!("error parsing value {v:?}: {e:?}"); - continue; + tracing::error!("error parsing value {value:?}: {err:?}"); } }; - self.insert_event(key); } if total.is_power_of_two() { tracing::debug!("endpoints read {}", total); } - *last_id = x.id; + *last_id = stream_id.id; } if return_when_finish && len <= self.config.default_batch_size { break; @@ -229,11 +235,11 @@ impl EndpointsCache { #[cfg(test)] mod tests { - use super::ControlPlaneEventKey; + use super::ControlPlaneEvent; #[test] - fn test() { - let s = "{\"branch_created\":null,\"endpoint_created\":{\"endpoint_id\":\"ep-rapid-thunder-w0qqw2q9\"},\"project_created\":null,\"type\":\"endpoint_created\"}"; - serde_json::from_str::(s).unwrap(); + fn test_parse_control_plane_event() { + let s = r#"{"branch_created":null,"endpoint_created":{"endpoint_id":"ep-rapid-thunder-w0qqw2q9"},"project_created":null,"type":"endpoint_created"}"#; + serde_json::from_str::(s).unwrap(); } } From ee68bbf6f5053d5ef7beb2fa711974866f66ca48 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arpad=20M=C3=BCller?= Date: Mon, 4 Nov 2024 21:01:18 +0100 Subject: [PATCH 08/52] Add tenant config option to allow timeline_offloading (#9598) Allow us to enable timeline offloading for single tenants without having to enable it for the entire pageserver. Part of #8088. --- control_plane/src/pageserver.rs | 10 ++++++++++ libs/pageserver_api/src/config.rs | 5 +++++ libs/pageserver_api/src/models.rs | 1 + pageserver/src/tenant.rs | 10 +++++++++- pageserver/src/tenant/config.rs | 8 ++++++++ test_runner/regress/test_attach_tenant_config.py | 1 + 6 files changed, 34 insertions(+), 1 deletion(-) diff --git a/control_plane/src/pageserver.rs b/control_plane/src/pageserver.rs index 8df0a714ec..db54965eb5 100644 --- a/control_plane/src/pageserver.rs +++ b/control_plane/src/pageserver.rs @@ -403,6 +403,11 @@ impl PageServerNode { lsn_lease_length_for_ts: settings .remove("lsn_lease_length_for_ts") .map(|x| x.to_string()), + timeline_offloading: settings + .remove("timeline_offloading") + .map(|x| x.parse::()) + .transpose() + .context("Failed to parse 'timeline_offloading' as bool")?, }; if !settings.is_empty() { bail!("Unrecognized tenant settings: {settings:?}") @@ -498,6 +503,11 @@ impl PageServerNode { lsn_lease_length_for_ts: settings .remove("lsn_lease_length_for_ts") .map(|x| x.to_string()), + timeline_offloading: settings + .remove("timeline_offloading") + .map(|x| x.parse::()) + .transpose() + .context("Failed to parse 'timeline_offloading' as bool")?, } }; diff --git a/libs/pageserver_api/src/config.rs b/libs/pageserver_api/src/config.rs index 6b2d6cf625..00cc426c3c 100644 --- a/libs/pageserver_api/src/config.rs +++ b/libs/pageserver_api/src/config.rs @@ -259,6 +259,10 @@ pub struct TenantConfigToml { /// Layers needed to reconstruct pages at LSN will not be GC-ed during this interval. #[serde(with = "humantime_serde")] pub lsn_lease_length_for_ts: Duration, + + /// Enable auto-offloading of timelines. + /// (either this flag or the pageserver-global one need to be set) + pub timeline_offloading: bool, } pub mod defaults { @@ -471,6 +475,7 @@ impl Default for TenantConfigToml { image_layer_creation_check_threshold: DEFAULT_IMAGE_LAYER_CREATION_CHECK_THRESHOLD, lsn_lease_length: LsnLease::DEFAULT_LENGTH, lsn_lease_length_for_ts: LsnLease::DEFAULT_LENGTH_FOR_TS, + timeline_offloading: false, } } } diff --git a/libs/pageserver_api/src/models.rs b/libs/pageserver_api/src/models.rs index 0a4992aea4..0dfa1ba817 100644 --- a/libs/pageserver_api/src/models.rs +++ b/libs/pageserver_api/src/models.rs @@ -310,6 +310,7 @@ pub struct TenantConfig { pub image_layer_creation_check_threshold: Option, pub lsn_lease_length: Option, pub lsn_lease_length_for_ts: Option, + pub timeline_offloading: Option, } /// The policy for the aux file storage. diff --git a/pageserver/src/tenant.rs b/pageserver/src/tenant.rs index 68f8f7e13c..d45c99a41b 100644 --- a/pageserver/src/tenant.rs +++ b/pageserver/src/tenant.rs @@ -2499,8 +2499,15 @@ impl Tenant { .iter() .any(|(_id, tl)| tl.get_ancestor_timeline_id() == Some(*timeline_id)) }; + let config_allows_offload = self.conf.timeline_offloading + || self + .tenant_conf + .load() + .tenant_conf + .timeline_offloading + .unwrap_or_default(); let can_offload = - can_offload && has_no_unoffloaded_children && self.conf.timeline_offloading; + can_offload && has_no_unoffloaded_children && config_allows_offload; if (is_active, can_offload) == (false, false) { None } else { @@ -4902,6 +4909,7 @@ pub(crate) mod harness { ), lsn_lease_length: Some(tenant_conf.lsn_lease_length), lsn_lease_length_for_ts: Some(tenant_conf.lsn_lease_length_for_ts), + timeline_offloading: Some(tenant_conf.timeline_offloading), } } } diff --git a/pageserver/src/tenant/config.rs b/pageserver/src/tenant/config.rs index ce686c89ef..4d6176bfd9 100644 --- a/pageserver/src/tenant/config.rs +++ b/pageserver/src/tenant/config.rs @@ -349,6 +349,10 @@ pub struct TenantConfOpt { #[serde(with = "humantime_serde")] #[serde(default)] pub lsn_lease_length_for_ts: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + #[serde(default)] + pub timeline_offloading: Option, } impl TenantConfOpt { @@ -411,6 +415,9 @@ impl TenantConfOpt { lsn_lease_length_for_ts: self .lsn_lease_length_for_ts .unwrap_or(global_conf.lsn_lease_length_for_ts), + timeline_offloading: self + .lazy_slru_download + .unwrap_or(global_conf.timeline_offloading), } } } @@ -464,6 +471,7 @@ impl From for models::TenantConfig { image_layer_creation_check_threshold: value.image_layer_creation_check_threshold, lsn_lease_length: value.lsn_lease_length.map(humantime), lsn_lease_length_for_ts: value.lsn_lease_length_for_ts.map(humantime), + timeline_offloading: value.timeline_offloading, } } } diff --git a/test_runner/regress/test_attach_tenant_config.py b/test_runner/regress/test_attach_tenant_config.py index 64de7626f4..7d19ba3b5d 100644 --- a/test_runner/regress/test_attach_tenant_config.py +++ b/test_runner/regress/test_attach_tenant_config.py @@ -174,6 +174,7 @@ def test_fully_custom_config(positive_env: NeonEnv): "image_layer_creation_check_threshold": 1, "lsn_lease_length": "1m", "lsn_lease_length_for_ts": "5s", + "timeline_offloading": True, } vps_http = env.storage_controller.pageserver_api() From 34812a6aaba2e2193b617cbb69ee4568aeca8274 Mon Sep 17 00:00:00 2001 From: Tristan Partin Date: Mon, 4 Nov 2024 15:52:01 -0600 Subject: [PATCH 09/52] Improve some typing related to performance testing for LR Signed-off-by: Tristan Partin --- test_runner/fixtures/neon_api.py | 11 +++++--- .../performance/test_logical_replication.py | 27 ++++++++++++------- 2 files changed, 25 insertions(+), 13 deletions(-) diff --git a/test_runner/fixtures/neon_api.py b/test_runner/fixtures/neon_api.py index 5934baccff..89c1f324b4 100644 --- a/test_runner/fixtures/neon_api.py +++ b/test_runner/fixtures/neon_api.py @@ -1,7 +1,7 @@ from __future__ import annotations import time -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING, cast, final import requests @@ -261,17 +261,22 @@ class NeonAPI: time.sleep(0.5) +@final class NeonApiEndpoint: def __init__(self, neon_api: NeonAPI, pg_version: PgVersion, project_id: Optional[str]): self.neon_api = neon_api + self.project_id: str + self.endpoint_id: str + self.connstr: str + if project_id is None: project = neon_api.create_project(pg_version) - neon_api.wait_for_operation_to_finish(project["project"]["id"]) + neon_api.wait_for_operation_to_finish(cast("str", project["project"]["id"])) self.project_id = project["project"]["id"] self.endpoint_id = project["endpoints"][0]["id"] self.connstr = project["connection_uris"][0]["connection_uri"] self.pgbench_env = connection_parameters_to_env( - project["connection_uris"][0]["connection_parameters"] + cast("dict[str, str]", project["connection_uris"][0]["connection_parameters"]) ) self.is_new = True else: diff --git a/test_runner/performance/test_logical_replication.py b/test_runner/performance/test_logical_replication.py index 8b2a296bdd..e62485905e 100644 --- a/test_runner/performance/test_logical_replication.py +++ b/test_runner/performance/test_logical_replication.py @@ -1,7 +1,7 @@ from __future__ import annotations import time -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast import psycopg2 import psycopg2.extras @@ -12,13 +12,16 @@ from fixtures.log_helper import log from fixtures.neon_fixtures import logical_replication_sync if TYPE_CHECKING: + from subprocess import Popen + from typing import AnyStr + from fixtures.benchmark_fixture import NeonBenchmarker from fixtures.neon_api import NeonApiEndpoint - from fixtures.neon_fixtures import NeonEnv, PgBin + from fixtures.neon_fixtures import NeonEnv, PgBin, VanillaPostgres @pytest.mark.timeout(1000) -def test_logical_replication(neon_simple_env: NeonEnv, pg_bin: PgBin, vanilla_pg): +def test_logical_replication(neon_simple_env: NeonEnv, pg_bin: PgBin, vanilla_pg: VanillaPostgres): env = neon_simple_env endpoint = env.endpoints.create_start("main") @@ -47,24 +50,28 @@ def test_logical_replication(neon_simple_env: NeonEnv, pg_bin: PgBin, vanilla_pg logical_replication_sync(vanilla_pg, endpoint) log.info(f"Sync with master took {time.time() - start} seconds") - sum_master = endpoint.safe_psql("select sum(abalance) from pgbench_accounts")[0][0] - sum_replica = vanilla_pg.safe_psql("select sum(abalance) from pgbench_accounts")[0][0] + sum_master = cast("int", endpoint.safe_psql("select sum(abalance) from pgbench_accounts")[0][0]) + sum_replica = cast( + "int", vanilla_pg.safe_psql("select sum(abalance) from pgbench_accounts")[0][0] + ) assert sum_master == sum_replica -def check_pgbench_still_running(pgbench, label=""): +def check_pgbench_still_running(pgbench: Popen[AnyStr], label: str = ""): rc = pgbench.poll() if rc is not None: raise RuntimeError(f"{label} pgbench terminated early with return code {rc}") -def measure_logical_replication_lag(sub_cur, pub_cur, timeout_sec=600): +def measure_logical_replication_lag( + sub_cur: psycopg2.cursor, pub_cur: psycopg2.cursor, timeout_sec: float = 600 +): start = time.time() pub_cur.execute("SELECT pg_current_wal_flush_lsn()") - pub_lsn = Lsn(pub_cur.fetchall()[0][0]) + pub_lsn = Lsn(cast("str", pub_cur.fetchall()[0][0])) while (time.time() - start) < timeout_sec: sub_cur.execute("SELECT latest_end_lsn FROM pg_catalog.pg_stat_subscription") - res = sub_cur.fetchall()[0][0] + res = cast("str", sub_cur.fetchall()[0][0]) if res: log.info(f"subscriber_lsn={res}") sub_lsn = Lsn(res) @@ -286,7 +293,7 @@ def test_snap_files( conn.autocommit = True with conn.cursor() as cur: cur.execute("SELECT rolsuper FROM pg_roles WHERE rolname = 'neondb_owner'") - is_super = cur.fetchall()[0][0] + is_super = cast("bool", cur.fetchall()[0][0]) assert is_super, "This benchmark won't work if we don't have superuser" pg_bin.run_capture(["pgbench", "-i", "-I", "dtGvp", "-s100"], env=env) From 1e16221f82b94fed3c37ce9e95583bb52deaa5d0 Mon Sep 17 00:00:00 2001 From: Tristan Partin Date: Mon, 4 Nov 2024 18:21:59 -0600 Subject: [PATCH 10/52] Update psycopg2 to latest version for complete PG 17 support Update the types to match. Changes the cursor import to match the C bindings[0]. Link: https://github.com/python/typeshed/issues/12578 [0] Signed-off-by: Tristan Partin --- poetry.lock | 155 +++++++++--------- pyproject.toml | 4 +- scripts/download_basebackup.py | 4 +- .../performance/test_logical_replication.py | 5 +- 4 files changed, 80 insertions(+), 88 deletions(-) diff --git a/poetry.lock b/poetry.lock index e06950cb52..d869761e8e 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.8.4 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand. [[package]] name = "aiohappyeyeballs" @@ -2106,83 +2106,78 @@ test = ["enum34", "ipaddress", "mock", "pywin32", "wmi"] [[package]] name = "psycopg2-binary" -version = "2.9.9" +version = "2.9.10" description = "psycopg2 - Python-PostgreSQL Database Adapter" optional = false -python-versions = ">=3.7" +python-versions = ">=3.8" files = [ - {file = "psycopg2-binary-2.9.9.tar.gz", hash = "sha256:7f01846810177d829c7692f1f5ada8096762d9172af1b1a28d4ab5b77c923c1c"}, - {file = "psycopg2_binary-2.9.9-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:c2470da5418b76232f02a2fcd2229537bb2d5a7096674ce61859c3229f2eb202"}, - {file = "psycopg2_binary-2.9.9-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:c6af2a6d4b7ee9615cbb162b0738f6e1fd1f5c3eda7e5da17861eacf4c717ea7"}, - {file = "psycopg2_binary-2.9.9-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:75723c3c0fbbf34350b46a3199eb50638ab22a0228f93fb472ef4d9becc2382b"}, - {file = "psycopg2_binary-2.9.9-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:83791a65b51ad6ee6cf0845634859d69a038ea9b03d7b26e703f94c7e93dbcf9"}, - {file = "psycopg2_binary-2.9.9-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:0ef4854e82c09e84cc63084a9e4ccd6d9b154f1dbdd283efb92ecd0b5e2b8c84"}, - {file = "psycopg2_binary-2.9.9-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ed1184ab8f113e8d660ce49a56390ca181f2981066acc27cf637d5c1e10ce46e"}, - {file = "psycopg2_binary-2.9.9-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:d2997c458c690ec2bc6b0b7ecbafd02b029b7b4283078d3b32a852a7ce3ddd98"}, - {file = "psycopg2_binary-2.9.9-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:b58b4710c7f4161b5e9dcbe73bb7c62d65670a87df7bcce9e1faaad43e715245"}, - {file = "psycopg2_binary-2.9.9-cp310-cp310-musllinux_1_1_ppc64le.whl", hash = "sha256:0c009475ee389757e6e34611d75f6e4f05f0cf5ebb76c6037508318e1a1e0d7e"}, - {file = "psycopg2_binary-2.9.9-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:8dbf6d1bc73f1d04ec1734bae3b4fb0ee3cb2a493d35ede9badbeb901fb40f6f"}, - {file = "psycopg2_binary-2.9.9-cp310-cp310-win32.whl", hash = "sha256:3f78fd71c4f43a13d342be74ebbc0666fe1f555b8837eb113cb7416856c79682"}, - {file = "psycopg2_binary-2.9.9-cp310-cp310-win_amd64.whl", hash = "sha256:876801744b0dee379e4e3c38b76fc89f88834bb15bf92ee07d94acd06ec890a0"}, - {file = "psycopg2_binary-2.9.9-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:ee825e70b1a209475622f7f7b776785bd68f34af6e7a46e2e42f27b659b5bc26"}, - {file = "psycopg2_binary-2.9.9-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:1ea665f8ce695bcc37a90ee52de7a7980be5161375d42a0b6c6abedbf0d81f0f"}, - {file = "psycopg2_binary-2.9.9-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:143072318f793f53819048fdfe30c321890af0c3ec7cb1dfc9cc87aa88241de2"}, - {file = "psycopg2_binary-2.9.9-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c332c8d69fb64979ebf76613c66b985414927a40f8defa16cf1bc028b7b0a7b0"}, - {file = "psycopg2_binary-2.9.9-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:f7fc5a5acafb7d6ccca13bfa8c90f8c51f13d8fb87d95656d3950f0158d3ce53"}, - {file = "psycopg2_binary-2.9.9-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:977646e05232579d2e7b9c59e21dbe5261f403a88417f6a6512e70d3f8a046be"}, - {file = "psycopg2_binary-2.9.9-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:b6356793b84728d9d50ead16ab43c187673831e9d4019013f1402c41b1db9b27"}, - {file = "psycopg2_binary-2.9.9-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:bc7bb56d04601d443f24094e9e31ae6deec9ccb23581f75343feebaf30423359"}, - {file = "psycopg2_binary-2.9.9-cp311-cp311-musllinux_1_1_ppc64le.whl", hash = "sha256:77853062a2c45be16fd6b8d6de2a99278ee1d985a7bd8b103e97e41c034006d2"}, - {file = "psycopg2_binary-2.9.9-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:78151aa3ec21dccd5cdef6c74c3e73386dcdfaf19bced944169697d7ac7482fc"}, - {file = "psycopg2_binary-2.9.9-cp311-cp311-win32.whl", hash = "sha256:dc4926288b2a3e9fd7b50dc6a1909a13bbdadfc67d93f3374d984e56f885579d"}, - {file = "psycopg2_binary-2.9.9-cp311-cp311-win_amd64.whl", hash = "sha256:b76bedd166805480ab069612119ea636f5ab8f8771e640ae103e05a4aae3e417"}, - {file = "psycopg2_binary-2.9.9-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:8532fd6e6e2dc57bcb3bc90b079c60de896d2128c5d9d6f24a63875a95a088cf"}, - {file = "psycopg2_binary-2.9.9-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:b0605eaed3eb239e87df0d5e3c6489daae3f7388d455d0c0b4df899519c6a38d"}, - {file = "psycopg2_binary-2.9.9-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8f8544b092a29a6ddd72f3556a9fcf249ec412e10ad28be6a0c0d948924f2212"}, - {file = "psycopg2_binary-2.9.9-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:2d423c8d8a3c82d08fe8af900ad5b613ce3632a1249fd6a223941d0735fce493"}, - {file = "psycopg2_binary-2.9.9-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:2e5afae772c00980525f6d6ecf7cbca55676296b580c0e6abb407f15f3706996"}, - {file = "psycopg2_binary-2.9.9-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6e6f98446430fdf41bd36d4faa6cb409f5140c1c2cf58ce0bbdaf16af7d3f119"}, - {file = "psycopg2_binary-2.9.9-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:c77e3d1862452565875eb31bdb45ac62502feabbd53429fdc39a1cc341d681ba"}, - {file = "psycopg2_binary-2.9.9-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:cb16c65dcb648d0a43a2521f2f0a2300f40639f6f8c1ecbc662141e4e3e1ee07"}, - {file = "psycopg2_binary-2.9.9-cp312-cp312-musllinux_1_1_ppc64le.whl", hash = "sha256:911dda9c487075abd54e644ccdf5e5c16773470a6a5d3826fda76699410066fb"}, - {file = "psycopg2_binary-2.9.9-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:57fede879f08d23c85140a360c6a77709113efd1c993923c59fde17aa27599fe"}, - {file = "psycopg2_binary-2.9.9-cp312-cp312-win32.whl", hash = "sha256:64cf30263844fa208851ebb13b0732ce674d8ec6a0c86a4e160495d299ba3c93"}, - {file = "psycopg2_binary-2.9.9-cp312-cp312-win_amd64.whl", hash = "sha256:81ff62668af011f9a48787564ab7eded4e9fb17a4a6a74af5ffa6a457400d2ab"}, - {file = "psycopg2_binary-2.9.9-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:2293b001e319ab0d869d660a704942c9e2cce19745262a8aba2115ef41a0a42a"}, - {file = "psycopg2_binary-2.9.9-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:03ef7df18daf2c4c07e2695e8cfd5ee7f748a1d54d802330985a78d2a5a6dca9"}, - {file = "psycopg2_binary-2.9.9-cp37-cp37m-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0a602ea5aff39bb9fac6308e9c9d82b9a35c2bf288e184a816002c9fae930b77"}, - {file = "psycopg2_binary-2.9.9-cp37-cp37m-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:8359bf4791968c5a78c56103702000105501adb557f3cf772b2c207284273984"}, - {file = "psycopg2_binary-2.9.9-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:275ff571376626195ab95a746e6a04c7df8ea34638b99fc11160de91f2fef503"}, - {file = "psycopg2_binary-2.9.9-cp37-cp37m-musllinux_1_1_aarch64.whl", hash = "sha256:f9b5571d33660d5009a8b3c25dc1db560206e2d2f89d3df1cb32d72c0d117d52"}, - {file = "psycopg2_binary-2.9.9-cp37-cp37m-musllinux_1_1_i686.whl", hash = "sha256:420f9bbf47a02616e8554e825208cb947969451978dceb77f95ad09c37791dae"}, - {file = "psycopg2_binary-2.9.9-cp37-cp37m-musllinux_1_1_ppc64le.whl", hash = "sha256:4154ad09dac630a0f13f37b583eae260c6aa885d67dfbccb5b02c33f31a6d420"}, - {file = "psycopg2_binary-2.9.9-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:a148c5d507bb9b4f2030a2025c545fccb0e1ef317393eaba42e7eabd28eb6041"}, - {file = "psycopg2_binary-2.9.9-cp37-cp37m-win32.whl", hash = "sha256:68fc1f1ba168724771e38bee37d940d2865cb0f562380a1fb1ffb428b75cb692"}, - {file = "psycopg2_binary-2.9.9-cp37-cp37m-win_amd64.whl", hash = "sha256:281309265596e388ef483250db3640e5f414168c5a67e9c665cafce9492eda2f"}, - {file = "psycopg2_binary-2.9.9-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:60989127da422b74a04345096c10d416c2b41bd7bf2a380eb541059e4e999980"}, - {file = "psycopg2_binary-2.9.9-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:246b123cc54bb5361588acc54218c8c9fb73068bf227a4a531d8ed56fa3ca7d6"}, - {file = "psycopg2_binary-2.9.9-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:34eccd14566f8fe14b2b95bb13b11572f7c7d5c36da61caf414d23b91fcc5d94"}, - {file = "psycopg2_binary-2.9.9-cp38-cp38-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:18d0ef97766055fec15b5de2c06dd8e7654705ce3e5e5eed3b6651a1d2a9a152"}, - {file = "psycopg2_binary-2.9.9-cp38-cp38-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:d3f82c171b4ccd83bbaf35aa05e44e690113bd4f3b7b6cc54d2219b132f3ae55"}, - {file = "psycopg2_binary-2.9.9-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ead20f7913a9c1e894aebe47cccf9dc834e1618b7aa96155d2091a626e59c972"}, - {file = "psycopg2_binary-2.9.9-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:ca49a8119c6cbd77375ae303b0cfd8c11f011abbbd64601167ecca18a87e7cdd"}, - {file = "psycopg2_binary-2.9.9-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:323ba25b92454adb36fa425dc5cf6f8f19f78948cbad2e7bc6cdf7b0d7982e59"}, - {file = "psycopg2_binary-2.9.9-cp38-cp38-musllinux_1_1_ppc64le.whl", hash = "sha256:1236ed0952fbd919c100bc839eaa4a39ebc397ed1c08a97fc45fee2a595aa1b3"}, - {file = "psycopg2_binary-2.9.9-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:729177eaf0aefca0994ce4cffe96ad3c75e377c7b6f4efa59ebf003b6d398716"}, - {file = "psycopg2_binary-2.9.9-cp38-cp38-win32.whl", hash = "sha256:804d99b24ad523a1fe18cc707bf741670332f7c7412e9d49cb5eab67e886b9b5"}, - {file = "psycopg2_binary-2.9.9-cp38-cp38-win_amd64.whl", hash = "sha256:a6cdcc3ede532f4a4b96000b6362099591ab4a3e913d70bcbac2b56c872446f7"}, - {file = "psycopg2_binary-2.9.9-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:72dffbd8b4194858d0941062a9766f8297e8868e1dd07a7b36212aaa90f49472"}, - {file = "psycopg2_binary-2.9.9-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:30dcc86377618a4c8f3b72418df92e77be4254d8f89f14b8e8f57d6d43603c0f"}, - {file = "psycopg2_binary-2.9.9-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:31a34c508c003a4347d389a9e6fcc2307cc2150eb516462a7a17512130de109e"}, - {file = "psycopg2_binary-2.9.9-cp39-cp39-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:15208be1c50b99203fe88d15695f22a5bed95ab3f84354c494bcb1d08557df67"}, - {file = "psycopg2_binary-2.9.9-cp39-cp39-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1873aade94b74715be2246321c8650cabf5a0d098a95bab81145ffffa4c13876"}, - {file = "psycopg2_binary-2.9.9-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3a58c98a7e9c021f357348867f537017057c2ed7f77337fd914d0bedb35dace7"}, - {file = "psycopg2_binary-2.9.9-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:4686818798f9194d03c9129a4d9a702d9e113a89cb03bffe08c6cf799e053291"}, - {file = "psycopg2_binary-2.9.9-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:ebdc36bea43063116f0486869652cb2ed7032dbc59fbcb4445c4862b5c1ecf7f"}, - {file = "psycopg2_binary-2.9.9-cp39-cp39-musllinux_1_1_ppc64le.whl", hash = "sha256:ca08decd2697fdea0aea364b370b1249d47336aec935f87b8bbfd7da5b2ee9c1"}, - {file = "psycopg2_binary-2.9.9-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:ac05fb791acf5e1a3e39402641827780fe44d27e72567a000412c648a85ba860"}, - {file = "psycopg2_binary-2.9.9-cp39-cp39-win32.whl", hash = "sha256:9dba73be7305b399924709b91682299794887cbbd88e38226ed9f6712eabee90"}, - {file = "psycopg2_binary-2.9.9-cp39-cp39-win_amd64.whl", hash = "sha256:f7ae5d65ccfbebdfa761585228eb4d0df3a8b15cfb53bd953e713e09fbb12957"}, + {file = "psycopg2-binary-2.9.10.tar.gz", hash = "sha256:4b3df0e6990aa98acda57d983942eff13d824135fe2250e6522edaa782a06de2"}, + {file = "psycopg2_binary-2.9.10-cp310-cp310-macosx_12_0_x86_64.whl", hash = "sha256:0ea8e3d0ae83564f2fc554955d327fa081d065c8ca5cc6d2abb643e2c9c1200f"}, + {file = "psycopg2_binary-2.9.10-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:3e9c76f0ac6f92ecfc79516a8034a544926430f7b080ec5a0537bca389ee0906"}, + {file = "psycopg2_binary-2.9.10-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2ad26b467a405c798aaa1458ba09d7e2b6e5f96b1ce0ac15d82fd9f95dc38a92"}, + {file = "psycopg2_binary-2.9.10-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:270934a475a0e4b6925b5f804e3809dd5f90f8613621d062848dd82f9cd62007"}, + {file = "psycopg2_binary-2.9.10-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:48b338f08d93e7be4ab2b5f1dbe69dc5e9ef07170fe1f86514422076d9c010d0"}, + {file = "psycopg2_binary-2.9.10-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7f4152f8f76d2023aac16285576a9ecd2b11a9895373a1f10fd9db54b3ff06b4"}, + {file = "psycopg2_binary-2.9.10-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:32581b3020c72d7a421009ee1c6bf4a131ef5f0a968fab2e2de0c9d2bb4577f1"}, + {file = "psycopg2_binary-2.9.10-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:2ce3e21dc3437b1d960521eca599d57408a695a0d3c26797ea0f72e834c7ffe5"}, + {file = "psycopg2_binary-2.9.10-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:e984839e75e0b60cfe75e351db53d6db750b00de45644c5d1f7ee5d1f34a1ce5"}, + {file = "psycopg2_binary-2.9.10-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:3c4745a90b78e51d9ba06e2088a2fe0c693ae19cc8cb051ccda44e8df8a6eb53"}, + {file = "psycopg2_binary-2.9.10-cp310-cp310-win32.whl", hash = "sha256:e5720a5d25e3b99cd0dc5c8a440570469ff82659bb09431c1439b92caf184d3b"}, + {file = "psycopg2_binary-2.9.10-cp310-cp310-win_amd64.whl", hash = "sha256:3c18f74eb4386bf35e92ab2354a12c17e5eb4d9798e4c0ad3a00783eae7cd9f1"}, + {file = "psycopg2_binary-2.9.10-cp311-cp311-macosx_12_0_x86_64.whl", hash = "sha256:04392983d0bb89a8717772a193cfaac58871321e3ec69514e1c4e0d4957b5aff"}, + {file = "psycopg2_binary-2.9.10-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:1a6784f0ce3fec4edc64e985865c17778514325074adf5ad8f80636cd029ef7c"}, + {file = "psycopg2_binary-2.9.10-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b5f86c56eeb91dc3135b3fd8a95dc7ae14c538a2f3ad77a19645cf55bab1799c"}, + {file = "psycopg2_binary-2.9.10-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:2b3d2491d4d78b6b14f76881905c7a8a8abcf974aad4a8a0b065273a0ed7a2cb"}, + {file = "psycopg2_binary-2.9.10-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:2286791ececda3a723d1910441c793be44625d86d1a4e79942751197f4d30341"}, + {file = "psycopg2_binary-2.9.10-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:512d29bb12608891e349af6a0cccedce51677725a921c07dba6342beaf576f9a"}, + {file = "psycopg2_binary-2.9.10-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:5a507320c58903967ef7384355a4da7ff3f28132d679aeb23572753cbf2ec10b"}, + {file = "psycopg2_binary-2.9.10-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:6d4fa1079cab9018f4d0bd2db307beaa612b0d13ba73b5c6304b9fe2fb441ff7"}, + {file = "psycopg2_binary-2.9.10-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:851485a42dbb0bdc1edcdabdb8557c09c9655dfa2ca0460ff210522e073e319e"}, + {file = "psycopg2_binary-2.9.10-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:35958ec9e46432d9076286dda67942ed6d968b9c3a6a2fd62b48939d1d78bf68"}, + {file = "psycopg2_binary-2.9.10-cp311-cp311-win32.whl", hash = "sha256:ecced182e935529727401b24d76634a357c71c9275b356efafd8a2a91ec07392"}, + {file = "psycopg2_binary-2.9.10-cp311-cp311-win_amd64.whl", hash = "sha256:ee0e8c683a7ff25d23b55b11161c2663d4b099770f6085ff0a20d4505778d6b4"}, + {file = "psycopg2_binary-2.9.10-cp312-cp312-macosx_12_0_x86_64.whl", hash = "sha256:880845dfe1f85d9d5f7c412efea7a08946a46894537e4e5d091732eb1d34d9a0"}, + {file = "psycopg2_binary-2.9.10-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:9440fa522a79356aaa482aa4ba500b65f28e5d0e63b801abf6aa152a29bd842a"}, + {file = "psycopg2_binary-2.9.10-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e3923c1d9870c49a2d44f795df0c889a22380d36ef92440ff618ec315757e539"}, + {file = "psycopg2_binary-2.9.10-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:7b2c956c028ea5de47ff3a8d6b3cc3330ab45cf0b7c3da35a2d6ff8420896526"}, + {file = "psycopg2_binary-2.9.10-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:f758ed67cab30b9a8d2833609513ce4d3bd027641673d4ebc9c067e4d208eec1"}, + {file = "psycopg2_binary-2.9.10-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8cd9b4f2cfab88ed4a9106192de509464b75a906462fb846b936eabe45c2063e"}, + {file = "psycopg2_binary-2.9.10-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:6dc08420625b5a20b53551c50deae6e231e6371194fa0651dbe0fb206452ae1f"}, + {file = "psycopg2_binary-2.9.10-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:d7cd730dfa7c36dbe8724426bf5612798734bff2d3c3857f36f2733f5bfc7c00"}, + {file = "psycopg2_binary-2.9.10-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:155e69561d54d02b3c3209545fb08938e27889ff5a10c19de8d23eb5a41be8a5"}, + {file = "psycopg2_binary-2.9.10-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:c3cc28a6fd5a4a26224007712e79b81dbaee2ffb90ff406256158ec4d7b52b47"}, + {file = "psycopg2_binary-2.9.10-cp312-cp312-win32.whl", hash = "sha256:ec8a77f521a17506a24a5f626cb2aee7850f9b69a0afe704586f63a464f3cd64"}, + {file = "psycopg2_binary-2.9.10-cp312-cp312-win_amd64.whl", hash = "sha256:18c5ee682b9c6dd3696dad6e54cc7ff3a1a9020df6a5c0f861ef8bfd338c3ca0"}, + {file = "psycopg2_binary-2.9.10-cp313-cp313-macosx_12_0_x86_64.whl", hash = "sha256:26540d4a9a4e2b096f1ff9cce51253d0504dca5a85872c7f7be23be5a53eb18d"}, + {file = "psycopg2_binary-2.9.10-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:e217ce4d37667df0bc1c397fdcd8de5e81018ef305aed9415c3b093faaeb10fb"}, + {file = "psycopg2_binary-2.9.10-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:245159e7ab20a71d989da00f280ca57da7641fa2cdcf71749c193cea540a74f7"}, + {file = "psycopg2_binary-2.9.10-cp313-cp313-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:3c4ded1a24b20021ebe677b7b08ad10bf09aac197d6943bfe6fec70ac4e4690d"}, + {file = "psycopg2_binary-2.9.10-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:3abb691ff9e57d4a93355f60d4f4c1dd2d68326c968e7db17ea96df3c023ef73"}, + {file = "psycopg2_binary-2.9.10-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8608c078134f0b3cbd9f89b34bd60a943b23fd33cc5f065e8d5f840061bd0673"}, + {file = "psycopg2_binary-2.9.10-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:230eeae2d71594103cd5b93fd29d1ace6420d0b86f4778739cb1a5a32f607d1f"}, + {file = "psycopg2_binary-2.9.10-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:bb89f0a835bcfc1d42ccd5f41f04870c1b936d8507c6df12b7737febc40f0909"}, + {file = "psycopg2_binary-2.9.10-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:f0c2d907a1e102526dd2986df638343388b94c33860ff3bbe1384130828714b1"}, + {file = "psycopg2_binary-2.9.10-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:f8157bed2f51db683f31306aa497311b560f2265998122abe1dce6428bd86567"}, + {file = "psycopg2_binary-2.9.10-cp38-cp38-macosx_12_0_x86_64.whl", hash = "sha256:eb09aa7f9cecb45027683bb55aebaaf45a0df8bf6de68801a6afdc7947bb09d4"}, + {file = "psycopg2_binary-2.9.10-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b73d6d7f0ccdad7bc43e6d34273f70d587ef62f824d7261c4ae9b8b1b6af90e8"}, + {file = "psycopg2_binary-2.9.10-cp38-cp38-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ce5ab4bf46a211a8e924d307c1b1fcda82368586a19d0a24f8ae166f5c784864"}, + {file = "psycopg2_binary-2.9.10-cp38-cp38-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:056470c3dc57904bbf63d6f534988bafc4e970ffd50f6271fc4ee7daad9498a5"}, + {file = "psycopg2_binary-2.9.10-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:73aa0e31fa4bb82578f3a6c74a73c273367727de397a7a0f07bd83cbea696baa"}, + {file = "psycopg2_binary-2.9.10-cp38-cp38-musllinux_1_2_aarch64.whl", hash = "sha256:8de718c0e1c4b982a54b41779667242bc630b2197948405b7bd8ce16bcecac92"}, + {file = "psycopg2_binary-2.9.10-cp38-cp38-musllinux_1_2_i686.whl", hash = "sha256:5c370b1e4975df846b0277b4deba86419ca77dbc25047f535b0bb03d1a544d44"}, + {file = "psycopg2_binary-2.9.10-cp38-cp38-musllinux_1_2_ppc64le.whl", hash = "sha256:ffe8ed017e4ed70f68b7b371d84b7d4a790368db9203dfc2d222febd3a9c8863"}, + {file = "psycopg2_binary-2.9.10-cp38-cp38-musllinux_1_2_x86_64.whl", hash = "sha256:8aecc5e80c63f7459a1a2ab2c64df952051df196294d9f739933a9f6687e86b3"}, + {file = "psycopg2_binary-2.9.10-cp39-cp39-macosx_12_0_x86_64.whl", hash = "sha256:7a813c8bdbaaaab1f078014b9b0b13f5de757e2b5d9be6403639b298a04d218b"}, + {file = "psycopg2_binary-2.9.10-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d00924255d7fc916ef66e4bf22f354a940c67179ad3fd7067d7a0a9c84d2fbfc"}, + {file = "psycopg2_binary-2.9.10-cp39-cp39-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:7559bce4b505762d737172556a4e6ea8a9998ecac1e39b5233465093e8cee697"}, + {file = "psycopg2_binary-2.9.10-cp39-cp39-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:e8b58f0a96e7a1e341fc894f62c1177a7c83febebb5ff9123b579418fdc8a481"}, + {file = "psycopg2_binary-2.9.10-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6b269105e59ac96aba877c1707c600ae55711d9dcd3fc4b5012e4af68e30c648"}, + {file = "psycopg2_binary-2.9.10-cp39-cp39-musllinux_1_2_aarch64.whl", hash = "sha256:79625966e176dc97ddabc142351e0409e28acf4660b88d1cf6adb876d20c490d"}, + {file = "psycopg2_binary-2.9.10-cp39-cp39-musllinux_1_2_i686.whl", hash = "sha256:8aabf1c1a04584c168984ac678a668094d831f152859d06e055288fa515e4d30"}, + {file = "psycopg2_binary-2.9.10-cp39-cp39-musllinux_1_2_ppc64le.whl", hash = "sha256:19721ac03892001ee8fdd11507e6a2e01f4e37014def96379411ca99d78aeb2c"}, + {file = "psycopg2_binary-2.9.10-cp39-cp39-musllinux_1_2_x86_64.whl", hash = "sha256:7f5d859928e635fa3ce3477704acee0f667b3a3d3e4bb109f2b18d4005f38287"}, + {file = "psycopg2_binary-2.9.10-cp39-cp39-win32.whl", hash = "sha256:3216ccf953b3f267691c90c6fe742e45d890d8272326b4a8b20850a03d05b7b8"}, + {file = "psycopg2_binary-2.9.10-cp39-cp39-win_amd64.whl", hash = "sha256:30e34c4e97964805f715206c7b789d54a78b70f3ff19fbe590104b71c45600e5"}, ] [[package]] @@ -3013,13 +3008,13 @@ files = [ [[package]] name = "types-psycopg2" -version = "2.9.21.10" +version = "2.9.21.20241019" description = "Typing stubs for psycopg2" optional = false -python-versions = "*" +python-versions = ">=3.8" files = [ - {file = "types-psycopg2-2.9.21.10.tar.gz", hash = "sha256:c2600892312ae1c34e12f145749795d93dc4eac3ef7dbf8a9c1bfd45385e80d7"}, - {file = "types_psycopg2-2.9.21.10-py3-none-any.whl", hash = "sha256:918224a0731a3650832e46633e720703b5beef7693a064e777d9748654fcf5e5"}, + {file = "types-psycopg2-2.9.21.20241019.tar.gz", hash = "sha256:bca89b988d2ebd19bcd08b177d22a877ea8b841decb10ed130afcf39404612fa"}, + {file = "types_psycopg2-2.9.21.20241019-py3-none-any.whl", hash = "sha256:44d091e67732d16a941baae48cd7b53bf91911bc36888652447cf1ef0c1fb3f6"}, ] [[package]] @@ -3489,4 +3484,4 @@ cffi = ["cffi (>=1.11)"] [metadata] lock-version = "2.0" python-versions = "^3.9" -content-hash = "13bfc7479aacfe051abb92252b8ddc2e0c429f4607b2d9d8c4b353d2f75c1927" +content-hash = "c656496f9fbb7c29b2df3143c1d72c95b5e121cb6340134c0b8d070f54a08508" diff --git a/pyproject.toml b/pyproject.toml index 92580ee156..9ea42bf46f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,7 +6,7 @@ package-mode = false [tool.poetry.dependencies] python = "^3.9" pytest = "^7.4.4" -psycopg2-binary = "^2.9.9" +psycopg2-binary = "^2.9.10" typing-extensions = "^4.6.1" PyJWT = {version = "^2.1.0", extras = ["crypto"]} requests = "^2.32.3" @@ -15,7 +15,7 @@ asyncpg = "^0.29.0" aiopg = "^1.4.0" Jinja2 = "^3.1.4" types-requests = "^2.31.0.0" -types-psycopg2 = "^2.9.21.10" +types-psycopg2 = "^2.9.21.20241019" boto3 = "^1.34.11" boto3-stubs = {extras = ["s3"], version = "^1.26.16"} moto = {extras = ["server"], version = "^5.0.6"} diff --git a/scripts/download_basebackup.py b/scripts/download_basebackup.py index f00ee87eb7..e23e4f99c3 100755 --- a/scripts/download_basebackup.py +++ b/scripts/download_basebackup.py @@ -23,9 +23,7 @@ def main(args: argparse.Namespace): psconn: PgConnection = psycopg2.connect(pageserver_connstr) psconn.autocommit = True - output = open(output_path, "wb") - - with psconn.cursor() as pscur: + with open(output_path, "wb", encoding="utf-8") as output, psconn.cursor() as pscur: pscur.copy_expert(f"basebackup {tenant_id} {timeline_id} {lsn}", output) diff --git a/test_runner/performance/test_logical_replication.py b/test_runner/performance/test_logical_replication.py index e62485905e..91d7e3446e 100644 --- a/test_runner/performance/test_logical_replication.py +++ b/test_runner/performance/test_logical_replication.py @@ -18,6 +18,7 @@ if TYPE_CHECKING: from fixtures.benchmark_fixture import NeonBenchmarker from fixtures.neon_api import NeonApiEndpoint from fixtures.neon_fixtures import NeonEnv, PgBin, VanillaPostgres + from psycopg2.extensions import cursor @pytest.mark.timeout(1000) @@ -63,9 +64,7 @@ def check_pgbench_still_running(pgbench: Popen[AnyStr], label: str = ""): raise RuntimeError(f"{label} pgbench terminated early with return code {rc}") -def measure_logical_replication_lag( - sub_cur: psycopg2.cursor, pub_cur: psycopg2.cursor, timeout_sec: float = 600 -): +def measure_logical_replication_lag(sub_cur: cursor, pub_cur: cursor, timeout_sec: float = 600): start = time.time() pub_cur.execute("SELECT pg_current_wal_flush_lsn()") pub_lsn = Lsn(cast("str", pub_cur.fetchall()[0][0])) From 8840f3858c829ccfd855c278dde6fdff6c60ff77 Mon Sep 17 00:00:00 2001 From: Erik Grinaker Date: Tue, 5 Nov 2024 13:16:55 +0100 Subject: [PATCH 11/52] pageserver: return 503 during tenant shutdown (#9635) ## Problem Tenant operations may return `409 Conflict` if the tenant is shutting down. This status code is not retried by the control plane, causing user-facing errors during pageserver restarts. Operations should instead return `503 Service Unavailable`, which may be retried for idempotent operations. ## Summary of changes Convert `GetActiveTenantError::WillNotBecomeActive(TenantState::Stopping)` to `ApiError::ShuttingDown` rather than `ApiError::Conflict`. This error is returned by `Tenant::wait_to_become_active` in most (all?) tenant/timeline-related HTTP routes. --- pageserver/src/http/routes.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pageserver/src/http/routes.rs b/pageserver/src/http/routes.rs index ef8efd3f27..72eb3e7ade 100644 --- a/pageserver/src/http/routes.rs +++ b/pageserver/src/http/routes.rs @@ -37,6 +37,7 @@ use pageserver_api::models::TenantShardLocation; use pageserver_api::models::TenantShardSplitRequest; use pageserver_api::models::TenantShardSplitResponse; use pageserver_api::models::TenantSorting; +use pageserver_api::models::TenantState; use pageserver_api::models::TimelineArchivalConfigRequest; use pageserver_api::models::TimelineCreateRequestMode; use pageserver_api::models::TimelinesInfoAndOffloaded; @@ -295,6 +296,9 @@ impl From for ApiError { GetActiveTenantError::Broken(reason) => { ApiError::InternalServerError(anyhow!("tenant is broken: {}", reason)) } + GetActiveTenantError::WillNotBecomeActive(TenantState::Stopping { .. }) => { + ApiError::ShuttingDown + } GetActiveTenantError::WillNotBecomeActive(_) => ApiError::Conflict(format!("{}", e)), GetActiveTenantError::Cancelled => ApiError::ShuttingDown, GetActiveTenantError::NotFound(gte) => gte.into(), From 70ae8c16da2578b6282c1825391bdf4508fa2feb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arpad=20M=C3=BCller?= Date: Tue, 5 Nov 2024 14:02:49 +0100 Subject: [PATCH 12/52] Construct models::TenantConfig only once (#9630) Since 5f83c9290b482dc90006c400dfc68e85a17af785/#1504 we've had duplication in construction of models::TenantConfig, where both constructs contained the same code. This PR removes one of the two locations to avoid the duplication. --- control_plane/src/pageserver.rs | 115 ++++---------------------------- 1 file changed, 14 insertions(+), 101 deletions(-) diff --git a/control_plane/src/pageserver.rs b/control_plane/src/pageserver.rs index db54965eb5..eab76e14c3 100644 --- a/control_plane/src/pageserver.rs +++ b/control_plane/src/pageserver.rs @@ -334,17 +334,20 @@ impl PageServerNode { checkpoint_distance: settings .remove("checkpoint_distance") .map(|x| x.parse::()) - .transpose()?, + .transpose() + .context("Failed to parse 'checkpoint_distance' as an integer")?, checkpoint_timeout: settings.remove("checkpoint_timeout").map(|x| x.to_string()), compaction_target_size: settings .remove("compaction_target_size") .map(|x| x.parse::()) - .transpose()?, + .transpose() + .context("Failed to parse 'compaction_target_size' as an integer")?, compaction_period: settings.remove("compaction_period").map(|x| x.to_string()), compaction_threshold: settings .remove("compaction_threshold") .map(|x| x.parse::()) - .transpose()?, + .transpose() + .context("Failed to parse 'compaction_threshold' as an integer")?, compaction_algorithm: settings .remove("compaction_algorithm") .map(serde_json::from_str) @@ -353,16 +356,19 @@ impl PageServerNode { gc_horizon: settings .remove("gc_horizon") .map(|x| x.parse::()) - .transpose()?, + .transpose() + .context("Failed to parse 'gc_horizon' as an integer")?, gc_period: settings.remove("gc_period").map(|x| x.to_string()), image_creation_threshold: settings .remove("image_creation_threshold") .map(|x| x.parse::()) - .transpose()?, + .transpose() + .context("Failed to parse 'image_creation_threshold' as non zero integer")?, image_layer_creation_check_threshold: settings .remove("image_layer_creation_check_threshold") .map(|x| x.parse::()) - .transpose()?, + .transpose() + .context("Failed to parse 'image_creation_check_threshold' as integer")?, pitr_interval: settings.remove("pitr_interval").map(|x| x.to_string()), walreceiver_connect_timeout: settings .remove("walreceiver_connect_timeout") @@ -419,102 +425,9 @@ impl PageServerNode { pub async fn tenant_config( &self, tenant_id: TenantId, - mut settings: HashMap<&str, &str>, + settings: HashMap<&str, &str>, ) -> anyhow::Result<()> { - let config = { - // Braces to make the diff easier to read - models::TenantConfig { - checkpoint_distance: settings - .remove("checkpoint_distance") - .map(|x| x.parse::()) - .transpose() - .context("Failed to parse 'checkpoint_distance' as an integer")?, - checkpoint_timeout: settings.remove("checkpoint_timeout").map(|x| x.to_string()), - compaction_target_size: settings - .remove("compaction_target_size") - .map(|x| x.parse::()) - .transpose() - .context("Failed to parse 'compaction_target_size' as an integer")?, - compaction_period: settings.remove("compaction_period").map(|x| x.to_string()), - compaction_threshold: settings - .remove("compaction_threshold") - .map(|x| x.parse::()) - .transpose() - .context("Failed to parse 'compaction_threshold' as an integer")?, - compaction_algorithm: settings - .remove("compactin_algorithm") - .map(serde_json::from_str) - .transpose() - .context("Failed to parse 'compaction_algorithm' json")?, - gc_horizon: settings - .remove("gc_horizon") - .map(|x| x.parse::()) - .transpose() - .context("Failed to parse 'gc_horizon' as an integer")?, - gc_period: settings.remove("gc_period").map(|x| x.to_string()), - image_creation_threshold: settings - .remove("image_creation_threshold") - .map(|x| x.parse::()) - .transpose() - .context("Failed to parse 'image_creation_threshold' as non zero integer")?, - image_layer_creation_check_threshold: settings - .remove("image_layer_creation_check_threshold") - .map(|x| x.parse::()) - .transpose() - .context("Failed to parse 'image_creation_check_threshold' as integer")?, - - pitr_interval: settings.remove("pitr_interval").map(|x| x.to_string()), - walreceiver_connect_timeout: settings - .remove("walreceiver_connect_timeout") - .map(|x| x.to_string()), - lagging_wal_timeout: settings - .remove("lagging_wal_timeout") - .map(|x| x.to_string()), - max_lsn_wal_lag: settings - .remove("max_lsn_wal_lag") - .map(|x| x.parse::()) - .transpose() - .context("Failed to parse 'max_lsn_wal_lag' as non zero integer")?, - eviction_policy: settings - .remove("eviction_policy") - .map(serde_json::from_str) - .transpose() - .context("Failed to parse 'eviction_policy' json")?, - min_resident_size_override: settings - .remove("min_resident_size_override") - .map(|x| x.parse::()) - .transpose() - .context("Failed to parse 'min_resident_size_override' as an integer")?, - evictions_low_residence_duration_metric_threshold: settings - .remove("evictions_low_residence_duration_metric_threshold") - .map(|x| x.to_string()), - heatmap_period: settings.remove("heatmap_period").map(|x| x.to_string()), - lazy_slru_download: settings - .remove("lazy_slru_download") - .map(|x| x.parse::()) - .transpose() - .context("Failed to parse 'lazy_slru_download' as bool")?, - timeline_get_throttle: settings - .remove("timeline_get_throttle") - .map(serde_json::from_str) - .transpose() - .context("parse `timeline_get_throttle` from json")?, - lsn_lease_length: settings.remove("lsn_lease_length").map(|x| x.to_string()), - lsn_lease_length_for_ts: settings - .remove("lsn_lease_length_for_ts") - .map(|x| x.to_string()), - timeline_offloading: settings - .remove("timeline_offloading") - .map(|x| x.parse::()) - .transpose() - .context("Failed to parse 'timeline_offloading' as bool")?, - } - }; - - if !settings.is_empty() { - bail!("Unrecognized tenant settings: {settings:?}") - } - + let config = Self::parse_config(settings)?; self.http_client .tenant_config(&models::TenantConfigRequest { tenant_id, config }) .await?; From e30f5fb922c44f47edddf4c0833f9a5fe95f33d0 Mon Sep 17 00:00:00 2001 From: John Spray Date: Tue, 5 Nov 2024 13:32:50 +0000 Subject: [PATCH 13/52] scrubber: remove AWS region assumption, tolerate negative max_project_size (#9636) ## Problem First issues noticed when trying to run scrubber find-garbage on Azure: - Azure staging contains projects with -1 set for max_project_size: apparently the control plane treats this as a signed field. - Scrubber code assumed that listing projects should filter to aws-$REGION. This is no longer needed (per comment in the code) because we know hit region-local APIs. This PR doesn't make it work all the way (`init_remote` still assumes S3), but these are necessary precursors. ## Summary of changes - Change max-project_size from unsigned to signed - Remove region filtering in favor of simply using the right region's API (which we already do) --- storage_scrubber/src/cloud_admin_api.rs | 6 +++--- storage_scrubber/src/garbage.rs | 4 +--- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/storage_scrubber/src/cloud_admin_api.rs b/storage_scrubber/src/cloud_admin_api.rs index 7b82a0b116..c9a62cd256 100644 --- a/storage_scrubber/src/cloud_admin_api.rs +++ b/storage_scrubber/src/cloud_admin_api.rs @@ -147,7 +147,7 @@ pub struct ProjectData { pub created_at: DateTime, pub updated_at: DateTime, pub pg_version: u32, - pub max_project_size: u64, + pub max_project_size: i64, pub remote_storage_size: u64, pub resident_size: u64, pub synthetic_storage_size: u64, @@ -261,7 +261,7 @@ impl CloudAdminApiClient { } } - pub async fn list_projects(&self, region_id: String) -> Result, Error> { + pub async fn list_projects(&self) -> Result, Error> { let _permit = self .request_limiter .acquire() @@ -318,7 +318,7 @@ impl CloudAdminApiClient { pagination_offset += response.data.len(); - result.extend(response.data.drain(..).filter(|t| t.region_id == region_id)); + result.append(&mut response.data); if pagination_offset >= response.total.unwrap_or(0) { break; diff --git a/storage_scrubber/src/garbage.rs b/storage_scrubber/src/garbage.rs index a0040ada08..863dbf960d 100644 --- a/storage_scrubber/src/garbage.rs +++ b/storage_scrubber/src/garbage.rs @@ -160,9 +160,7 @@ async fn find_garbage_inner( // Build a set of console-known tenants, for quickly eliminating known-active tenants without having // to issue O(N) console API requests. let console_projects: HashMap = cloud_admin_api_client - // FIXME: we can't just assume that all console's region ids are aws-. This hack - // will go away when we are talking to Control Plane APIs, which are per-region. - .list_projects(format!("aws-{}", bucket_config.region)) + .list_projects() .await? .into_iter() .map(|t| (t.tenant, t)) From 2f1a56c8f9c6666fb1972636be1f6c92801d2772 Mon Sep 17 00:00:00 2001 From: Ivan Efremov Date: Tue, 5 Nov 2024 17:33:41 +0200 Subject: [PATCH 14/52] proxy: Unify local and remote conn pool client structures (#9604) Unify client, EndpointConnPool and DbUserConnPool for remote and local conn. - Use new ClientDataEnum for additional client data. - Add ClientInnerCommon client structure. - Remove Client and EndpointConnPool code from local_conn_pool.rs --- proxy/src/serverless/backend.rs | 6 +- proxy/src/serverless/conn_pool.rs | 61 ++-- proxy/src/serverless/conn_pool_lib.rs | 289 +++++++++++----- proxy/src/serverless/http_conn_pool.rs | 20 +- proxy/src/serverless/local_conn_pool.rs | 418 ++++++------------------ proxy/src/serverless/sql_over_http.rs | 7 +- 6 files changed, 341 insertions(+), 460 deletions(-) diff --git a/proxy/src/serverless/backend.rs b/proxy/src/serverless/backend.rs index c89e0f0232..c2b0de1876 100644 --- a/proxy/src/serverless/backend.rs +++ b/proxy/src/serverless/backend.rs @@ -14,7 +14,7 @@ use tracing::{debug, info}; use super::conn_pool::poll_client; use super::conn_pool_lib::{Client, ConnInfo, GlobalConnPool}; use super::http_conn_pool::{self, poll_http2_client, Send}; -use super::local_conn_pool::{self, LocalClient, LocalConnPool, EXT_NAME, EXT_SCHEMA, EXT_VERSION}; +use super::local_conn_pool::{self, LocalConnPool, EXT_NAME, EXT_SCHEMA, EXT_VERSION}; use crate::auth::backend::local::StaticAuthRules; use crate::auth::backend::{ComputeCredentials, ComputeUserInfo}; use crate::auth::{self, check_peer_addr_is_in_list, AuthError}; @@ -205,7 +205,7 @@ impl PoolingBackend { conn_info: ConnInfo, ) -> Result, HttpConnError> { info!("pool: looking for an existing connection"); - if let Some(client) = self.http_conn_pool.get(ctx, &conn_info) { + if let Ok(Some(client)) = self.http_conn_pool.get(ctx, &conn_info) { return Ok(client); } @@ -248,7 +248,7 @@ impl PoolingBackend { &self, ctx: &RequestMonitoring, conn_info: ConnInfo, - ) -> Result, HttpConnError> { + ) -> Result, HttpConnError> { if let Some(client) = self.local_pool.get(ctx, &conn_info)? { return Ok(client); } diff --git a/proxy/src/serverless/conn_pool.rs b/proxy/src/serverless/conn_pool.rs index 7fa3357b5b..1845603bf7 100644 --- a/proxy/src/serverless/conn_pool.rs +++ b/proxy/src/serverless/conn_pool.rs @@ -18,7 +18,9 @@ use { std::{sync::atomic, time::Duration}, }; -use super::conn_pool_lib::{Client, ClientInnerExt, ConnInfo, GlobalConnPool}; +use super::conn_pool_lib::{ + Client, ClientDataEnum, ClientInnerCommon, ClientInnerExt, ConnInfo, GlobalConnPool, +}; use crate::context::RequestMonitoring; use crate::control_plane::messages::MetricsAuxInfo; use crate::metrics::Metrics; @@ -152,53 +154,30 @@ pub(crate) fn poll_client( } .instrument(span)); - let inner = ClientInnerRemote { + let inner = ClientInnerCommon { inner: client, - session: tx, - cancel, aux, conn_id, + data: ClientDataEnum::Remote(ClientDataRemote { + session: tx, + cancel, + }), }; + Client::new(inner, conn_info, pool_clone) } -pub(crate) struct ClientInnerRemote { - inner: C, +pub(crate) struct ClientDataRemote { session: tokio::sync::watch::Sender, cancel: CancellationToken, - aux: MetricsAuxInfo, - conn_id: uuid::Uuid, } -impl ClientInnerRemote { - pub(crate) fn inner_mut(&mut self) -> &mut C { - &mut self.inner - } - - pub(crate) fn inner(&self) -> &C { - &self.inner - } - - pub(crate) fn session(&mut self) -> &mut tokio::sync::watch::Sender { +impl ClientDataRemote { + pub fn session(&mut self) -> &mut tokio::sync::watch::Sender { &mut self.session } - pub(crate) fn aux(&self) -> &MetricsAuxInfo { - &self.aux - } - - pub(crate) fn get_conn_id(&self) -> uuid::Uuid { - self.conn_id - } - - pub(crate) fn is_closed(&self) -> bool { - self.inner.is_closed() - } -} - -impl Drop for ClientInnerRemote { - fn drop(&mut self) { - // on client drop, tell the conn to shut down + pub fn cancel(&mut self) { self.cancel.cancel(); } } @@ -228,15 +207,13 @@ mod tests { } } - fn create_inner() -> ClientInnerRemote { + fn create_inner() -> ClientInnerCommon { create_inner_with(MockClient::new(false)) } - fn create_inner_with(client: MockClient) -> ClientInnerRemote { - ClientInnerRemote { + fn create_inner_with(client: MockClient) -> ClientInnerCommon { + ClientInnerCommon { inner: client, - session: tokio::sync::watch::Sender::new(uuid::Uuid::new_v4()), - cancel: CancellationToken::new(), aux: MetricsAuxInfo { endpoint_id: (&EndpointId::from("endpoint")).into(), project_id: (&ProjectId::from("project")).into(), @@ -244,6 +221,10 @@ mod tests { cold_start_info: crate::control_plane::messages::ColdStartInfo::Warm, }, conn_id: uuid::Uuid::new_v4(), + data: ClientDataEnum::Remote(ClientDataRemote { + session: tokio::sync::watch::Sender::new(uuid::Uuid::new_v4()), + cancel: CancellationToken::new(), + }), } } @@ -280,7 +261,7 @@ mod tests { { let mut client = Client::new(create_inner(), conn_info.clone(), ep_pool.clone()); assert_eq!(0, pool.get_global_connections_count()); - client.inner_mut().1.discard(); + client.inner().1.discard(); // Discard should not add the connection from the pool. assert_eq!(0, pool.get_global_connections_count()); } diff --git a/proxy/src/serverless/conn_pool_lib.rs b/proxy/src/serverless/conn_pool_lib.rs index 8830cddf0c..00a8ac4768 100644 --- a/proxy/src/serverless/conn_pool_lib.rs +++ b/proxy/src/serverless/conn_pool_lib.rs @@ -11,10 +11,13 @@ use tokio_postgres::ReadyForQueryStatus; use tracing::{debug, info, Span}; use super::backend::HttpConnError; -use super::conn_pool::ClientInnerRemote; +use super::conn_pool::ClientDataRemote; +use super::http_conn_pool::ClientDataHttp; +use super::local_conn_pool::ClientDataLocal; use crate::auth::backend::ComputeUserInfo; use crate::context::RequestMonitoring; use crate::control_plane::messages::ColdStartInfo; +use crate::control_plane::messages::MetricsAuxInfo; use crate::metrics::{HttpEndpointPoolsGuard, Metrics}; use crate::types::{DbName, EndpointCacheKey, RoleName}; use crate::usage_metrics::{Ids, MetricCounter, USAGE_METRICS}; @@ -41,8 +44,46 @@ impl ConnInfo { } } +pub(crate) enum ClientDataEnum { + Remote(ClientDataRemote), + Local(ClientDataLocal), + #[allow(dead_code)] + Http(ClientDataHttp), +} + +pub(crate) struct ClientInnerCommon { + pub(crate) inner: C, + pub(crate) aux: MetricsAuxInfo, + pub(crate) conn_id: uuid::Uuid, + pub(crate) data: ClientDataEnum, // custom client data like session, key, jti +} + +impl Drop for ClientInnerCommon { + fn drop(&mut self) { + match &mut self.data { + ClientDataEnum::Remote(remote_data) => { + remote_data.cancel(); + } + ClientDataEnum::Local(local_data) => { + local_data.cancel(); + } + ClientDataEnum::Http(_http_data) => (), + } + } +} + +impl ClientInnerCommon { + pub(crate) fn get_conn_id(&self) -> uuid::Uuid { + self.conn_id + } + + pub(crate) fn get_data(&mut self) -> &mut ClientDataEnum { + &mut self.data + } +} + pub(crate) struct ConnPoolEntry { - pub(crate) conn: ClientInnerRemote, + pub(crate) conn: ClientInnerCommon, pub(crate) _last_access: std::time::Instant, } @@ -55,10 +96,33 @@ pub(crate) struct EndpointConnPool { _guard: HttpEndpointPoolsGuard<'static>, global_connections_count: Arc, global_pool_size_max_conns: usize, + pool_name: String, } impl EndpointConnPool { - fn get_conn_entry(&mut self, db_user: (DbName, RoleName)) -> Option> { + pub(crate) fn new( + hmap: HashMap<(DbName, RoleName), DbUserConnPool>, + tconns: usize, + max_conns_per_endpoint: usize, + global_connections_count: Arc, + max_total_conns: usize, + pname: String, + ) -> Self { + Self { + pools: hmap, + total_conns: tconns, + max_conns: max_conns_per_endpoint, + _guard: Metrics::get().proxy.http_endpoint_pools.guard(), + global_connections_count, + global_pool_size_max_conns: max_total_conns, + pool_name: pname, + } + } + + pub(crate) fn get_conn_entry( + &mut self, + db_user: (DbName, RoleName), + ) -> Option> { let Self { pools, total_conns, @@ -84,9 +148,10 @@ impl EndpointConnPool { .. } = self; if let Some(pool) = pools.get_mut(&db_user) { - let old_len = pool.conns.len(); - pool.conns.retain(|conn| conn.conn.get_conn_id() != conn_id); - let new_len = pool.conns.len(); + let old_len = pool.get_conns().len(); + pool.get_conns() + .retain(|conn| conn.conn.get_conn_id() != conn_id); + let new_len = pool.get_conns().len(); let removed = old_len - new_len; if removed > 0 { global_connections_count.fetch_sub(removed, atomic::Ordering::Relaxed); @@ -103,11 +168,26 @@ impl EndpointConnPool { } } - pub(crate) fn put(pool: &RwLock, conn_info: &ConnInfo, client: ClientInnerRemote) { - let conn_id = client.get_conn_id(); + pub(crate) fn get_name(&self) -> &str { + &self.pool_name + } - if client.is_closed() { - info!(%conn_id, "pool: throwing away connection '{conn_info}' because connection is closed"); + pub(crate) fn get_pool(&self, db_user: (DbName, RoleName)) -> Option<&DbUserConnPool> { + self.pools.get(&db_user) + } + + pub(crate) fn get_pool_mut( + &mut self, + db_user: (DbName, RoleName), + ) -> Option<&mut DbUserConnPool> { + self.pools.get_mut(&db_user) + } + + pub(crate) fn put(pool: &RwLock, conn_info: &ConnInfo, client: ClientInnerCommon) { + let conn_id = client.get_conn_id(); + let pool_name = pool.read().get_name().to_string(); + if client.inner.is_closed() { + info!(%conn_id, "{}: throwing away connection '{conn_info}' because connection is closed", pool_name); return; } @@ -118,7 +198,7 @@ impl EndpointConnPool { .load(atomic::Ordering::Relaxed) >= global_max_conn { - info!(%conn_id, "pool: throwing away connection '{conn_info}' because pool is full"); + info!(%conn_id, "{}: throwing away connection '{conn_info}' because pool is full", pool_name); return; } @@ -130,13 +210,13 @@ impl EndpointConnPool { if pool.total_conns < pool.max_conns { let pool_entries = pool.pools.entry(conn_info.db_and_user()).or_default(); - pool_entries.conns.push(ConnPoolEntry { + pool_entries.get_conns().push(ConnPoolEntry { conn: client, _last_access: std::time::Instant::now(), }); returned = true; - per_db_size = pool_entries.conns.len(); + per_db_size = pool_entries.get_conns().len(); pool.total_conns += 1; pool.global_connections_count @@ -153,9 +233,9 @@ impl EndpointConnPool { // do logging outside of the mutex if returned { - info!(%conn_id, "pool: returning connection '{conn_info}' back to the pool, total_conns={total_conns}, for this (db, user)={per_db_size}"); + info!(%conn_id, "{pool_name}: returning connection '{conn_info}' back to the pool, total_conns={total_conns}, for this (db, user)={per_db_size}"); } else { - info!(%conn_id, "pool: throwing away connection '{conn_info}' because pool is full, total_conns={total_conns}"); + info!(%conn_id, "{pool_name}: throwing away connection '{conn_info}' because pool is full, total_conns={total_conns}"); } } } @@ -176,19 +256,39 @@ impl Drop for EndpointConnPool { pub(crate) struct DbUserConnPool { pub(crate) conns: Vec>, + pub(crate) initialized: Option, // a bit ugly, exists only for local pools } impl Default for DbUserConnPool { fn default() -> Self { - Self { conns: Vec::new() } + Self { + conns: Vec::new(), + initialized: None, + } } } -impl DbUserConnPool { +pub(crate) trait DbUserConn: Default { + fn set_initialized(&mut self); + fn is_initialized(&self) -> bool; + fn clear_closed_clients(&mut self, conns: &mut usize) -> usize; + fn get_conn_entry(&mut self, conns: &mut usize) -> (Option>, usize); + fn get_conns(&mut self) -> &mut Vec>; +} + +impl DbUserConn for DbUserConnPool { + fn set_initialized(&mut self) { + self.initialized = Some(true); + } + + fn is_initialized(&self) -> bool { + self.initialized.unwrap_or(false) + } + fn clear_closed_clients(&mut self, conns: &mut usize) -> usize { let old_len = self.conns.len(); - self.conns.retain(|conn| !conn.conn.is_closed()); + self.conns.retain(|conn| !conn.conn.inner.is_closed()); let new_len = self.conns.len(); let removed = old_len - new_len; @@ -196,10 +296,7 @@ impl DbUserConnPool { removed } - pub(crate) fn get_conn_entry( - &mut self, - conns: &mut usize, - ) -> (Option>, usize) { + fn get_conn_entry(&mut self, conns: &mut usize) -> (Option>, usize) { let mut removed = self.clear_closed_clients(conns); let conn = self.conns.pop(); if conn.is_some() { @@ -215,6 +312,10 @@ impl DbUserConnPool { (conn, removed) } + + fn get_conns(&mut self) -> &mut Vec> { + &mut self.conns + } } pub(crate) struct GlobalConnPool { @@ -278,6 +379,60 @@ impl GlobalConnPool { self.config.pool_options.idle_timeout } + pub(crate) fn get( + self: &Arc, + ctx: &RequestMonitoring, + conn_info: &ConnInfo, + ) -> Result>, HttpConnError> { + let mut client: Option> = None; + let Some(endpoint) = conn_info.endpoint_cache_key() else { + return Ok(None); + }; + + let endpoint_pool = self.get_or_create_endpoint_pool(&endpoint); + if let Some(entry) = endpoint_pool + .write() + .get_conn_entry(conn_info.db_and_user()) + { + client = Some(entry.conn); + } + let endpoint_pool = Arc::downgrade(&endpoint_pool); + + // ok return cached connection if found and establish a new one otherwise + if let Some(mut client) = client { + if client.inner.is_closed() { + info!("pool: cached connection '{conn_info}' is closed, opening a new one"); + return Ok(None); + } + tracing::Span::current() + .record("conn_id", tracing::field::display(client.get_conn_id())); + tracing::Span::current().record( + "pid", + tracing::field::display(client.inner.get_process_id()), + ); + info!( + cold_start_info = ColdStartInfo::HttpPoolHit.as_str(), + "pool: reusing connection '{conn_info}'" + ); + + match client.get_data() { + ClientDataEnum::Local(data) => { + data.session().send(ctx.session_id())?; + } + + ClientDataEnum::Remote(data) => { + data.session().send(ctx.session_id())?; + } + ClientDataEnum::Http(_) => (), + } + + ctx.set_cold_start_info(ColdStartInfo::HttpPoolHit); + ctx.success(); + return Ok(Some(Client::new(client, conn_info.clone(), endpoint_pool))); + } + Ok(None) + } + pub(crate) fn shutdown(&self) { // drops all strong references to endpoint-pools self.global_pool.clear(); @@ -374,6 +529,7 @@ impl GlobalConnPool { _guard: Metrics::get().proxy.http_endpoint_pools.guard(), global_connections_count: self.global_connections_count.clone(), global_pool_size_max_conns: self.config.pool_options.max_total_conns, + pool_name: String::from("remote"), })); // find or create a pool for this endpoint @@ -400,55 +556,23 @@ impl GlobalConnPool { pool } +} - pub(crate) fn get( - self: &Arc, - ctx: &RequestMonitoring, - conn_info: &ConnInfo, - ) -> Result>, HttpConnError> { - let mut client: Option> = None; - let Some(endpoint) = conn_info.endpoint_cache_key() else { - return Ok(None); - }; +pub(crate) struct Client { + span: Span, + inner: Option>, + conn_info: ConnInfo, + pool: Weak>>, +} - let endpoint_pool = self.get_or_create_endpoint_pool(&endpoint); - if let Some(entry) = endpoint_pool - .write() - .get_conn_entry(conn_info.db_and_user()) - { - client = Some(entry.conn); - } - let endpoint_pool = Arc::downgrade(&endpoint_pool); - - // ok return cached connection if found and establish a new one otherwise - if let Some(mut client) = client { - if client.is_closed() { - info!("pool: cached connection '{conn_info}' is closed, opening a new one"); - return Ok(None); - } - tracing::Span::current() - .record("conn_id", tracing::field::display(client.get_conn_id())); - tracing::Span::current().record( - "pid", - tracing::field::display(client.inner().get_process_id()), - ); - info!( - cold_start_info = ColdStartInfo::HttpPoolHit.as_str(), - "pool: reusing connection '{conn_info}'" - ); - - client.session().send(ctx.session_id())?; - ctx.set_cold_start_info(ColdStartInfo::HttpPoolHit); - ctx.success(); - return Ok(Some(Client::new(client, conn_info.clone(), endpoint_pool))); - } - Ok(None) - } +pub(crate) struct Discard<'a, C: ClientInnerExt> { + conn_info: &'a ConnInfo, + pool: &'a mut Weak>>, } impl Client { pub(crate) fn new( - inner: ClientInnerRemote, + inner: ClientInnerCommon, conn_info: ConnInfo, pool: Weak>>, ) -> Self { @@ -460,7 +584,18 @@ impl Client { } } - pub(crate) fn inner_mut(&mut self) -> (&mut C, Discard<'_, C>) { + pub(crate) fn client_inner(&mut self) -> (&mut ClientInnerCommon, Discard<'_, C>) { + let Self { + inner, + pool, + conn_info, + span: _, + } = self; + let inner_m = inner.as_mut().expect("client inner should not be removed"); + (inner_m, Discard { conn_info, pool }) + } + + pub(crate) fn inner(&mut self) -> (&mut C, Discard<'_, C>) { let Self { inner, pool, @@ -468,12 +603,11 @@ impl Client { span: _, } = self; let inner = inner.as_mut().expect("client inner should not be removed"); - let inner_ref = inner.inner_mut(); - (inner_ref, Discard { conn_info, pool }) + (&mut inner.inner, Discard { conn_info, pool }) } pub(crate) fn metrics(&self) -> Arc { - let aux = &self.inner.as_ref().unwrap().aux(); + let aux = &self.inner.as_ref().unwrap().aux; USAGE_METRICS.register(Ids { endpoint_id: aux.endpoint_id, branch_id: aux.branch_id, @@ -498,13 +632,6 @@ impl Client { } } -pub(crate) struct Client { - span: Span, - inner: Option>, - conn_info: ConnInfo, - pool: Weak>>, -} - impl Drop for Client { fn drop(&mut self) { if let Some(drop) = self.do_drop() { @@ -517,10 +644,11 @@ impl Deref for Client { type Target = C; fn deref(&self) -> &Self::Target { - self.inner + &self + .inner .as_ref() .expect("client inner should not be removed") - .inner() + .inner } } @@ -539,11 +667,6 @@ impl ClientInnerExt for tokio_postgres::Client { } } -pub(crate) struct Discard<'a, C: ClientInnerExt> { - conn_info: &'a ConnInfo, - pool: &'a mut Weak>>, -} - impl Discard<'_, C> { pub(crate) fn check_idle(&mut self, status: ReadyForQueryStatus) { let conn_info = &self.conn_info; diff --git a/proxy/src/serverless/http_conn_pool.rs b/proxy/src/serverless/http_conn_pool.rs index b92ae31310..56be70abec 100644 --- a/proxy/src/serverless/http_conn_pool.rs +++ b/proxy/src/serverless/http_conn_pool.rs @@ -7,9 +7,11 @@ use hyper::client::conn::http2; use hyper_util::rt::{TokioExecutor, TokioIo}; use parking_lot::RwLock; use rand::Rng; +use std::result::Result::Ok; use tokio::net::TcpStream; use tracing::{debug, error, info, info_span, Instrument}; +use super::backend::HttpConnError; use super::conn_pool_lib::{ClientInnerExt, ConnInfo}; use crate::context::RequestMonitoring; use crate::control_plane::messages::{ColdStartInfo, MetricsAuxInfo}; @@ -28,6 +30,8 @@ pub(crate) struct ConnPoolEntry { aux: MetricsAuxInfo, } +pub(crate) struct ClientDataHttp(); + // Per-endpoint connection pool // Number of open connections is limited by the `max_conns_per_endpoint`. pub(crate) struct EndpointConnPool { @@ -206,14 +210,22 @@ impl GlobalConnPool { } } + #[expect(unused_results)] pub(crate) fn get( self: &Arc, ctx: &RequestMonitoring, conn_info: &ConnInfo, - ) -> Option> { - let endpoint = conn_info.endpoint_cache_key()?; + ) -> Result>, HttpConnError> { + let result: Result>, HttpConnError>; + let Some(endpoint) = conn_info.endpoint_cache_key() else { + result = Ok(None); + return result; + }; let endpoint_pool = self.get_or_create_endpoint_pool(&endpoint); - let client = endpoint_pool.write().get_conn_entry()?; + let Some(client) = endpoint_pool.write().get_conn_entry() else { + result = Ok(None); + return result; + }; tracing::Span::current().record("conn_id", tracing::field::display(client.conn_id)); info!( @@ -222,7 +234,7 @@ impl GlobalConnPool { ); ctx.set_cold_start_info(ColdStartInfo::HttpPoolHit); ctx.success(); - Some(Client::new(client.conn, client.aux)) + Ok(Some(Client::new(client.conn, client.aux))) } fn get_or_create_endpoint_pool( diff --git a/proxy/src/serverless/local_conn_pool.rs b/proxy/src/serverless/local_conn_pool.rs index 064e7db7b3..99d4329f88 100644 --- a/proxy/src/serverless/local_conn_pool.rs +++ b/proxy/src/serverless/local_conn_pool.rs @@ -11,7 +11,8 @@ use std::collections::HashMap; use std::pin::pin; -use std::sync::{Arc, Weak}; +use std::sync::atomic::AtomicUsize; +use std::sync::Arc; use std::task::{ready, Poll}; use std::time::Duration; @@ -26,177 +27,42 @@ use signature::Signer; use tokio::time::Instant; use tokio_postgres::tls::NoTlsStream; use tokio_postgres::types::ToSql; -use tokio_postgres::{AsyncMessage, ReadyForQueryStatus, Socket}; +use tokio_postgres::{AsyncMessage, Socket}; use tokio_util::sync::CancellationToken; -use tracing::{error, info, info_span, warn, Instrument, Span}; +use tracing::{error, info, info_span, warn, Instrument}; use super::backend::HttpConnError; -use super::conn_pool_lib::{ClientInnerExt, ConnInfo}; +use super::conn_pool_lib::{ + Client, ClientDataEnum, ClientInnerCommon, ClientInnerExt, ConnInfo, DbUserConn, + EndpointConnPool, +}; use crate::context::RequestMonitoring; use crate::control_plane::messages::{ColdStartInfo, MetricsAuxInfo}; use crate::metrics::Metrics; -use crate::types::{DbName, RoleName}; -use crate::usage_metrics::{Ids, MetricCounter, USAGE_METRICS}; pub(crate) const EXT_NAME: &str = "pg_session_jwt"; pub(crate) const EXT_VERSION: &str = "0.1.2"; pub(crate) const EXT_SCHEMA: &str = "auth"; -struct ConnPoolEntry { - conn: ClientInner, - _last_access: std::time::Instant, +pub(crate) struct ClientDataLocal { + session: tokio::sync::watch::Sender, + cancel: CancellationToken, + key: SigningKey, + jti: u64, } -// Per-endpoint connection pool, (dbname, username) -> DbUserConnPool -// Number of open connections is limited by the `max_conns_per_endpoint`. -pub(crate) struct EndpointConnPool { - pools: HashMap<(DbName, RoleName), DbUserConnPool>, - total_conns: usize, - max_conns: usize, - global_pool_size_max_conns: usize, -} - -impl EndpointConnPool { - fn get_conn_entry(&mut self, db_user: (DbName, RoleName)) -> Option> { - let Self { - pools, total_conns, .. - } = self; - pools - .get_mut(&db_user) - .and_then(|pool_entries| pool_entries.get_conn_entry(total_conns)) +impl ClientDataLocal { + pub fn session(&mut self) -> &mut tokio::sync::watch::Sender { + &mut self.session } - fn remove_client(&mut self, db_user: (DbName, RoleName), conn_id: uuid::Uuid) -> bool { - let Self { - pools, total_conns, .. - } = self; - if let Some(pool) = pools.get_mut(&db_user) { - let old_len = pool.conns.len(); - pool.conns.retain(|conn| conn.conn.conn_id != conn_id); - let new_len = pool.conns.len(); - let removed = old_len - new_len; - if removed > 0 { - Metrics::get() - .proxy - .http_pool_opened_connections - .get_metric() - .dec_by(removed as i64); - } - *total_conns -= removed; - removed > 0 - } else { - false - } - } - - fn put(pool: &RwLock, conn_info: &ConnInfo, client: ClientInner) { - let conn_id = client.conn_id; - - if client.is_closed() { - info!(%conn_id, "local_pool: throwing away connection '{conn_info}' because connection is closed"); - return; - } - let global_max_conn = pool.read().global_pool_size_max_conns; - if pool.read().total_conns >= global_max_conn { - info!(%conn_id, "local_pool: throwing away connection '{conn_info}' because pool is full"); - return; - } - - // return connection to the pool - let mut returned = false; - let mut per_db_size = 0; - let total_conns = { - let mut pool = pool.write(); - - if pool.total_conns < pool.max_conns { - let pool_entries = pool.pools.entry(conn_info.db_and_user()).or_default(); - pool_entries.conns.push(ConnPoolEntry { - conn: client, - _last_access: std::time::Instant::now(), - }); - - returned = true; - per_db_size = pool_entries.conns.len(); - - pool.total_conns += 1; - Metrics::get() - .proxy - .http_pool_opened_connections - .get_metric() - .inc(); - } - - pool.total_conns - }; - - // do logging outside of the mutex - if returned { - info!(%conn_id, "local_pool: returning connection '{conn_info}' back to the pool, total_conns={total_conns}, for this (db, user)={per_db_size}"); - } else { - info!(%conn_id, "local_pool: throwing away connection '{conn_info}' because pool is full, total_conns={total_conns}"); - } - } -} - -impl Drop for EndpointConnPool { - fn drop(&mut self) { - if self.total_conns > 0 { - Metrics::get() - .proxy - .http_pool_opened_connections - .get_metric() - .dec_by(self.total_conns as i64); - } - } -} - -pub(crate) struct DbUserConnPool { - conns: Vec>, - - // true if we have definitely installed the extension and - // granted the role access to the auth schema. - initialized: bool, -} - -impl Default for DbUserConnPool { - fn default() -> Self { - Self { - conns: Vec::new(), - initialized: false, - } - } -} - -impl DbUserConnPool { - fn clear_closed_clients(&mut self, conns: &mut usize) -> usize { - let old_len = self.conns.len(); - - self.conns.retain(|conn| !conn.conn.is_closed()); - - let new_len = self.conns.len(); - let removed = old_len - new_len; - *conns -= removed; - removed - } - - fn get_conn_entry(&mut self, conns: &mut usize) -> Option> { - let mut removed = self.clear_closed_clients(conns); - let conn = self.conns.pop(); - if conn.is_some() { - *conns -= 1; - removed += 1; - } - Metrics::get() - .proxy - .http_pool_opened_connections - .get_metric() - .dec_by(removed as i64); - conn + pub fn cancel(&mut self) { + self.cancel.cancel(); } } pub(crate) struct LocalConnPool { - global_pool: RwLock>, + global_pool: Arc>>, config: &'static crate::config::HttpConfig, } @@ -204,12 +70,14 @@ pub(crate) struct LocalConnPool { impl LocalConnPool { pub(crate) fn new(config: &'static crate::config::HttpConfig) -> Arc { Arc::new(Self { - global_pool: RwLock::new(EndpointConnPool { - pools: HashMap::new(), - total_conns: 0, - max_conns: config.pool_options.max_conns_per_endpoint, - global_pool_size_max_conns: config.pool_options.max_total_conns, - }), + global_pool: Arc::new(RwLock::new(EndpointConnPool::new( + HashMap::new(), + 0, + config.pool_options.max_conns_per_endpoint, + Arc::new(AtomicUsize::new(0)), + config.pool_options.max_total_conns, + String::from("local_pool"), + ))), config, }) } @@ -222,7 +90,7 @@ impl LocalConnPool { self: &Arc, ctx: &RequestMonitoring, conn_info: &ConnInfo, - ) -> Result>, HttpConnError> { + ) -> Result>, HttpConnError> { let client = self .global_pool .write() @@ -230,12 +98,14 @@ impl LocalConnPool { .map(|entry| entry.conn); // ok return cached connection if found and establish a new one otherwise - if let Some(client) = client { - if client.is_closed() { + if let Some(mut client) = client { + if client.inner.is_closed() { info!("local_pool: cached connection '{conn_info}' is closed, opening a new one"); return Ok(None); } - tracing::Span::current().record("conn_id", tracing::field::display(client.conn_id)); + + tracing::Span::current() + .record("conn_id", tracing::field::display(client.get_conn_id())); tracing::Span::current().record( "pid", tracing::field::display(client.inner.get_process_id()), @@ -244,47 +114,59 @@ impl LocalConnPool { cold_start_info = ColdStartInfo::HttpPoolHit.as_str(), "local_pool: reusing connection '{conn_info}'" ); - client.session.send(ctx.session_id())?; + + match client.get_data() { + ClientDataEnum::Local(data) => { + data.session().send(ctx.session_id())?; + } + + ClientDataEnum::Remote(data) => { + data.session().send(ctx.session_id())?; + } + ClientDataEnum::Http(_) => (), + } + ctx.set_cold_start_info(ColdStartInfo::HttpPoolHit); ctx.success(); - return Ok(Some(LocalClient::new( + + return Ok(Some(Client::new( client, conn_info.clone(), - Arc::downgrade(self), + Arc::downgrade(&self.global_pool), ))); } Ok(None) } pub(crate) fn initialized(self: &Arc, conn_info: &ConnInfo) -> bool { - self.global_pool - .read() - .pools - .get(&conn_info.db_and_user()) - .map_or(false, |pool| pool.initialized) + if let Some(pool) = self.global_pool.read().get_pool(conn_info.db_and_user()) { + return pool.is_initialized(); + } + false } pub(crate) fn set_initialized(self: &Arc, conn_info: &ConnInfo) { - self.global_pool + if let Some(pool) = self + .global_pool .write() - .pools - .entry(conn_info.db_and_user()) - .or_default() - .initialized = true; + .get_pool_mut(conn_info.db_and_user()) + { + pool.set_initialized(); + } } } #[allow(clippy::too_many_arguments)] -pub(crate) fn poll_client( - global_pool: Arc>, +pub(crate) fn poll_client( + global_pool: Arc>, ctx: &RequestMonitoring, conn_info: ConnInfo, - client: tokio_postgres::Client, + client: C, mut connection: tokio_postgres::Connection, key: SigningKey, conn_id: uuid::Uuid, aux: MetricsAuxInfo, -) -> LocalClient { +) -> Client { let conn_gauge = Metrics::get().proxy.db_connections.guard(ctx.protocol()); let mut session_id = ctx.session_id(); let (tx, mut rx) = tokio::sync::watch::channel(session_id); @@ -377,111 +259,47 @@ pub(crate) fn poll_client( } .instrument(span)); - let inner = ClientInner { + let inner = ClientInnerCommon { inner: client, - session: tx, - cancel, aux, conn_id, - key, - jti: 0, + data: ClientDataEnum::Local(ClientDataLocal { + session: tx, + cancel, + key, + jti: 0, + }), }; - LocalClient::new(inner, conn_info, pool_clone) + + Client::new( + inner, + conn_info, + Arc::downgrade(&pool_clone.upgrade().unwrap().global_pool), + ) } -pub(crate) struct ClientInner { - inner: C, - session: tokio::sync::watch::Sender, - cancel: CancellationToken, - aux: MetricsAuxInfo, - conn_id: uuid::Uuid, - - // needed for pg_session_jwt state - key: SigningKey, - jti: u64, -} - -impl Drop for ClientInner { - fn drop(&mut self) { - // on client drop, tell the conn to shut down - self.cancel.cancel(); - } -} - -impl ClientInner { - pub(crate) fn is_closed(&self) -> bool { - self.inner.is_closed() - } -} - -impl ClientInner { +impl ClientInnerCommon { pub(crate) async fn set_jwt_session(&mut self, payload: &[u8]) -> Result<(), HttpConnError> { - self.jti += 1; - let token = resign_jwt(&self.key, payload, self.jti)?; + if let ClientDataEnum::Local(local_data) = &mut self.data { + local_data.jti += 1; + let token = resign_jwt(&local_data.key, payload, local_data.jti)?; - // initiates the auth session - self.inner.simple_query("discard all").await?; - self.inner - .query( - "select auth.jwt_session_init($1)", - &[&token as &(dyn ToSql + Sync)], - ) - .await?; + // initiates the auth session + self.inner.simple_query("discard all").await?; + self.inner + .query( + "select auth.jwt_session_init($1)", + &[&token as &(dyn ToSql + Sync)], + ) + .await?; - let pid = self.inner.get_process_id(); - info!(pid, jti = self.jti, "user session state init"); - - Ok(()) - } -} - -pub(crate) struct LocalClient { - span: Span, - inner: Option>, - conn_info: ConnInfo, - pool: Weak>, -} - -pub(crate) struct Discard<'a, C: ClientInnerExt> { - conn_info: &'a ConnInfo, - pool: &'a mut Weak>, -} - -impl LocalClient { - pub(self) fn new( - inner: ClientInner, - conn_info: ConnInfo, - pool: Weak>, - ) -> Self { - Self { - inner: Some(inner), - span: Span::current(), - conn_info, - pool, + let pid = self.inner.get_process_id(); + info!(pid, jti = local_data.jti, "user session state init"); + Ok(()) + } else { + panic!("unexpected client data type"); } } - - pub(crate) fn client_inner(&mut self) -> (&mut ClientInner, Discard<'_, C>) { - let Self { - inner, - pool, - conn_info, - span: _, - } = self; - let inner_m = inner.as_mut().expect("client inner should not be removed"); - (inner_m, Discard { conn_info, pool }) - } - - pub(crate) fn inner(&mut self) -> (&mut C, Discard<'_, C>) { - let Self { - inner, - pool, - conn_info, - span: _, - } = self; - let inner = inner.as_mut().expect("client inner should not be removed"); - (&mut inner.inner, Discard { conn_info, pool }) - } } /// implements relatively efficient in-place json object key upserting @@ -547,58 +365,6 @@ fn sign_jwt(sk: &SigningKey, payload: &[u8]) -> String { jwt } -impl LocalClient { - pub(crate) fn metrics(&self) -> Arc { - let aux = &self.inner.as_ref().unwrap().aux; - USAGE_METRICS.register(Ids { - endpoint_id: aux.endpoint_id, - branch_id: aux.branch_id, - }) - } - - fn do_drop(&mut self) -> Option> { - let conn_info = self.conn_info.clone(); - let client = self - .inner - .take() - .expect("client inner should not be removed"); - if let Some(conn_pool) = std::mem::take(&mut self.pool).upgrade() { - let current_span = self.span.clone(); - // return connection to the pool - return Some(move || { - let _span = current_span.enter(); - EndpointConnPool::put(&conn_pool.global_pool, &conn_info, client); - }); - } - None - } -} - -impl Drop for LocalClient { - fn drop(&mut self) { - if let Some(drop) = self.do_drop() { - tokio::task::spawn_blocking(drop); - } - } -} - -impl Discard<'_, C> { - pub(crate) fn check_idle(&mut self, status: ReadyForQueryStatus) { - let conn_info = &self.conn_info; - if status != ReadyForQueryStatus::Idle && std::mem::take(self.pool).strong_count() > 0 { - info!( - "local_pool: throwing away connection '{conn_info}' because connection is not idle" - ); - } - } - pub(crate) fn discard(&mut self) { - let conn_info = &self.conn_info; - if std::mem::take(self.pool).strong_count() > 0 { - info!("local_pool: throwing away connection '{conn_info}' because connection is potentially in a broken state"); - } - } -} - #[cfg(test)] mod tests { use p256::ecdsa::SigningKey; diff --git a/proxy/src/serverless/sql_over_http.rs b/proxy/src/serverless/sql_over_http.rs index 0713c27d65..f0975617d4 100644 --- a/proxy/src/serverless/sql_over_http.rs +++ b/proxy/src/serverless/sql_over_http.rs @@ -31,7 +31,6 @@ use super::conn_pool_lib::{self, ConnInfo}; use super::error::HttpCodeError; use super::http_util::json_response; use super::json::{json_to_pg_text, pg_text_row_to_json, JsonConversionError}; -use super::local_conn_pool; use crate::auth::backend::{ComputeCredentialKeys, ComputeUserInfo}; use crate::auth::{endpoint_sni, ComputeUserInfoParseError}; use crate::config::{AuthenticationConfig, HttpConfig, ProxyConfig, TlsConfig}; @@ -1052,12 +1051,12 @@ async fn query_to_json( enum Client { Remote(conn_pool_lib::Client), - Local(local_conn_pool::LocalClient), + Local(conn_pool_lib::Client), } enum Discard<'a> { Remote(conn_pool_lib::Discard<'a, tokio_postgres::Client>), - Local(local_conn_pool::Discard<'a, tokio_postgres::Client>), + Local(conn_pool_lib::Discard<'a, tokio_postgres::Client>), } impl Client { @@ -1071,7 +1070,7 @@ impl Client { fn inner(&mut self) -> (&mut tokio_postgres::Client, Discard<'_>) { match self { Client::Remote(client) => { - let (c, d) = client.inner_mut(); + let (c, d) = client.inner(); (c, Discard::Remote(d)) } Client::Local(local_client) => { From babfeb70ba3c8036210a4ff2e8e33dd9ecde0dc4 Mon Sep 17 00:00:00 2001 From: Erik Grinaker Date: Tue, 5 Nov 2024 18:05:30 +0100 Subject: [PATCH 15/52] safekeeper: don't allocate send buffers on stack (#9644) ## Problem While experimenting with `MAX_SEND_SIZE` for benchmarking, I saw stack overflows when increasing it to 1 MB. Turns out a few buffers of this size are stack-allocated rather than heap-allocated. Even at the default 128 KB size, that's a bit large to allocate on the stack. ## Summary of changes Heap-allocate buffers of size `MAX_SEND_SIZE`. --- safekeeper/src/copy_timeline.rs | 2 +- safekeeper/src/debug_dump.rs | 2 +- safekeeper/src/send_wal.rs | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/safekeeper/src/copy_timeline.rs b/safekeeper/src/copy_timeline.rs index 1bf0cc668f..07fa98212f 100644 --- a/safekeeper/src/copy_timeline.rs +++ b/safekeeper/src/copy_timeline.rs @@ -172,7 +172,7 @@ async fn copy_disk_segments( ) -> Result<()> { let mut wal_reader = tli.get_walreader(start_lsn).await?; - let mut buf = [0u8; MAX_SEND_SIZE]; + let mut buf = vec![0u8; MAX_SEND_SIZE]; let first_segment = start_lsn.segment_number(wal_seg_size); let last_segment = end_lsn.segment_number(wal_seg_size); diff --git a/safekeeper/src/debug_dump.rs b/safekeeper/src/debug_dump.rs index 125f5af7f3..a2d0c49768 100644 --- a/safekeeper/src/debug_dump.rs +++ b/safekeeper/src/debug_dump.rs @@ -383,7 +383,7 @@ pub async fn calculate_digest( let mut wal_reader = tli.get_walreader(request.from_lsn).await?; let mut hasher = Sha256::new(); - let mut buf = [0u8; MAX_SEND_SIZE]; + let mut buf = vec![0u8; MAX_SEND_SIZE]; let mut bytes_left = (request.until_lsn.0 - request.from_lsn.0) as usize; while bytes_left > 0 { diff --git a/safekeeper/src/send_wal.rs b/safekeeper/src/send_wal.rs index 6d677f405a..6d94ff98b1 100644 --- a/safekeeper/src/send_wal.rs +++ b/safekeeper/src/send_wal.rs @@ -467,7 +467,7 @@ impl SafekeeperPostgresHandler { end_watch, ws_guard: ws_guard.clone(), wal_reader, - send_buf: [0; MAX_SEND_SIZE], + send_buf: vec![0u8; MAX_SEND_SIZE], }; let mut reply_reader = ReplyReader { reader, @@ -548,7 +548,7 @@ struct WalSender<'a, IO> { ws_guard: Arc, wal_reader: WalReader, // buffer for readling WAL into to send it - send_buf: [u8; MAX_SEND_SIZE], + send_buf: Vec, } const POLL_STATE_TIMEOUT: Duration = Duration::from_secs(1); From fcde40d60035ef5187615b9864321041cf7bc024 Mon Sep 17 00:00:00 2001 From: Conrad Ludgate Date: Tue, 5 Nov 2024 17:23:00 +0000 Subject: [PATCH 16/52] [proxy] use the proxy protocol v2 command to silence some logs (#9620) The PROXY Protocol V2 offers a "command" concept. It can be of two different values. "Local" and "Proxy". The spec suggests that "Local" be used for health-checks. We can thus use this to silence logging for such health checks such as those from NLB. This additionally refactors the flow to be a bit more type-safe, self documenting and using zerocopy deser. --- Cargo.lock | 3 + Cargo.toml | 1 + proxy/Cargo.toml | 1 + proxy/src/console_redirect_proxy.rs | 19 +- proxy/src/protocol2.rs | 319 +++++++++++++++++----------- proxy/src/proxy/mod.rs | 19 +- proxy/src/serverless/mod.rs | 15 +- workspace_hack/Cargo.toml | 2 + 8 files changed, 238 insertions(+), 141 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 484769bd16..9c2a0b455e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4365,6 +4365,7 @@ dependencies = [ "walkdir", "workspace_hack", "x509-parser", + "zerocopy", ] [[package]] @@ -7374,6 +7375,7 @@ dependencies = [ "tracing", "tracing-core", "url", + "zerocopy", "zeroize", "zstd", "zstd-safe", @@ -7446,6 +7448,7 @@ version = "0.7.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1c4061bedbb353041c12f413700357bec76df2c7e2ca8e4df8bac24c6bf68e3d" dependencies = [ + "byteorder", "zerocopy-derive", ] diff --git a/Cargo.toml b/Cargo.toml index e5f7719e7f..8207726caa 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -196,6 +196,7 @@ walkdir = "2.3.2" rustls-native-certs = "0.8" x509-parser = "0.16" whoami = "1.5.1" +zerocopy = { version = "0.7", features = ["derive"] } ## TODO replace this with tracing env_logger = "0.10" diff --git a/proxy/Cargo.toml b/proxy/Cargo.toml index 2580b1cf8a..bc4da95a91 100644 --- a/proxy/Cargo.toml +++ b/proxy/Cargo.toml @@ -98,6 +98,7 @@ rustls-native-certs.workspace = true x509-parser.workspace = true postgres-protocol.workspace = true redis.workspace = true +zerocopy.workspace = true # jwt stuff jose-jwa = "0.1.2" diff --git a/proxy/src/console_redirect_proxy.rs b/proxy/src/console_redirect_proxy.rs index 243ef07854..cc456f3667 100644 --- a/proxy/src/console_redirect_proxy.rs +++ b/proxy/src/console_redirect_proxy.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use futures::TryFutureExt; use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt}; use tokio_util::sync::CancellationToken; -use tracing::{error, info, Instrument}; +use tracing::{debug, error, info, Instrument}; use crate::auth::backend::ConsoleRedirectBackend; use crate::cancellation::{CancellationHandlerMain, CancellationHandlerMainInternal}; @@ -11,7 +11,7 @@ use crate::config::{ProxyConfig, ProxyProtocolV2}; use crate::context::RequestMonitoring; use crate::error::ReportableError; use crate::metrics::{Metrics, NumClientConnectionsGuard}; -use crate::protocol2::{read_proxy_protocol, ConnectionInfo}; +use crate::protocol2::{read_proxy_protocol, ConnectHeader, ConnectionInfo}; use crate::proxy::connect_compute::{connect_to_compute, TcpMechanism}; use crate::proxy::handshake::{handshake, HandshakeData}; use crate::proxy::passthrough::ProxyPassthrough; @@ -49,7 +49,7 @@ pub async fn task_main( let session_id = uuid::Uuid::new_v4(); let cancellation_handler = Arc::clone(&cancellation_handler); - tracing::info!(protocol = "tcp", %session_id, "accepted new TCP connection"); + debug!(protocol = "tcp", %session_id, "accepted new TCP connection"); connections.spawn(async move { let (socket, peer_addr) = match read_proxy_protocol(socket).await { @@ -57,16 +57,21 @@ pub async fn task_main( error!("per-client task finished with an error: {e:#}"); return; } - Ok((_socket, None)) if config.proxy_protocol_v2 == ProxyProtocolV2::Required => { + // our load balancers will not send any more data. let's just exit immediately + Ok((_socket, ConnectHeader::Local)) => { + debug!("healthcheck received"); + return; + } + Ok((_socket, ConnectHeader::Missing)) if config.proxy_protocol_v2 == ProxyProtocolV2::Required => { error!("missing required proxy protocol header"); return; } - Ok((_socket, Some(_))) if config.proxy_protocol_v2 == ProxyProtocolV2::Rejected => { + Ok((_socket, ConnectHeader::Proxy(_))) if config.proxy_protocol_v2 == ProxyProtocolV2::Rejected => { error!("proxy protocol header not supported"); return; } - Ok((socket, Some(info))) => (socket, info), - Ok((socket, None)) => (socket, ConnectionInfo{ addr: peer_addr, extra: None }), + Ok((socket, ConnectHeader::Proxy(info))) => (socket, info), + Ok((socket, ConnectHeader::Missing)) => (socket, ConnectionInfo{ addr: peer_addr, extra: None }), }; match socket.inner.set_nodelay(true) { diff --git a/proxy/src/protocol2.rs b/proxy/src/protocol2.rs index d1084ca2ff..33a5eb5e1e 100644 --- a/proxy/src/protocol2.rs +++ b/proxy/src/protocol2.rs @@ -11,6 +11,7 @@ use bytes::{Buf, Bytes, BytesMut}; use pin_project_lite::pin_project; use strum_macros::FromRepr; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, ReadBuf}; +use zerocopy::{FromBytes, FromZeroes}; pin_project! { /// A chained [`AsyncRead`] with [`AsyncWrite`] passthrough @@ -57,16 +58,31 @@ impl AsyncWrite for ChainRW { } /// Proxy Protocol Version 2 Header -const HEADER: [u8; 12] = [ +const SIGNATURE: [u8; 12] = [ 0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A, ]; +const LOCAL_V2: u8 = 0x20; +const PROXY_V2: u8 = 0x21; + +const TCP_OVER_IPV4: u8 = 0x11; +const UDP_OVER_IPV4: u8 = 0x12; +const TCP_OVER_IPV6: u8 = 0x21; +const UDP_OVER_IPV6: u8 = 0x22; + #[derive(PartialEq, Eq, Clone, Debug)] pub struct ConnectionInfo { pub addr: SocketAddr, pub extra: Option, } +#[derive(PartialEq, Eq, Clone, Debug)] +pub enum ConnectHeader { + Missing, + Local, + Proxy(ConnectionInfo), +} + impl fmt::Display for ConnectionInfo { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match &self.extra { @@ -89,96 +105,31 @@ pub enum ConnectionInfoExtra { pub(crate) async fn read_proxy_protocol( mut read: T, -) -> std::io::Result<(ChainRW, Option)> { +) -> std::io::Result<(ChainRW, ConnectHeader)> { let mut buf = BytesMut::with_capacity(128); - while buf.len() < 16 { + let header = loop { let bytes_read = read.read_buf(&mut buf).await?; - // exit for bad header - let len = usize::min(buf.len(), HEADER.len()); - if buf[..len] != HEADER[..len] { - return Ok((ChainRW { inner: read, buf }, None)); + // exit for bad header signature + let len = usize::min(buf.len(), SIGNATURE.len()); + if buf[..len] != SIGNATURE[..len] { + return Ok((ChainRW { inner: read, buf }, ConnectHeader::Missing)); } // if no more bytes available then exit if bytes_read == 0 { - return Ok((ChainRW { inner: read, buf }, None)); + return Ok((ChainRW { inner: read, buf }, ConnectHeader::Missing)); }; - } - let header = buf.split_to(16); - - // The next byte (the 13th one) is the protocol version and command. - // The highest four bits contains the version. As of this specification, it must - // always be sent as \x2 and the receiver must only accept this value. - let vc = header[12]; - let version = vc >> 4; - let command = vc & 0b1111; - if version != 2 { - return Err(io::Error::new( - io::ErrorKind::Other, - "invalid proxy protocol version. expected version 2", - )); - } - match command { - // the connection was established on purpose by the proxy - // without being relayed. The connection endpoints are the sender and the - // receiver. Such connections exist when the proxy sends health-checks to the - // server. The receiver must accept this connection as valid and must use the - // real connection endpoints and discard the protocol block including the - // family which is ignored. - 0 => {} - // the connection was established on behalf of another node, - // and reflects the original connection endpoints. The receiver must then use - // the information provided in the protocol block to get original the address. - 1 => {} - // other values are unassigned and must not be emitted by senders. Receivers - // must drop connections presenting unexpected values here. - _ => { - return Err(io::Error::new( - io::ErrorKind::Other, - "invalid proxy protocol command. expected local (0) or proxy (1)", - )) + // check if we have enough bytes to continue + if let Some(header) = buf.try_get::() { + break header; } }; - // The 14th byte contains the transport protocol and address family. The highest 4 - // bits contain the address family, the lowest 4 bits contain the protocol. - let ft = header[13]; - let address_length = match ft { - // - \x11 : TCP over IPv4 : the forwarded connection uses TCP over the AF_INET - // protocol family. Address length is 2*4 + 2*2 = 12 bytes. - // - \x12 : UDP over IPv4 : the forwarded connection uses UDP over the AF_INET - // protocol family. Address length is 2*4 + 2*2 = 12 bytes. - 0x11 | 0x12 => 12, - // - \x21 : TCP over IPv6 : the forwarded connection uses TCP over the AF_INET6 - // protocol family. Address length is 2*16 + 2*2 = 36 bytes. - // - \x22 : UDP over IPv6 : the forwarded connection uses UDP over the AF_INET6 - // protocol family. Address length is 2*16 + 2*2 = 36 bytes. - 0x21 | 0x22 => 36, - // unspecified or unix stream. ignore the addresses - _ => 0, - }; + let remaining_length = usize::from(header.len.get()); - // The 15th and 16th bytes is the address length in bytes in network endian order. - // It is used so that the receiver knows how many address bytes to skip even when - // it does not implement the presented protocol. Thus the length of the protocol - // header in bytes is always exactly 16 + this value. When a sender presents a - // LOCAL connection, it should not present any address so it sets this field to - // zero. Receivers MUST always consider this field to skip the appropriate number - // of bytes and must not assume zero is presented for LOCAL connections. When a - // receiver accepts an incoming connection showing an UNSPEC address family or - // protocol, it may or may not decide to log the address information if present. - let remaining_length = u16::from_be_bytes(header[14..16].try_into().unwrap()); - if remaining_length < address_length { - return Err(io::Error::new( - io::ErrorKind::Other, - "invalid proxy protocol length. not enough to fit requested IP addresses", - )); - } - drop(header); - - while buf.len() < remaining_length as usize { + while buf.len() < remaining_length { if read.read_buf(&mut buf).await? == 0 { return Err(io::Error::new( io::ErrorKind::UnexpectedEof, @@ -186,36 +137,69 @@ pub(crate) async fn read_proxy_protocol( )); } } + let payload = buf.split_to(remaining_length); - // Starting from the 17th byte, addresses are presented in network byte order. - // The address order is always the same : - // - source layer 3 address in network byte order - // - destination layer 3 address in network byte order - // - source layer 4 address if any, in network byte order (port) - // - destination layer 4 address if any, in network byte order (port) - let mut header = buf.split_to(usize::from(remaining_length)); - let mut addr = header.split_to(usize::from(address_length)); - let socket = match addr.len() { - 12 => { - let src_addr = Ipv4Addr::from_bits(addr.get_u32()); - let _dst_addr = Ipv4Addr::from_bits(addr.get_u32()); - let src_port = addr.get_u16(); - let _dst_port = addr.get_u16(); - Some(SocketAddr::from((src_addr, src_port))) + let res = process_proxy_payload(header, payload)?; + Ok((ChainRW { inner: read, buf }, res)) +} + +fn process_proxy_payload( + header: ProxyProtocolV2Header, + mut payload: BytesMut, +) -> std::io::Result { + match header.version_and_command { + // the connection was established on purpose by the proxy + // without being relayed. The connection endpoints are the sender and the + // receiver. Such connections exist when the proxy sends health-checks to the + // server. The receiver must accept this connection as valid and must use the + // real connection endpoints and discard the protocol block including the + // family which is ignored. + LOCAL_V2 => return Ok(ConnectHeader::Local), + // the connection was established on behalf of another node, + // and reflects the original connection endpoints. The receiver must then use + // the information provided in the protocol block to get original the address. + PROXY_V2 => {} + // other values are unassigned and must not be emitted by senders. Receivers + // must drop connections presenting unexpected values here. + #[rustfmt::skip] // https://github.com/rust-lang/rustfmt/issues/6384 + _ => return Err(io::Error::new( + io::ErrorKind::Other, + format!( + "invalid proxy protocol command 0x{:02X}. expected local (0x20) or proxy (0x21)", + header.version_and_command + ), + )), + }; + + let size_err = + "invalid proxy protocol length. payload not large enough to fit requested IP addresses"; + let addr = match header.protocol_and_family { + TCP_OVER_IPV4 | UDP_OVER_IPV4 => { + let addr = payload + .try_get::() + .ok_or_else(|| io::Error::new(io::ErrorKind::Other, size_err))?; + + SocketAddr::from((addr.src_addr.get(), addr.src_port.get())) } - 36 => { - let src_addr = Ipv6Addr::from_bits(addr.get_u128()); - let _dst_addr = Ipv6Addr::from_bits(addr.get_u128()); - let src_port = addr.get_u16(); - let _dst_port = addr.get_u16(); - Some(SocketAddr::from((src_addr, src_port))) + TCP_OVER_IPV6 | UDP_OVER_IPV6 => { + let addr = payload + .try_get::() + .ok_or_else(|| io::Error::new(io::ErrorKind::Other, size_err))?; + + SocketAddr::from((addr.src_addr.get(), addr.src_port.get())) + } + // unspecified or unix stream. ignore the addresses + _ => { + return Err(io::Error::new( + io::ErrorKind::Other, + "invalid proxy protocol address family/transport protocol.", + )) } - _ => None, }; let mut extra = None; - while let Some(mut tlv) = read_tlv(&mut header) { + while let Some(mut tlv) = read_tlv(&mut payload) { match Pp2Kind::from_repr(tlv.kind) { Some(Pp2Kind::Aws) => { if tlv.value.is_empty() { @@ -259,9 +243,7 @@ pub(crate) async fn read_proxy_protocol( } } - let conn_info = socket.map(|addr| ConnectionInfo { addr, extra }); - - Ok((ChainRW { inner: read, buf }, conn_info)) + Ok(ConnectHeader::Proxy(ConnectionInfo { addr, extra })) } #[derive(FromRepr, Debug, Copy, Clone)] @@ -337,27 +319,93 @@ struct Tlv { } fn read_tlv(b: &mut BytesMut) -> Option { - if b.len() < 3 { - return None; - } - let kind = b.get_u8(); - let len = usize::from(b.get_u16()); + let tlv_header = b.try_get::()?; + let len = usize::from(tlv_header.len.get()); if b.len() < len { return None; } - let value = b.split_to(len).freeze(); - Some(Tlv { kind, value }) + Some(Tlv { + kind: tlv_header.kind, + value: b.split_to(len).freeze(), + }) +} + +trait BufExt: Sized { + fn try_get(&mut self) -> Option; +} +impl BufExt for BytesMut { + fn try_get(&mut self) -> Option { + let res = T::read_from_prefix(self)?; + self.advance(size_of::()); + Some(res) + } +} + +#[derive(FromBytes, FromZeroes, Copy, Clone)] +#[repr(C)] +struct ProxyProtocolV2Header { + signature: [u8; 12], + version_and_command: u8, + protocol_and_family: u8, + len: zerocopy::byteorder::network_endian::U16, +} + +#[derive(FromBytes, FromZeroes, Copy, Clone)] +#[repr(C)] +struct ProxyProtocolV2HeaderV4 { + src_addr: NetworkEndianIpv4, + dst_addr: NetworkEndianIpv4, + src_port: zerocopy::byteorder::network_endian::U16, + dst_port: zerocopy::byteorder::network_endian::U16, +} + +#[derive(FromBytes, FromZeroes, Copy, Clone)] +#[repr(C)] +struct ProxyProtocolV2HeaderV6 { + src_addr: NetworkEndianIpv6, + dst_addr: NetworkEndianIpv6, + src_port: zerocopy::byteorder::network_endian::U16, + dst_port: zerocopy::byteorder::network_endian::U16, +} + +#[derive(FromBytes, FromZeroes, Copy, Clone)] +#[repr(C)] +struct TlvHeader { + kind: u8, + len: zerocopy::byteorder::network_endian::U16, +} + +#[derive(FromBytes, FromZeroes, Copy, Clone)] +#[repr(transparent)] +struct NetworkEndianIpv4(zerocopy::byteorder::network_endian::U32); +impl NetworkEndianIpv4 { + #[inline] + fn get(self) -> Ipv4Addr { + Ipv4Addr::from_bits(self.0.get()) + } +} + +#[derive(FromBytes, FromZeroes, Copy, Clone)] +#[repr(transparent)] +struct NetworkEndianIpv6(zerocopy::byteorder::network_endian::U128); +impl NetworkEndianIpv6 { + #[inline] + fn get(self) -> Ipv6Addr { + Ipv6Addr::from_bits(self.0.get()) + } } #[cfg(test)] mod tests { use tokio::io::AsyncReadExt; - use crate::protocol2::read_proxy_protocol; + use crate::protocol2::{ + read_proxy_protocol, ConnectHeader, LOCAL_V2, PROXY_V2, TCP_OVER_IPV4, UDP_OVER_IPV6, + }; #[tokio::test] async fn test_ipv4() { - let header = super::HEADER + let header = super::SIGNATURE // Proxy command, IPV4 | TCP .chain([(2 << 4) | 1, (1 << 4) | 1].as_slice()) // 12 + 3 bytes @@ -384,15 +432,17 @@ mod tests { assert_eq!(bytes, extra_data); - let info = info.unwrap(); + let ConnectHeader::Proxy(info) = info else { + panic!() + }; assert_eq!(info.addr, ([127, 0, 0, 1], 65535).into()); } #[tokio::test] async fn test_ipv6() { - let header = super::HEADER + let header = super::SIGNATURE // Proxy command, IPV6 | UDP - .chain([(2 << 4) | 1, (2 << 4) | 2].as_slice()) + .chain([PROXY_V2, UDP_OVER_IPV6].as_slice()) // 36 + 3 bytes .chain([0, 39].as_slice()) // src ip @@ -417,7 +467,9 @@ mod tests { assert_eq!(bytes, extra_data); - let info = info.unwrap(); + let ConnectHeader::Proxy(info) = info else { + panic!() + }; assert_eq!( info.addr, ([15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0], 257).into() @@ -433,7 +485,7 @@ mod tests { let mut bytes = vec![]; read.read_to_end(&mut bytes).await.unwrap(); assert_eq!(bytes, data); - assert_eq!(info, None); + assert_eq!(info, ConnectHeader::Missing); } #[tokio::test] @@ -445,7 +497,7 @@ mod tests { let mut bytes = vec![]; read.read_to_end(&mut bytes).await.unwrap(); assert_eq!(bytes, data); - assert_eq!(info, None); + assert_eq!(info, ConnectHeader::Missing); } #[tokio::test] @@ -454,9 +506,9 @@ mod tests { let tlv_len = (tlv.len() as u16).to_be_bytes(); let len = (12 + 3 + tlv.len() as u16).to_be_bytes(); - let header = super::HEADER + let header = super::SIGNATURE // Proxy command, Inet << 4 | Stream - .chain([(2 << 4) | 1, (1 << 4) | 1].as_slice()) + .chain([PROXY_V2, TCP_OVER_IPV4].as_slice()) // 12 + 3 bytes .chain(len.as_slice()) // src ip @@ -483,7 +535,30 @@ mod tests { assert_eq!(bytes, extra_data); - let info = info.unwrap(); + let ConnectHeader::Proxy(info) = info else { + panic!() + }; assert_eq!(info.addr, ([55, 56, 57, 58], 65535).into()); } + + #[tokio::test] + async fn test_local() { + let len = 0u16.to_be_bytes(); + let header = super::SIGNATURE + .chain([LOCAL_V2, 0x00].as_slice()) + .chain(len.as_slice()); + + let extra_data = [0xaa; 256]; + + let (mut read, info) = read_proxy_protocol(header.chain(extra_data.as_slice())) + .await + .unwrap(); + + let mut bytes = vec![]; + read.read_to_end(&mut bytes).await.unwrap(); + + assert_eq!(bytes, extra_data); + + let ConnectHeader::Local = info else { panic!() }; + } } diff --git a/proxy/src/proxy/mod.rs b/proxy/src/proxy/mod.rs index 922646d889..17721c23d5 100644 --- a/proxy/src/proxy/mod.rs +++ b/proxy/src/proxy/mod.rs @@ -19,7 +19,7 @@ use smol_str::{format_smolstr, SmolStr}; use thiserror::Error; use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt}; use tokio_util::sync::CancellationToken; -use tracing::{error, info, warn, Instrument}; +use tracing::{debug, error, info, warn, Instrument}; use self::connect_compute::{connect_to_compute, TcpMechanism}; use self::passthrough::ProxyPassthrough; @@ -28,7 +28,7 @@ use crate::config::{ProxyConfig, ProxyProtocolV2, TlsConfig}; use crate::context::RequestMonitoring; use crate::error::ReportableError; use crate::metrics::{Metrics, NumClientConnectionsGuard}; -use crate::protocol2::{read_proxy_protocol, ConnectionInfo}; +use crate::protocol2::{read_proxy_protocol, ConnectHeader, ConnectionInfo}; use crate::proxy::handshake::{handshake, HandshakeData}; use crate::rate_limiter::EndpointRateLimiter; use crate::stream::{PqStream, Stream}; @@ -83,7 +83,7 @@ pub async fn task_main( let session_id = uuid::Uuid::new_v4(); let cancellation_handler = Arc::clone(&cancellation_handler); - tracing::info!(protocol = "tcp", %session_id, "accepted new TCP connection"); + debug!(protocol = "tcp", %session_id, "accepted new TCP connection"); let endpoint_rate_limiter2 = endpoint_rate_limiter.clone(); connections.spawn(async move { @@ -92,16 +92,21 @@ pub async fn task_main( warn!("per-client task finished with an error: {e:#}"); return; } - Ok((_socket, None)) if config.proxy_protocol_v2 == ProxyProtocolV2::Required => { + // our load balancers will not send any more data. let's just exit immediately + Ok((_socket, ConnectHeader::Local)) => { + debug!("healthcheck received"); + return; + } + Ok((_socket, ConnectHeader::Missing)) if config.proxy_protocol_v2 == ProxyProtocolV2::Required => { warn!("missing required proxy protocol header"); return; } - Ok((_socket, Some(_))) if config.proxy_protocol_v2 == ProxyProtocolV2::Rejected => { + Ok((_socket, ConnectHeader::Proxy(_))) if config.proxy_protocol_v2 == ProxyProtocolV2::Rejected => { warn!("proxy protocol header not supported"); return; } - Ok((socket, Some(info))) => (socket, info), - Ok((socket, None)) => (socket, ConnectionInfo { addr: peer_addr, extra: None }), + Ok((socket, ConnectHeader::Proxy(info))) => (socket, info), + Ok((socket, ConnectHeader::Missing)) => (socket, ConnectionInfo { addr: peer_addr, extra: None }), }; match socket.inner.set_nodelay(true) { diff --git a/proxy/src/serverless/mod.rs b/proxy/src/serverless/mod.rs index 4b60ddf60f..cf758855fa 100644 --- a/proxy/src/serverless/mod.rs +++ b/proxy/src/serverless/mod.rs @@ -47,7 +47,7 @@ use crate::cancellation::CancellationHandlerMain; use crate::config::{ProxyConfig, ProxyProtocolV2}; use crate::context::RequestMonitoring; use crate::metrics::Metrics; -use crate::protocol2::{read_proxy_protocol, ChainRW, ConnectionInfo}; +use crate::protocol2::{read_proxy_protocol, ChainRW, ConnectHeader, ConnectionInfo}; use crate::proxy::run_until_cancelled; use crate::rate_limiter::EndpointRateLimiter; use crate::serverless::backend::PoolingBackend; @@ -251,16 +251,21 @@ async fn connection_startup( }; let conn_info = match peer { - None if config.proxy_protocol_v2 == ProxyProtocolV2::Required => { + // our load balancers will not send any more data. let's just exit immediately + ConnectHeader::Local => { + tracing::debug!("healthcheck received"); + return None; + } + ConnectHeader::Missing if config.proxy_protocol_v2 == ProxyProtocolV2::Required => { tracing::warn!("missing required proxy protocol header"); return None; } - Some(_) if config.proxy_protocol_v2 == ProxyProtocolV2::Rejected => { + ConnectHeader::Proxy(_) if config.proxy_protocol_v2 == ProxyProtocolV2::Rejected => { tracing::warn!("proxy protocol header not supported"); return None; } - Some(info) => info, - None => ConnectionInfo { + ConnectHeader::Proxy(info) => info, + ConnectHeader::Missing => ConnectionInfo { addr: peer_addr, extra: None, }, diff --git a/workspace_hack/Cargo.toml b/workspace_hack/Cargo.toml index 8d83d9d9e2..71ebab4119 100644 --- a/workspace_hack/Cargo.toml +++ b/workspace_hack/Cargo.toml @@ -88,6 +88,7 @@ tower = { version = "0.4", default-features = false, features = ["balance", "buf tracing = { version = "0.1", features = ["log"] } tracing-core = { version = "0.1" } url = { version = "2", features = ["serde"] } +zerocopy = { version = "0.7", features = ["derive", "simd"] } zeroize = { version = "1", features = ["derive", "serde"] } zstd = { version = "0.13" } zstd-safe = { version = "7", default-features = false, features = ["arrays", "legacy", "std", "zdict_builder"] } @@ -126,6 +127,7 @@ serde = { version = "1", features = ["alloc", "derive"] } syn = { version = "2", features = ["extra-traits", "fold", "full", "visit", "visit-mut"] } time-macros = { version = "0.2", default-features = false, features = ["formatting", "parsing", "serde"] } toml_edit = { version = "0.22", features = ["serde"] } +zerocopy = { version = "0.7", features = ["derive", "simd"] } zstd = { version = "0.13" } zstd-safe = { version = "7", default-features = false, features = ["arrays", "legacy", "std", "zdict_builder"] } zstd-sys = { version = "2", default-features = false, features = ["legacy", "std", "zdict_builder"] } From 754d2950a37935224620a0c8db0b2d311ae0c226 Mon Sep 17 00:00:00 2001 From: Folke Behrens Date: Tue, 5 Nov 2024 22:32:33 +0100 Subject: [PATCH 17/52] proxy: Revert ControlPlaneEvent back to struct (#9649) Due to neondatabase/cloud#19815 we need to be more tolerant when reading events. --- proxy/src/cache/endpoints.rs | 77 ++++++++++++++++++++++-------------- 1 file changed, 47 insertions(+), 30 deletions(-) diff --git a/proxy/src/cache/endpoints.rs b/proxy/src/cache/endpoints.rs index 400c76291e..a488d8a9cd 100644 --- a/proxy/src/cache/endpoints.rs +++ b/proxy/src/cache/endpoints.rs @@ -19,26 +19,28 @@ use crate::rate_limiter::GlobalRateLimiter; use crate::redis::connection_with_credentials_provider::ConnectionWithCredentialsProvider; use crate::types::EndpointId; -#[allow(clippy::enum_variant_names)] -#[derive(Deserialize, Debug, Clone)] -#[serde(tag = "type", rename_all(deserialize = "snake_case"))] -enum ControlPlaneEvent { - EndpointCreated { endpoint_created: EndpointCreated }, - BranchCreated { branch_created: BranchCreated }, - ProjectCreated { project_created: ProjectCreated }, +// TODO: this could be an enum, but events in Redis need to be fixed first. +// ProjectCreated was sent with type:branch_created. So we ignore type. +#[derive(Deserialize, Debug, Clone, PartialEq)] +struct ControlPlaneEvent { + endpoint_created: Option, + branch_created: Option, + project_created: Option, + #[serde(rename = "type")] + _type: Option, } -#[derive(Deserialize, Debug, Clone)] +#[derive(Deserialize, Debug, Clone, PartialEq)] struct EndpointCreated { endpoint_id: String, } -#[derive(Deserialize, Debug, Clone)] +#[derive(Deserialize, Debug, Clone, PartialEq)] struct BranchCreated { branch_id: String, } -#[derive(Deserialize, Debug, Clone)] +#[derive(Deserialize, Debug, Clone, PartialEq)] struct ProjectCreated { project_id: String, } @@ -104,24 +106,28 @@ impl EndpointsCache { } fn insert_event(&self, event: ControlPlaneEvent) { - let counter = match event { - ControlPlaneEvent::EndpointCreated { endpoint_created } => { - self.endpoints - .insert(EndpointIdInt::from(&endpoint_created.endpoint_id.into())); - RedisEventsCount::EndpointCreated - } - ControlPlaneEvent::BranchCreated { branch_created } => { - self.branches - .insert(BranchIdInt::from(&branch_created.branch_id.into())); - RedisEventsCount::BranchCreated - } - ControlPlaneEvent::ProjectCreated { project_created } => { - self.projects - .insert(ProjectIdInt::from(&project_created.project_id.into())); - RedisEventsCount::ProjectCreated - } - }; - Metrics::get().proxy.redis_events_count.inc(counter); + if let Some(endpoint_created) = event.endpoint_created { + self.endpoints + .insert(EndpointIdInt::from(&endpoint_created.endpoint_id.into())); + Metrics::get() + .proxy + .redis_events_count + .inc(RedisEventsCount::EndpointCreated); + } else if let Some(branch_created) = event.branch_created { + self.branches + .insert(BranchIdInt::from(&branch_created.branch_id.into())); + Metrics::get() + .proxy + .redis_events_count + .inc(RedisEventsCount::BranchCreated); + } else if let Some(project_created) = event.project_created { + self.projects + .insert(ProjectIdInt::from(&project_created.project_id.into())); + Metrics::get() + .proxy + .redis_events_count + .inc(RedisEventsCount::ProjectCreated); + } } pub async fn do_read( @@ -235,11 +241,22 @@ impl EndpointsCache { #[cfg(test)] mod tests { - use super::ControlPlaneEvent; + use super::*; #[test] fn test_parse_control_plane_event() { let s = r#"{"branch_created":null,"endpoint_created":{"endpoint_id":"ep-rapid-thunder-w0qqw2q9"},"project_created":null,"type":"endpoint_created"}"#; - serde_json::from_str::(s).unwrap(); + + assert_eq!( + serde_json::from_str::(s).unwrap(), + ControlPlaneEvent { + endpoint_created: Some(EndpointCreated { + endpoint_id: "ep-rapid-thunder-w0qqw2q9".into() + }), + branch_created: None, + project_created: None, + _type: Some("endpoint_created".into()), + } + ); } } From ebc43efebc0adb99043d6b2cd6600abf4a588810 Mon Sep 17 00:00:00 2001 From: Folke Behrens Date: Tue, 5 Nov 2024 23:03:53 +0100 Subject: [PATCH 18/52] proxy: Refactor cplane types (#9643) The overall idea of the PR is to rename a few types to make their purpose more clear, reduce abstraction where not needed, and move types to to more better suited modules. --- proxy/src/auth/backend/console_redirect.rs | 3 +- proxy/src/auth/backend/mod.rs | 41 +- proxy/src/bin/proxy.rs | 13 +- proxy/src/compute.rs | 2 +- proxy/src/config.rs | 10 +- .../{provider => client}/mock.rs | 19 +- proxy/src/control_plane/client/mod.rs | 281 +++++++++ .../{provider => client}/neon.rs | 44 +- proxy/src/control_plane/errors.rs | 216 +++++++ proxy/src/control_plane/messages.rs | 10 +- proxy/src/control_plane/mod.rs | 127 +++- proxy/src/control_plane/provider/mod.rs | 588 ------------------ proxy/src/proxy/tests/mod.rs | 34 +- proxy/src/proxy/wake_compute.rs | 2 +- proxy/src/serverless/backend.rs | 2 +- 15 files changed, 708 insertions(+), 684 deletions(-) rename proxy/src/control_plane/{provider => client}/mock.rs (94%) create mode 100644 proxy/src/control_plane/client/mod.rs rename proxy/src/control_plane/{provider => client}/neon.rs (93%) create mode 100644 proxy/src/control_plane/errors.rs delete mode 100644 proxy/src/control_plane/provider/mod.rs diff --git a/proxy/src/auth/backend/console_redirect.rs b/proxy/src/auth/backend/console_redirect.rs index 255e1fed54..1df59b4893 100644 --- a/proxy/src/auth/backend/console_redirect.rs +++ b/proxy/src/auth/backend/console_redirect.rs @@ -9,8 +9,7 @@ use super::ComputeCredentialKeys; use crate::cache::Cached; use crate::config::AuthenticationConfig; use crate::context::RequestMonitoring; -use crate::control_plane::provider::NodeInfo; -use crate::control_plane::{self, CachedNodeInfo}; +use crate::control_plane::{self, CachedNodeInfo, NodeInfo}; use crate::error::{ReportableError, UserFacingError}; use crate::proxy::connect_compute::ComputeConnectBackend; use crate::stream::PqStream; diff --git a/proxy/src/auth/backend/mod.rs b/proxy/src/auth/backend/mod.rs index 0eb68e6412..ccd5cf3071 100644 --- a/proxy/src/auth/backend/mod.rs +++ b/proxy/src/auth/backend/mod.rs @@ -21,11 +21,11 @@ use crate::auth::{self, validate_password_and_exchange, AuthError, ComputeUserIn use crate::cache::Cached; use crate::config::AuthenticationConfig; use crate::context::RequestMonitoring; +use crate::control_plane::client::ControlPlaneClient; use crate::control_plane::errors::GetAuthInfoError; -use crate::control_plane::provider::{ - CachedAllowedIps, CachedNodeInfo, CachedRoleSecret, ControlPlaneBackend, +use crate::control_plane::{ + self, AuthSecret, CachedAllowedIps, CachedNodeInfo, CachedRoleSecret, ControlPlaneApi, }; -use crate::control_plane::{self, Api, AuthSecret}; use crate::intern::EndpointIdInt; use crate::metrics::Metrics; use crate::proxy::connect_compute::ComputeConnectBackend; @@ -62,42 +62,26 @@ impl std::ops::Deref for MaybeOwned<'_, T> { /// backends which require them for the authentication process. pub enum Backend<'a, T> { /// Cloud API (V2). - ControlPlane(MaybeOwned<'a, ControlPlaneBackend>, T), + ControlPlane(MaybeOwned<'a, ControlPlaneClient>, T), /// Local proxy uses configured auth credentials and does not wake compute Local(MaybeOwned<'a, LocalBackend>), } -#[cfg(test)] -pub(crate) trait TestBackend: Send + Sync + 'static { - fn wake_compute(&self) -> Result; - fn get_allowed_ips_and_secret( - &self, - ) -> Result<(CachedAllowedIps, Option), control_plane::errors::GetAuthInfoError>; - fn dyn_clone(&self) -> Box; -} - -#[cfg(test)] -impl Clone for Box { - fn clone(&self) -> Self { - TestBackend::dyn_clone(&**self) - } -} - impl std::fmt::Display for Backend<'_, ()> { fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { Self::ControlPlane(api, ()) => match &**api { - ControlPlaneBackend::Management(endpoint) => fmt - .debug_tuple("ControlPlane::Management") + ControlPlaneClient::Neon(endpoint) => fmt + .debug_tuple("ControlPlane::Neon") .field(&endpoint.url()) .finish(), #[cfg(any(test, feature = "testing"))] - ControlPlaneBackend::PostgresMock(endpoint) => fmt + ControlPlaneClient::PostgresMock(endpoint) => fmt .debug_tuple("ControlPlane::PostgresMock") .field(&endpoint.url()) .finish(), #[cfg(test)] - ControlPlaneBackend::Test(_) => fmt.debug_tuple("ControlPlane::Test").finish(), + ControlPlaneClient::Test(_) => fmt.debug_tuple("ControlPlane::Test").finish(), }, Self::Local(_) => fmt.debug_tuple("Local").finish(), } @@ -282,7 +266,7 @@ impl AuthenticationConfig { /// All authentication flows will emit an AuthenticationOk message if successful. async fn auth_quirks( ctx: &RequestMonitoring, - api: &impl control_plane::Api, + api: &impl control_plane::ControlPlaneApi, user_info: ComputeUserInfoMaybeEndpoint, client: &mut stream::PqStream>, allow_cleartext: bool, @@ -499,12 +483,12 @@ mod tests { use std::time::Duration; use bytes::BytesMut; + use control_plane::AuthSecret; use fallible_iterator::FallibleIterator; use once_cell::sync::Lazy; use postgres_protocol::authentication::sasl::{ChannelBinding, ScramSha256}; use postgres_protocol::message::backend::Message as PgMessage; use postgres_protocol::message::frontend; - use provider::AuthSecret; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWriteExt}; use super::jwt::JwkCache; @@ -513,8 +497,7 @@ mod tests { use crate::auth::{ComputeUserInfoMaybeEndpoint, IpPattern}; use crate::config::AuthenticationConfig; use crate::context::RequestMonitoring; - use crate::control_plane::provider::{self, CachedAllowedIps, CachedRoleSecret}; - use crate::control_plane::{self, CachedNodeInfo}; + use crate::control_plane::{self, CachedAllowedIps, CachedNodeInfo, CachedRoleSecret}; use crate::proxy::NeonOptions; use crate::rate_limiter::{EndpointRateLimiter, RateBucketInfo}; use crate::scram::threadpool::ThreadPool; @@ -526,7 +509,7 @@ mod tests { secret: AuthSecret, } - impl control_plane::Api for Auth { + impl control_plane::ControlPlaneApi for Auth { async fn get_role_secret( &self, _ctx: &RequestMonitoring, diff --git a/proxy/src/bin/proxy.rs b/proxy/src/bin/proxy.rs index 82c259efc8..3179a929c4 100644 --- a/proxy/src/bin/proxy.rs +++ b/proxy/src/bin/proxy.rs @@ -513,7 +513,7 @@ async fn main() -> anyhow::Result<()> { } if let Either::Left(auth::Backend::ControlPlane(api, _)) = &auth_backend { - if let proxy::control_plane::provider::ControlPlaneBackend::Management(api) = &**api { + if let proxy::control_plane::client::ControlPlaneClient::Neon(api) = &**api { match (redis_notifications_client, regional_redis_client.clone()) { (None, None) => {} (client1, client2) => { @@ -732,13 +732,13 @@ fn build_auth_backend( RateBucketInfo::validate(&mut wake_compute_rps_limit)?; let wake_compute_endpoint_rate_limiter = Arc::new(WakeComputeRateLimiter::new(wake_compute_rps_limit)); - let api = control_plane::provider::neon::Api::new( + let api = control_plane::client::neon::NeonControlPlaneClient::new( endpoint, caches, locks, wake_compute_endpoint_rate_limiter, ); - let api = control_plane::provider::ControlPlaneBackend::Management(api); + let api = control_plane::client::ControlPlaneClient::Neon(api); let auth_backend = auth::Backend::ControlPlane(MaybeOwned::Owned(api), ()); let config = Box::leak(Box::new(auth_backend)); @@ -749,8 +749,11 @@ fn build_auth_backend( #[cfg(feature = "testing")] AuthBackendType::Postgres => { let url = args.auth_endpoint.parse()?; - let api = control_plane::provider::mock::Api::new(url, !args.is_private_access_proxy); - let api = control_plane::provider::ControlPlaneBackend::PostgresMock(api); + let api = control_plane::client::mock::MockControlPlane::new( + url, + !args.is_private_access_proxy, + ); + let api = control_plane::client::ControlPlaneClient::PostgresMock(api); let auth_backend = auth::Backend::ControlPlane(MaybeOwned::Owned(api), ()); diff --git a/proxy/src/compute.rs b/proxy/src/compute.rs index 65b6dd215b..d397fc5160 100644 --- a/proxy/src/compute.rs +++ b/proxy/src/compute.rs @@ -19,9 +19,9 @@ use tracing::{error, info, warn}; use crate::auth::parse_endpoint_param; use crate::cancellation::CancelClosure; use crate::context::RequestMonitoring; +use crate::control_plane::client::ApiLockError; use crate::control_plane::errors::WakeComputeError; use crate::control_plane::messages::MetricsAuxInfo; -use crate::control_plane::provider::ApiLockError; use crate::error::{ReportableError, UserFacingError}; use crate::metrics::{Metrics, NumDbConnectionsGuard}; use crate::proxy::neon_option; diff --git a/proxy/src/config.rs b/proxy/src/config.rs index 2870e100b7..f63d7e45aa 100644 --- a/proxy/src/config.rs +++ b/proxy/src/config.rs @@ -366,7 +366,7 @@ pub struct EndpointCacheConfig { } impl EndpointCacheConfig { - /// Default options for [`crate::control_plane::provider::NodeInfoCache`]. + /// Default options for [`crate::control_plane::NodeInfoCache`]. /// Notice that by default the limiter is empty, which means that cache is disabled. pub const CACHE_DEFAULT_OPTIONS: &'static str = "initial_batch_size=1000,default_batch_size=10,xread_timeout=5m,stream_name=controlPlane,disable_cache=true,limiter_info=1000@1s,retry_interval=1s"; @@ -441,7 +441,7 @@ pub struct CacheOptions { } impl CacheOptions { - /// Default options for [`crate::control_plane::provider::NodeInfoCache`]. + /// Default options for [`crate::control_plane::NodeInfoCache`]. pub const CACHE_DEFAULT_OPTIONS: &'static str = "size=4000,ttl=4m"; /// Parse cache options passed via cmdline. @@ -497,7 +497,7 @@ pub struct ProjectInfoCacheOptions { } impl ProjectInfoCacheOptions { - /// Default options for [`crate::control_plane::provider::NodeInfoCache`]. + /// Default options for [`crate::control_plane::NodeInfoCache`]. pub const CACHE_DEFAULT_OPTIONS: &'static str = "size=10000,ttl=4m,max_roles=10,gc_interval=60m"; @@ -616,9 +616,9 @@ pub struct ConcurrencyLockOptions { } impl ConcurrencyLockOptions { - /// Default options for [`crate::control_plane::provider::ApiLocks`]. + /// Default options for [`crate::control_plane::client::ApiLocks`]. pub const DEFAULT_OPTIONS_WAKE_COMPUTE_LOCK: &'static str = "permits=0"; - /// Default options for [`crate::control_plane::provider::ApiLocks`]. + /// Default options for [`crate::control_plane::client::ApiLocks`]. pub const DEFAULT_OPTIONS_CONNECT_COMPUTE_LOCK: &'static str = "shards=64,permits=100,epoch=10m,timeout=10ms"; diff --git a/proxy/src/control_plane/provider/mock.rs b/proxy/src/control_plane/client/mock.rs similarity index 94% rename from proxy/src/control_plane/provider/mock.rs rename to proxy/src/control_plane/client/mock.rs index 75a242d8d3..fd333d2aac 100644 --- a/proxy/src/control_plane/provider/mock.rs +++ b/proxy/src/control_plane/client/mock.rs @@ -9,16 +9,17 @@ use tokio_postgres::config::SslMode; use tokio_postgres::Client; use tracing::{error, info, info_span, warn, Instrument}; -use super::errors::{ApiError, GetAuthInfoError, WakeComputeError}; -use super::{AuthInfo, AuthSecret, CachedNodeInfo, NodeInfo}; use crate::auth::backend::jwt::AuthRule; use crate::auth::backend::ComputeUserInfo; use crate::auth::IpPattern; use crate::cache::Cached; use crate::context::RequestMonitoring; -use crate::control_plane::errors::GetEndpointJwksError; +use crate::control_plane::client::{CachedAllowedIps, CachedRoleSecret}; +use crate::control_plane::errors::{ + ControlPlaneError, GetAuthInfoError, GetEndpointJwksError, WakeComputeError, +}; use crate::control_plane::messages::MetricsAuxInfo; -use crate::control_plane::provider::{CachedAllowedIps, CachedRoleSecret}; +use crate::control_plane::{AuthInfo, AuthSecret, CachedNodeInfo, NodeInfo}; use crate::error::io_error; use crate::intern::RoleNameInt; use crate::types::{BranchId, EndpointId, ProjectId, RoleName}; @@ -31,25 +32,25 @@ enum MockApiError { PasswordNotSet(tokio_postgres::Error), } -impl From for ApiError { +impl From for ControlPlaneError { fn from(e: MockApiError) -> Self { io_error(e).into() } } -impl From for ApiError { +impl From for ControlPlaneError { fn from(e: tokio_postgres::Error) -> Self { io_error(e).into() } } #[derive(Clone)] -pub struct Api { +pub struct MockControlPlane { endpoint: ApiUrl, ip_allowlist_check_enabled: bool, } -impl Api { +impl MockControlPlane { pub fn new(endpoint: ApiUrl, ip_allowlist_check_enabled: bool) -> Self { Self { endpoint, @@ -201,7 +202,7 @@ async fn get_execute_postgres_query( Ok(Some(entry)) } -impl super::Api for Api { +impl super::ControlPlaneApi for MockControlPlane { #[tracing::instrument(skip_all)] async fn get_role_secret( &self, diff --git a/proxy/src/control_plane/client/mod.rs b/proxy/src/control_plane/client/mod.rs new file mode 100644 index 0000000000..e388d8a538 --- /dev/null +++ b/proxy/src/control_plane/client/mod.rs @@ -0,0 +1,281 @@ +#[cfg(any(test, feature = "testing"))] +pub mod mock; +pub mod neon; + +use std::hash::Hash; +use std::sync::Arc; +use std::time::Duration; + +use dashmap::DashMap; +use tokio::time::Instant; +use tracing::info; + +use crate::auth::backend::jwt::{AuthRule, FetchAuthRules, FetchAuthRulesError}; +use crate::auth::backend::ComputeUserInfo; +use crate::cache::endpoints::EndpointsCache; +use crate::cache::project_info::ProjectInfoCacheImpl; +use crate::config::{CacheOptions, EndpointCacheConfig, ProjectInfoCacheOptions}; +use crate::context::RequestMonitoring; +use crate::control_plane::{ + errors, CachedAllowedIps, CachedNodeInfo, CachedRoleSecret, ControlPlaneApi, NodeInfoCache, +}; +use crate::error::ReportableError; +use crate::metrics::ApiLockMetrics; +use crate::rate_limiter::{DynamicLimiter, Outcome, RateLimiterConfig, Token}; +use crate::types::EndpointId; + +#[non_exhaustive] +#[derive(Clone)] +pub enum ControlPlaneClient { + /// Current Management API (V2). + Neon(neon::NeonControlPlaneClient), + /// Local mock control plane. + #[cfg(any(test, feature = "testing"))] + PostgresMock(mock::MockControlPlane), + /// Internal testing + #[cfg(test)] + #[allow(private_interfaces)] + Test(Box), +} + +impl ControlPlaneApi for ControlPlaneClient { + async fn get_role_secret( + &self, + ctx: &RequestMonitoring, + user_info: &ComputeUserInfo, + ) -> Result { + match self { + Self::Neon(api) => api.get_role_secret(ctx, user_info).await, + #[cfg(any(test, feature = "testing"))] + Self::PostgresMock(api) => api.get_role_secret(ctx, user_info).await, + #[cfg(test)] + Self::Test(_) => { + unreachable!("this function should never be called in the test backend") + } + } + } + + async fn get_allowed_ips_and_secret( + &self, + ctx: &RequestMonitoring, + user_info: &ComputeUserInfo, + ) -> Result<(CachedAllowedIps, Option), errors::GetAuthInfoError> { + match self { + Self::Neon(api) => api.get_allowed_ips_and_secret(ctx, user_info).await, + #[cfg(any(test, feature = "testing"))] + Self::PostgresMock(api) => api.get_allowed_ips_and_secret(ctx, user_info).await, + #[cfg(test)] + Self::Test(api) => api.get_allowed_ips_and_secret(), + } + } + + async fn get_endpoint_jwks( + &self, + ctx: &RequestMonitoring, + endpoint: EndpointId, + ) -> Result, errors::GetEndpointJwksError> { + match self { + Self::Neon(api) => api.get_endpoint_jwks(ctx, endpoint).await, + #[cfg(any(test, feature = "testing"))] + Self::PostgresMock(api) => api.get_endpoint_jwks(ctx, endpoint).await, + #[cfg(test)] + Self::Test(_api) => Ok(vec![]), + } + } + + async fn wake_compute( + &self, + ctx: &RequestMonitoring, + user_info: &ComputeUserInfo, + ) -> Result { + match self { + Self::Neon(api) => api.wake_compute(ctx, user_info).await, + #[cfg(any(test, feature = "testing"))] + Self::PostgresMock(api) => api.wake_compute(ctx, user_info).await, + #[cfg(test)] + Self::Test(api) => api.wake_compute(), + } + } +} + +#[cfg(test)] +pub(crate) trait TestControlPlaneClient: Send + Sync + 'static { + fn wake_compute(&self) -> Result; + + fn get_allowed_ips_and_secret( + &self, + ) -> Result<(CachedAllowedIps, Option), errors::GetAuthInfoError>; + + fn dyn_clone(&self) -> Box; +} + +#[cfg(test)] +impl Clone for Box { + fn clone(&self) -> Self { + TestControlPlaneClient::dyn_clone(&**self) + } +} + +/// Various caches for [`control_plane`](super). +pub struct ApiCaches { + /// Cache for the `wake_compute` API method. + pub(crate) node_info: NodeInfoCache, + /// Cache which stores project_id -> endpoint_ids mapping. + pub project_info: Arc, + /// List of all valid endpoints. + pub endpoints_cache: Arc, +} + +impl ApiCaches { + pub fn new( + wake_compute_cache_config: CacheOptions, + project_info_cache_config: ProjectInfoCacheOptions, + endpoint_cache_config: EndpointCacheConfig, + ) -> Self { + Self { + node_info: NodeInfoCache::new( + "node_info_cache", + wake_compute_cache_config.size, + wake_compute_cache_config.ttl, + true, + ), + project_info: Arc::new(ProjectInfoCacheImpl::new(project_info_cache_config)), + endpoints_cache: Arc::new(EndpointsCache::new(endpoint_cache_config)), + } + } +} + +/// Various caches for [`control_plane`](super). +pub struct ApiLocks { + name: &'static str, + node_locks: DashMap>, + config: RateLimiterConfig, + timeout: Duration, + epoch: std::time::Duration, + metrics: &'static ApiLockMetrics, +} + +#[derive(Debug, thiserror::Error)] +pub(crate) enum ApiLockError { + #[error("timeout acquiring resource permit")] + TimeoutError(#[from] tokio::time::error::Elapsed), +} + +impl ReportableError for ApiLockError { + fn get_error_kind(&self) -> crate::error::ErrorKind { + match self { + ApiLockError::TimeoutError(_) => crate::error::ErrorKind::RateLimit, + } + } +} + +impl ApiLocks { + pub fn new( + name: &'static str, + config: RateLimiterConfig, + shards: usize, + timeout: Duration, + epoch: std::time::Duration, + metrics: &'static ApiLockMetrics, + ) -> prometheus::Result { + Ok(Self { + name, + node_locks: DashMap::with_shard_amount(shards), + config, + timeout, + epoch, + metrics, + }) + } + + pub(crate) async fn get_permit(&self, key: &K) -> Result { + if self.config.initial_limit == 0 { + return Ok(WakeComputePermit { + permit: Token::disabled(), + }); + } + 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.metrics.semaphores_registered.inc(); + DynamicLimiter::new(self.config) + }) + .clone() + } + }; + let permit = semaphore.acquire_timeout(self.timeout).await; + + self.metrics + .semaphore_acquire_seconds + .observe(now.elapsed().as_secs_f64()); + info!("acquired permit {:?}", now.elapsed().as_secs_f64()); + Ok(WakeComputePermit { permit: permit? }) + } + + pub async fn garbage_collect_worker(&self) { + if self.config.initial_limit == 0 { + return; + } + let mut interval = + tokio::time::interval(self.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.metrics.reclamation_lag_seconds.start_timer(); + let count = lock + .extract_if(|_, semaphore| Arc::strong_count(semaphore.get_mut()) == 1) + .count(); + drop(lock); + self.metrics.semaphores_unregistered.inc_by(count as u64); + timer.observe(); + } + } + } +} + +pub(crate) struct WakeComputePermit { + permit: Token, +} + +impl WakeComputePermit { + pub(crate) fn should_check_cache(&self) -> bool { + !self.permit.is_disabled() + } + pub(crate) fn release(self, outcome: Outcome) { + self.permit.release(outcome); + } + pub(crate) fn release_result(self, res: Result) -> Result { + match res { + Ok(_) => self.release(Outcome::Success), + Err(_) => self.release(Outcome::Overload), + } + res + } +} + +impl FetchAuthRules for ControlPlaneClient { + async fn fetch_auth_rules( + &self, + ctx: &RequestMonitoring, + endpoint: EndpointId, + ) -> Result, FetchAuthRulesError> { + self.get_endpoint_jwks(ctx, endpoint) + .await + .map_err(FetchAuthRulesError::GetEndpointJwks) + } +} diff --git a/proxy/src/control_plane/provider/neon.rs b/proxy/src/control_plane/client/neon.rs similarity index 93% rename from proxy/src/control_plane/provider/neon.rs rename to proxy/src/control_plane/client/neon.rs index 8ea91d7875..6c67d2df96 100644 --- a/proxy/src/control_plane/provider/neon.rs +++ b/proxy/src/control_plane/client/neon.rs @@ -10,18 +10,20 @@ use tokio::time::Instant; use tokio_postgres::config::SslMode; use tracing::{debug, info, info_span, warn, Instrument}; -use super::super::messages::{ControlPlaneError, GetRoleSecret, WakeCompute}; -use super::errors::{ApiError, GetAuthInfoError, WakeComputeError}; -use super::{ - ApiCaches, ApiLocks, AuthInfo, AuthSecret, CachedAllowedIps, CachedNodeInfo, CachedRoleSecret, - NodeInfo, -}; +use super::super::messages::{ControlPlaneErrorMessage, GetRoleSecret, WakeCompute}; use crate::auth::backend::jwt::AuthRule; use crate::auth::backend::ComputeUserInfo; use crate::cache::Cached; use crate::context::RequestMonitoring; -use crate::control_plane::errors::GetEndpointJwksError; +use crate::control_plane::caches::ApiCaches; +use crate::control_plane::errors::{ + ControlPlaneError, GetAuthInfoError, GetEndpointJwksError, WakeComputeError, +}; +use crate::control_plane::locks::ApiLocks; use crate::control_plane::messages::{ColdStartInfo, EndpointJwksResponse, Reason}; +use crate::control_plane::{ + AuthInfo, AuthSecret, CachedAllowedIps, CachedNodeInfo, CachedRoleSecret, NodeInfo, +}; use crate::metrics::{CacheOutcome, Metrics}; use crate::rate_limiter::WakeComputeRateLimiter; use crate::types::{EndpointCacheKey, EndpointId}; @@ -30,7 +32,7 @@ use crate::{compute, http, scram}; const X_REQUEST_ID: HeaderName = HeaderName::from_static("x-request-id"); #[derive(Clone)] -pub struct Api { +pub struct NeonControlPlaneClient { endpoint: http::Endpoint, pub caches: &'static ApiCaches, pub(crate) locks: &'static ApiLocks, @@ -39,7 +41,7 @@ pub struct Api { jwt: Arc, } -impl Api { +impl NeonControlPlaneClient { /// Construct an API object containing the auth parameters. pub fn new( endpoint: http::Endpoint, @@ -256,7 +258,7 @@ impl Api { } } -impl super::Api for Api { +impl super::ControlPlaneApi for NeonControlPlaneClient { #[tracing::instrument(skip_all)] async fn get_role_secret( &self, @@ -356,7 +358,7 @@ impl super::Api for Api { let (cached, info) = cached.take_value(); let info = info.map_err(|c| { info!(key = &*key, "found cached wake_compute error"); - WakeComputeError::ApiError(ApiError::ControlPlane(Box::new(*c))) + WakeComputeError::ControlPlane(ControlPlaneError::Message(Box::new(*c))) })?; debug!(key = &*key, "found cached compute node info"); @@ -403,9 +405,11 @@ impl super::Api for Api { Ok(cached.map(|()| node)) } Err(err) => match err { - WakeComputeError::ApiError(ApiError::ControlPlane(err)) => { + WakeComputeError::ControlPlane(ControlPlaneError::Message(err)) => { let Some(status) = &err.status else { - return Err(WakeComputeError::ApiError(ApiError::ControlPlane(err))); + return Err(WakeComputeError::ControlPlane(ControlPlaneError::Message( + err, + ))); }; let reason = status @@ -415,7 +419,9 @@ impl super::Api for Api { // if we can retry this error, do not cache it. if reason.can_retry() { - return Err(WakeComputeError::ApiError(ApiError::ControlPlane(err))); + return Err(WakeComputeError::ControlPlane(ControlPlaneError::Message( + err, + ))); } // at this point, we should only have quota errors. @@ -430,7 +436,9 @@ impl super::Api for Api { Duration::from_secs(30), ); - Err(WakeComputeError::ApiError(ApiError::ControlPlane(err))) + Err(WakeComputeError::ControlPlane(ControlPlaneError::Message( + err, + ))) } err => return Err(err), }, @@ -441,7 +449,7 @@ impl super::Api for Api { /// Parse http response body, taking status code into account. async fn parse_body serde::Deserialize<'a>>( response: http::Response, -) -> Result { +) -> Result { let status = response.status(); if status.is_success() { // We shouldn't log raw body because it may contain secrets. @@ -456,7 +464,7 @@ async fn parse_body serde::Deserialize<'a>>( // as the fact that the request itself has failed. let mut body = serde_json::from_slice(&s).unwrap_or_else(|e| { warn!("failed to parse error body: {e}"); - ControlPlaneError { + ControlPlaneErrorMessage { error: "reason unclear (malformed error message)".into(), http_status_code: status, status: None, @@ -465,7 +473,7 @@ async fn parse_body serde::Deserialize<'a>>( body.http_status_code = status; warn!("console responded with an error ({status}): {body:?}"); - Err(ApiError::ControlPlane(Box::new(body))) + Err(ControlPlaneError::Message(Box::new(body))) } fn parse_host_port(input: &str) -> Option<(&str, u16)> { diff --git a/proxy/src/control_plane/errors.rs b/proxy/src/control_plane/errors.rs new file mode 100644 index 0000000000..d6f565e34a --- /dev/null +++ b/proxy/src/control_plane/errors.rs @@ -0,0 +1,216 @@ +use thiserror::Error; + +use crate::control_plane::client::ApiLockError; +use crate::control_plane::messages::{self, ControlPlaneErrorMessage, Reason}; +use crate::error::{io_error, ErrorKind, ReportableError, UserFacingError}; +use crate::proxy::retry::CouldRetry; + +/// A go-to error message which doesn't leak any detail. +pub(crate) const REQUEST_FAILED: &str = "Console request failed"; + +/// Common console API error. +#[derive(Debug, Error)] +pub(crate) enum ControlPlaneError { + /// Error returned by the console itself. + #[error("{REQUEST_FAILED} with {0}")] + Message(Box), + + /// Various IO errors like broken pipe or malformed payload. + #[error("{REQUEST_FAILED}: {0}")] + Transport(#[from] std::io::Error), +} + +impl ControlPlaneError { + /// Returns HTTP status code if it's the reason for failure. + pub(crate) fn get_reason(&self) -> messages::Reason { + match self { + ControlPlaneError::Message(e) => e.get_reason(), + ControlPlaneError::Transport(_) => messages::Reason::Unknown, + } + } +} + +impl UserFacingError for ControlPlaneError { + fn to_string_client(&self) -> String { + match self { + // To minimize risks, only select errors are forwarded to users. + ControlPlaneError::Message(c) => c.get_user_facing_message(), + ControlPlaneError::Transport(_) => REQUEST_FAILED.to_owned(), + } + } +} + +impl ReportableError for ControlPlaneError { + fn get_error_kind(&self) -> crate::error::ErrorKind { + match self { + ControlPlaneError::Message(e) => match e.get_reason() { + Reason::RoleProtected => ErrorKind::User, + Reason::ResourceNotFound => ErrorKind::User, + Reason::ProjectNotFound => ErrorKind::User, + Reason::EndpointNotFound => ErrorKind::User, + Reason::BranchNotFound => ErrorKind::User, + Reason::RateLimitExceeded => ErrorKind::ServiceRateLimit, + Reason::NonDefaultBranchComputeTimeExceeded => ErrorKind::Quota, + Reason::ActiveTimeQuotaExceeded => ErrorKind::Quota, + Reason::ComputeTimeQuotaExceeded => ErrorKind::Quota, + Reason::WrittenDataQuotaExceeded => ErrorKind::Quota, + Reason::DataTransferQuotaExceeded => ErrorKind::Quota, + Reason::LogicalSizeQuotaExceeded => ErrorKind::Quota, + Reason::ConcurrencyLimitReached => ErrorKind::ControlPlane, + Reason::LockAlreadyTaken => ErrorKind::ControlPlane, + Reason::RunningOperations => ErrorKind::ControlPlane, + Reason::ActiveEndpointsLimitExceeded => ErrorKind::ControlPlane, + Reason::Unknown => ErrorKind::ControlPlane, + }, + ControlPlaneError::Transport(_) => crate::error::ErrorKind::ControlPlane, + } + } +} + +impl CouldRetry for ControlPlaneError { + fn could_retry(&self) -> bool { + match self { + // retry some transport errors + Self::Transport(io) => io.could_retry(), + Self::Message(e) => e.could_retry(), + } + } +} + +impl From for ControlPlaneError { + fn from(e: reqwest::Error) -> Self { + io_error(e).into() + } +} + +impl From for ControlPlaneError { + fn from(e: reqwest_middleware::Error) -> Self { + io_error(e).into() + } +} + +#[derive(Debug, Error)] +pub(crate) enum GetAuthInfoError { + // We shouldn't include the actual secret here. + #[error("Console responded with a malformed auth secret")] + BadSecret, + + #[error(transparent)] + ApiError(ControlPlaneError), +} + +// This allows more useful interactions than `#[from]`. +impl> From for GetAuthInfoError { + fn from(e: E) -> Self { + Self::ApiError(e.into()) + } +} + +impl UserFacingError for GetAuthInfoError { + fn to_string_client(&self) -> String { + match self { + // We absolutely should not leak any secrets! + Self::BadSecret => REQUEST_FAILED.to_owned(), + // However, API might return a meaningful error. + Self::ApiError(e) => e.to_string_client(), + } + } +} + +impl ReportableError for GetAuthInfoError { + fn get_error_kind(&self) -> crate::error::ErrorKind { + match self { + Self::BadSecret => crate::error::ErrorKind::ControlPlane, + Self::ApiError(_) => crate::error::ErrorKind::ControlPlane, + } + } +} + +#[derive(Debug, Error)] +pub(crate) enum WakeComputeError { + #[error("Console responded with a malformed compute address: {0}")] + BadComputeAddress(Box), + + #[error(transparent)] + ControlPlane(ControlPlaneError), + + #[error("Too many connections attempts")] + TooManyConnections, + + #[error("error acquiring resource permit: {0}")] + TooManyConnectionAttempts(#[from] ApiLockError), +} + +// This allows more useful interactions than `#[from]`. +impl> From for WakeComputeError { + fn from(e: E) -> Self { + Self::ControlPlane(e.into()) + } +} + +impl UserFacingError for WakeComputeError { + fn to_string_client(&self) -> String { + match self { + // We shouldn't show user the address even if it's broken. + // Besides, user is unlikely to care about this detail. + Self::BadComputeAddress(_) => REQUEST_FAILED.to_owned(), + // However, control plane might return a meaningful error. + Self::ControlPlane(e) => e.to_string_client(), + + Self::TooManyConnections => self.to_string(), + + Self::TooManyConnectionAttempts(_) => { + "Failed to acquire permit to connect to the database. Too many database connection attempts are currently ongoing.".to_owned() + } + } + } +} + +impl ReportableError for WakeComputeError { + fn get_error_kind(&self) -> crate::error::ErrorKind { + match self { + Self::BadComputeAddress(_) => crate::error::ErrorKind::ControlPlane, + Self::ControlPlane(e) => e.get_error_kind(), + Self::TooManyConnections => crate::error::ErrorKind::RateLimit, + Self::TooManyConnectionAttempts(e) => e.get_error_kind(), + } + } +} + +impl CouldRetry for WakeComputeError { + fn could_retry(&self) -> bool { + match self { + Self::BadComputeAddress(_) => false, + Self::ControlPlane(e) => e.could_retry(), + Self::TooManyConnections => false, + Self::TooManyConnectionAttempts(_) => false, + } + } +} + +#[derive(Debug, Error)] +pub enum GetEndpointJwksError { + #[error("endpoint not found")] + EndpointNotFound, + + #[error("failed to build control plane request: {0}")] + RequestBuild(#[source] reqwest::Error), + + #[error("failed to send control plane request: {0}")] + RequestExecute(#[source] reqwest_middleware::Error), + + #[error(transparent)] + ControlPlane(#[from] ControlPlaneError), + + #[cfg(any(test, feature = "testing"))] + #[error(transparent)] + TokioPostgres(#[from] tokio_postgres::Error), + + #[cfg(any(test, feature = "testing"))] + #[error(transparent)] + ParseUrl(#[from] url::ParseError), + + #[cfg(any(test, feature = "testing"))] + #[error(transparent)] + TaskJoin(#[from] tokio::task::JoinError), +} diff --git a/proxy/src/control_plane/messages.rs b/proxy/src/control_plane/messages.rs index 13a54145b1..75c932e6ab 100644 --- a/proxy/src/control_plane/messages.rs +++ b/proxy/src/control_plane/messages.rs @@ -10,14 +10,14 @@ use crate::proxy::retry::CouldRetry; /// Generic error response with human-readable description. /// Note that we can't always present it to user as is. #[derive(Debug, Deserialize, Clone)] -pub(crate) struct ControlPlaneError { +pub(crate) struct ControlPlaneErrorMessage { pub(crate) error: Box, #[serde(skip)] pub(crate) http_status_code: http::StatusCode, pub(crate) status: Option, } -impl ControlPlaneError { +impl ControlPlaneErrorMessage { pub(crate) fn get_reason(&self) -> Reason { self.status .as_ref() @@ -26,7 +26,7 @@ impl ControlPlaneError { } pub(crate) fn get_user_facing_message(&self) -> String { - use super::provider::errors::REQUEST_FAILED; + use super::errors::REQUEST_FAILED; self.status .as_ref() .and_then(|s| s.details.user_facing_message.as_ref()) @@ -51,7 +51,7 @@ impl ControlPlaneError { } } -impl Display for ControlPlaneError { +impl Display for ControlPlaneErrorMessage { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { let msg: &str = self .status @@ -62,7 +62,7 @@ impl Display for ControlPlaneError { } } -impl CouldRetry for ControlPlaneError { +impl CouldRetry for ControlPlaneErrorMessage { fn could_retry(&self) -> bool { // If the error message does not have a status, // the error is unknown and probably should not retry automatically diff --git a/proxy/src/control_plane/mod.rs b/proxy/src/control_plane/mod.rs index 87d8e781aa..70607ac0d0 100644 --- a/proxy/src/control_plane/mod.rs +++ b/proxy/src/control_plane/mod.rs @@ -5,18 +5,137 @@ pub mod messages; /// Wrappers for console APIs and their mocks. -pub mod provider; -pub(crate) use provider::{errors, Api, AuthSecret, CachedNodeInfo, NodeInfo}; +pub mod client; + +pub(crate) mod errors; + +use std::sync::Arc; +use std::time::Duration; + +use crate::auth::backend::jwt::AuthRule; +use crate::auth::backend::{ComputeCredentialKeys, ComputeUserInfo}; +use crate::auth::IpPattern; +use crate::cache::project_info::ProjectInfoCacheImpl; +use crate::cache::{Cached, TimedLru}; +use crate::context::RequestMonitoring; +use crate::control_plane::messages::{ControlPlaneErrorMessage, MetricsAuxInfo}; +use crate::intern::ProjectIdInt; +use crate::types::{EndpointCacheKey, EndpointId}; +use crate::{compute, scram}; /// Various cache-related types. pub mod caches { - pub use super::provider::ApiCaches; + pub use super::client::ApiCaches; } /// Various cache-related types. pub mod locks { - pub use super::provider::ApiLocks; + pub use super::client::ApiLocks; } /// Console's management API. pub mod mgmt; + +/// Auth secret which is managed by the cloud. +#[derive(Clone, Eq, PartialEq, Debug)] +pub(crate) enum AuthSecret { + #[cfg(any(test, feature = "testing"))] + /// Md5 hash of user's password. + Md5([u8; 16]), + + /// [SCRAM](crate::scram) authentication info. + Scram(scram::ServerSecret), +} + +#[derive(Default)] +pub(crate) struct AuthInfo { + pub(crate) secret: Option, + /// List of IP addresses allowed for the autorization. + pub(crate) allowed_ips: Vec, + /// Project ID. This is used for cache invalidation. + pub(crate) project_id: Option, +} + +/// Info for establishing a connection to a compute node. +/// This is what we get after auth succeeded, but not before! +#[derive(Clone)] +pub(crate) struct NodeInfo { + /// Compute node connection params. + /// It's sad that we have to clone this, but this will improve + /// once we migrate to a bespoke connection logic. + pub(crate) config: compute::ConnCfg, + + /// Labels for proxy's metrics. + pub(crate) aux: MetricsAuxInfo, + + /// Whether we should accept self-signed certificates (for testing) + pub(crate) allow_self_signed_compute: bool, +} + +impl NodeInfo { + pub(crate) async fn connect( + &self, + ctx: &RequestMonitoring, + timeout: Duration, + ) -> Result { + self.config + .connect( + ctx, + self.allow_self_signed_compute, + self.aux.clone(), + timeout, + ) + .await + } + pub(crate) fn reuse_settings(&mut self, other: Self) { + self.allow_self_signed_compute = other.allow_self_signed_compute; + self.config.reuse_password(other.config); + } + + pub(crate) fn set_keys(&mut self, keys: &ComputeCredentialKeys) { + match keys { + #[cfg(any(test, feature = "testing"))] + ComputeCredentialKeys::Password(password) => self.config.password(password), + ComputeCredentialKeys::AuthKeys(auth_keys) => self.config.auth_keys(*auth_keys), + ComputeCredentialKeys::JwtPayload(_) | ComputeCredentialKeys::None => &mut self.config, + }; + } +} + +pub(crate) type NodeInfoCache = + TimedLru>>; +pub(crate) type CachedNodeInfo = Cached<&'static NodeInfoCache, NodeInfo>; +pub(crate) type CachedRoleSecret = Cached<&'static ProjectInfoCacheImpl, Option>; +pub(crate) type CachedAllowedIps = Cached<&'static ProjectInfoCacheImpl, Arc>>; + +/// This will allocate per each call, but the http requests alone +/// already require a few allocations, so it should be fine. +pub(crate) trait ControlPlaneApi { + /// Get the client's auth secret for authentication. + /// Returns option because user not found situation is special. + /// We still have to mock the scram to avoid leaking information that user doesn't exist. + async fn get_role_secret( + &self, + ctx: &RequestMonitoring, + user_info: &ComputeUserInfo, + ) -> Result; + + async fn get_allowed_ips_and_secret( + &self, + ctx: &RequestMonitoring, + user_info: &ComputeUserInfo, + ) -> Result<(CachedAllowedIps, Option), errors::GetAuthInfoError>; + + async fn get_endpoint_jwks( + &self, + ctx: &RequestMonitoring, + endpoint: EndpointId, + ) -> Result, errors::GetEndpointJwksError>; + + /// Wake up the compute node and return the corresponding connection info. + async fn wake_compute( + &self, + ctx: &RequestMonitoring, + user_info: &ComputeUserInfo, + ) -> Result; +} diff --git a/proxy/src/control_plane/provider/mod.rs b/proxy/src/control_plane/provider/mod.rs deleted file mode 100644 index 49e57b6b7e..0000000000 --- a/proxy/src/control_plane/provider/mod.rs +++ /dev/null @@ -1,588 +0,0 @@ -#[cfg(any(test, feature = "testing"))] -pub mod mock; -pub mod neon; - -use std::hash::Hash; -use std::sync::Arc; -use std::time::Duration; - -use dashmap::DashMap; -use tokio::time::Instant; -use tracing::info; - -use super::messages::{ControlPlaneError, MetricsAuxInfo}; -use crate::auth::backend::jwt::{AuthRule, FetchAuthRules, FetchAuthRulesError}; -use crate::auth::backend::{ComputeCredentialKeys, ComputeUserInfo}; -use crate::auth::IpPattern; -use crate::cache::endpoints::EndpointsCache; -use crate::cache::project_info::ProjectInfoCacheImpl; -use crate::cache::{Cached, TimedLru}; -use crate::config::{CacheOptions, EndpointCacheConfig, ProjectInfoCacheOptions}; -use crate::context::RequestMonitoring; -use crate::error::ReportableError; -use crate::intern::ProjectIdInt; -use crate::metrics::ApiLockMetrics; -use crate::rate_limiter::{DynamicLimiter, Outcome, RateLimiterConfig, Token}; -use crate::types::{EndpointCacheKey, EndpointId}; -use crate::{compute, scram}; - -pub(crate) mod errors { - use thiserror::Error; - - use super::ApiLockError; - use crate::control_plane::messages::{self, ControlPlaneError, Reason}; - use crate::error::{io_error, ErrorKind, ReportableError, UserFacingError}; - use crate::proxy::retry::CouldRetry; - - /// A go-to error message which doesn't leak any detail. - pub(crate) const REQUEST_FAILED: &str = "Console request failed"; - - /// Common console API error. - #[derive(Debug, Error)] - pub(crate) enum ApiError { - /// Error returned by the console itself. - #[error("{REQUEST_FAILED} with {0}")] - ControlPlane(Box), - - /// Various IO errors like broken pipe or malformed payload. - #[error("{REQUEST_FAILED}: {0}")] - Transport(#[from] std::io::Error), - } - - impl ApiError { - /// Returns HTTP status code if it's the reason for failure. - pub(crate) fn get_reason(&self) -> messages::Reason { - match self { - ApiError::ControlPlane(e) => e.get_reason(), - ApiError::Transport(_) => messages::Reason::Unknown, - } - } - } - - impl UserFacingError for ApiError { - fn to_string_client(&self) -> String { - match self { - // To minimize risks, only select errors are forwarded to users. - ApiError::ControlPlane(c) => c.get_user_facing_message(), - ApiError::Transport(_) => REQUEST_FAILED.to_owned(), - } - } - } - - impl ReportableError for ApiError { - fn get_error_kind(&self) -> crate::error::ErrorKind { - match self { - ApiError::ControlPlane(e) => match e.get_reason() { - Reason::RoleProtected => ErrorKind::User, - Reason::ResourceNotFound => ErrorKind::User, - Reason::ProjectNotFound => ErrorKind::User, - Reason::EndpointNotFound => ErrorKind::User, - Reason::BranchNotFound => ErrorKind::User, - Reason::RateLimitExceeded => ErrorKind::ServiceRateLimit, - Reason::NonDefaultBranchComputeTimeExceeded => ErrorKind::Quota, - Reason::ActiveTimeQuotaExceeded => ErrorKind::Quota, - Reason::ComputeTimeQuotaExceeded => ErrorKind::Quota, - Reason::WrittenDataQuotaExceeded => ErrorKind::Quota, - Reason::DataTransferQuotaExceeded => ErrorKind::Quota, - Reason::LogicalSizeQuotaExceeded => ErrorKind::Quota, - Reason::ConcurrencyLimitReached => ErrorKind::ControlPlane, - Reason::LockAlreadyTaken => ErrorKind::ControlPlane, - Reason::RunningOperations => ErrorKind::ControlPlane, - Reason::ActiveEndpointsLimitExceeded => ErrorKind::ControlPlane, - Reason::Unknown => ErrorKind::ControlPlane, - }, - ApiError::Transport(_) => crate::error::ErrorKind::ControlPlane, - } - } - } - - impl CouldRetry for ApiError { - fn could_retry(&self) -> bool { - match self { - // retry some transport errors - Self::Transport(io) => io.could_retry(), - Self::ControlPlane(e) => e.could_retry(), - } - } - } - - impl From for ApiError { - fn from(e: reqwest::Error) -> Self { - io_error(e).into() - } - } - - impl From for ApiError { - fn from(e: reqwest_middleware::Error) -> Self { - io_error(e).into() - } - } - - #[derive(Debug, Error)] - pub(crate) enum GetAuthInfoError { - // We shouldn't include the actual secret here. - #[error("Console responded with a malformed auth secret")] - BadSecret, - - #[error(transparent)] - ApiError(ApiError), - } - - // This allows more useful interactions than `#[from]`. - impl> From for GetAuthInfoError { - fn from(e: E) -> Self { - Self::ApiError(e.into()) - } - } - - impl UserFacingError for GetAuthInfoError { - fn to_string_client(&self) -> String { - match self { - // We absolutely should not leak any secrets! - Self::BadSecret => REQUEST_FAILED.to_owned(), - // However, API might return a meaningful error. - Self::ApiError(e) => e.to_string_client(), - } - } - } - - impl ReportableError for GetAuthInfoError { - fn get_error_kind(&self) -> crate::error::ErrorKind { - match self { - Self::BadSecret => crate::error::ErrorKind::ControlPlane, - Self::ApiError(_) => crate::error::ErrorKind::ControlPlane, - } - } - } - - #[derive(Debug, Error)] - pub(crate) enum WakeComputeError { - #[error("Console responded with a malformed compute address: {0}")] - BadComputeAddress(Box), - - #[error(transparent)] - ApiError(ApiError), - - #[error("Too many connections attempts")] - TooManyConnections, - - #[error("error acquiring resource permit: {0}")] - TooManyConnectionAttempts(#[from] ApiLockError), - } - - // This allows more useful interactions than `#[from]`. - impl> From for WakeComputeError { - fn from(e: E) -> Self { - Self::ApiError(e.into()) - } - } - - impl UserFacingError for WakeComputeError { - fn to_string_client(&self) -> String { - match self { - // We shouldn't show user the address even if it's broken. - // Besides, user is unlikely to care about this detail. - Self::BadComputeAddress(_) => REQUEST_FAILED.to_owned(), - // However, API might return a meaningful error. - Self::ApiError(e) => e.to_string_client(), - - Self::TooManyConnections => self.to_string(), - - Self::TooManyConnectionAttempts(_) => { - "Failed to acquire permit to connect to the database. Too many database connection attempts are currently ongoing.".to_owned() - } - } - } - } - - impl ReportableError for WakeComputeError { - fn get_error_kind(&self) -> crate::error::ErrorKind { - match self { - Self::BadComputeAddress(_) => crate::error::ErrorKind::ControlPlane, - Self::ApiError(e) => e.get_error_kind(), - Self::TooManyConnections => crate::error::ErrorKind::RateLimit, - Self::TooManyConnectionAttempts(e) => e.get_error_kind(), - } - } - } - - impl CouldRetry for WakeComputeError { - fn could_retry(&self) -> bool { - match self { - Self::BadComputeAddress(_) => false, - Self::ApiError(e) => e.could_retry(), - Self::TooManyConnections => false, - Self::TooManyConnectionAttempts(_) => false, - } - } - } - - #[derive(Debug, Error)] - pub enum GetEndpointJwksError { - #[error("endpoint not found")] - EndpointNotFound, - - #[error("failed to build control plane request: {0}")] - RequestBuild(#[source] reqwest::Error), - - #[error("failed to send control plane request: {0}")] - RequestExecute(#[source] reqwest_middleware::Error), - - #[error(transparent)] - ControlPlane(#[from] ApiError), - - #[cfg(any(test, feature = "testing"))] - #[error(transparent)] - TokioPostgres(#[from] tokio_postgres::Error), - - #[cfg(any(test, feature = "testing"))] - #[error(transparent)] - ParseUrl(#[from] url::ParseError), - - #[cfg(any(test, feature = "testing"))] - #[error(transparent)] - TaskJoin(#[from] tokio::task::JoinError), - } -} - -/// Auth secret which is managed by the cloud. -#[derive(Clone, Eq, PartialEq, Debug)] -pub(crate) enum AuthSecret { - #[cfg(any(test, feature = "testing"))] - /// Md5 hash of user's password. - Md5([u8; 16]), - - /// [SCRAM](crate::scram) authentication info. - Scram(scram::ServerSecret), -} - -#[derive(Default)] -pub(crate) struct AuthInfo { - pub(crate) secret: Option, - /// List of IP addresses allowed for the autorization. - pub(crate) allowed_ips: Vec, - /// Project ID. This is used for cache invalidation. - pub(crate) project_id: Option, -} - -/// Info for establishing a connection to a compute node. -/// This is what we get after auth succeeded, but not before! -#[derive(Clone)] -pub(crate) struct NodeInfo { - /// Compute node connection params. - /// It's sad that we have to clone this, but this will improve - /// once we migrate to a bespoke connection logic. - pub(crate) config: compute::ConnCfg, - - /// Labels for proxy's metrics. - pub(crate) aux: MetricsAuxInfo, - - /// Whether we should accept self-signed certificates (for testing) - pub(crate) allow_self_signed_compute: bool, -} - -impl NodeInfo { - pub(crate) async fn connect( - &self, - ctx: &RequestMonitoring, - timeout: Duration, - ) -> Result { - self.config - .connect( - ctx, - self.allow_self_signed_compute, - self.aux.clone(), - timeout, - ) - .await - } - pub(crate) fn reuse_settings(&mut self, other: Self) { - self.allow_self_signed_compute = other.allow_self_signed_compute; - self.config.reuse_password(other.config); - } - - pub(crate) fn set_keys(&mut self, keys: &ComputeCredentialKeys) { - match keys { - #[cfg(any(test, feature = "testing"))] - ComputeCredentialKeys::Password(password) => self.config.password(password), - ComputeCredentialKeys::AuthKeys(auth_keys) => self.config.auth_keys(*auth_keys), - ComputeCredentialKeys::JwtPayload(_) | ComputeCredentialKeys::None => &mut self.config, - }; - } -} - -pub(crate) type NodeInfoCache = - TimedLru>>; -pub(crate) type CachedNodeInfo = Cached<&'static NodeInfoCache, NodeInfo>; -pub(crate) type CachedRoleSecret = Cached<&'static ProjectInfoCacheImpl, Option>; -pub(crate) type CachedAllowedIps = Cached<&'static ProjectInfoCacheImpl, Arc>>; - -/// This will allocate per each call, but the http requests alone -/// already require a few allocations, so it should be fine. -pub(crate) trait Api { - /// Get the client's auth secret for authentication. - /// Returns option because user not found situation is special. - /// We still have to mock the scram to avoid leaking information that user doesn't exist. - async fn get_role_secret( - &self, - ctx: &RequestMonitoring, - user_info: &ComputeUserInfo, - ) -> Result; - - async fn get_allowed_ips_and_secret( - &self, - ctx: &RequestMonitoring, - user_info: &ComputeUserInfo, - ) -> Result<(CachedAllowedIps, Option), errors::GetAuthInfoError>; - - async fn get_endpoint_jwks( - &self, - ctx: &RequestMonitoring, - endpoint: EndpointId, - ) -> Result, errors::GetEndpointJwksError>; - - /// Wake up the compute node and return the corresponding connection info. - async fn wake_compute( - &self, - ctx: &RequestMonitoring, - user_info: &ComputeUserInfo, - ) -> Result; -} - -#[non_exhaustive] -#[derive(Clone)] -pub enum ControlPlaneBackend { - /// Current Management API (V2). - Management(neon::Api), - /// Local mock control plane. - #[cfg(any(test, feature = "testing"))] - PostgresMock(mock::Api), - /// Internal testing - #[cfg(test)] - #[allow(private_interfaces)] - Test(Box), -} - -impl Api for ControlPlaneBackend { - async fn get_role_secret( - &self, - ctx: &RequestMonitoring, - user_info: &ComputeUserInfo, - ) -> Result { - match self { - Self::Management(api) => api.get_role_secret(ctx, user_info).await, - #[cfg(any(test, feature = "testing"))] - Self::PostgresMock(api) => api.get_role_secret(ctx, user_info).await, - #[cfg(test)] - Self::Test(_) => { - unreachable!("this function should never be called in the test backend") - } - } - } - - async fn get_allowed_ips_and_secret( - &self, - ctx: &RequestMonitoring, - user_info: &ComputeUserInfo, - ) -> Result<(CachedAllowedIps, Option), errors::GetAuthInfoError> { - match self { - Self::Management(api) => api.get_allowed_ips_and_secret(ctx, user_info).await, - #[cfg(any(test, feature = "testing"))] - Self::PostgresMock(api) => api.get_allowed_ips_and_secret(ctx, user_info).await, - #[cfg(test)] - Self::Test(api) => api.get_allowed_ips_and_secret(), - } - } - - async fn get_endpoint_jwks( - &self, - ctx: &RequestMonitoring, - endpoint: EndpointId, - ) -> Result, errors::GetEndpointJwksError> { - match self { - Self::Management(api) => api.get_endpoint_jwks(ctx, endpoint).await, - #[cfg(any(test, feature = "testing"))] - Self::PostgresMock(api) => api.get_endpoint_jwks(ctx, endpoint).await, - #[cfg(test)] - Self::Test(_api) => Ok(vec![]), - } - } - - async fn wake_compute( - &self, - ctx: &RequestMonitoring, - user_info: &ComputeUserInfo, - ) -> Result { - match self { - Self::Management(api) => api.wake_compute(ctx, user_info).await, - #[cfg(any(test, feature = "testing"))] - Self::PostgresMock(api) => api.wake_compute(ctx, user_info).await, - #[cfg(test)] - Self::Test(api) => api.wake_compute(), - } - } -} - -/// Various caches for [`control_plane`](super). -pub struct ApiCaches { - /// Cache for the `wake_compute` API method. - pub(crate) node_info: NodeInfoCache, - /// Cache which stores project_id -> endpoint_ids mapping. - pub project_info: Arc, - /// List of all valid endpoints. - pub endpoints_cache: Arc, -} - -impl ApiCaches { - pub fn new( - wake_compute_cache_config: CacheOptions, - project_info_cache_config: ProjectInfoCacheOptions, - endpoint_cache_config: EndpointCacheConfig, - ) -> Self { - Self { - node_info: NodeInfoCache::new( - "node_info_cache", - wake_compute_cache_config.size, - wake_compute_cache_config.ttl, - true, - ), - project_info: Arc::new(ProjectInfoCacheImpl::new(project_info_cache_config)), - endpoints_cache: Arc::new(EndpointsCache::new(endpoint_cache_config)), - } - } -} - -/// Various caches for [`control_plane`](super). -pub struct ApiLocks { - name: &'static str, - node_locks: DashMap>, - config: RateLimiterConfig, - timeout: Duration, - epoch: std::time::Duration, - metrics: &'static ApiLockMetrics, -} - -#[derive(Debug, thiserror::Error)] -pub(crate) enum ApiLockError { - #[error("timeout acquiring resource permit")] - TimeoutError(#[from] tokio::time::error::Elapsed), -} - -impl ReportableError for ApiLockError { - fn get_error_kind(&self) -> crate::error::ErrorKind { - match self { - ApiLockError::TimeoutError(_) => crate::error::ErrorKind::RateLimit, - } - } -} - -impl ApiLocks { - pub fn new( - name: &'static str, - config: RateLimiterConfig, - shards: usize, - timeout: Duration, - epoch: std::time::Duration, - metrics: &'static ApiLockMetrics, - ) -> prometheus::Result { - Ok(Self { - name, - node_locks: DashMap::with_shard_amount(shards), - config, - timeout, - epoch, - metrics, - }) - } - - pub(crate) async fn get_permit(&self, key: &K) -> Result { - if self.config.initial_limit == 0 { - return Ok(WakeComputePermit { - permit: Token::disabled(), - }); - } - 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.metrics.semaphores_registered.inc(); - DynamicLimiter::new(self.config) - }) - .clone() - } - }; - let permit = semaphore.acquire_timeout(self.timeout).await; - - self.metrics - .semaphore_acquire_seconds - .observe(now.elapsed().as_secs_f64()); - info!("acquired permit {:?}", now.elapsed().as_secs_f64()); - Ok(WakeComputePermit { permit: permit? }) - } - - pub async fn garbage_collect_worker(&self) { - if self.config.initial_limit == 0 { - return; - } - let mut interval = - tokio::time::interval(self.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.metrics.reclamation_lag_seconds.start_timer(); - let count = lock - .extract_if(|_, semaphore| Arc::strong_count(semaphore.get_mut()) == 1) - .count(); - drop(lock); - self.metrics.semaphores_unregistered.inc_by(count as u64); - timer.observe(); - } - } - } -} - -pub(crate) struct WakeComputePermit { - permit: Token, -} - -impl WakeComputePermit { - pub(crate) fn should_check_cache(&self) -> bool { - !self.permit.is_disabled() - } - pub(crate) fn release(self, outcome: Outcome) { - self.permit.release(outcome); - } - pub(crate) fn release_result(self, res: Result) -> Result { - match res { - Ok(_) => self.release(Outcome::Success), - Err(_) => self.release(Outcome::Overload), - } - res - } -} - -impl FetchAuthRules for ControlPlaneBackend { - async fn fetch_auth_rules( - &self, - ctx: &RequestMonitoring, - endpoint: EndpointId, - ) -> Result, FetchAuthRulesError> { - self.get_endpoint_jwks(ctx, endpoint) - .await - .map_err(FetchAuthRulesError::GetEndpointJwks) - } -} diff --git a/proxy/src/proxy/tests/mod.rs b/proxy/src/proxy/tests/mod.rs index abb0599d08..be821925b5 100644 --- a/proxy/src/proxy/tests/mod.rs +++ b/proxy/src/proxy/tests/mod.rs @@ -20,14 +20,14 @@ use super::connect_compute::ConnectMechanism; use super::retry::CouldRetry; use super::*; use crate::auth::backend::{ - ComputeCredentialKeys, ComputeCredentials, ComputeUserInfo, MaybeOwned, TestBackend, + ComputeCredentialKeys, ComputeCredentials, ComputeUserInfo, MaybeOwned, }; use crate::config::{CertResolver, RetryConfig}; -use crate::control_plane::messages::{ControlPlaneError, Details, MetricsAuxInfo, Status}; -use crate::control_plane::provider::{ - CachedAllowedIps, CachedRoleSecret, ControlPlaneBackend, NodeInfoCache, +use crate::control_plane::client::{ControlPlaneClient, TestControlPlaneClient}; +use crate::control_plane::messages::{ControlPlaneErrorMessage, Details, MetricsAuxInfo, Status}; +use crate::control_plane::{ + self, CachedAllowedIps, CachedNodeInfo, CachedRoleSecret, NodeInfo, NodeInfoCache, }; -use crate::control_plane::{self, CachedNodeInfo, NodeInfo}; use crate::error::ErrorKind; use crate::types::{BranchId, EndpointId, ProjectId}; use crate::{sasl, scram}; @@ -490,7 +490,7 @@ impl ConnectMechanism for TestConnectMechanism { fn update_connect_config(&self, _conf: &mut compute::ConnCfg) {} } -impl TestBackend for TestConnectMechanism { +impl TestControlPlaneClient for TestConnectMechanism { fn wake_compute(&self) -> Result { let mut counter = self.counter.lock().unwrap(); let action = self.sequence[*counter]; @@ -498,18 +498,19 @@ impl TestBackend for TestConnectMechanism { match action { ConnectAction::Wake => Ok(helper_create_cached_node_info(self.cache)), ConnectAction::WakeFail => { - let err = - control_plane::errors::ApiError::ControlPlane(Box::new(ControlPlaneError { + let err = control_plane::errors::ControlPlaneError::Message(Box::new( + ControlPlaneErrorMessage { http_status_code: StatusCode::BAD_REQUEST, error: "TEST".into(), status: None, - })); + }, + )); assert!(!err.could_retry()); - Err(control_plane::errors::WakeComputeError::ApiError(err)) + Err(control_plane::errors::WakeComputeError::ControlPlane(err)) } ConnectAction::WakeRetry => { - let err = - control_plane::errors::ApiError::ControlPlane(Box::new(ControlPlaneError { + let err = control_plane::errors::ControlPlaneError::Message(Box::new( + ControlPlaneErrorMessage { http_status_code: StatusCode::BAD_REQUEST, error: "TEST".into(), status: Some(Status { @@ -523,9 +524,10 @@ impl TestBackend for TestConnectMechanism { user_facing_message: None, }, }), - })); + }, + )); assert!(err.could_retry()); - Err(control_plane::errors::WakeComputeError::ApiError(err)) + Err(control_plane::errors::WakeComputeError::ControlPlane(err)) } x => panic!("expecting action {x:?}, wake_compute is called instead"), } @@ -538,7 +540,7 @@ impl TestBackend for TestConnectMechanism { unimplemented!("not used in tests") } - fn dyn_clone(&self) -> Box { + fn dyn_clone(&self) -> Box { Box::new(self.clone()) } } @@ -562,7 +564,7 @@ fn helper_create_connect_info( mechanism: &TestConnectMechanism, ) -> auth::Backend<'static, ComputeCredentials> { let user_info = auth::Backend::ControlPlane( - MaybeOwned::Owned(ControlPlaneBackend::Test(Box::new(mechanism.clone()))), + MaybeOwned::Owned(ControlPlaneClient::Test(Box::new(mechanism.clone()))), ComputeCredentials { info: ComputeUserInfo { endpoint: "endpoint".into(), diff --git a/proxy/src/proxy/wake_compute.rs b/proxy/src/proxy/wake_compute.rs index 4e61094264..f9f46bb66c 100644 --- a/proxy/src/proxy/wake_compute.rs +++ b/proxy/src/proxy/wake_compute.rs @@ -4,7 +4,7 @@ use super::connect_compute::ComputeConnectBackend; use crate::config::RetryConfig; use crate::context::RequestMonitoring; use crate::control_plane::errors::WakeComputeError; -use crate::control_plane::provider::CachedNodeInfo; +use crate::control_plane::CachedNodeInfo; use crate::error::ReportableError; use crate::metrics::{ ConnectOutcome, ConnectionFailuresBreakdownGroup, Metrics, RetriesMetricGroup, RetryType, diff --git a/proxy/src/serverless/backend.rs b/proxy/src/serverless/backend.rs index c2b0de1876..7fc5bd236d 100644 --- a/proxy/src/serverless/backend.rs +++ b/proxy/src/serverless/backend.rs @@ -24,9 +24,9 @@ use crate::compute_ctl::{ }; use crate::config::ProxyConfig; use crate::context::RequestMonitoring; +use crate::control_plane::client::ApiLockError; use crate::control_plane::errors::{GetAuthInfoError, WakeComputeError}; use crate::control_plane::locks::ApiLocks; -use crate::control_plane::provider::ApiLockError; use crate::control_plane::CachedNodeInfo; use crate::error::{ErrorKind, ReportableError, UserFacingError}; use crate::intern::EndpointIdInt; From 5d8284c7fec4b8c58cd66a1b6219fb09d08f504b Mon Sep 17 00:00:00 2001 From: Folke Behrens Date: Wed, 6 Nov 2024 11:27:55 +0100 Subject: [PATCH 19/52] proxy: Read cplane JWT with clap arg (#9654) --- proxy/Cargo.toml | 2 +- proxy/src/bin/proxy.rs | 9 +++++++++ proxy/src/control_plane/client/neon.rs | 4 +--- workspace_hack/Cargo.toml | 4 ++-- 4 files changed, 13 insertions(+), 6 deletions(-) diff --git a/proxy/Cargo.toml b/proxy/Cargo.toml index bc4da95a91..efd336dbea 100644 --- a/proxy/Cargo.toml +++ b/proxy/Cargo.toml @@ -23,7 +23,7 @@ bstr.workspace = true bytes = { workspace = true, features = ["serde"] } camino.workspace = true chrono.workspace = true -clap.workspace = true +clap = { workspace = true, features = ["derive", "env"] } compute_api.workspace = true consumption_metrics.workspace = true dashmap.workspace = true diff --git a/proxy/src/bin/proxy.rs b/proxy/src/bin/proxy.rs index 3179a929c4..efb3747829 100644 --- a/proxy/src/bin/proxy.rs +++ b/proxy/src/bin/proxy.rs @@ -92,6 +92,14 @@ struct ProxyCliArgs { default_value = "http://localhost:3000/authenticate_proxy_request/" )] auth_endpoint: String, + /// JWT used to connect to control plane. + #[clap( + long, + value_name = "JWT", + default_value = "", + env = "NEON_PROXY_TO_CONTROLPLANE_TOKEN" + )] + control_plane_token: Arc, /// if this is not local proxy, this toggles whether we accept jwt or passwords for http #[clap(long, default_value_t = false, value_parser = clap::builder::BoolishValueParser::new(), action = clap::ArgAction::Set)] is_auth_broker: bool, @@ -734,6 +742,7 @@ fn build_auth_backend( Arc::new(WakeComputeRateLimiter::new(wake_compute_rps_limit)); let api = control_plane::client::neon::NeonControlPlaneClient::new( endpoint, + args.control_plane_token.clone(), caches, locks, wake_compute_endpoint_rate_limiter, diff --git a/proxy/src/control_plane/client/neon.rs b/proxy/src/control_plane/client/neon.rs index 6c67d2df96..1588e50423 100644 --- a/proxy/src/control_plane/client/neon.rs +++ b/proxy/src/control_plane/client/neon.rs @@ -45,13 +45,11 @@ impl NeonControlPlaneClient { /// Construct an API object containing the auth parameters. pub fn new( endpoint: http::Endpoint, + jwt: Arc, caches: &'static ApiCaches, locks: &'static ApiLocks, wake_compute_endpoint_rate_limiter: Arc, ) -> Self { - let jwt = std::env::var("NEON_PROXY_TO_CONTROLPLANE_TOKEN") - .unwrap_or_default() - .into(); Self { endpoint, caches, diff --git a/workspace_hack/Cargo.toml b/workspace_hack/Cargo.toml index 71ebab4119..02deecd385 100644 --- a/workspace_hack/Cargo.toml +++ b/workspace_hack/Cargo.toml @@ -24,8 +24,8 @@ base64ct = { version = "1", default-features = false, features = ["std"] } bytes = { version = "1", features = ["serde"] } camino = { version = "1", default-features = false, features = ["serde1"] } chrono = { version = "0.4", default-features = false, features = ["clock", "serde", "wasmbind"] } -clap = { version = "4", features = ["derive", "string"] } -clap_builder = { version = "4", default-features = false, features = ["color", "help", "std", "string", "suggestions", "usage"] } +clap = { version = "4", features = ["derive", "env", "string"] } +clap_builder = { version = "4", default-features = false, features = ["color", "env", "help", "std", "string", "suggestions", "usage"] } crypto-bigint = { version = "0.5", features = ["generic-array", "zeroize"] } der = { version = "0.7", default-features = false, features = ["oid", "pem", "std"] } deranged = { version = "0.3", default-features = false, features = ["powerfmt", "serde", "std"] } From bdd492b1d871271c94462736e438e8d38b1fa997 Mon Sep 17 00:00:00 2001 From: Folke Behrens Date: Wed, 6 Nov 2024 12:03:38 +0100 Subject: [PATCH 20/52] proxy: Replace "web(auth)" with "console redirect" everywhere (#9655) --- proxy/src/auth/backend/console_redirect.rs | 20 +++++++++++--------- proxy/src/auth/backend/mod.rs | 4 ++-- proxy/src/auth/mod.rs | 6 +++--- proxy/src/bin/local_proxy.rs | 2 +- proxy/src/bin/proxy.rs | 21 +++++++++++---------- proxy/src/compute.rs | 4 ++-- proxy/src/config.rs | 4 ++-- proxy/src/context/mod.rs | 2 +- proxy/src/context/parquet.rs | 2 +- proxy/src/control_plane/messages.rs | 2 +- proxy/src/control_plane/mgmt.rs | 8 ++++---- 11 files changed, 39 insertions(+), 36 deletions(-) diff --git a/proxy/src/auth/backend/console_redirect.rs b/proxy/src/auth/backend/console_redirect.rs index 1df59b4893..e25dc3d45e 100644 --- a/proxy/src/auth/backend/console_redirect.rs +++ b/proxy/src/auth/backend/console_redirect.rs @@ -16,7 +16,7 @@ use crate::stream::PqStream; use crate::{auth, compute, waiters}; #[derive(Debug, Error)] -pub(crate) enum WebAuthError { +pub(crate) enum ConsoleRedirectError { #[error(transparent)] WaiterRegister(#[from] waiters::RegisterError), @@ -32,13 +32,13 @@ pub struct ConsoleRedirectBackend { console_uri: reqwest::Url, } -impl UserFacingError for WebAuthError { +impl UserFacingError for ConsoleRedirectError { fn to_string_client(&self) -> String { "Internal error".to_string() } } -impl ReportableError for WebAuthError { +impl ReportableError for ConsoleRedirectError { fn get_error_kind(&self) -> crate::error::ErrorKind { match self { Self::WaiterRegister(_) => crate::error::ErrorKind::Service, @@ -103,7 +103,7 @@ async fn authenticate( link_uri: &reqwest::Url, client: &mut PqStream, ) -> auth::Result { - ctx.set_auth_method(crate::context::AuthMethod::Web); + ctx.set_auth_method(crate::context::AuthMethod::ConsoleRedirect); // registering waiter can fail if we get unlucky with rng. // just try again. @@ -116,7 +116,7 @@ async fn authenticate( } }; - let span = info_span!("web", psql_session_id = &psql_session_id); + let span = info_span!("console_redirect", psql_session_id = &psql_session_id); let greeting = hello_message(link_uri, &psql_session_id); // Give user a URL to spawn a new database. @@ -127,14 +127,16 @@ async fn authenticate( .write_message(&Be::NoticeResponse(&greeting)) .await?; - // Wait for web console response (see `mgmt`). + // Wait for console response via control plane (see `mgmt`). info!(parent: &span, "waiting for console's reply..."); - let db_info = tokio::time::timeout(auth_config.webauth_confirmation_timeout, waiter) + let db_info = tokio::time::timeout(auth_config.console_redirect_confirmation_timeout, waiter) .await .map_err(|_elapsed| { - auth::AuthError::confirmation_timeout(auth_config.webauth_confirmation_timeout.into()) + auth::AuthError::confirmation_timeout( + auth_config.console_redirect_confirmation_timeout.into(), + ) })? - .map_err(WebAuthError::from)?; + .map_err(ConsoleRedirectError::from)?; if auth_config.ip_allowlist_check_enabled { if let Some(allowed_ips) = &db_info.allowed_ips { diff --git a/proxy/src/auth/backend/mod.rs b/proxy/src/auth/backend/mod.rs index ccd5cf3071..242fe99de2 100644 --- a/proxy/src/auth/backend/mod.rs +++ b/proxy/src/auth/backend/mod.rs @@ -9,7 +9,7 @@ use std::sync::Arc; use std::time::Duration; pub use console_redirect::ConsoleRedirectBackend; -pub(crate) use console_redirect::WebAuthError; +pub(crate) use console_redirect::ConsoleRedirectError; use ipnet::{Ipv4Net, Ipv6Net}; use local::LocalBackend; use tokio::io::{AsyncRead, AsyncWrite}; @@ -560,7 +560,7 @@ mod tests { ip_allowlist_check_enabled: true, is_auth_broker: false, accept_jwts: false, - webauth_confirmation_timeout: std::time::Duration::from_secs(5), + console_redirect_confirmation_timeout: std::time::Duration::from_secs(5), }); async fn read_message(r: &mut (impl AsyncRead + Unpin), b: &mut BytesMut) -> PgMessage { diff --git a/proxy/src/auth/mod.rs b/proxy/src/auth/mod.rs index 2bd7a2da3d..0198cc306e 100644 --- a/proxy/src/auth/mod.rs +++ b/proxy/src/auth/mod.rs @@ -32,7 +32,7 @@ pub(crate) type Result = std::result::Result; #[derive(Debug, Error)] pub(crate) enum AuthError { #[error(transparent)] - Web(#[from] backend::WebAuthError), + ConsoleRedirect(#[from] backend::ConsoleRedirectError), #[error(transparent)] GetAuthInfo(#[from] control_plane::errors::GetAuthInfoError), @@ -115,7 +115,7 @@ impl AuthError { impl UserFacingError for AuthError { fn to_string_client(&self) -> String { match self { - Self::Web(e) => e.to_string_client(), + Self::ConsoleRedirect(e) => e.to_string_client(), Self::GetAuthInfo(e) => e.to_string_client(), Self::Sasl(e) => e.to_string_client(), Self::PasswordFailed(_) => self.to_string(), @@ -135,7 +135,7 @@ impl UserFacingError for AuthError { impl ReportableError for AuthError { fn get_error_kind(&self) -> crate::error::ErrorKind { match self { - Self::Web(e) => e.get_error_kind(), + Self::ConsoleRedirect(e) => e.get_error_kind(), Self::GetAuthInfo(e) => e.get_error_kind(), Self::Sasl(e) => e.get_error_kind(), Self::PasswordFailed(_) => crate::error::ErrorKind::User, diff --git a/proxy/src/bin/local_proxy.rs b/proxy/src/bin/local_proxy.rs index df3628465f..fbdb1dec15 100644 --- a/proxy/src/bin/local_proxy.rs +++ b/proxy/src/bin/local_proxy.rs @@ -281,7 +281,7 @@ fn build_config(args: &LocalProxyCliArgs) -> anyhow::Result<&'static ProxyConfig ip_allowlist_check_enabled: true, is_auth_broker: false, accept_jwts: true, - webauth_confirmation_timeout: Duration::ZERO, + console_redirect_confirmation_timeout: Duration::ZERO, }, proxy_protocol_v2: config::ProxyProtocolV2::Rejected, handshake_timeout: Duration::from_secs(10), diff --git a/proxy/src/bin/proxy.rs b/proxy/src/bin/proxy.rs index efb3747829..fda5b25961 100644 --- a/proxy/src/bin/proxy.rs +++ b/proxy/src/bin/proxy.rs @@ -51,11 +51,11 @@ static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; #[derive(Clone, Debug, ValueEnum)] enum AuthBackendType { - Console, - // clap only shows the name, not the alias, in usage text. - // TODO: swap name/alias and deprecate "link" - #[value(name("link"), alias("web"))] - Web, + #[value(name("console"), alias("cplane"))] + ControlPlane, + + #[value(name("link"), alias("control-redirect"))] + ConsoleRedirect, #[cfg(feature = "testing")] Postgres, @@ -71,7 +71,7 @@ struct ProxyCliArgs { /// listen for incoming client connections on ip:port #[clap(short, long, default_value = "127.0.0.1:4432")] proxy: String, - #[clap(value_enum, long, default_value_t = AuthBackendType::Web)] + #[clap(value_enum, long, default_value_t = AuthBackendType::ConsoleRedirect)] auth_backend: AuthBackendType, /// listen for management callback connection on ip:port #[clap(short, long, default_value = "127.0.0.1:7000")] @@ -82,7 +82,7 @@ struct ProxyCliArgs { /// listen for incoming wss connections on ip:port #[clap(long)] wss: Option, - /// redirect unauthenticated users to the given uri in case of web auth + /// redirect unauthenticated users to the given uri in case of console redirect auth #[clap(short, long, default_value = "http://localhost:3000/psql_session/")] uri: String, /// cloud API endpoint for authenticating users @@ -231,6 +231,7 @@ struct ProxyCliArgs { proxy_protocol_v2: ProxyProtocolV2, /// Time the proxy waits for the webauth session to be confirmed by the control plane. + // TODO: rename to `console_redirect_confirmation_timeout`. #[clap(long, default_value = "2m", value_parser = humantime::parse_duration)] webauth_confirmation_timeout: std::time::Duration, } @@ -667,7 +668,7 @@ fn build_config(args: &ProxyCliArgs) -> anyhow::Result<&'static ProxyConfig> { ip_allowlist_check_enabled: !args.is_private_access_proxy, is_auth_broker: args.is_auth_broker, accept_jwts: args.is_auth_broker, - webauth_confirmation_timeout: args.webauth_confirmation_timeout, + console_redirect_confirmation_timeout: args.webauth_confirmation_timeout, }; let config = ProxyConfig { @@ -698,7 +699,7 @@ fn build_auth_backend( args: &ProxyCliArgs, ) -> anyhow::Result, &'static ConsoleRedirectBackend>> { match &args.auth_backend { - AuthBackendType::Console => { + AuthBackendType::ControlPlane => { let wake_compute_cache_config: CacheOptions = args.wake_compute_cache.parse()?; let project_info_cache_config: ProjectInfoCacheOptions = args.project_info_cache.parse()?; @@ -771,7 +772,7 @@ fn build_auth_backend( Ok(Either::Left(config)) } - AuthBackendType::Web => { + AuthBackendType::ConsoleRedirect => { let url = args.uri.parse()?; let backend = ConsoleRedirectBackend::new(url); diff --git a/proxy/src/compute.rs b/proxy/src/compute.rs index d397fc5160..ca4a348ed8 100644 --- a/proxy/src/compute.rs +++ b/proxy/src/compute.rs @@ -135,13 +135,13 @@ impl ConnCfg { /// Apply startup message params to the connection config. pub(crate) fn set_startup_params(&mut self, params: &StartupMessageParams) { // Only set `user` if it's not present in the config. - // Web auth flow takes username from the console's response. + // Console redirect auth flow takes username from the console's response. if let (None, Some(user)) = (self.get_user(), params.get("user")) { self.user(user); } // Only set `dbname` if it's not present in the config. - // Web auth flow takes dbname from the console's response. + // Console redirect auth flow takes dbname from the console's response. if let (None, Some(dbname)) = (self.get_dbname(), params.get("database")) { self.dbname(dbname); } diff --git a/proxy/src/config.rs b/proxy/src/config.rs index f63d7e45aa..b048c9d389 100644 --- a/proxy/src/config.rs +++ b/proxy/src/config.rs @@ -78,7 +78,7 @@ pub struct AuthenticationConfig { pub jwks_cache: JwkCache, pub is_auth_broker: bool, pub accept_jwts: bool, - pub webauth_confirmation_timeout: tokio::time::Duration, + pub console_redirect_confirmation_timeout: tokio::time::Duration, } impl TlsConfig { @@ -271,7 +271,7 @@ impl CertResolver { // auth-broker does not use SNI and instead uses the Neon-Connection-String header. // Auth broker has the subdomain `apiauth` we need to remove for the purposes of validating the Neon-Connection-String. // - // Console Web proxy does not use any wildcard domains and does not need any certificate selection or conn string + // Console Redirect proxy does not use any wildcard domains and does not need any certificate selection or conn string // validation, so let's we can continue with any common-name let common_name = if let Some(s) = common_name.strip_prefix("CN=*.") { s.to_string() diff --git a/proxy/src/context/mod.rs b/proxy/src/context/mod.rs index 2a6c9c5969..6cf99c0c97 100644 --- a/proxy/src/context/mod.rs +++ b/proxy/src/context/mod.rs @@ -75,7 +75,7 @@ struct RequestMonitoringInner { #[derive(Clone, Debug)] pub(crate) enum AuthMethod { // aka passwordless, fka link - Web, + ConsoleRedirect, ScramSha256, ScramSha256Plus, Cleartext, diff --git a/proxy/src/context/parquet.rs b/proxy/src/context/parquet.rs index adbb74c8e5..4112de646f 100644 --- a/proxy/src/context/parquet.rs +++ b/proxy/src/context/parquet.rs @@ -134,7 +134,7 @@ impl From<&RequestMonitoringInner> for RequestData { .as_ref() .and_then(|options| serde_json::to_string(&Options { options }).ok()), auth_method: value.auth_method.as_ref().map(|x| match x { - super::AuthMethod::Web => "web", + super::AuthMethod::ConsoleRedirect => "console_redirect", super::AuthMethod::ScramSha256 => "scram_sha_256", super::AuthMethod::ScramSha256Plus => "scram_sha_256_plus", super::AuthMethod::Cleartext => "cleartext", diff --git a/proxy/src/control_plane/messages.rs b/proxy/src/control_plane/messages.rs index 75c932e6ab..8762ba874b 100644 --- a/proxy/src/control_plane/messages.rs +++ b/proxy/src/control_plane/messages.rs @@ -245,7 +245,7 @@ pub(crate) struct WakeCompute { pub(crate) aux: MetricsAuxInfo, } -/// Async response which concludes the web auth flow. +/// Async response which concludes the console redirect auth flow. /// Also known as `kickResponse` in the console. #[derive(Debug, Deserialize)] pub(crate) struct KickSession<'a> { diff --git a/proxy/src/control_plane/mgmt.rs b/proxy/src/control_plane/mgmt.rs index 5ac3acd28a..2f7359240d 100644 --- a/proxy/src/control_plane/mgmt.rs +++ b/proxy/src/control_plane/mgmt.rs @@ -24,8 +24,8 @@ pub(crate) fn notify(psql_session_id: &str, msg: ComputeReady) -> Result<(), wai CPLANE_WAITERS.notify(psql_session_id, msg) } -/// Console management API listener task. -/// It spawns console response handlers needed for the web auth. +/// Management API listener task. +/// It spawns management response handlers needed for the console redirect auth flow. pub async fn task_main(listener: TcpListener) -> anyhow::Result { scopeguard::defer! { info!("mgmt has shut down"); @@ -43,13 +43,13 @@ pub async fn task_main(listener: TcpListener) -> anyhow::Result { tokio::task::spawn( async move { - info!("serving a new console management API connection"); + info!("serving a new management API connection"); // these might be long running connections, have a separate logging for cancelling // on shutdown and other ways of stopping. let cancelled = scopeguard::guard(tracing::Span::current(), |span| { let _e = span.entered(); - info!("console management API task cancelled"); + info!("management API task cancelled"); }); if let Err(e) = handle_connection(socket).await { From 4dfa0c221b1ea4164e4260da8f62bcd98e816920 Mon Sep 17 00:00:00 2001 From: Vlad Lazar Date: Wed, 6 Nov 2024 14:10:32 +0000 Subject: [PATCH 21/52] pageserver: ingest pre-serialized batches of values (#9579) ## Problem https://github.com/neondatabase/neon/pull/9524 split the decoding and interpretation step from ingestion. The output of the first phase is a `wal_decoder::models::InterpretedWalRecord`. Before this patch set that struct contained a list of `Value` instances. We wish to lift the decoding and interpretation step to the safekeeper, but it would be nice if the safekeeper gave us a batch containing the raw data instead of actual values. ## Summary of changes Main goal here is to make `InterpretedWalRecord` hold a raw buffer which contains pre-serialized Values. For this we do: 1. Add a `SerializedValueBatch` type. This is `inmemory_layer::SerializedBatch` with some extra functionality for extension, observing values for shard 0 and tests. 2. Replace `inmemory_layer::SerializedBatch` with `SerializedValueBatch` 3. Make `DatadirModification` maintain a `SerializedValueBatch`. ### `DatadirModification` changes `DatadirModification` now maintains a `SerializedValueBatch` and extends it as new WAL records come in (to avoid flushing to disk on every record). In turn, this cascaded into a number of modifications to `DatadirModification`: 1. Replace `pending_data_pages` and `pending_zero_data_pages` with `pending_data_batch`. 2. Removal of `pending_zero_data_pages` and its cousin `on_wal_record_end` 3. Rename `pending_bytes` to `pending_metadata_bytes` since this is what it tracks now. 4. Adapting of various utility methods like `len`, `approx_pending_bytes` and `has_dirty_data_pages`. Removal of `pending_zero_data_pages` and the optimisation associated with it ((1) and (2)) deserves more detail. Previously all zero data pages went through `pending_zero_data_pages`. We wrote zero data pages when filling gaps caused by relation extension (case A) and when handling special wal records (case B). If it happened that the same WAL record contained a non zero write for an entry in `pending_zero_data_pages` we skipped the zero write. Case A: We handle this differently now. When ingesting the `SerialiezdValueBatch` associated with one PG WAL record, we identify the gaps and fill the them in one go. Essentially, we move from a per key process (gaps were filled after each new key), and replace it with a per record process. Hence, the optimisation is not required anymore. Case B: When the handling of a special record needs to zero out a key, it just adds that to the current batch. I inspected the code, and I don't think the optimisation kicked in here. --- libs/wal_decoder/Cargo.toml | 2 +- libs/wal_decoder/src/decoder.rs | 113 +-- libs/wal_decoder/src/lib.rs | 1 + libs/wal_decoder/src/models.rs | 16 +- libs/wal_decoder/src/serialized_batch.rs | 862 ++++++++++++++++++ pageserver/benches/bench_ingest.rs | 6 +- pageserver/src/pgdatadir_mapping.rs | 361 ++++++-- .../tenant/storage_layer/inmemory_layer.rs | 116 +-- pageserver/src/tenant/timeline.rs | 35 +- .../walreceiver/walreceiver_connection.rs | 14 +- pageserver/src/walingest.rs | 84 +- 11 files changed, 1247 insertions(+), 363 deletions(-) create mode 100644 libs/wal_decoder/src/serialized_batch.rs diff --git a/libs/wal_decoder/Cargo.toml b/libs/wal_decoder/Cargo.toml index 3f80f8fcdb..c8c0f4c990 100644 --- a/libs/wal_decoder/Cargo.toml +++ b/libs/wal_decoder/Cargo.toml @@ -5,7 +5,7 @@ edition.workspace = true license.workspace = true [features] -testing = [] +testing = ["pageserver_api/testing"] [dependencies] anyhow.workspace = true diff --git a/libs/wal_decoder/src/decoder.rs b/libs/wal_decoder/src/decoder.rs index 780fce3d69..684718d220 100644 --- a/libs/wal_decoder/src/decoder.rs +++ b/libs/wal_decoder/src/decoder.rs @@ -2,15 +2,13 @@ //! raw bytes which represent a raw Postgres WAL record. use crate::models::*; -use bytes::{Buf, Bytes, BytesMut}; -use pageserver_api::key::rel_block_to_key; -use pageserver_api::record::NeonWalRecord; +use crate::serialized_batch::SerializedValueBatch; +use bytes::{Buf, Bytes}; use pageserver_api::reltag::{RelTag, SlruKind}; use pageserver_api::shard::ShardIdentity; -use pageserver_api::value::Value; +use postgres_ffi::pg_constants; use postgres_ffi::relfile_utils::VISIBILITYMAP_FORKNUM; use postgres_ffi::walrecord::*; -use postgres_ffi::{page_is_new, page_set_lsn, pg_constants, BLCKSZ}; use utils::lsn::Lsn; impl InterpretedWalRecord { @@ -21,11 +19,12 @@ impl InterpretedWalRecord { pub fn from_bytes_filtered( buf: Bytes, shard: &ShardIdentity, - lsn: Lsn, + record_end_lsn: Lsn, pg_version: u32, ) -> anyhow::Result { let mut decoded = DecodedWALRecord::default(); decode_wal_record(buf, &mut decoded, pg_version)?; + let xid = decoded.xl_xid; let flush_uncommitted = if decoded.is_dbase_create_copy(pg_version) { FlushUncommittedRecords::Yes @@ -33,96 +32,20 @@ impl InterpretedWalRecord { FlushUncommittedRecords::No }; - let metadata_record = MetadataRecord::from_decoded(&decoded, lsn, pg_version)?; - - let mut blocks = Vec::default(); - for blk in decoded.blocks.iter() { - let rel = RelTag { - spcnode: blk.rnode_spcnode, - dbnode: blk.rnode_dbnode, - relnode: blk.rnode_relnode, - forknum: blk.forknum, - }; - - let key = rel_block_to_key(rel, blk.blkno); - - if !key.is_valid_key_on_write_path() { - anyhow::bail!("Unsupported key decoded at LSN {}: {}", lsn, key); - } - - let key_is_local = shard.is_key_local(&key); - - tracing::debug!( - lsn=%lsn, - key=%key, - "ingest: shard decision {}", - if !key_is_local { "drop" } else { "keep" }, - ); - - if !key_is_local { - if shard.is_shard_zero() { - // Shard 0 tracks relation sizes. Although we will not store this block, we will observe - // its blkno in case it implicitly extends a relation. - blocks.push((key.to_compact(), None)); - } - - continue; - } - - // Instead of storing full-page-image WAL record, - // it is better to store extracted image: we can skip wal-redo - // in this case. Also some FPI records may contain multiple (up to 32) pages, - // so them have to be copied multiple times. - // - let value = if blk.apply_image - && blk.has_image - && decoded.xl_rmid == pg_constants::RM_XLOG_ID - && (decoded.xl_info == pg_constants::XLOG_FPI - || decoded.xl_info == pg_constants::XLOG_FPI_FOR_HINT) - // compression of WAL is not yet supported: fall back to storing the original WAL record - && !postgres_ffi::bkpimage_is_compressed(blk.bimg_info, pg_version) - // do not materialize null pages because them most likely be soon replaced with real data - && blk.bimg_len != 0 - { - // Extract page image from FPI record - let img_len = blk.bimg_len as usize; - let img_offs = blk.bimg_offset as usize; - let mut image = BytesMut::with_capacity(BLCKSZ as usize); - // TODO(vlad): skip the copy - image.extend_from_slice(&decoded.record[img_offs..img_offs + img_len]); - - if blk.hole_length != 0 { - let tail = image.split_off(blk.hole_offset as usize); - image.resize(image.len() + blk.hole_length as usize, 0u8); - image.unsplit(tail); - } - // - // Match the logic of XLogReadBufferForRedoExtended: - // The page may be uninitialized. If so, we can't set the LSN because - // that would corrupt the page. - // - if !page_is_new(&image) { - page_set_lsn(&mut image, lsn) - } - assert_eq!(image.len(), BLCKSZ as usize); - - Value::Image(image.freeze()) - } else { - Value::WalRecord(NeonWalRecord::Postgres { - will_init: blk.will_init || blk.apply_image, - rec: decoded.record.clone(), - }) - }; - - blocks.push((key.to_compact(), Some(value))); - } + let metadata_record = MetadataRecord::from_decoded(&decoded, record_end_lsn, pg_version)?; + let batch = SerializedValueBatch::from_decoded_filtered( + decoded, + shard, + record_end_lsn, + pg_version, + )?; Ok(InterpretedWalRecord { metadata_record, - blocks, - lsn, + batch, + end_lsn: record_end_lsn, flush_uncommitted, - xid: decoded.xl_xid, + xid, }) } } @@ -130,7 +53,7 @@ impl InterpretedWalRecord { impl MetadataRecord { fn from_decoded( decoded: &DecodedWALRecord, - lsn: Lsn, + record_end_lsn: Lsn, pg_version: u32, ) -> anyhow::Result> { // Note: this doesn't actually copy the bytes since @@ -151,7 +74,7 @@ impl MetadataRecord { Ok(None) } pg_constants::RM_CLOG_ID => Self::decode_clog_record(&mut buf, decoded, pg_version), - pg_constants::RM_XACT_ID => Self::decode_xact_record(&mut buf, decoded, lsn), + pg_constants::RM_XACT_ID => Self::decode_xact_record(&mut buf, decoded, record_end_lsn), pg_constants::RM_MULTIXACT_ID => { Self::decode_multixact_record(&mut buf, decoded, pg_version) } @@ -163,7 +86,7 @@ impl MetadataRecord { // // Alternatively, one can make the checkpoint part of the subscription protocol // to the pageserver. This should work fine, but can be done at a later point. - pg_constants::RM_XLOG_ID => Self::decode_xlog_record(&mut buf, decoded, lsn), + pg_constants::RM_XLOG_ID => Self::decode_xlog_record(&mut buf, decoded, record_end_lsn), pg_constants::RM_LOGICALMSG_ID => { Self::decode_logical_message_record(&mut buf, decoded) } diff --git a/libs/wal_decoder/src/lib.rs b/libs/wal_decoder/src/lib.rs index 05349d17c9..a8a26956e6 100644 --- a/libs/wal_decoder/src/lib.rs +++ b/libs/wal_decoder/src/lib.rs @@ -1,2 +1,3 @@ pub mod decoder; pub mod models; +pub mod serialized_batch; diff --git a/libs/wal_decoder/src/models.rs b/libs/wal_decoder/src/models.rs index 92b66fcefd..5d90eeb69c 100644 --- a/libs/wal_decoder/src/models.rs +++ b/libs/wal_decoder/src/models.rs @@ -2,7 +2,8 @@ //! ready for the pageserver to interpret. They are derived from the original //! WAL records, so that each struct corresponds closely to one WAL record of //! a specific kind. They contain the same information as the original WAL records, -//! just decoded into structs and fields for easier access. +//! but the values are already serialized in a [`SerializedValueBatch`], which +//! is the format that the pageserver is expecting them in. //! //! The ingestion code uses these structs to help with parsing the WAL records, //! and it splits them into a stream of modifications to the key-value pairs that @@ -25,9 +26,7 @@ //! |--> write to KV store within the pageserver use bytes::Bytes; -use pageserver_api::key::CompactKey; use pageserver_api::reltag::{RelTag, SlruKind}; -use pageserver_api::value::Value; use postgres_ffi::walrecord::{ XlMultiXactCreate, XlMultiXactTruncate, XlRelmapUpdate, XlReploriginDrop, XlReploriginSet, XlSmgrTruncate, XlXactParsedRecord, @@ -35,6 +34,8 @@ use postgres_ffi::walrecord::{ use postgres_ffi::{Oid, TransactionId}; use utils::lsn::Lsn; +use crate::serialized_batch::SerializedValueBatch; + pub enum FlushUncommittedRecords { Yes, No, @@ -45,12 +46,11 @@ pub struct InterpretedWalRecord { /// Optional metadata record - may cause writes to metadata keys /// in the storage engine pub metadata_record: Option, - /// Images or deltas for blocks modified in the original WAL record. - /// The [`Value`] is optional to avoid sending superfluous data to - /// shard 0 for relation size tracking. - pub blocks: Vec<(CompactKey, Option)>, + /// A pre-serialized batch along with the required metadata for ingestion + /// by the pageserver + pub batch: SerializedValueBatch, /// Byte offset within WAL for the end of the original PG WAL record - pub lsn: Lsn, + pub end_lsn: Lsn, /// Whether to flush all uncommitted modifications to the storage engine /// before ingesting this record. This is currently only used for legacy PG /// database creations which read pages from a template database. Such WAL diff --git a/libs/wal_decoder/src/serialized_batch.rs b/libs/wal_decoder/src/serialized_batch.rs new file mode 100644 index 0000000000..8f33291023 --- /dev/null +++ b/libs/wal_decoder/src/serialized_batch.rs @@ -0,0 +1,862 @@ +//! This module implements batch type for serialized [`pageserver_api::value::Value`] +//! instances. Each batch contains a raw buffer (serialized values) +//! and a list of metadata for each (key, LSN) tuple present in the batch. +//! +//! Such batches are created from decoded PG wal records and ingested +//! by the pageserver by writing directly to the ephemeral file. + +use std::collections::BTreeSet; + +use bytes::{Bytes, BytesMut}; +use pageserver_api::key::rel_block_to_key; +use pageserver_api::keyspace::KeySpace; +use pageserver_api::record::NeonWalRecord; +use pageserver_api::reltag::RelTag; +use pageserver_api::shard::ShardIdentity; +use pageserver_api::{key::CompactKey, value::Value}; +use postgres_ffi::walrecord::{DecodedBkpBlock, DecodedWALRecord}; +use postgres_ffi::{page_is_new, page_set_lsn, pg_constants, BLCKSZ}; +use utils::bin_ser::BeSer; +use utils::lsn::Lsn; + +use pageserver_api::key::Key; + +static ZERO_PAGE: Bytes = Bytes::from_static(&[0u8; BLCKSZ as usize]); + +/// Accompanying metadata for the batch +/// A value may be serialized and stored into the batch or just "observed". +/// Shard 0 currently "observes" all values in order to accurately track +/// relation sizes. In the case of "observed" values, we only need to know +/// the key and LSN, so two types of metadata are supported to save on network +/// bandwidth. +pub enum ValueMeta { + Serialized(SerializedValueMeta), + Observed(ObservedValueMeta), +} + +impl ValueMeta { + pub fn key(&self) -> CompactKey { + match self { + Self::Serialized(ser) => ser.key, + Self::Observed(obs) => obs.key, + } + } + + pub fn lsn(&self) -> Lsn { + match self { + Self::Serialized(ser) => ser.lsn, + Self::Observed(obs) => obs.lsn, + } + } +} + +/// Wrapper around [`ValueMeta`] that implements ordering by +/// (key, LSN) tuples +struct OrderedValueMeta(ValueMeta); + +impl Ord for OrderedValueMeta { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + (self.0.key(), self.0.lsn()).cmp(&(other.0.key(), other.0.lsn())) + } +} + +impl PartialOrd for OrderedValueMeta { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl PartialEq for OrderedValueMeta { + fn eq(&self, other: &Self) -> bool { + (self.0.key(), self.0.lsn()) == (other.0.key(), other.0.lsn()) + } +} + +impl Eq for OrderedValueMeta {} + +/// Metadata for a [`Value`] serialized into the batch. +pub struct SerializedValueMeta { + pub key: CompactKey, + pub lsn: Lsn, + /// Starting offset of the value for the (key, LSN) tuple + /// in [`SerializedValueBatch::raw`] + pub batch_offset: u64, + pub len: usize, + pub will_init: bool, +} + +/// Metadata for a [`Value`] observed by the batch +pub struct ObservedValueMeta { + pub key: CompactKey, + pub lsn: Lsn, +} + +/// Batch of serialized [`Value`]s. +pub struct SerializedValueBatch { + /// [`Value`]s serialized in EphemeralFile's native format, + /// ready for disk write by the pageserver + pub raw: Vec, + + /// Metadata to make sense of the bytes in [`Self::raw`] + /// and represent "observed" values. + /// + /// Invariant: Metadata entries for any given key are ordered + /// by LSN. Note that entries for a key do not have to be contiguous. + pub metadata: Vec, + + /// The highest LSN of any value in the batch + pub max_lsn: Lsn, + + /// Number of values encoded by [`Self::raw`] + pub len: usize, +} + +impl Default for SerializedValueBatch { + fn default() -> Self { + Self { + raw: Default::default(), + metadata: Default::default(), + max_lsn: Lsn(0), + len: 0, + } + } +} + +impl SerializedValueBatch { + /// Build a batch of serialized values from a decoded PG WAL record + /// + /// The batch will only contain values for keys targeting the specifiec + /// shard. Shard 0 is a special case, where any keys that don't belong to + /// it are "observed" by the batch (i.e. present in [`SerializedValueBatch::metadata`], + /// but absent from the raw buffer [`SerializedValueBatch::raw`]). + pub(crate) fn from_decoded_filtered( + decoded: DecodedWALRecord, + shard: &ShardIdentity, + record_end_lsn: Lsn, + pg_version: u32, + ) -> anyhow::Result { + // First determine how big the buffer needs to be and allocate it up-front. + // This duplicates some of the work below, but it's empirically much faster. + let estimated_buffer_size = Self::estimate_buffer_size(&decoded, shard, pg_version); + let mut buf = Vec::::with_capacity(estimated_buffer_size); + + let mut metadata: Vec = Vec::with_capacity(decoded.blocks.len()); + let mut max_lsn: Lsn = Lsn(0); + let mut len: usize = 0; + for blk in decoded.blocks.iter() { + let relative_off = buf.len() as u64; + + let rel = RelTag { + spcnode: blk.rnode_spcnode, + dbnode: blk.rnode_dbnode, + relnode: blk.rnode_relnode, + forknum: blk.forknum, + }; + + let key = rel_block_to_key(rel, blk.blkno); + + if !key.is_valid_key_on_write_path() { + anyhow::bail!("Unsupported key decoded at LSN {}: {}", record_end_lsn, key); + } + + let key_is_local = shard.is_key_local(&key); + + tracing::debug!( + lsn=%record_end_lsn, + key=%key, + "ingest: shard decision {}", + if !key_is_local { "drop" } else { "keep" }, + ); + + if !key_is_local { + if shard.is_shard_zero() { + // Shard 0 tracks relation sizes. Although we will not store this block, we will observe + // its blkno in case it implicitly extends a relation. + metadata.push(ValueMeta::Observed(ObservedValueMeta { + key: key.to_compact(), + lsn: record_end_lsn, + })) + } + + continue; + } + + // Instead of storing full-page-image WAL record, + // it is better to store extracted image: we can skip wal-redo + // in this case. Also some FPI records may contain multiple (up to 32) pages, + // so them have to be copied multiple times. + // + let val = if Self::block_is_image(&decoded, blk, pg_version) { + // Extract page image from FPI record + let img_len = blk.bimg_len as usize; + let img_offs = blk.bimg_offset as usize; + let mut image = BytesMut::with_capacity(BLCKSZ as usize); + // TODO(vlad): skip the copy + image.extend_from_slice(&decoded.record[img_offs..img_offs + img_len]); + + if blk.hole_length != 0 { + let tail = image.split_off(blk.hole_offset as usize); + image.resize(image.len() + blk.hole_length as usize, 0u8); + image.unsplit(tail); + } + // + // Match the logic of XLogReadBufferForRedoExtended: + // The page may be uninitialized. If so, we can't set the LSN because + // that would corrupt the page. + // + if !page_is_new(&image) { + page_set_lsn(&mut image, record_end_lsn) + } + assert_eq!(image.len(), BLCKSZ as usize); + + Value::Image(image.freeze()) + } else { + Value::WalRecord(NeonWalRecord::Postgres { + will_init: blk.will_init || blk.apply_image, + rec: decoded.record.clone(), + }) + }; + + val.ser_into(&mut buf) + .expect("Writing into in-memory buffer is infallible"); + + let val_ser_size = buf.len() - relative_off as usize; + + metadata.push(ValueMeta::Serialized(SerializedValueMeta { + key: key.to_compact(), + lsn: record_end_lsn, + batch_offset: relative_off, + len: val_ser_size, + will_init: val.will_init(), + })); + max_lsn = std::cmp::max(max_lsn, record_end_lsn); + len += 1; + } + + if cfg!(any(debug_assertions, test)) { + let batch = Self { + raw: buf, + metadata, + max_lsn, + len, + }; + + batch.validate_lsn_order(); + + return Ok(batch); + } + + Ok(Self { + raw: buf, + metadata, + max_lsn, + len, + }) + } + + /// Look into the decoded PG WAL record and determine + /// roughly how large the buffer for serialized values needs to be. + fn estimate_buffer_size( + decoded: &DecodedWALRecord, + shard: &ShardIdentity, + pg_version: u32, + ) -> usize { + let mut estimate: usize = 0; + + for blk in decoded.blocks.iter() { + let rel = RelTag { + spcnode: blk.rnode_spcnode, + dbnode: blk.rnode_dbnode, + relnode: blk.rnode_relnode, + forknum: blk.forknum, + }; + + let key = rel_block_to_key(rel, blk.blkno); + + if !shard.is_key_local(&key) { + continue; + } + + if Self::block_is_image(decoded, blk, pg_version) { + // 4 bytes for the Value::Image discriminator + // 8 bytes for encoding the size of the buffer + // BLCKSZ for the raw image + estimate += (4 + 8 + BLCKSZ) as usize; + } else { + // 4 bytes for the Value::WalRecord discriminator + // 4 bytes for the NeonWalRecord::Postgres discriminator + // 1 bytes for NeonWalRecord::Postgres::will_init + // 8 bytes for encoding the size of the buffer + // length of the raw record + estimate += 8 + 1 + 8 + decoded.record.len(); + } + } + + estimate + } + + fn block_is_image(decoded: &DecodedWALRecord, blk: &DecodedBkpBlock, pg_version: u32) -> bool { + blk.apply_image + && blk.has_image + && decoded.xl_rmid == pg_constants::RM_XLOG_ID + && (decoded.xl_info == pg_constants::XLOG_FPI + || decoded.xl_info == pg_constants::XLOG_FPI_FOR_HINT) + // compression of WAL is not yet supported: fall back to storing the original WAL record + && !postgres_ffi::bkpimage_is_compressed(blk.bimg_info, pg_version) + // do not materialize null pages because them most likely be soon replaced with real data + && blk.bimg_len != 0 + } + + /// Encode a list of values and metadata into a serialized batch + /// + /// This is used by the pageserver ingest code to conveniently generate + /// batches for metadata writes. + pub fn from_values(batch: Vec<(CompactKey, Lsn, usize, Value)>) -> Self { + // Pre-allocate a big flat buffer to write into. This should be large but not huge: it is soft-limited in practice by + // [`crate::pgdatadir_mapping::DatadirModification::MAX_PENDING_BYTES`] + let buffer_size = batch.iter().map(|i| i.2).sum::(); + let mut buf = Vec::::with_capacity(buffer_size); + + let mut metadata: Vec = Vec::with_capacity(batch.len()); + let mut max_lsn: Lsn = Lsn(0); + let len = batch.len(); + for (key, lsn, val_ser_size, val) in batch { + let relative_off = buf.len() as u64; + + val.ser_into(&mut buf) + .expect("Writing into in-memory buffer is infallible"); + + metadata.push(ValueMeta::Serialized(SerializedValueMeta { + key, + lsn, + batch_offset: relative_off, + len: val_ser_size, + will_init: val.will_init(), + })); + max_lsn = std::cmp::max(max_lsn, lsn); + } + + // Assert that we didn't do any extra allocations while building buffer. + debug_assert!(buf.len() <= buffer_size); + + if cfg!(any(debug_assertions, test)) { + let batch = Self { + raw: buf, + metadata, + max_lsn, + len, + }; + + batch.validate_lsn_order(); + + return batch; + } + + Self { + raw: buf, + metadata, + max_lsn, + len, + } + } + + /// Add one value to the batch + /// + /// This is used by the pageserver ingest code to include metadata block + /// updates for a single key. + pub fn put(&mut self, key: CompactKey, value: Value, lsn: Lsn) { + let relative_off = self.raw.len() as u64; + value.ser_into(&mut self.raw).unwrap(); + + let val_ser_size = self.raw.len() - relative_off as usize; + self.metadata + .push(ValueMeta::Serialized(SerializedValueMeta { + key, + lsn, + batch_offset: relative_off, + len: val_ser_size, + will_init: value.will_init(), + })); + + self.max_lsn = std::cmp::max(self.max_lsn, lsn); + self.len += 1; + + if cfg!(any(debug_assertions, test)) { + self.validate_lsn_order(); + } + } + + /// Extend with the contents of another batch + /// + /// One batch is generated for each decoded PG WAL record. + /// They are then merged to accumulate reasonably sized writes. + pub fn extend(&mut self, mut other: SerializedValueBatch) { + let extend_batch_start_offset = self.raw.len() as u64; + + self.raw.extend(other.raw); + + // Shift the offsets in the batch we are extending with + other.metadata.iter_mut().for_each(|meta| match meta { + ValueMeta::Serialized(ser) => { + ser.batch_offset += extend_batch_start_offset; + if cfg!(debug_assertions) { + let value_end = ser.batch_offset + ser.len as u64; + assert!((value_end as usize) <= self.raw.len()); + } + } + ValueMeta::Observed(_) => {} + }); + self.metadata.extend(other.metadata); + + self.max_lsn = std::cmp::max(self.max_lsn, other.max_lsn); + + self.len += other.len; + + if cfg!(any(debug_assertions, test)) { + self.validate_lsn_order(); + } + } + + /// Add zero images for the (key, LSN) tuples specified + /// + /// PG versions below 16 do not zero out pages before extending + /// a relation and may leave gaps. Such gaps need to be identified + /// by the pageserver ingest logic and get patched up here. + /// + /// Note that this function does not validate that the gaps have been + /// identified correctly (it does not know relation sizes), so it's up + /// to the call-site to do it properly. + pub fn zero_gaps(&mut self, gaps: Vec<(KeySpace, Lsn)>) { + // Implementation note: + // + // Values within [`SerializedValueBatch::raw`] do not have any ordering requirements, + // but the metadata entries should be ordered properly (see + // [`SerializedValueBatch::metadata`]). + // + // Exploiting this observation we do: + // 1. Drain all the metadata entries into an ordered set. + // The use of a BTreeSet keyed by (Key, Lsn) relies on the observation that Postgres never + // includes more than one update to the same block in the same WAL record. + // 2. For each (key, LSN) gap tuple, append a zero image to the raw buffer + // and add an index entry to the ordered metadata set. + // 3. Drain the ordered set back into a metadata vector + + let mut ordered_metas = self + .metadata + .drain(..) + .map(OrderedValueMeta) + .collect::>(); + for (keyspace, lsn) in gaps { + self.max_lsn = std::cmp::max(self.max_lsn, lsn); + + for gap_range in keyspace.ranges { + let mut key = gap_range.start; + while key != gap_range.end { + let relative_off = self.raw.len() as u64; + + // TODO(vlad): Can we be cheeky and write only one zero image, and + // make all index entries requiring a zero page point to it? + // Alternatively, we can change the index entry format to represent zero pages + // without writing them at all. + Value::Image(ZERO_PAGE.clone()) + .ser_into(&mut self.raw) + .unwrap(); + let val_ser_size = self.raw.len() - relative_off as usize; + + ordered_metas.insert(OrderedValueMeta(ValueMeta::Serialized( + SerializedValueMeta { + key: key.to_compact(), + lsn, + batch_offset: relative_off, + len: val_ser_size, + will_init: true, + }, + ))); + + self.len += 1; + + key = key.next(); + } + } + } + + self.metadata = ordered_metas.into_iter().map(|ord| ord.0).collect(); + + if cfg!(any(debug_assertions, test)) { + self.validate_lsn_order(); + } + } + + /// Checks if the batch is empty + /// + /// A batch is empty when it contains no serialized values. + /// Note that it may still contain observed values. + pub fn is_empty(&self) -> bool { + let empty = self.raw.is_empty(); + + if cfg!(debug_assertions) && empty { + assert!(self + .metadata + .iter() + .all(|meta| matches!(meta, ValueMeta::Observed(_)))); + } + + empty + } + + /// Returns the number of values serialized in the batch + pub fn len(&self) -> usize { + self.len + } + + /// Returns the size of the buffer wrapped by the batch + pub fn buffer_size(&self) -> usize { + self.raw.len() + } + + pub fn updates_key(&self, key: &Key) -> bool { + self.metadata.iter().any(|meta| match meta { + ValueMeta::Serialized(ser) => key.to_compact() == ser.key, + ValueMeta::Observed(_) => false, + }) + } + + pub fn validate_lsn_order(&self) { + use std::collections::HashMap; + + let mut last_seen_lsn_per_key: HashMap = HashMap::default(); + + for meta in self.metadata.iter() { + let lsn = meta.lsn(); + let key = meta.key(); + + if let Some(prev_lsn) = last_seen_lsn_per_key.insert(key, lsn) { + assert!( + lsn >= prev_lsn, + "Ordering violated by {}: {} < {}", + Key::from_compact(key), + lsn, + prev_lsn + ); + } + } + } +} + +#[cfg(all(test, feature = "testing"))] +mod tests { + use super::*; + + fn validate_batch( + batch: &SerializedValueBatch, + values: &[(CompactKey, Lsn, usize, Value)], + gaps: Option<&Vec<(KeySpace, Lsn)>>, + ) { + // Invariant 1: The metadata for a given entry in the batch + // is correct and can be used to deserialize back to the original value. + for (key, lsn, size, value) in values.iter() { + let meta = batch + .metadata + .iter() + .find(|meta| (meta.key(), meta.lsn()) == (*key, *lsn)) + .unwrap(); + let meta = match meta { + ValueMeta::Serialized(ser) => ser, + ValueMeta::Observed(_) => unreachable!(), + }; + + assert_eq!(meta.len, *size); + assert_eq!(meta.will_init, value.will_init()); + + let start = meta.batch_offset as usize; + let end = meta.batch_offset as usize + meta.len; + let value_from_batch = Value::des(&batch.raw[start..end]).unwrap(); + assert_eq!(&value_from_batch, value); + } + + let mut expected_buffer_size: usize = values.iter().map(|(_, _, size, _)| size).sum(); + let mut gap_pages_count: usize = 0; + + // Invariant 2: Zero pages were added for identified gaps and their metadata + // is correct. + if let Some(gaps) = gaps { + for (gap_keyspace, lsn) in gaps { + for gap_range in &gap_keyspace.ranges { + let mut gap_key = gap_range.start; + while gap_key != gap_range.end { + let meta = batch + .metadata + .iter() + .find(|meta| (meta.key(), meta.lsn()) == (gap_key.to_compact(), *lsn)) + .unwrap(); + let meta = match meta { + ValueMeta::Serialized(ser) => ser, + ValueMeta::Observed(_) => unreachable!(), + }; + + let zero_value = Value::Image(ZERO_PAGE.clone()); + let zero_value_size = zero_value.serialized_size().unwrap() as usize; + + assert_eq!(meta.len, zero_value_size); + assert_eq!(meta.will_init, zero_value.will_init()); + + let start = meta.batch_offset as usize; + let end = meta.batch_offset as usize + meta.len; + let value_from_batch = Value::des(&batch.raw[start..end]).unwrap(); + assert_eq!(value_from_batch, zero_value); + + gap_pages_count += 1; + expected_buffer_size += zero_value_size; + gap_key = gap_key.next(); + } + } + } + } + + // Invariant 3: The length of the batch is equal to the number + // of values inserted, plus the number of gap pages. This extends + // to the raw buffer size. + assert_eq!(batch.len(), values.len() + gap_pages_count); + assert_eq!(expected_buffer_size, batch.buffer_size()); + + // Invariant 4: Metadata entries for any given key are sorted in LSN order. + batch.validate_lsn_order(); + } + + #[test] + fn test_creation_from_values() { + const LSN: Lsn = Lsn(0x10); + let key = Key::from_hex("110000000033333333444444445500000001").unwrap(); + + let values = vec![ + ( + key.to_compact(), + LSN, + Value::WalRecord(NeonWalRecord::wal_append("foo")), + ), + ( + key.next().to_compact(), + LSN, + Value::WalRecord(NeonWalRecord::wal_append("bar")), + ), + ( + key.to_compact(), + Lsn(LSN.0 + 0x10), + Value::WalRecord(NeonWalRecord::wal_append("baz")), + ), + ( + key.next().next().to_compact(), + LSN, + Value::WalRecord(NeonWalRecord::wal_append("taz")), + ), + ]; + + let values = values + .into_iter() + .map(|(key, lsn, value)| (key, lsn, value.serialized_size().unwrap() as usize, value)) + .collect::>(); + let batch = SerializedValueBatch::from_values(values.clone()); + + validate_batch(&batch, &values, None); + + assert!(!batch.is_empty()); + } + + #[test] + fn test_put() { + const LSN: Lsn = Lsn(0x10); + let key = Key::from_hex("110000000033333333444444445500000001").unwrap(); + + let values = vec![ + ( + key.to_compact(), + LSN, + Value::WalRecord(NeonWalRecord::wal_append("foo")), + ), + ( + key.next().to_compact(), + LSN, + Value::WalRecord(NeonWalRecord::wal_append("bar")), + ), + ]; + + let mut values = values + .into_iter() + .map(|(key, lsn, value)| (key, lsn, value.serialized_size().unwrap() as usize, value)) + .collect::>(); + let mut batch = SerializedValueBatch::from_values(values.clone()); + + validate_batch(&batch, &values, None); + + let value = ( + key.to_compact(), + Lsn(LSN.0 + 0x10), + Value::WalRecord(NeonWalRecord::wal_append("baz")), + ); + let serialized_size = value.2.serialized_size().unwrap() as usize; + let value = (value.0, value.1, serialized_size, value.2); + values.push(value.clone()); + batch.put(value.0, value.3, value.1); + + validate_batch(&batch, &values, None); + + let value = ( + key.next().next().to_compact(), + LSN, + Value::WalRecord(NeonWalRecord::wal_append("taz")), + ); + let serialized_size = value.2.serialized_size().unwrap() as usize; + let value = (value.0, value.1, serialized_size, value.2); + values.push(value.clone()); + batch.put(value.0, value.3, value.1); + + validate_batch(&batch, &values, None); + } + + #[test] + fn test_extension() { + const LSN: Lsn = Lsn(0x10); + let key = Key::from_hex("110000000033333333444444445500000001").unwrap(); + + let values = vec![ + ( + key.to_compact(), + LSN, + Value::WalRecord(NeonWalRecord::wal_append("foo")), + ), + ( + key.next().to_compact(), + LSN, + Value::WalRecord(NeonWalRecord::wal_append("bar")), + ), + ( + key.next().next().to_compact(), + LSN, + Value::WalRecord(NeonWalRecord::wal_append("taz")), + ), + ]; + + let mut values = values + .into_iter() + .map(|(key, lsn, value)| (key, lsn, value.serialized_size().unwrap() as usize, value)) + .collect::>(); + let mut batch = SerializedValueBatch::from_values(values.clone()); + + let other_values = vec![ + ( + key.to_compact(), + Lsn(LSN.0 + 0x10), + Value::WalRecord(NeonWalRecord::wal_append("foo")), + ), + ( + key.next().to_compact(), + Lsn(LSN.0 + 0x10), + Value::WalRecord(NeonWalRecord::wal_append("bar")), + ), + ( + key.next().next().to_compact(), + Lsn(LSN.0 + 0x10), + Value::WalRecord(NeonWalRecord::wal_append("taz")), + ), + ]; + + let other_values = other_values + .into_iter() + .map(|(key, lsn, value)| (key, lsn, value.serialized_size().unwrap() as usize, value)) + .collect::>(); + let other_batch = SerializedValueBatch::from_values(other_values.clone()); + + values.extend(other_values); + batch.extend(other_batch); + + validate_batch(&batch, &values, None); + } + + #[test] + fn test_gap_zeroing() { + const LSN: Lsn = Lsn(0x10); + let rel_foo_base_key = Key::from_hex("110000000033333333444444445500000001").unwrap(); + + let rel_bar_base_key = { + let mut key = rel_foo_base_key; + key.field4 += 1; + key + }; + + let values = vec![ + ( + rel_foo_base_key.to_compact(), + LSN, + Value::WalRecord(NeonWalRecord::wal_append("foo1")), + ), + ( + rel_foo_base_key.add(1).to_compact(), + LSN, + Value::WalRecord(NeonWalRecord::wal_append("foo2")), + ), + ( + rel_foo_base_key.add(5).to_compact(), + LSN, + Value::WalRecord(NeonWalRecord::wal_append("foo3")), + ), + ( + rel_foo_base_key.add(1).to_compact(), + Lsn(LSN.0 + 0x10), + Value::WalRecord(NeonWalRecord::wal_append("foo4")), + ), + ( + rel_foo_base_key.add(10).to_compact(), + Lsn(LSN.0 + 0x10), + Value::WalRecord(NeonWalRecord::wal_append("foo5")), + ), + ( + rel_foo_base_key.add(11).to_compact(), + Lsn(LSN.0 + 0x10), + Value::WalRecord(NeonWalRecord::wal_append("foo6")), + ), + ( + rel_foo_base_key.add(12).to_compact(), + Lsn(LSN.0 + 0x10), + Value::WalRecord(NeonWalRecord::wal_append("foo7")), + ), + ( + rel_bar_base_key.to_compact(), + LSN, + Value::WalRecord(NeonWalRecord::wal_append("bar1")), + ), + ( + rel_bar_base_key.add(4).to_compact(), + LSN, + Value::WalRecord(NeonWalRecord::wal_append("bar2")), + ), + ]; + + let values = values + .into_iter() + .map(|(key, lsn, value)| (key, lsn, value.serialized_size().unwrap() as usize, value)) + .collect::>(); + + let mut batch = SerializedValueBatch::from_values(values.clone()); + + let gaps = vec![ + ( + KeySpace { + ranges: vec![ + rel_foo_base_key.add(2)..rel_foo_base_key.add(5), + rel_bar_base_key.add(1)..rel_bar_base_key.add(4), + ], + }, + LSN, + ), + ( + KeySpace { + ranges: vec![rel_foo_base_key.add(6)..rel_foo_base_key.add(10)], + }, + Lsn(LSN.0 + 0x10), + ), + ]; + + batch.zero_gaps(gaps.clone()); + validate_batch(&batch, &values, Some(&gaps)); + } +} diff --git a/pageserver/benches/bench_ingest.rs b/pageserver/benches/bench_ingest.rs index 0a1ad9cd6b..f6b2a8e031 100644 --- a/pageserver/benches/bench_ingest.rs +++ b/pageserver/benches/bench_ingest.rs @@ -9,7 +9,6 @@ use pageserver::{ l0_flush::{L0FlushConfig, L0FlushGlobalState}, page_cache, task_mgr::TaskKind, - tenant::storage_layer::inmemory_layer::SerializedBatch, tenant::storage_layer::InMemoryLayer, virtual_file, }; @@ -18,6 +17,7 @@ use utils::{ bin_ser::BeSer, id::{TenantId, TimelineId}, }; +use wal_decoder::serialized_batch::SerializedValueBatch; // A very cheap hash for generating non-sequential keys. fn murmurhash32(mut h: u32) -> u32 { @@ -102,13 +102,13 @@ async fn ingest( batch.push((key.to_compact(), lsn, data_ser_size, data.clone())); if batch.len() >= BATCH_SIZE { let this_batch = std::mem::take(&mut batch); - let serialized = SerializedBatch::from_values(this_batch).unwrap(); + let serialized = SerializedValueBatch::from_values(this_batch); layer.put_batch(serialized, &ctx).await?; } } if !batch.is_empty() { let this_batch = std::mem::take(&mut batch); - let serialized = SerializedBatch::from_values(this_batch).unwrap(); + let serialized = SerializedValueBatch::from_values(this_batch); layer.put_batch(serialized, &ctx).await?; } layer.freeze(lsn + 1).await; diff --git a/pageserver/src/pgdatadir_mapping.rs b/pageserver/src/pgdatadir_mapping.rs index dc2dc08b53..7b106569a4 100644 --- a/pageserver/src/pgdatadir_mapping.rs +++ b/pageserver/src/pgdatadir_mapping.rs @@ -24,6 +24,7 @@ use pageserver_api::key::{ use pageserver_api::keyspace::SparseKeySpace; use pageserver_api::record::NeonWalRecord; use pageserver_api::reltag::{BlockNumber, RelTag, SlruKind}; +use pageserver_api::shard::ShardIdentity; use pageserver_api::value::Value; use postgres_ffi::relfile_utils::{FSM_FORKNUM, VISIBILITYMAP_FORKNUM}; use postgres_ffi::BLCKSZ; @@ -38,6 +39,7 @@ use tracing::{debug, trace, warn}; use utils::bin_ser::DeserializeError; use utils::pausable_failpoint; use utils::{bin_ser::BeSer, lsn::Lsn}; +use wal_decoder::serialized_batch::SerializedValueBatch; /// Max delta records appended to the AUX_FILES_KEY (for aux v1). The write path will write a full image once this threshold is reached. pub const MAX_AUX_FILE_DELTAS: usize = 1024; @@ -170,12 +172,11 @@ impl Timeline { tline: self, pending_lsns: Vec::new(), pending_metadata_pages: HashMap::new(), - pending_data_pages: Vec::new(), - pending_zero_data_pages: Default::default(), + pending_data_batch: None, pending_deletions: Vec::new(), pending_nblocks: 0, pending_directory_entries: Vec::new(), - pending_bytes: 0, + pending_metadata_bytes: 0, lsn, } } @@ -1025,21 +1026,14 @@ pub struct DatadirModification<'a> { /// Data writes, ready to be flushed into an ephemeral layer. See [`Self::is_data_key`] for /// which keys are stored here. - pending_data_pages: Vec<(CompactKey, Lsn, usize, Value)>, - - // Sometimes during ingest, for example when extending a relation, we would like to write a zero page. However, - // if we encounter a write from postgres in the same wal record, we will drop this entry. - // - // Unlike other 'pending' fields, this does not last until the next call to commit(): it is flushed - // at the end of each wal record, and all these writes implicitly are at lsn Self::lsn - pending_zero_data_pages: HashSet, + pending_data_batch: Option, /// For special "directory" keys that store key-value maps, track the size of the map /// if it was updated in this modification. pending_directory_entries: Vec<(DirectoryKind, usize)>, - /// An **approximation** of how large our EphemeralFile write will be when committed. - pending_bytes: usize, + /// An **approximation** of how many metadata bytes will be written to the EphemeralFile. + pending_metadata_bytes: usize, } impl<'a> DatadirModification<'a> { @@ -1054,11 +1048,17 @@ impl<'a> DatadirModification<'a> { } pub(crate) fn approx_pending_bytes(&self) -> usize { - self.pending_bytes + self.pending_data_batch + .as_ref() + .map_or(0, |b| b.buffer_size()) + + self.pending_metadata_bytes } - pub(crate) fn has_dirty_data_pages(&self) -> bool { - (!self.pending_data_pages.is_empty()) || (!self.pending_zero_data_pages.is_empty()) + pub(crate) fn has_dirty_data(&self) -> bool { + !self + .pending_data_batch + .as_ref() + .map_or(true, |b| b.is_empty()) } /// Set the current lsn @@ -1070,9 +1070,6 @@ impl<'a> DatadirModification<'a> { self.lsn ); - // If we are advancing LSN, then state from previous wal record should have been flushed. - assert!(self.pending_zero_data_pages.is_empty()); - if lsn > self.lsn { self.pending_lsns.push(self.lsn); self.lsn = lsn; @@ -1147,6 +1144,107 @@ impl<'a> DatadirModification<'a> { Ok(()) } + /// Creates a relation if it is not already present. + /// Returns the current size of the relation + pub(crate) async fn create_relation_if_required( + &mut self, + rel: RelTag, + ctx: &RequestContext, + ) -> Result { + // Get current size and put rel creation if rel doesn't exist + // + // NOTE: we check the cache first even though get_rel_exists and get_rel_size would + // check the cache too. This is because eagerly checking the cache results in + // less work overall and 10% better performance. It's more work on cache miss + // but cache miss is rare. + if let Some(nblocks) = self.tline.get_cached_rel_size(&rel, self.get_lsn()) { + Ok(nblocks) + } else if !self + .tline + .get_rel_exists(rel, Version::Modified(self), ctx) + .await? + { + // create it with 0 size initially, the logic below will extend it + self.put_rel_creation(rel, 0, ctx) + .await + .context("Relation Error")?; + Ok(0) + } else { + self.tline + .get_rel_size(rel, Version::Modified(self), ctx) + .await + } + } + + /// Given a block number for a relation (which represents a newly written block), + /// the previous block count of the relation, and the shard info, find the gaps + /// that were created by the newly written block if any. + fn find_gaps( + rel: RelTag, + blkno: u32, + previous_nblocks: u32, + shard: &ShardIdentity, + ) -> Option { + let mut key = rel_block_to_key(rel, blkno); + let mut gap_accum = None; + + for gap_blkno in previous_nblocks..blkno { + key.field6 = gap_blkno; + + if shard.get_shard_number(&key) != shard.number { + continue; + } + + gap_accum + .get_or_insert_with(KeySpaceAccum::new) + .add_key(key); + } + + gap_accum.map(|accum| accum.to_keyspace()) + } + + pub async fn ingest_batch( + &mut self, + mut batch: SerializedValueBatch, + // TODO(vlad): remove this argument and replace the shard check with is_key_local + shard: &ShardIdentity, + ctx: &RequestContext, + ) -> anyhow::Result<()> { + let mut gaps_at_lsns = Vec::default(); + + for meta in batch.metadata.iter() { + let (rel, blkno) = Key::from_compact(meta.key()).to_rel_block()?; + let new_nblocks = blkno + 1; + + let old_nblocks = self.create_relation_if_required(rel, ctx).await?; + if new_nblocks > old_nblocks { + self.put_rel_extend(rel, new_nblocks, ctx).await?; + } + + if let Some(gaps) = Self::find_gaps(rel, blkno, old_nblocks, shard) { + gaps_at_lsns.push((gaps, meta.lsn())); + } + } + + if !gaps_at_lsns.is_empty() { + batch.zero_gaps(gaps_at_lsns); + } + + match self.pending_data_batch.as_mut() { + Some(pending_batch) => { + pending_batch.extend(batch); + } + None if !batch.is_empty() => { + self.pending_data_batch = Some(batch); + } + None => { + // Nothing to initialize the batch with + } + } + + Ok(()) + } + /// Put a new page version that can be constructed from a WAL record /// /// NOTE: this will *not* implicitly extend the relation, if the page is beyond the @@ -1229,8 +1327,13 @@ impl<'a> DatadirModification<'a> { self.lsn ); } - self.pending_zero_data_pages.insert(key.to_compact()); - self.pending_bytes += ZERO_PAGE.len(); + + let batch = self + .pending_data_batch + .get_or_insert_with(SerializedValueBatch::default); + + batch.put(key.to_compact(), Value::Image(ZERO_PAGE.clone()), self.lsn); + Ok(()) } @@ -1248,17 +1351,14 @@ impl<'a> DatadirModification<'a> { self.lsn ); } - self.pending_zero_data_pages.insert(key.to_compact()); - self.pending_bytes += ZERO_PAGE.len(); - Ok(()) - } - /// Call this at the end of each WAL record. - pub(crate) fn on_record_end(&mut self) { - let pending_zero_data_pages = std::mem::take(&mut self.pending_zero_data_pages); - for key in pending_zero_data_pages { - self.put_data(key, Value::Image(ZERO_PAGE.clone())); - } + let batch = self + .pending_data_batch + .get_or_insert_with(SerializedValueBatch::default); + + batch.put(key.to_compact(), Value::Image(ZERO_PAGE.clone()), self.lsn); + + Ok(()) } /// Store a relmapper file (pg_filenode.map) in the repository @@ -1750,12 +1850,17 @@ impl<'a> DatadirModification<'a> { let mut writer = self.tline.writer().await; // Flush relation and SLRU data blocks, keep metadata. - let pending_data_pages = std::mem::take(&mut self.pending_data_pages); + if let Some(batch) = self.pending_data_batch.take() { + tracing::debug!( + "Flushing batch with max_lsn={}. Last record LSN is {}", + batch.max_lsn, + self.tline.get_last_record_lsn() + ); - // This bails out on first error without modifying pending_updates. - // That's Ok, cf this function's doc comment. - writer.put_batch(pending_data_pages, ctx).await?; - self.pending_bytes = 0; + // This bails out on first error without modifying pending_updates. + // That's Ok, cf this function's doc comment. + writer.put_batch(batch, ctx).await?; + } if pending_nblocks != 0 { writer.update_current_logical_size(pending_nblocks * i64::from(BLCKSZ)); @@ -1775,9 +1880,6 @@ impl<'a> DatadirModification<'a> { /// All the modifications in this atomic update are stamped by the specified LSN. /// pub async fn commit(&mut self, ctx: &RequestContext) -> anyhow::Result<()> { - // Commit should never be called mid-wal-record - assert!(self.pending_zero_data_pages.is_empty()); - let mut writer = self.tline.writer().await; let pending_nblocks = self.pending_nblocks; @@ -1785,21 +1887,49 @@ impl<'a> DatadirModification<'a> { // Ordering: the items in this batch do not need to be in any global order, but values for // a particular Key must be in Lsn order relative to one another. InMemoryLayer relies on - // this to do efficient updates to its index. - let mut write_batch = std::mem::take(&mut self.pending_data_pages); + // this to do efficient updates to its index. See [`wal_decoder::serialized_batch`] for + // more details. - write_batch.extend( - self.pending_metadata_pages + let metadata_batch = { + let pending_meta = self + .pending_metadata_pages .drain() .flat_map(|(key, values)| { values .into_iter() .map(move |(lsn, value_size, value)| (key, lsn, value_size, value)) - }), - ); + }) + .collect::>(); - if !write_batch.is_empty() { - writer.put_batch(write_batch, ctx).await?; + if pending_meta.is_empty() { + None + } else { + Some(SerializedValueBatch::from_values(pending_meta)) + } + }; + + let data_batch = self.pending_data_batch.take(); + + let maybe_batch = match (data_batch, metadata_batch) { + (Some(mut data), Some(metadata)) => { + data.extend(metadata); + Some(data) + } + (Some(data), None) => Some(data), + (None, Some(metadata)) => Some(metadata), + (None, None) => None, + }; + + if let Some(batch) = maybe_batch { + tracing::debug!( + "Flushing batch with max_lsn={}. Last record LSN is {}", + batch.max_lsn, + self.tline.get_last_record_lsn() + ); + + // This bails out on first error without modifying pending_updates. + // That's Ok, cf this function's doc comment. + writer.put_batch(batch, ctx).await?; } if !self.pending_deletions.is_empty() { @@ -1809,6 +1939,9 @@ impl<'a> DatadirModification<'a> { self.pending_lsns.push(self.lsn); for pending_lsn in self.pending_lsns.drain(..) { + // TODO(vlad): pretty sure the comment below is not valid anymore + // and we can call finish write with the latest LSN + // // Ideally, we should be able to call writer.finish_write() only once // with the highest LSN. However, the last_record_lsn variable in the // timeline keeps track of the latest LSN and the immediate previous LSN @@ -1824,14 +1957,14 @@ impl<'a> DatadirModification<'a> { writer.update_directory_entries_count(kind, count as u64); } - self.pending_bytes = 0; + self.pending_metadata_bytes = 0; Ok(()) } pub(crate) fn len(&self) -> usize { self.pending_metadata_pages.len() - + self.pending_data_pages.len() + + self.pending_data_batch.as_ref().map_or(0, |b| b.len()) + self.pending_deletions.len() } @@ -1873,11 +2006,10 @@ impl<'a> DatadirModification<'a> { // modifications before ingesting DB create operations, which are the only kind that reads // data pages during ingest. if cfg!(debug_assertions) { - for (dirty_key, _, _, _) in &self.pending_data_pages { - debug_assert!(&key.to_compact() != dirty_key); - } - - debug_assert!(!self.pending_zero_data_pages.contains(&key.to_compact())) + assert!(!self + .pending_data_batch + .as_ref() + .map_or(false, |b| b.updates_key(&key))); } } @@ -1895,18 +2027,10 @@ impl<'a> DatadirModification<'a> { } fn put_data(&mut self, key: CompactKey, val: Value) { - let val_serialized_size = val.serialized_size().unwrap() as usize; - - // If this page was previously zero'd in the same WalRecord, then drop the previous zero page write. This - // is an optimization that avoids persisting both the zero page generated by us (e.g. during a relation extend), - // and the subsequent postgres-originating write - if self.pending_zero_data_pages.remove(&key) { - self.pending_bytes -= ZERO_PAGE.len(); - } - - self.pending_bytes += val_serialized_size; - self.pending_data_pages - .push((key, self.lsn, val_serialized_size, val)) + let batch = self + .pending_data_batch + .get_or_insert_with(SerializedValueBatch::default); + batch.put(key, val, self.lsn); } fn put_metadata(&mut self, key: CompactKey, val: Value) { @@ -1914,10 +2038,10 @@ impl<'a> DatadirModification<'a> { // Replace the previous value if it exists at the same lsn if let Some((last_lsn, last_value_ser_size, last_value)) = values.last_mut() { if *last_lsn == self.lsn { - // Update the pending_bytes contribution from this entry, and update the serialized size in place - self.pending_bytes -= *last_value_ser_size; + // Update the pending_metadata_bytes contribution from this entry, and update the serialized size in place + self.pending_metadata_bytes -= *last_value_ser_size; *last_value_ser_size = val.serialized_size().unwrap() as usize; - self.pending_bytes += *last_value_ser_size; + self.pending_metadata_bytes += *last_value_ser_size; // Use the latest value, this replaces any earlier write to the same (key,lsn), such as much // have been generated by synthesized zero page writes prior to the first real write to a page. @@ -1927,8 +2051,12 @@ impl<'a> DatadirModification<'a> { } let val_serialized_size = val.serialized_size().unwrap() as usize; - self.pending_bytes += val_serialized_size; + self.pending_metadata_bytes += val_serialized_size; values.push((self.lsn, val_serialized_size, val)); + + if key == CHECKPOINT_KEY.to_compact() { + tracing::debug!("Checkpoint key added to pending with size {val_serialized_size}"); + } } fn delete(&mut self, key_range: Range) { @@ -2037,7 +2165,11 @@ static ZERO_PAGE: Bytes = Bytes::from_static(&[0u8; BLCKSZ as usize]); #[cfg(test)] mod tests { use hex_literal::hex; - use utils::id::TimelineId; + use pageserver_api::{models::ShardParameters, shard::ShardStripeSize}; + use utils::{ + id::TimelineId, + shard::{ShardCount, ShardNumber}, + }; use super::*; @@ -2091,6 +2223,93 @@ mod tests { Ok(()) } + #[test] + fn gap_finding() { + let rel = RelTag { + spcnode: 1663, + dbnode: 208101, + relnode: 2620, + forknum: 0, + }; + let base_blkno = 1; + + let base_key = rel_block_to_key(rel, base_blkno); + let before_base_key = rel_block_to_key(rel, base_blkno - 1); + + let shard = ShardIdentity::unsharded(); + + let mut previous_nblocks = 0; + for i in 0..10 { + let crnt_blkno = base_blkno + i; + let gaps = DatadirModification::find_gaps(rel, crnt_blkno, previous_nblocks, &shard); + + previous_nblocks = crnt_blkno + 1; + + if i == 0 { + // The first block we write is 1, so we should find the gap. + assert_eq!(gaps.unwrap(), KeySpace::single(before_base_key..base_key)); + } else { + assert!(gaps.is_none()); + } + } + + // This is an update to an already existing block. No gaps here. + let update_blkno = 5; + let gaps = DatadirModification::find_gaps(rel, update_blkno, previous_nblocks, &shard); + assert!(gaps.is_none()); + + // This is an update past the current end block. + let after_gap_blkno = 20; + let gaps = DatadirModification::find_gaps(rel, after_gap_blkno, previous_nblocks, &shard); + + let gap_start_key = rel_block_to_key(rel, previous_nblocks); + let after_gap_key = rel_block_to_key(rel, after_gap_blkno); + assert_eq!( + gaps.unwrap(), + KeySpace::single(gap_start_key..after_gap_key) + ); + } + + #[test] + fn sharded_gap_finding() { + let rel = RelTag { + spcnode: 1663, + dbnode: 208101, + relnode: 2620, + forknum: 0, + }; + + let first_blkno = 6; + + // This shard will get the even blocks + let shard = ShardIdentity::from_params( + ShardNumber(0), + &ShardParameters { + count: ShardCount(2), + stripe_size: ShardStripeSize(1), + }, + ); + + // Only keys belonging to this shard are considered as gaps. + let mut previous_nblocks = 0; + let gaps = + DatadirModification::find_gaps(rel, first_blkno, previous_nblocks, &shard).unwrap(); + assert!(!gaps.ranges.is_empty()); + for gap_range in gaps.ranges { + let mut k = gap_range.start; + while k != gap_range.end { + assert_eq!(shard.get_shard_number(&k), shard.number); + k = k.next(); + } + } + + previous_nblocks = first_blkno; + + let update_blkno = 2; + let gaps = DatadirModification::find_gaps(rel, update_blkno, previous_nblocks, &shard); + assert!(gaps.is_none()); + } + /* fn assert_current_logical_size(timeline: &DatadirTimeline, lsn: Lsn) { let incremental = timeline.get_current_logical_size(); diff --git a/pageserver/src/tenant/storage_layer/inmemory_layer.rs b/pageserver/src/tenant/storage_layer/inmemory_layer.rs index df448a0963..2ce26ed2eb 100644 --- a/pageserver/src/tenant/storage_layer/inmemory_layer.rs +++ b/pageserver/src/tenant/storage_layer/inmemory_layer.rs @@ -12,7 +12,7 @@ use crate::tenant::timeline::GetVectoredError; use crate::tenant::PageReconstructError; use crate::virtual_file::owned_buffers_io::io_buf_ext::IoBufExt; use crate::{l0_flush, page_cache}; -use anyhow::{anyhow, Context, Result}; +use anyhow::{anyhow, Result}; use camino::Utf8PathBuf; use pageserver_api::key::CompactKey; use pageserver_api::key::Key; @@ -25,6 +25,7 @@ use std::sync::{Arc, OnceLock}; use std::time::Instant; use tracing::*; use utils::{bin_ser::BeSer, id::TimelineId, lsn::Lsn, vec_map::VecMap}; +use wal_decoder::serialized_batch::{SerializedValueBatch, SerializedValueMeta, ValueMeta}; // avoid binding to Write (conflicts with std::io::Write) // while being able to use std::fmt::Write's methods use crate::metrics::TIMELINE_EPHEMERAL_BYTES; @@ -452,6 +453,7 @@ impl InMemoryLayer { len, will_init, } = index_entry.unpack(); + reads.entry(key).or_default().push(ValueRead { entry_lsn: *entry_lsn, read: vectored_dio_read::LogicalRead::new( @@ -513,68 +515,6 @@ impl InMemoryLayer { } } -/// Offset of a particular Value within a serialized batch. -struct SerializedBatchOffset { - key: CompactKey, - lsn: Lsn, - // TODO: separate type when we start serde-serializing this value, to avoid coupling - // in-memory representation to serialization format. - index_entry: IndexEntry, -} - -pub struct SerializedBatch { - /// Blobs serialized in EphemeralFile's native format, ready for passing to [`EphemeralFile::write_raw`]. - pub(crate) raw: Vec, - - /// Index of values in [`Self::raw`], using offsets relative to the start of the buffer. - offsets: Vec, - - /// The highest LSN of any value in the batch - pub(crate) max_lsn: Lsn, -} - -impl SerializedBatch { - pub fn from_values(batch: Vec<(CompactKey, Lsn, usize, Value)>) -> anyhow::Result { - // Pre-allocate a big flat buffer to write into. This should be large but not huge: it is soft-limited in practice by - // [`crate::pgdatadir_mapping::DatadirModification::MAX_PENDING_BYTES`] - let buffer_size = batch.iter().map(|i| i.2).sum::(); - let mut cursor = std::io::Cursor::new(Vec::::with_capacity(buffer_size)); - - let mut offsets: Vec = Vec::with_capacity(batch.len()); - let mut max_lsn: Lsn = Lsn(0); - for (key, lsn, val_ser_size, val) in batch { - let relative_off = cursor.position(); - - val.ser_into(&mut cursor) - .expect("Writing into in-memory buffer is infallible"); - - offsets.push(SerializedBatchOffset { - key, - lsn, - index_entry: IndexEntry::new(IndexEntryNewArgs { - base_offset: 0, - batch_offset: relative_off, - len: val_ser_size, - will_init: val.will_init(), - }) - .context("higher-level code ensures that values are within supported ranges")?, - }); - max_lsn = std::cmp::max(max_lsn, lsn); - } - - let buffer = cursor.into_inner(); - - // Assert that we didn't do any extra allocations while building buffer. - debug_assert!(buffer.len() <= buffer_size); - - Ok(Self { - raw: buffer, - offsets, - max_lsn, - }) - } -} - fn inmem_layer_display(mut f: impl Write, start_lsn: Lsn, end_lsn: Lsn) -> std::fmt::Result { write!(f, "inmem-{:016X}-{:016X}", start_lsn.0, end_lsn.0) } @@ -642,7 +582,7 @@ impl InMemoryLayer { /// TODO: it can be made retryable if we aborted the process on EphemeralFile write errors. pub async fn put_batch( &self, - serialized_batch: SerializedBatch, + serialized_batch: SerializedValueBatch, ctx: &RequestContext, ) -> anyhow::Result<()> { let mut inner = self.inner.write().await; @@ -650,27 +590,13 @@ impl InMemoryLayer { let base_offset = inner.file.len(); - let SerializedBatch { + let SerializedValueBatch { raw, - mut offsets, + metadata, max_lsn: _, + len: _, } = serialized_batch; - // Add the base_offset to the batch's index entries which are relative to the batch start. - for offset in &mut offsets { - let IndexEntryUnpacked { - will_init, - len, - pos, - } = offset.index_entry.unpack(); - offset.index_entry = IndexEntry::new(IndexEntryNewArgs { - base_offset, - batch_offset: pos, - len: len.into_usize(), - will_init, - })?; - } - // Write the batch to the file inner.file.write_raw(&raw, ctx).await?; let new_size = inner.file.len(); @@ -683,12 +609,28 @@ impl InMemoryLayer { assert_eq!(new_size, expected_new_len); // Update the index with the new entries - for SerializedBatchOffset { - key, - lsn, - index_entry, - } in offsets - { + for meta in metadata { + let SerializedValueMeta { + key, + lsn, + batch_offset, + len, + will_init, + } = match meta { + ValueMeta::Serialized(ser) => ser, + ValueMeta::Observed(_) => { + continue; + } + }; + + // Add the base_offset to the batch's index entries which are relative to the batch start. + let index_entry = IndexEntry::new(IndexEntryNewArgs { + base_offset, + batch_offset, + len, + will_init, + })?; + let vec_map = inner.index.entry(key).or_default(); let old = vec_map.append_or_update_last(lsn, index_entry).unwrap().0; if old.is_some() { diff --git a/pageserver/src/tenant/timeline.rs b/pageserver/src/tenant/timeline.rs index 12919866a3..ee823beca8 100644 --- a/pageserver/src/tenant/timeline.rs +++ b/pageserver/src/tenant/timeline.rs @@ -24,8 +24,8 @@ use offload::OffloadError; use once_cell::sync::Lazy; use pageserver_api::{ key::{ - CompactKey, KEY_SIZE, METADATA_KEY_BEGIN_PREFIX, METADATA_KEY_END_PREFIX, - NON_INHERITED_RANGE, NON_INHERITED_SPARSE_RANGE, + KEY_SIZE, METADATA_KEY_BEGIN_PREFIX, METADATA_KEY_END_PREFIX, NON_INHERITED_RANGE, + NON_INHERITED_SPARSE_RANGE, }, keyspace::{KeySpaceAccum, KeySpaceRandomAccum, SparseKeyPartitioning}, models::{ @@ -49,6 +49,7 @@ use utils::{ fs_ext, pausable_failpoint, sync::gate::{Gate, GateGuard}, }; +use wal_decoder::serialized_batch::SerializedValueBatch; use std::sync::atomic::Ordering as AtomicOrdering; use std::sync::{Arc, Mutex, RwLock, Weak}; @@ -131,7 +132,6 @@ use crate::task_mgr::TaskKind; use crate::tenant::gc_result::GcResult; use crate::ZERO_PAGE; use pageserver_api::key::Key; -use pageserver_api::value::Value; use self::delete::DeleteTimelineFlow; pub(super) use self::eviction_task::EvictionTaskTenantState; @@ -141,9 +141,7 @@ use self::logical_size::LogicalSize; use self::walreceiver::{WalReceiver, WalReceiverConf}; use super::{ - config::TenantConf, - storage_layer::{inmemory_layer, LayerVisibilityHint}, - upload_queue::NotInitialized, + config::TenantConf, storage_layer::LayerVisibilityHint, upload_queue::NotInitialized, MaybeOffloaded, }; use super::{debug_assert_current_span_has_tenant_and_timeline_id, AttachedTenantConf}; @@ -157,6 +155,9 @@ use super::{ GcError, }; +#[cfg(test)] +use pageserver_api::value::Value; + #[derive(Debug, PartialEq, Eq, Clone, Copy)] pub(crate) enum FlushLoopState { NotStarted, @@ -5736,23 +5737,22 @@ impl<'a> TimelineWriter<'a> { /// Put a batch of keys at the specified Lsns. pub(crate) async fn put_batch( &mut self, - batch: Vec<(CompactKey, Lsn, usize, Value)>, + batch: SerializedValueBatch, ctx: &RequestContext, ) -> anyhow::Result<()> { if batch.is_empty() { return Ok(()); } - let serialized_batch = inmemory_layer::SerializedBatch::from_values(batch)?; - let batch_max_lsn = serialized_batch.max_lsn; - let buf_size: u64 = serialized_batch.raw.len() as u64; + let batch_max_lsn = batch.max_lsn; + let buf_size: u64 = batch.buffer_size() as u64; let action = self.get_open_layer_action(batch_max_lsn, buf_size); let layer = self .handle_open_layer_action(batch_max_lsn, action, ctx) .await?; - let res = layer.put_batch(serialized_batch, ctx).await; + let res = layer.put_batch(batch, ctx).await; if res.is_ok() { // Update the current size only when the entire write was ok. @@ -5787,11 +5787,14 @@ impl<'a> TimelineWriter<'a> { ); } let val_ser_size = value.serialized_size().unwrap() as usize; - self.put_batch( - vec![(key.to_compact(), lsn, val_ser_size, value.clone())], - ctx, - ) - .await + let batch = SerializedValueBatch::from_values(vec![( + key.to_compact(), + lsn, + val_ser_size, + value.clone(), + )]); + + self.put_batch(batch, ctx).await } pub(crate) async fn delete_batch( diff --git a/pageserver/src/tenant/timeline/walreceiver/walreceiver_connection.rs b/pageserver/src/tenant/timeline/walreceiver/walreceiver_connection.rs index eb19fb691f..34bf959058 100644 --- a/pageserver/src/tenant/timeline/walreceiver/walreceiver_connection.rs +++ b/pageserver/src/tenant/timeline/walreceiver/walreceiver_connection.rs @@ -331,11 +331,11 @@ pub(super) async fn handle_walreceiver_connection( Ok(()) } - while let Some((lsn, recdata)) = waldecoder.poll_decode()? { + while let Some((record_end_lsn, recdata)) = waldecoder.poll_decode()? { // It is important to deal with the aligned records as lsn in getPage@LSN is // aligned and can be several bytes bigger. Without this alignment we are // at risk of hitting a deadlock. - if !lsn.is_aligned() { + if !record_end_lsn.is_aligned() { return Err(WalReceiverError::Other(anyhow!("LSN not aligned"))); } @@ -343,7 +343,7 @@ pub(super) async fn handle_walreceiver_connection( let interpreted = InterpretedWalRecord::from_bytes_filtered( recdata, modification.tline.get_shard_identity(), - lsn, + record_end_lsn, modification.tline.pg_version, )?; @@ -366,9 +366,11 @@ pub(super) async fn handle_walreceiver_connection( let ingested = walingest .ingest_record(interpreted, &mut modification, &ctx) .await - .with_context(|| format!("could not ingest record at {lsn}"))?; + .with_context(|| { + format!("could not ingest record at {record_end_lsn}") + })?; if !ingested { - tracing::debug!("ingest: filtered out record @ LSN {lsn}"); + tracing::debug!("ingest: filtered out record @ LSN {record_end_lsn}"); WAL_INGEST.records_filtered.inc(); filtered_records += 1; } @@ -378,7 +380,7 @@ pub(super) async fn handle_walreceiver_connection( // to timeout the tests. fail_point!("walreceiver-after-ingest"); - last_rec_lsn = lsn; + last_rec_lsn = record_end_lsn; // Commit every ingest_batch_size records. Even if we filtered out // all records, we still need to call commit to advance the LSN. diff --git a/pageserver/src/walingest.rs b/pageserver/src/walingest.rs index 84353970b7..c3ccd8a2e4 100644 --- a/pageserver/src/walingest.rs +++ b/pageserver/src/walingest.rs @@ -28,14 +28,13 @@ use std::time::Duration; use std::time::Instant; use std::time::SystemTime; -use pageserver_api::key::Key; use pageserver_api::shard::ShardIdentity; use postgres_ffi::fsm_logical_to_physical; use postgres_ffi::walrecord::*; use postgres_ffi::{dispatch_pgversion, enum_pgversion, enum_pgversion_dispatch, TimestampTz}; use wal_decoder::models::*; -use anyhow::{bail, Context, Result}; +use anyhow::{bail, Result}; use bytes::{Buf, Bytes}; use tracing::*; use utils::failpoint_support; @@ -51,7 +50,6 @@ use crate::ZERO_PAGE; use pageserver_api::key::rel_block_to_key; use pageserver_api::record::NeonWalRecord; use pageserver_api::reltag::{BlockNumber, RelTag, SlruKind}; -use pageserver_api::value::Value; use postgres_ffi::pg_constants; use postgres_ffi::relfile_utils::{FSM_FORKNUM, INIT_FORKNUM, MAIN_FORKNUM, VISIBILITYMAP_FORKNUM}; use postgres_ffi::TransactionId; @@ -156,12 +154,12 @@ impl WalIngest { WAL_INGEST.records_received.inc(); let prev_len = modification.len(); - modification.set_lsn(interpreted.lsn)?; + modification.set_lsn(interpreted.end_lsn)?; if matches!(interpreted.flush_uncommitted, FlushUncommittedRecords::Yes) { // Records of this type should always be preceded by a commit(), as they // rely on reading data pages back from the Timeline. - assert!(!modification.has_dirty_data_pages()); + assert!(!modification.has_dirty_data()); } assert!(!self.checkpoint_modified); @@ -275,28 +273,9 @@ impl WalIngest { } } - // Iterate through all the key value pairs provided in the interpreted block - // and update the modification currently in-flight to include them. - for (compact_key, maybe_value) in interpreted.blocks.into_iter() { - let (rel, blk) = Key::from_compact(compact_key).to_rel_block()?; - match maybe_value { - Some(Value::Image(img)) => { - self.put_rel_page_image(modification, rel, blk, img, ctx) - .await?; - } - Some(Value::WalRecord(rec)) => { - self.put_rel_wal_record(modification, rel, blk, rec, ctx) - .await?; - } - None => { - // Shard 0 tracks relation sizes. We will observe - // its blkno in case it implicitly extends a relation. - assert!(self.shard.is_shard_zero()); - self.observe_decoded_block(modification, rel, blk, ctx) - .await?; - } - } - } + modification + .ingest_batch(interpreted.batch, &self.shard, ctx) + .await?; // If checkpoint data was updated, store the new version in the repository if self.checkpoint_modified { @@ -310,8 +289,6 @@ impl WalIngest { // until commit() is called to flush the data into the repository and update // the latest LSN. - modification.on_record_end(); - Ok(modification.len() > prev_len) } @@ -334,17 +311,6 @@ impl WalIngest { Ok((epoch as u64) << 32 | xid as u64) } - /// Do not store this block, but observe it for the purposes of updating our relation size state. - async fn observe_decoded_block( - &mut self, - modification: &mut DatadirModification<'_>, - rel: RelTag, - blkno: BlockNumber, - ctx: &RequestContext, - ) -> Result<(), PageReconstructError> { - self.handle_rel_extend(modification, rel, blkno, ctx).await - } - async fn ingest_clear_vm_bits( &mut self, clear_vm_bits: ClearVmBits, @@ -1248,6 +1214,7 @@ impl WalIngest { Ok(()) } + #[cfg(test)] async fn put_rel_page_image( &mut self, modification: &mut DatadirModification<'_>, @@ -1297,36 +1264,7 @@ impl WalIngest { let new_nblocks = blknum + 1; // Check if the relation exists. We implicitly create relations on first // record. - // TODO: would be nice if to be more explicit about it - - // Get current size and put rel creation if rel doesn't exist - // - // NOTE: we check the cache first even though get_rel_exists and get_rel_size would - // check the cache too. This is because eagerly checking the cache results in - // less work overall and 10% better performance. It's more work on cache miss - // but cache miss is rare. - let old_nblocks = if let Some(nblocks) = modification - .tline - .get_cached_rel_size(&rel, modification.get_lsn()) - { - nblocks - } else if !modification - .tline - .get_rel_exists(rel, Version::Modified(modification), ctx) - .await? - { - // create it with 0 size initially, the logic below will extend it - modification - .put_rel_creation(rel, 0, ctx) - .await - .context("Relation Error")?; - 0 - } else { - modification - .tline - .get_rel_size(rel, Version::Modified(modification), ctx) - .await? - }; + let old_nblocks = modification.create_relation_if_required(rel, ctx).await?; if new_nblocks > old_nblocks { //info!("extending {} {} to {}", rel, old_nblocks, new_nblocks); @@ -1553,25 +1491,21 @@ mod tests { walingest .put_rel_page_image(&mut m, TESTREL_A, 0, test_img("foo blk 0 at 2"), &ctx) .await?; - m.on_record_end(); m.commit(&ctx).await?; let mut m = tline.begin_modification(Lsn(0x30)); walingest .put_rel_page_image(&mut m, TESTREL_A, 0, test_img("foo blk 0 at 3"), &ctx) .await?; - m.on_record_end(); m.commit(&ctx).await?; let mut m = tline.begin_modification(Lsn(0x40)); walingest .put_rel_page_image(&mut m, TESTREL_A, 1, test_img("foo blk 1 at 4"), &ctx) .await?; - m.on_record_end(); m.commit(&ctx).await?; let mut m = tline.begin_modification(Lsn(0x50)); walingest .put_rel_page_image(&mut m, TESTREL_A, 2, test_img("foo blk 2 at 5"), &ctx) .await?; - m.on_record_end(); m.commit(&ctx).await?; assert_current_logical_size(&tline, Lsn(0x50)); @@ -1713,7 +1647,6 @@ mod tests { walingest .put_rel_page_image(&mut m, TESTREL_A, 1, test_img("foo blk 1"), &ctx) .await?; - m.on_record_end(); m.commit(&ctx).await?; assert_eq!( tline @@ -1739,7 +1672,6 @@ mod tests { walingest .put_rel_page_image(&mut m, TESTREL_A, 1500, test_img("foo blk 1500"), &ctx) .await?; - m.on_record_end(); m.commit(&ctx).await?; assert_eq!( tline From d182ff294c85f01e77bd0cf4e5345e3357986f23 Mon Sep 17 00:00:00 2001 From: John Spray Date: Wed, 6 Nov 2024 15:14:43 +0000 Subject: [PATCH 22/52] storcon: respect tenant scheduling policy in drain/fill (#9657) ## Problem Pinning a tenant by setting Pause scheduling policy doesn't work because drain/fill code moves the tenant around during deploys. Closes: https://github.com/neondatabase/neon/issues/9612 ## Summary of changes - In drain, only move a tenant if it is in Active or Essential mode - In fill, only move a tenant if it is in Active mode. The asymmetry is a bit annoying, but it faithfully respects the purposes of the modes: Essential is meant to endeavor to keep the tenant available, which means it needs to be drained but doesn't need to be migrated during fills. --- storage_controller/src/drain_utils.rs | 16 +++++++++++++++- storage_controller/src/service.rs | 10 ++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/storage_controller/src/drain_utils.rs b/storage_controller/src/drain_utils.rs index dea1f04649..47f4276ff2 100644 --- a/storage_controller/src/drain_utils.rs +++ b/storage_controller/src/drain_utils.rs @@ -3,7 +3,7 @@ use std::{ sync::Arc, }; -use pageserver_api::controller_api::NodeSchedulingPolicy; +use pageserver_api::controller_api::{NodeSchedulingPolicy, ShardSchedulingPolicy}; use utils::{id::NodeId, shard::TenantShardId}; use crate::{ @@ -98,6 +98,20 @@ impl TenantShardDrain { return None; } + // Only tenants with a normal (Active) scheduling policy are proactively moved + // around during a node drain. Shards which have been manually configured to a different + // policy are only rescheduled by manual intervention. + match tenant_shard.get_scheduling_policy() { + ShardSchedulingPolicy::Active | ShardSchedulingPolicy::Essential => { + // A migration during drain is classed as 'essential' because it is required to + // uphold our availability goals for the tenant: this shard is elegible for migration. + } + ShardSchedulingPolicy::Pause | ShardSchedulingPolicy::Stop => { + // If we have been asked to avoid rescheduling this shard, then do not migrate it during a drain + return None; + } + } + match scheduler.node_preferred(tenant_shard.intent.get_secondary()) { Some(node) => Some(node), None => { diff --git a/storage_controller/src/service.rs b/storage_controller/src/service.rs index 3f6cbfef59..e3a147bc06 100644 --- a/storage_controller/src/service.rs +++ b/storage_controller/src/service.rs @@ -6721,6 +6721,16 @@ impl Service { .tenants .iter_mut() .filter_map(|(tid, tenant_shard)| { + if !matches!( + tenant_shard.get_scheduling_policy(), + ShardSchedulingPolicy::Active + ) { + // Only include tenants in fills if they have a normal (Active) scheduling policy. We + // even exclude Essential, because moving to fill a node is not essential to keeping this + // tenant available. + return None; + } + if tenant_shard.intent.get_secondary().contains(&node_id) { if let Some(primary) = tenant_shard.intent.get_attached() { return Some((*primary, *tid)); From 73bdc9a2d0d3731c0191d0c2063039985825f26f Mon Sep 17 00:00:00 2001 From: Conrad Ludgate Date: Wed, 6 Nov 2024 17:40:40 +0000 Subject: [PATCH 23/52] [proxy]: minor changes to endpoint-cache handling (#9666) I think I meant to make these changes over 6 months ago. alas, better late than never. 1. should_reject doesn't eagerly intern the endpoint string 2. Rate limiter uses a std Mutex instead of a tokio Mutex. 3. Recently I introduced a `-local-proxy` endpoint suffix. I forgot to add this to normalize. 4. Random but a small cleanup making the ControlPlaneEvent deser directly to the interned strings. --- proxy/src/cache/endpoints.rs | 74 ++++++++++++++++---------- proxy/src/control_plane/client/neon.rs | 2 - proxy/src/types.rs | 29 +++++----- 3 files changed, 59 insertions(+), 46 deletions(-) diff --git a/proxy/src/cache/endpoints.rs b/proxy/src/cache/endpoints.rs index a488d8a9cd..07769e053c 100644 --- a/proxy/src/cache/endpoints.rs +++ b/proxy/src/cache/endpoints.rs @@ -1,13 +1,12 @@ use std::convert::Infallible; use std::future::pending; use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use dashmap::DashSet; use redis::streams::{StreamReadOptions, StreamReadReply}; use redis::{AsyncCommands, FromRedisValue, Value}; use serde::Deserialize; -use tokio::sync::Mutex; use tokio_util::sync::CancellationToken; use tracing::info; @@ -32,17 +31,17 @@ struct ControlPlaneEvent { #[derive(Deserialize, Debug, Clone, PartialEq)] struct EndpointCreated { - endpoint_id: String, + endpoint_id: EndpointIdInt, } #[derive(Deserialize, Debug, Clone, PartialEq)] struct BranchCreated { - branch_id: String, + branch_id: BranchIdInt, } #[derive(Deserialize, Debug, Clone, PartialEq)] struct ProjectCreated { - project_id: String, + project_id: ProjectIdInt, } impl TryFrom<&Value> for ControlPlaneEvent { @@ -76,53 +75,72 @@ impl EndpointsCache { } } - pub(crate) async fn is_valid(&self, ctx: &RequestMonitoring, endpoint: &EndpointId) -> bool { + pub(crate) fn is_valid(&self, ctx: &RequestMonitoring, endpoint: &EndpointId) -> bool { if !self.ready.load(Ordering::Acquire) { + // the endpoint cache is not yet fully initialised. return true; } - let rejected = self.should_reject(endpoint); - ctx.set_rejected(rejected); - info!(?rejected, "check endpoint is valid, disabled cache"); - // If cache is disabled, just collect the metrics and return or - // If the limiter allows, we don't need to check the cache. - if self.config.disable_cache || self.limiter.lock().await.check() { + + if !self.should_reject(endpoint) { + ctx.set_rejected(false); return true; } - !rejected + + // report that we might want to reject this endpoint + ctx.set_rejected(true); + + // If cache is disabled, just collect the metrics and return. + if self.config.disable_cache { + return true; + } + + // If the limiter allows, we can pretend like it's valid + // (incase it is, due to redis channel lag). + if self.limiter.lock().unwrap().check() { + return true; + } + + // endpoint not found, and there's too much load. + false } fn should_reject(&self, endpoint: &EndpointId) -> bool { if endpoint.is_endpoint() { - !self.endpoints.contains(&EndpointIdInt::from(endpoint)) + let Some(endpoint) = EndpointIdInt::get(endpoint) else { + // if we haven't interned this endpoint, it's not in the cache. + return true; + }; + !self.endpoints.contains(&endpoint) } else if endpoint.is_branch() { - !self - .branches - .contains(&BranchIdInt::from(&endpoint.as_branch())) + let Some(branch) = BranchIdInt::get(endpoint) else { + // if we haven't interned this branch, it's not in the cache. + return true; + }; + !self.branches.contains(&branch) } else { - !self - .projects - .contains(&ProjectIdInt::from(&endpoint.as_project())) + let Some(project) = ProjectIdInt::get(endpoint) else { + // if we haven't interned this project, it's not in the cache. + return true; + }; + !self.projects.contains(&project) } } fn insert_event(&self, event: ControlPlaneEvent) { if let Some(endpoint_created) = event.endpoint_created { - self.endpoints - .insert(EndpointIdInt::from(&endpoint_created.endpoint_id.into())); + self.endpoints.insert(endpoint_created.endpoint_id); Metrics::get() .proxy .redis_events_count .inc(RedisEventsCount::EndpointCreated); } else if let Some(branch_created) = event.branch_created { - self.branches - .insert(BranchIdInt::from(&branch_created.branch_id.into())); + self.branches.insert(branch_created.branch_id); Metrics::get() .proxy .redis_events_count .inc(RedisEventsCount::BranchCreated); } else if let Some(project_created) = event.project_created { - self.projects - .insert(ProjectIdInt::from(&project_created.project_id.into())); + self.projects.insert(project_created.project_id); Metrics::get() .proxy .redis_events_count @@ -247,11 +265,13 @@ mod tests { fn test_parse_control_plane_event() { let s = r#"{"branch_created":null,"endpoint_created":{"endpoint_id":"ep-rapid-thunder-w0qqw2q9"},"project_created":null,"type":"endpoint_created"}"#; + let endpoint_id: EndpointId = "ep-rapid-thunder-w0qqw2q9".into(); + assert_eq!( serde_json::from_str::(s).unwrap(), ControlPlaneEvent { endpoint_created: Some(EndpointCreated { - endpoint_id: "ep-rapid-thunder-w0qqw2q9".into() + endpoint_id: endpoint_id.into(), }), branch_created: None, project_created: None, diff --git a/proxy/src/control_plane/client/neon.rs b/proxy/src/control_plane/client/neon.rs index 1588e50423..26ff4e1402 100644 --- a/proxy/src/control_plane/client/neon.rs +++ b/proxy/src/control_plane/client/neon.rs @@ -72,7 +72,6 @@ impl NeonControlPlaneClient { .caches .endpoints_cache .is_valid(ctx, &user_info.endpoint.normalize()) - .await { info!("endpoint is not valid, skipping the request"); return Ok(AuthInfo::default()); @@ -145,7 +144,6 @@ impl NeonControlPlaneClient { .caches .endpoints_cache .is_valid(ctx, &endpoint.normalize()) - .await { return Err(GetEndpointJwksError::EndpointNotFound); } diff --git a/proxy/src/types.rs b/proxy/src/types.rs index b0408a51d1..6e0bd61c94 100644 --- a/proxy/src/types.rs +++ b/proxy/src/types.rs @@ -64,24 +64,28 @@ macro_rules! smol_str_wrapper { } const POOLER_SUFFIX: &str = "-pooler"; +pub(crate) const LOCAL_PROXY_SUFFIX: &str = "-local-proxy"; impl EndpointId { #[must_use] - pub fn normalize(&self) -> Self { + fn normalize_str(&self) -> &str { if let Some(stripped) = self.as_ref().strip_suffix(POOLER_SUFFIX) { - stripped.into() + stripped + } else if let Some(stripped) = self.as_ref().strip_suffix(LOCAL_PROXY_SUFFIX) { + stripped } else { - self.clone() + self } } + #[must_use] + pub fn normalize(&self) -> Self { + self.normalize_str().into() + } + #[must_use] pub fn normalize_intern(&self) -> EndpointIdInt { - if let Some(stripped) = self.as_ref().strip_suffix(POOLER_SUFFIX) { - EndpointIdTag::get_interner().get_or_intern(stripped) - } else { - self.into() - } + EndpointIdTag::get_interner().get_or_intern(self.normalize_str()) } } @@ -110,13 +114,4 @@ impl EndpointId { pub(crate) fn is_branch(&self) -> bool { self.0.starts_with("br-") } - // pub(crate) fn is_project(&self) -> bool { - // !self.is_endpoint() && !self.is_branch() - // } - pub(crate) fn as_branch(&self) -> BranchId { - BranchId(self.0.clone()) - } - pub(crate) fn as_project(&self) -> ProjectId { - ProjectId(self.0.clone()) - } } From 1d3559d4bc0d9be3f657fe15de5acf55e19d5c0a Mon Sep 17 00:00:00 2001 From: "Alex Chi Z." <4198311+skyzh@users.noreply.github.com> Date: Wed, 6 Nov 2024 13:17:02 -0500 Subject: [PATCH 24/52] feat(pageserver): add fast path for sparse keyspace read (#9631) In https://github.com/neondatabase/neon/issues/9441, the tenant has a lot of aux keys spread in multiple aux files. The perf tool shows that a significant amount of time is spent on remove_overlapping_keys. For sparse keyspaces, we don't need to report missing key errors anyways, and it's very likely that we will need to read all layers intersecting with the key range. Therefore, this patch adds a new fast path for sparse keyspace reads that we do not track `unmapped_keyspace` in a fine-grained way. We only modify it when we find an image layer. In debug mode, it was ~5min to read the aux files for a dump of the tenant, and now it's only 8s, that's a 60x speedup. ## Summary of changes * Do not add sparse keys into `keys_done` so that remove_overlapping does nothing. * Allow `ValueReconstructSituation::Complete` to be updated again in `ValuesReconstructState::update_key` for sparse keyspaces. --------- Signed-off-by: Alex Chi Z --- pageserver/src/tenant/storage_layer.rs | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/pageserver/src/tenant/storage_layer.rs b/pageserver/src/tenant/storage_layer.rs index 8f4219bbbc..9e3a25cbbc 100644 --- a/pageserver/src/tenant/storage_layer.rs +++ b/pageserver/src/tenant/storage_layer.rs @@ -12,7 +12,7 @@ pub mod merge_iterator; use crate::context::{AccessStatsBehavior, RequestContext}; use bytes::Bytes; -use pageserver_api::key::Key; +use pageserver_api::key::{Key, NON_INHERITED_SPARSE_RANGE}; use pageserver_api::keyspace::{KeySpace, KeySpaceRandomAccum}; use pageserver_api::record::NeonWalRecord; use pageserver_api::value::Value; @@ -196,6 +196,9 @@ impl ValuesReconstructState { /// Returns true if this was the last value needed for the key and false otherwise. /// /// If the key is done after the update, mark it as such. + /// + /// If the key is in the sparse keyspace (i.e., aux files), we do not track them in + /// `key_done`. pub(crate) fn update_key( &mut self, key: &Key, @@ -206,10 +209,18 @@ impl ValuesReconstructState { .keys .entry(*key) .or_insert(Ok(VectoredValueReconstructState::default())); - + let is_sparse_key = NON_INHERITED_SPARSE_RANGE.contains(key); if let Ok(state) = state { let key_done = match state.situation { - ValueReconstructSituation::Complete => unreachable!(), + ValueReconstructSituation::Complete => { + if is_sparse_key { + // Sparse keyspace might be visited multiple times because + // we don't track unmapped keyspaces. + return ValueReconstructSituation::Complete; + } else { + unreachable!() + } + } ValueReconstructSituation::Continue => match value { Value::Image(img) => { state.img = Some((lsn, img)); @@ -234,7 +245,9 @@ impl ValuesReconstructState { if key_done && state.situation == ValueReconstructSituation::Continue { state.situation = ValueReconstructSituation::Complete; - self.keys_done.add_key(*key); + if !is_sparse_key { + self.keys_done.add_key(*key); + } } state.situation From 93123f2623e5a8df90b0814aba3bf14edac27351 Mon Sep 17 00:00:00 2001 From: Tristan Partin Date: Wed, 6 Nov 2024 13:28:23 -0600 Subject: [PATCH 25/52] Rename compute_backpressure_throttling_ms to compute_backpressure_throttling_seconds This is in line with the Prometheus guidance[0]. We also haven't started using this metric, so renaming is essentially free. Link: https://prometheus.io/docs/practices/naming/ [0] Signed-off-by: Tristan Partin --- compute/etc/neon_collector.jsonnet | 2 +- .../etc/sql_exporter/compute_backpressure_throttling_ms.sql | 1 - ...nnet => compute_backpressure_throttling_seconds.libsonnet} | 4 ++-- .../sql_exporter/compute_backpressure_throttling_seconds.sql | 1 + 4 files changed, 4 insertions(+), 4 deletions(-) delete mode 100644 compute/etc/sql_exporter/compute_backpressure_throttling_ms.sql rename compute/etc/sql_exporter/{compute_backpressure_throttling_ms.libsonnet => compute_backpressure_throttling_seconds.libsonnet} (71%) create mode 100644 compute/etc/sql_exporter/compute_backpressure_throttling_seconds.sql diff --git a/compute/etc/neon_collector.jsonnet b/compute/etc/neon_collector.jsonnet index e73fb132ee..c6fa645b41 100644 --- a/compute/etc/neon_collector.jsonnet +++ b/compute/etc/neon_collector.jsonnet @@ -3,7 +3,7 @@ metrics: [ import 'sql_exporter/checkpoints_req.libsonnet', import 'sql_exporter/checkpoints_timed.libsonnet', - import 'sql_exporter/compute_backpressure_throttling_ms.libsonnet', + import 'sql_exporter/compute_backpressure_throttling_seconds.libsonnet', import 'sql_exporter/compute_current_lsn.libsonnet', import 'sql_exporter/compute_logical_snapshot_files.libsonnet', import 'sql_exporter/compute_receive_lsn.libsonnet', diff --git a/compute/etc/sql_exporter/compute_backpressure_throttling_ms.sql b/compute/etc/sql_exporter/compute_backpressure_throttling_ms.sql deleted file mode 100644 index 1fa62d38a4..0000000000 --- a/compute/etc/sql_exporter/compute_backpressure_throttling_ms.sql +++ /dev/null @@ -1 +0,0 @@ -SELECT neon.backpressure_throttling_time() AS throttled; diff --git a/compute/etc/sql_exporter/compute_backpressure_throttling_ms.libsonnet b/compute/etc/sql_exporter/compute_backpressure_throttling_seconds.libsonnet similarity index 71% rename from compute/etc/sql_exporter/compute_backpressure_throttling_ms.libsonnet rename to compute/etc/sql_exporter/compute_backpressure_throttling_seconds.libsonnet index b25bb73d0f..02c803cfa6 100644 --- a/compute/etc/sql_exporter/compute_backpressure_throttling_ms.libsonnet +++ b/compute/etc/sql_exporter/compute_backpressure_throttling_seconds.libsonnet @@ -1,10 +1,10 @@ { - metric_name: 'compute_backpressure_throttling_ms', + metric_name: 'compute_backpressure_throttling_seconds', type: 'gauge', help: 'Time compute has spent throttled', key_labels: null, values: [ 'throttled', ], - query: importstr 'sql_exporter/compute_backpressure_throttling_ms.sql', + query: importstr 'sql_exporter/compute_backpressure_throttling_seconds.sql', } diff --git a/compute/etc/sql_exporter/compute_backpressure_throttling_seconds.sql b/compute/etc/sql_exporter/compute_backpressure_throttling_seconds.sql new file mode 100644 index 0000000000..459c586d18 --- /dev/null +++ b/compute/etc/sql_exporter/compute_backpressure_throttling_seconds.sql @@ -0,0 +1 @@ +SELECT neon.backpressure_throttling_time()::float8 / 1000 AS throttled; From 11fc1a4c12e435215e44f5f0d8a424335f7018f6 Mon Sep 17 00:00:00 2001 From: Yuchen Liang <70461588+yliang412@users.noreply.github.com> Date: Wed, 6 Nov 2024 15:18:21 -0500 Subject: [PATCH 26/52] fix(test): use layer map dump in `test_readonly_node_gc` to validate layers protected by leases (#9551) Fixes #9518. ## Problem After removing the assertion `layers_removed == 0` in #9506, we could miss breakage if we solely rely on the successful execution of the `SELECT` query to check if lease is properly protecting layers. Details listed in #9518. Also, in integration tests, we sometimes run into the race condition where getpage request comes before the lease get renewed (item 2 of #8817), even if compute_ctl sends a lease renewal as soon as it sees a `/configure` API calls that updates the `pageserver_connstr`. In this case, we would observe a getpage request error stating that we `tried to request a page version that was garbage collected` (as we seen in [Allure Report](https://neon-github-public-dev.s3.amazonaws.com/reports/pr-8613/11550393107/index.html#suites/3ccffb1d100105b98aed3dc19b717917/d1a1ba47bc180493)). ## Summary of changes - Use layer map dump to verify if the lease protects what it claimed: Record all historical layers that has `start_lsn <= lease_lsn` before and after running timeline gc. This is the same check as https://github.com/neondatabase/neon/blob/ad79f4246030583e5af418cb087bf7582266accc/pageserver/src/tenant/timeline.rs#L5025-L5027 The set recorded after GC should contain every layer in the set recorded before GC. - Wait until log contains another successful lease request before running the `SELECT` query after GC. We argued in #8817 that the bad request can only exist within a short period after migration/restart, and our test shows that as long as a lease renewal is done before the first getpage request sent after reconfiguration, we will not have bad request. Signed-off-by: Yuchen Liang --- test_runner/regress/test_readonly_node.py | 71 ++++++++++++++++++++--- 1 file changed, 64 insertions(+), 7 deletions(-) diff --git a/test_runner/regress/test_readonly_node.py b/test_runner/regress/test_readonly_node.py index 8151160477..f257f0853b 100644 --- a/test_runner/regress/test_readonly_node.py +++ b/test_runner/regress/test_readonly_node.py @@ -1,19 +1,22 @@ from __future__ import annotations import time +from typing import Union import pytest -from fixtures.common_types import Lsn +from fixtures.common_types import Lsn, TenantId, TenantShardId, TimelineId from fixtures.log_helper import log from fixtures.neon_fixtures import ( Endpoint, + LogCursor, NeonEnv, NeonEnvBuilder, last_flush_lsn_upload, tenant_get_shards, ) +from fixtures.pageserver.http import PageserverHttpClient from fixtures.pageserver.utils import wait_for_last_record_lsn -from fixtures.utils import query_scalar +from fixtures.utils import query_scalar, wait_until # @@ -169,23 +172,63 @@ def test_readonly_node_gc(neon_env_builder: NeonEnvBuilder): ) return last_flush_lsn - def trigger_gc_and_select(env: NeonEnv, ep_static: Endpoint, ctx: str): + def get_layers_protected_by_lease( + ps_http: PageserverHttpClient, + tenant_id: Union[TenantId, TenantShardId], + timeline_id: TimelineId, + lease_lsn: Lsn, + ) -> set[str]: + """Get all layers whose start_lsn is less than or equal to the lease lsn.""" + layer_map_info = ps_http.layer_map_info(tenant_id, timeline_id) + return set( + x.layer_file_name + for x in layer_map_info.historic_layers + if Lsn(x.lsn_start) <= lease_lsn + ) + + def trigger_gc_and_select( + env: NeonEnv, + ep_static: Endpoint, + lease_lsn: Lsn, + ctx: str, + offset: None | LogCursor = None, + ) -> LogCursor: """ Trigger GC manually on all pageservers. Then run an `SELECT` query. """ for shard, ps in tenant_get_shards(env, env.initial_tenant): client = ps.http_client() + layers_guarded_before_gc = get_layers_protected_by_lease( + client, shard, env.initial_timeline, lease_lsn=lsn + ) gc_result = client.timeline_gc(shard, env.initial_timeline, 0) + layers_guarded_after_gc = get_layers_protected_by_lease( + client, shard, env.initial_timeline, lease_lsn=lsn + ) + # Note: cannot assert on `layers_removed` here because it could be layers - # not guarded by the lease. Rely on successful execution of the query instead. + # not guarded by the lease. Instead, use layer map dump. + assert layers_guarded_before_gc.issubset( + layers_guarded_after_gc + ), "Layers guarded by lease before GC should not be removed" + log.info(f"{gc_result=}") + # wait for lease renewal before running query. + _, offset = wait_until( + 20, + 0.5, + lambda: ep_static.assert_log_contains( + "lsn_lease_bg_task.*Request succeeded", offset=offset + ), + ) with ep_static.cursor() as cur: # Following query should succeed if pages are properly guarded by leases. cur.execute("SELECT count(*) FROM t0") assert cur.fetchone() == (ROW_COUNT,) log.info(f"`SELECT` query succeed after GC, {ctx=}") + return offset # Insert some records on main branch with env.endpoints.create_start("main") as ep_main: @@ -213,7 +256,9 @@ def test_readonly_node_gc(neon_env_builder: NeonEnvBuilder): generate_updates_on_main(env, ep_main, 3, end=100) - trigger_gc_and_select(env, ep_static, ctx="Before pageservers restart") + offset = trigger_gc_and_select( + env, ep_static, lease_lsn=lsn, ctx="Before pageservers restart" + ) # Trigger Pageserver restarts for ps in env.pageservers: @@ -222,7 +267,13 @@ def test_readonly_node_gc(neon_env_builder: NeonEnvBuilder): time.sleep(LSN_LEASE_LENGTH / 2) ps.start() - trigger_gc_and_select(env, ep_static, ctx="After pageservers restart") + trigger_gc_and_select( + env, + ep_static, + lease_lsn=lsn, + ctx="After pageservers restart", + offset=offset, + ) # Reconfigure pageservers env.pageservers[0].stop() @@ -231,7 +282,13 @@ def test_readonly_node_gc(neon_env_builder: NeonEnvBuilder): ) env.storage_controller.reconcile_until_idle() - trigger_gc_and_select(env, ep_static, ctx="After putting pageserver 0 offline") + trigger_gc_and_select( + env, + ep_static, + lease_lsn=lsn, + ctx="After putting pageserver 0 offline", + offset=offset, + ) # Do some update so we can increment latest_gc_cutoff generate_updates_on_main(env, ep_main, i, end=100) From 2a95a51a0df6602aa444fd07842e4b347f970ac9 Mon Sep 17 00:00:00 2001 From: "Alex Chi Z." <4198311+skyzh@users.noreply.github.com> Date: Wed, 6 Nov 2024 15:41:01 -0500 Subject: [PATCH 27/52] refactor(pageserver): better pageservice command parsing (#9597) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit close https://github.com/neondatabase/neon/issues/9460 ## Summary of changes A full rewrite of pagestream cmdline parsing to make it more robust and readable. --------- Signed-off-by: Alex Chi Z Co-authored-by: Arpad Müller --- pageserver/src/page_service.rs | 682 ++++++++++++++++++++++++--------- 1 file changed, 504 insertions(+), 178 deletions(-) diff --git a/pageserver/src/page_service.rs b/pageserver/src/page_service.rs index aed8a87851..f07474df6a 100644 --- a/pageserver/src/page_service.rs +++ b/pageserver/src/page_service.rs @@ -1,10 +1,11 @@ //! The Page Service listens for client connections and serves their GetPage@LSN //! requests. -use anyhow::Context; +use anyhow::{bail, Context}; use async_compression::tokio::write::GzipEncoder; use bytes::Buf; use futures::FutureExt; +use itertools::Itertools; use once_cell::sync::OnceCell; use pageserver_api::models::TenantState; use pageserver_api::models::{ @@ -1221,6 +1222,222 @@ impl PageServerHandler { } } +/// `basebackup tenant timeline [lsn] [--gzip] [--replica]` +#[derive(Debug, Clone, Eq, PartialEq)] +struct BaseBackupCmd { + tenant_id: TenantId, + timeline_id: TimelineId, + lsn: Option, + gzip: bool, + replica: bool, +} + +/// `fullbackup tenant timeline [lsn] [prev_lsn]` +#[derive(Debug, Clone, Eq, PartialEq)] +struct FullBackupCmd { + tenant_id: TenantId, + timeline_id: TimelineId, + lsn: Option, + prev_lsn: Option, +} + +/// `pagestream_v2 tenant timeline` +#[derive(Debug, Clone, Eq, PartialEq)] +struct PageStreamCmd { + tenant_id: TenantId, + timeline_id: TimelineId, +} + +/// `lease lsn tenant timeline lsn` +#[derive(Debug, Clone, Eq, PartialEq)] +struct LeaseLsnCmd { + tenant_shard_id: TenantShardId, + timeline_id: TimelineId, + lsn: Lsn, +} + +#[derive(Debug, Clone, Eq, PartialEq)] +enum PageServiceCmd { + Set, + PageStream(PageStreamCmd), + BaseBackup(BaseBackupCmd), + FullBackup(FullBackupCmd), + LeaseLsn(LeaseLsnCmd), +} + +impl PageStreamCmd { + fn parse(query: &str) -> anyhow::Result { + let parameters = query.split_whitespace().collect_vec(); + if parameters.len() != 2 { + bail!( + "invalid number of parameters for pagestream command: {}", + query + ); + } + let tenant_id = TenantId::from_str(parameters[0]) + .with_context(|| format!("Failed to parse tenant id from {}", parameters[0]))?; + let timeline_id = TimelineId::from_str(parameters[1]) + .with_context(|| format!("Failed to parse timeline id from {}", parameters[1]))?; + Ok(Self { + tenant_id, + timeline_id, + }) + } +} + +impl FullBackupCmd { + fn parse(query: &str) -> anyhow::Result { + let parameters = query.split_whitespace().collect_vec(); + if parameters.len() < 2 || parameters.len() > 4 { + bail!( + "invalid number of parameters for basebackup command: {}", + query + ); + } + let tenant_id = TenantId::from_str(parameters[0]) + .with_context(|| format!("Failed to parse tenant id from {}", parameters[0]))?; + let timeline_id = TimelineId::from_str(parameters[1]) + .with_context(|| format!("Failed to parse timeline id from {}", parameters[1]))?; + // The caller is responsible for providing correct lsn and prev_lsn. + let lsn = if let Some(lsn_str) = parameters.get(2) { + Some( + Lsn::from_str(lsn_str) + .with_context(|| format!("Failed to parse Lsn from {lsn_str}"))?, + ) + } else { + None + }; + let prev_lsn = if let Some(prev_lsn_str) = parameters.get(3) { + Some( + Lsn::from_str(prev_lsn_str) + .with_context(|| format!("Failed to parse Lsn from {prev_lsn_str}"))?, + ) + } else { + None + }; + Ok(Self { + tenant_id, + timeline_id, + lsn, + prev_lsn, + }) + } +} + +impl BaseBackupCmd { + fn parse(query: &str) -> anyhow::Result { + let parameters = query.split_whitespace().collect_vec(); + if parameters.len() < 2 { + bail!( + "invalid number of parameters for basebackup command: {}", + query + ); + } + let tenant_id = TenantId::from_str(parameters[0]) + .with_context(|| format!("Failed to parse tenant id from {}", parameters[0]))?; + let timeline_id = TimelineId::from_str(parameters[1]) + .with_context(|| format!("Failed to parse timeline id from {}", parameters[1]))?; + let lsn; + let flags_parse_from; + if let Some(maybe_lsn) = parameters.get(2) { + if *maybe_lsn == "latest" { + lsn = None; + flags_parse_from = 3; + } else if maybe_lsn.starts_with("--") { + lsn = None; + flags_parse_from = 2; + } else { + lsn = Some( + Lsn::from_str(maybe_lsn) + .with_context(|| format!("Failed to parse lsn from {maybe_lsn}"))?, + ); + flags_parse_from = 3; + } + } else { + lsn = None; + flags_parse_from = 2; + } + + let mut gzip = false; + let mut replica = false; + + for ¶m in ¶meters[flags_parse_from..] { + match param { + "--gzip" => { + if gzip { + bail!("duplicate parameter for basebackup command: {param}") + } + gzip = true + } + "--replica" => { + if replica { + bail!("duplicate parameter for basebackup command: {param}") + } + replica = true + } + _ => bail!("invalid parameter for basebackup command: {param}"), + } + } + Ok(Self { + tenant_id, + timeline_id, + lsn, + gzip, + replica, + }) + } +} + +impl LeaseLsnCmd { + fn parse(query: &str) -> anyhow::Result { + let parameters = query.split_whitespace().collect_vec(); + if parameters.len() != 3 { + bail!( + "invalid number of parameters for lease lsn command: {}", + query + ); + } + let tenant_shard_id = TenantShardId::from_str(parameters[0]) + .with_context(|| format!("Failed to parse tenant id from {}", parameters[0]))?; + let timeline_id = TimelineId::from_str(parameters[1]) + .with_context(|| format!("Failed to parse timeline id from {}", parameters[1]))?; + let lsn = Lsn::from_str(parameters[2]) + .with_context(|| format!("Failed to parse lsn from {}", parameters[2]))?; + Ok(Self { + tenant_shard_id, + timeline_id, + lsn, + }) + } +} + +impl PageServiceCmd { + fn parse(query: &str) -> anyhow::Result { + let query = query.trim(); + let Some((cmd, other)) = query.split_once(' ') else { + bail!("cannot parse query: {query}") + }; + match cmd.to_ascii_lowercase().as_str() { + "pagestream_v2" => Ok(Self::PageStream(PageStreamCmd::parse(other)?)), + "basebackup" => Ok(Self::BaseBackup(BaseBackupCmd::parse(other)?)), + "fullbackup" => Ok(Self::FullBackup(FullBackupCmd::parse(other)?)), + "lease" => { + let Some((cmd2, other)) = other.split_once(' ') else { + bail!("invalid lease command: {cmd}"); + }; + let cmd2 = cmd2.to_ascii_lowercase(); + if cmd2 == "lsn" { + Ok(Self::LeaseLsn(LeaseLsnCmd::parse(other)?)) + } else { + bail!("invalid lease command: {cmd}"); + } + } + "set" => Ok(Self::Set), + _ => Err(anyhow::anyhow!("unsupported command {cmd} in {query}")), + } + } +} + impl postgres_backend::Handler for PageServerHandler where IO: AsyncRead + AsyncWrite + Send + Sync + Unpin, @@ -1277,206 +1494,137 @@ where fail::fail_point!("ps::connection-start::process-query"); let ctx = self.connection_ctx.attached_child(); - debug!("process query {query_string:?}"); - let parts = query_string.split_whitespace().collect::>(); - if let Some(params) = parts.strip_prefix(&["pagestream_v2"]) { - if params.len() != 2 { - return Err(QueryError::Other(anyhow::anyhow!( - "invalid param number for pagestream command" - ))); - } - let tenant_id = TenantId::from_str(params[0]) - .with_context(|| format!("Failed to parse tenant id from {}", params[0]))?; - let timeline_id = TimelineId::from_str(params[1]) - .with_context(|| format!("Failed to parse timeline id from {}", params[1]))?; - - tracing::Span::current() - .record("tenant_id", field::display(tenant_id)) - .record("timeline_id", field::display(timeline_id)); - - self.check_permission(Some(tenant_id))?; - - COMPUTE_COMMANDS_COUNTERS - .for_command(ComputeCommandKind::PageStreamV2) - .inc(); - - self.handle_pagerequests( - pgb, + debug!("process query {query_string}"); + let query = PageServiceCmd::parse(query_string)?; + match query { + PageServiceCmd::PageStream(PageStreamCmd { tenant_id, timeline_id, - PagestreamProtocolVersion::V2, - ctx, - ) - .await?; - } else if let Some(params) = parts.strip_prefix(&["basebackup"]) { - if params.len() < 2 { - return Err(QueryError::Other(anyhow::anyhow!( - "invalid param number for basebackup command" - ))); + }) => { + tracing::Span::current() + .record("tenant_id", field::display(tenant_id)) + .record("timeline_id", field::display(timeline_id)); + + self.check_permission(Some(tenant_id))?; + + COMPUTE_COMMANDS_COUNTERS + .for_command(ComputeCommandKind::PageStreamV2) + .inc(); + + self.handle_pagerequests( + pgb, + tenant_id, + timeline_id, + PagestreamProtocolVersion::V2, + ctx, + ) + .await?; } + PageServiceCmd::BaseBackup(BaseBackupCmd { + tenant_id, + timeline_id, + lsn, + gzip, + replica, + }) => { + tracing::Span::current() + .record("tenant_id", field::display(tenant_id)) + .record("timeline_id", field::display(timeline_id)); - let tenant_id = TenantId::from_str(params[0]) - .with_context(|| format!("Failed to parse tenant id from {}", params[0]))?; - let timeline_id = TimelineId::from_str(params[1]) - .with_context(|| format!("Failed to parse timeline id from {}", params[1]))?; + self.check_permission(Some(tenant_id))?; - tracing::Span::current() - .record("tenant_id", field::display(tenant_id)) - .record("timeline_id", field::display(timeline_id)); - - self.check_permission(Some(tenant_id))?; - - COMPUTE_COMMANDS_COUNTERS - .for_command(ComputeCommandKind::Basebackup) - .inc(); - - let mut lsn = None; - let mut replica = false; - let mut gzip = false; - for param in ¶ms[2..] { - if param.starts_with("--") { - match *param { - "--gzip" => gzip = true, - "--replica" => replica = true, - _ => { - return Err(QueryError::Other(anyhow::anyhow!( - "Unknown parameter {param}", - ))) - } - } - } else { - lsn = Some( - Lsn::from_str(param) - .with_context(|| format!("Failed to parse Lsn from {param}"))?, - ); + COMPUTE_COMMANDS_COUNTERS + .for_command(ComputeCommandKind::Basebackup) + .inc(); + let metric_recording = metrics::BASEBACKUP_QUERY_TIME.start_recording(&ctx); + let res = async { + self.handle_basebackup_request( + pgb, + tenant_id, + timeline_id, + lsn, + None, + false, + gzip, + replica, + &ctx, + ) + .await?; + pgb.write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?; + Result::<(), QueryError>::Ok(()) } + .await; + metric_recording.observe(&res); + res?; } + // same as basebackup, but result includes relational data as well + PageServiceCmd::FullBackup(FullBackupCmd { + tenant_id, + timeline_id, + lsn, + prev_lsn, + }) => { + tracing::Span::current() + .record("tenant_id", field::display(tenant_id)) + .record("timeline_id", field::display(timeline_id)); - let metric_recording = metrics::BASEBACKUP_QUERY_TIME.start_recording(&ctx); - let res = async { + self.check_permission(Some(tenant_id))?; + + COMPUTE_COMMANDS_COUNTERS + .for_command(ComputeCommandKind::Fullbackup) + .inc(); + + // Check that the timeline exists self.handle_basebackup_request( pgb, tenant_id, timeline_id, lsn, - None, + prev_lsn, + true, + false, false, - gzip, - replica, &ctx, ) .await?; pgb.write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?; - Result::<(), QueryError>::Ok(()) } - .await; - metric_recording.observe(&res); - res?; - } - // same as basebackup, but result includes relational data as well - else if let Some(params) = parts.strip_prefix(&["fullbackup"]) { - if params.len() < 2 { - return Err(QueryError::Other(anyhow::anyhow!( - "invalid param number for fullbackup command" - ))); + PageServiceCmd::Set => { + // important because psycopg2 executes "SET datestyle TO 'ISO'" + // on connect + pgb.write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?; } - - let tenant_id = TenantId::from_str(params[0]) - .with_context(|| format!("Failed to parse tenant id from {}", params[0]))?; - let timeline_id = TimelineId::from_str(params[1]) - .with_context(|| format!("Failed to parse timeline id from {}", params[1]))?; - - tracing::Span::current() - .record("tenant_id", field::display(tenant_id)) - .record("timeline_id", field::display(timeline_id)); - - // The caller is responsible for providing correct lsn and prev_lsn. - let lsn = if let Some(lsn_str) = params.get(2) { - Some( - Lsn::from_str(lsn_str) - .with_context(|| format!("Failed to parse Lsn from {lsn_str}"))?, - ) - } else { - None - }; - let prev_lsn = if let Some(prev_lsn_str) = params.get(3) { - Some( - Lsn::from_str(prev_lsn_str) - .with_context(|| format!("Failed to parse Lsn from {prev_lsn_str}"))?, - ) - } else { - None - }; - - self.check_permission(Some(tenant_id))?; - - COMPUTE_COMMANDS_COUNTERS - .for_command(ComputeCommandKind::Fullbackup) - .inc(); - - // Check that the timeline exists - self.handle_basebackup_request( - pgb, - tenant_id, + PageServiceCmd::LeaseLsn(LeaseLsnCmd { + tenant_shard_id, timeline_id, lsn, - prev_lsn, - true, - false, - false, - &ctx, - ) - .await?; - pgb.write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?; - } else if query_string.to_ascii_lowercase().starts_with("set ") { - // important because psycopg2 executes "SET datestyle TO 'ISO'" - // on connect - pgb.write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?; - } else if query_string.starts_with("lease lsn ") { - let params = &parts[2..]; - if params.len() != 3 { - return Err(QueryError::Other(anyhow::anyhow!( - "invalid param number {} for lease lsn command", - params.len() - ))); + }) => { + tracing::Span::current() + .record("tenant_id", field::display(tenant_shard_id)) + .record("timeline_id", field::display(timeline_id)); + + self.check_permission(Some(tenant_shard_id.tenant_id))?; + + COMPUTE_COMMANDS_COUNTERS + .for_command(ComputeCommandKind::LeaseLsn) + .inc(); + + match self + .handle_make_lsn_lease(pgb, tenant_shard_id, timeline_id, lsn, &ctx) + .await + { + Ok(()) => { + pgb.write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))? + } + Err(e) => { + error!("error obtaining lsn lease for {lsn}: {e:?}"); + pgb.write_message_noflush(&BeMessage::ErrorResponse( + &e.to_string(), + Some(e.pg_error_code()), + ))? + } + }; } - - let tenant_shard_id = TenantShardId::from_str(params[0]) - .with_context(|| format!("Failed to parse tenant id from {}", params[0]))?; - let timeline_id = TimelineId::from_str(params[1]) - .with_context(|| format!("Failed to parse timeline id from {}", params[1]))?; - - tracing::Span::current() - .record("tenant_id", field::display(tenant_shard_id)) - .record("timeline_id", field::display(timeline_id)); - - self.check_permission(Some(tenant_shard_id.tenant_id))?; - - COMPUTE_COMMANDS_COUNTERS - .for_command(ComputeCommandKind::LeaseLsn) - .inc(); - - // The caller is responsible for providing correct lsn. - let lsn = Lsn::from_str(params[2]) - .with_context(|| format!("Failed to parse Lsn from {}", params[2]))?; - - match self - .handle_make_lsn_lease(pgb, tenant_shard_id, timeline_id, lsn, &ctx) - .await - { - Ok(()) => pgb.write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?, - Err(e) => { - error!("error obtaining lsn lease for {lsn}: {e:?}"); - pgb.write_message_noflush(&BeMessage::ErrorResponse( - &e.to_string(), - Some(e.pg_error_code()), - ))? - } - }; - } else { - return Err(QueryError::Other(anyhow::anyhow!( - "unknown command {query_string}" - ))); } Ok(()) @@ -1525,3 +1673,181 @@ fn set_tracing_field_shard_id(timeline: &Timeline) { ); debug_assert_current_span_has_tenant_and_timeline_id(); } + +#[cfg(test)] +mod tests { + use utils::shard::ShardCount; + + use super::*; + + #[test] + fn pageservice_cmd_parse() { + let tenant_id = TenantId::generate(); + let timeline_id = TimelineId::generate(); + let cmd = + PageServiceCmd::parse(&format!("pagestream_v2 {tenant_id} {timeline_id}")).unwrap(); + assert_eq!( + cmd, + PageServiceCmd::PageStream(PageStreamCmd { + tenant_id, + timeline_id + }) + ); + let cmd = PageServiceCmd::parse(&format!("basebackup {tenant_id} {timeline_id}")).unwrap(); + assert_eq!( + cmd, + PageServiceCmd::BaseBackup(BaseBackupCmd { + tenant_id, + timeline_id, + lsn: None, + gzip: false, + replica: false + }) + ); + let cmd = + PageServiceCmd::parse(&format!("basebackup {tenant_id} {timeline_id} --gzip")).unwrap(); + assert_eq!( + cmd, + PageServiceCmd::BaseBackup(BaseBackupCmd { + tenant_id, + timeline_id, + lsn: None, + gzip: true, + replica: false + }) + ); + let cmd = + PageServiceCmd::parse(&format!("basebackup {tenant_id} {timeline_id} latest")).unwrap(); + assert_eq!( + cmd, + PageServiceCmd::BaseBackup(BaseBackupCmd { + tenant_id, + timeline_id, + lsn: None, + gzip: false, + replica: false + }) + ); + let cmd = PageServiceCmd::parse(&format!("basebackup {tenant_id} {timeline_id} 0/16ABCDE")) + .unwrap(); + assert_eq!( + cmd, + PageServiceCmd::BaseBackup(BaseBackupCmd { + tenant_id, + timeline_id, + lsn: Some(Lsn::from_str("0/16ABCDE").unwrap()), + gzip: false, + replica: false + }) + ); + let cmd = PageServiceCmd::parse(&format!( + "basebackup {tenant_id} {timeline_id} --replica --gzip" + )) + .unwrap(); + assert_eq!( + cmd, + PageServiceCmd::BaseBackup(BaseBackupCmd { + tenant_id, + timeline_id, + lsn: None, + gzip: true, + replica: true + }) + ); + let cmd = PageServiceCmd::parse(&format!( + "basebackup {tenant_id} {timeline_id} 0/16ABCDE --replica --gzip" + )) + .unwrap(); + assert_eq!( + cmd, + PageServiceCmd::BaseBackup(BaseBackupCmd { + tenant_id, + timeline_id, + lsn: Some(Lsn::from_str("0/16ABCDE").unwrap()), + gzip: true, + replica: true + }) + ); + let cmd = PageServiceCmd::parse(&format!("fullbackup {tenant_id} {timeline_id}")).unwrap(); + assert_eq!( + cmd, + PageServiceCmd::FullBackup(FullBackupCmd { + tenant_id, + timeline_id, + lsn: None, + prev_lsn: None + }) + ); + let cmd = PageServiceCmd::parse(&format!( + "fullbackup {tenant_id} {timeline_id} 0/16ABCDE 0/16ABCDF" + )) + .unwrap(); + assert_eq!( + cmd, + PageServiceCmd::FullBackup(FullBackupCmd { + tenant_id, + timeline_id, + lsn: Some(Lsn::from_str("0/16ABCDE").unwrap()), + prev_lsn: Some(Lsn::from_str("0/16ABCDF").unwrap()), + }) + ); + let tenant_shard_id = TenantShardId::unsharded(tenant_id); + let cmd = PageServiceCmd::parse(&format!( + "lease lsn {tenant_shard_id} {timeline_id} 0/16ABCDE" + )) + .unwrap(); + assert_eq!( + cmd, + PageServiceCmd::LeaseLsn(LeaseLsnCmd { + tenant_shard_id, + timeline_id, + lsn: Lsn::from_str("0/16ABCDE").unwrap(), + }) + ); + let tenant_shard_id = TenantShardId::split(&tenant_shard_id, ShardCount(8))[1]; + let cmd = PageServiceCmd::parse(&format!( + "lease lsn {tenant_shard_id} {timeline_id} 0/16ABCDE" + )) + .unwrap(); + assert_eq!( + cmd, + PageServiceCmd::LeaseLsn(LeaseLsnCmd { + tenant_shard_id, + timeline_id, + lsn: Lsn::from_str("0/16ABCDE").unwrap(), + }) + ); + let cmd = PageServiceCmd::parse("set a = b").unwrap(); + assert_eq!(cmd, PageServiceCmd::Set); + let cmd = PageServiceCmd::parse("SET foo").unwrap(); + assert_eq!(cmd, PageServiceCmd::Set); + } + + #[test] + fn pageservice_cmd_err_handling() { + let tenant_id = TenantId::generate(); + let timeline_id = TimelineId::generate(); + let cmd = PageServiceCmd::parse("unknown_command"); + assert!(cmd.is_err()); + let cmd = PageServiceCmd::parse("pagestream_v2"); + assert!(cmd.is_err()); + let cmd = PageServiceCmd::parse(&format!("pagestream_v2 {tenant_id}xxx")); + assert!(cmd.is_err()); + let cmd = PageServiceCmd::parse(&format!("pagestream_v2 {tenant_id}xxx {timeline_id}xxx")); + assert!(cmd.is_err()); + let cmd = PageServiceCmd::parse(&format!( + "basebackup {tenant_id} {timeline_id} --gzip --gzip" + )); + assert!(cmd.is_err()); + let cmd = PageServiceCmd::parse(&format!( + "basebackup {tenant_id} {timeline_id} --gzip --unknown" + )); + assert!(cmd.is_err()); + let cmd = PageServiceCmd::parse(&format!( + "basebackup {tenant_id} {timeline_id} --gzip 0/16ABCDE" + )); + assert!(cmd.is_err()); + let cmd = PageServiceCmd::parse(&format!("lease {tenant_id} {timeline_id} gzip 0/16ABCDE")); + assert!(cmd.is_err()); + } +} From 011c0a175f91d3f636acb11dc8c7613b64ccb6cc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arpad=20M=C3=BCller?= Date: Thu, 7 Nov 2024 01:53:58 +0100 Subject: [PATCH 28/52] Support copying layers in detach_ancestor from before shard splits (#9669) We need to use the shard associated with the layer file, not the shard associated with our current tenant shard ID. Due to shard splits, the shard IDs can refer to older files. close https://github.com/neondatabase/neon/issues/9667 --- .../src/tenant/remote_timeline_client.rs | 6 +-- .../src/tenant/timeline/detach_ancestor.rs | 14 +++++-- .../regress/test_timeline_detach_ancestor.py | 41 ++++++++++++++++--- 3 files changed, 49 insertions(+), 12 deletions(-) diff --git a/pageserver/src/tenant/remote_timeline_client.rs b/pageserver/src/tenant/remote_timeline_client.rs index 03ec18c882..0aa8d61036 100644 --- a/pageserver/src/tenant/remote_timeline_client.rs +++ b/pageserver/src/tenant/remote_timeline_client.rs @@ -1445,7 +1445,7 @@ impl RemoteTimelineClient { let remote_path = remote_layer_path( &self.tenant_shard_id.tenant_id, &self.timeline_id, - self.tenant_shard_id.to_index(), + uploaded.metadata().shard, &uploaded.layer_desc().layer_name(), uploaded.metadata().generation, ); @@ -1486,7 +1486,7 @@ impl RemoteTimelineClient { &adopted .get_timeline_id() .expect("Source timeline should be alive"), - self.tenant_shard_id.to_index(), + adopted.metadata().shard, &adopted.layer_desc().layer_name(), adopted.metadata().generation, ); @@ -1494,7 +1494,7 @@ impl RemoteTimelineClient { let target_remote_path = remote_layer_path( &self.tenant_shard_id.tenant_id, &self.timeline_id, - self.tenant_shard_id.to_index(), + adopted_as.metadata().shard, &adopted_as.layer_desc().layer_name(), adopted_as.metadata().generation, ); diff --git a/pageserver/src/tenant/timeline/detach_ancestor.rs b/pageserver/src/tenant/timeline/detach_ancestor.rs index b4c0ab0329..f8bc4352e2 100644 --- a/pageserver/src/tenant/timeline/detach_ancestor.rs +++ b/pageserver/src/tenant/timeline/detach_ancestor.rs @@ -12,7 +12,7 @@ use crate::{ virtual_file::{MaybeFatalIo, VirtualFile}, }; use anyhow::Context; -use pageserver_api::models::detach_ancestor::AncestorDetached; +use pageserver_api::{models::detach_ancestor::AncestorDetached, shard::ShardIdentity}; use tokio::sync::Semaphore; use tokio_util::sync::CancellationToken; use tracing::Instrument; @@ -376,8 +376,14 @@ pub(super) async fn prepare( tasks.spawn( async move { let _permit = limiter.acquire().await; - let owned = - remote_copy(&adopted, &timeline, timeline.generation, &timeline.cancel).await?; + let owned = remote_copy( + &adopted, + &timeline, + timeline.generation, + timeline.shard_identity, + &timeline.cancel, + ) + .await?; tracing::info!(layer=%owned, "remote copied"); Ok(owned) } @@ -629,6 +635,7 @@ async fn remote_copy( adopted: &Layer, adoptee: &Arc, generation: Generation, + shard_identity: ShardIdentity, cancel: &CancellationToken, ) -> Result { // depending if Layer::keep_resident we could hardlink @@ -636,6 +643,7 @@ async fn remote_copy( let mut metadata = adopted.metadata(); debug_assert!(metadata.generation <= generation); metadata.generation = generation; + metadata.shard = shard_identity.shard_index(); let owned = crate::tenant::storage_layer::Layer::for_evicted( adoptee.conf, diff --git a/test_runner/regress/test_timeline_detach_ancestor.py b/test_runner/regress/test_timeline_detach_ancestor.py index ed47f9432b..0e8519e07b 100644 --- a/test_runner/regress/test_timeline_detach_ancestor.py +++ b/test_runner/regress/test_timeline_detach_ancestor.py @@ -16,6 +16,7 @@ from fixtures.neon_fixtures import ( NeonEnvBuilder, PgBin, flush_ep_to_pageserver, + last_flush_lsn_upload, wait_for_last_flush_lsn, ) from fixtures.pageserver.http import HistoricLayerInfo, PageserverApiException @@ -576,27 +577,49 @@ def test_compaction_induced_by_detaches_in_history( assert_pageserver_backups_equal(fullbackup_before, fullbackup_after, set()) -@pytest.mark.parametrize("sharded", [True, False]) +@pytest.mark.parametrize("shards_initial_after", [(1, 1), (2, 2), (1, 4)]) def test_timeline_ancestor_detach_idempotent_success( - neon_env_builder: NeonEnvBuilder, sharded: bool + neon_env_builder: NeonEnvBuilder, shards_initial_after: tuple[int, int] ): - shards = 2 if sharded else 1 + shards_initial = shards_initial_after[0] + shards_after = shards_initial_after[1] - neon_env_builder.num_pageservers = shards - env = neon_env_builder.init_start(initial_tenant_shard_count=shards if sharded else None) + neon_env_builder.num_pageservers = shards_after + env = neon_env_builder.init_start( + initial_tenant_shard_count=shards_initial if shards_initial > 1 else None, + initial_tenant_conf={ + # small checkpointing and compaction targets to ensure we generate many upload operations + "checkpoint_distance": 512 * 1024, + "compaction_threshold": 1, + "compaction_target_size": 512 * 1024, + # disable background compaction and GC. We invoke it manually when we want it to happen. + "gc_period": "0s", + "compaction_period": "0s", + }, + ) pageservers = dict((int(p.id), p) for p in env.pageservers) for ps in pageservers.values(): ps.allowed_errors.extend(SHUTDOWN_ALLOWED_ERRORS) - if sharded: + if shards_after > 1: # FIXME: should this be in the neon_env_builder.init_start? env.storage_controller.reconcile_until_idle() client = env.storage_controller.pageserver_api() else: client = env.pageserver.http_client() + # Write some data so that we have some layers to copy + with env.endpoints.create_start("main", tenant_id=env.initial_tenant) as endpoint: + endpoint.safe_psql_many( + [ + "CREATE TABLE foo(key serial primary key, t text default 'data_content')", + "INSERT INTO foo SELECT FROM generate_series(1,1024)", + ] + ) + last_flush_lsn_upload(env, endpoint, env.initial_tenant, env.initial_timeline) + first_branch = env.create_branch("first_branch") _ = env.create_branch("second_branch", ancestor_branch_name="first_branch") @@ -607,6 +630,12 @@ def test_timeline_ancestor_detach_idempotent_success( reparented1 = env.create_branch("first_reparented", ancestor_branch_name="main") reparented2 = env.create_branch("second_reparented", ancestor_branch_name="main") + if shards_after > shards_initial: + # Do a shard split + # This is a reproducer for https://github.com/neondatabase/neon/issues/9667 + env.storage_controller.tenant_shard_split(env.initial_tenant, shards_after) + env.storage_controller.reconcile_until_idle() + first_reparenting_response = client.detach_ancestor(env.initial_tenant, first_branch) assert set(first_reparenting_response) == {reparented1, reparented2} From e1d0b73824dfd42895202c666c24779023d4c4b0 Mon Sep 17 00:00:00 2001 From: Cheng Chen Date: Wed, 6 Nov 2024 20:41:18 -0800 Subject: [PATCH 29/52] chore(compute): Bump pg_mooncake to the latest version --- compute/compute-node.Dockerfile | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/compute/compute-node.Dockerfile b/compute/compute-node.Dockerfile index 30126de56c..f070f66c0a 100644 --- a/compute/compute-node.Dockerfile +++ b/compute/compute-node.Dockerfile @@ -1151,8 +1151,8 @@ COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/ # The topmost commit in the `neon` branch at the time of writing this # https://github.com/Mooncake-Labs/pg_mooncake/commits/neon/ -# https://github.com/Mooncake-Labs/pg_mooncake/commit/568b5a82b5fc16136bdf4ca5aac3e0cc261ab48d -ENV PG_MOONCAKE_VERSION=568b5a82b5fc16136bdf4ca5aac3e0cc261ab48d +# https://github.com/Mooncake-Labs/pg_mooncake/commit/077c92c452bb6896a7b7776ee95f039984f076af +ENV PG_MOONCAKE_VERSION=077c92c452bb6896a7b7776ee95f039984f076af ENV PATH="/usr/local/pgsql/bin/:$PATH" RUN case "${PG_VERSION}" in \ From d6aa26a533844671a3b48e03ed9bb8c2345bb137 Mon Sep 17 00:00:00 2001 From: Erik Grinaker Date: Thu, 7 Nov 2024 11:38:39 +0100 Subject: [PATCH 30/52] postgres_ffi: make `WalGenerator` generic over record generator (#9614) ## Problem Benchmarks need more control over the WAL generated by `WalGenerator`. In particular, they need to vary the size of logical messages. ## Summary of changes * Make `WalGenerator` generic over `RecordGenerator`, which constructs WAL records. * Add `LogicalMessageGenerator` which emits logical messages, with a configurable payload. * Minor tweaks and code reorganization. There are no changes to the core logic or emitted WAL. --- libs/postgres_ffi/src/wal_generator.rs | 236 +++++++++++------- libs/postgres_ffi/src/xlog_utils.rs | 14 +- .../tests/walproposer_sim/walproposer_disk.rs | 8 +- 3 files changed, 156 insertions(+), 102 deletions(-) diff --git a/libs/postgres_ffi/src/wal_generator.rs b/libs/postgres_ffi/src/wal_generator.rs index 97968c269b..dc679eea33 100644 --- a/libs/postgres_ffi/src/wal_generator.rs +++ b/libs/postgres_ffi/src/wal_generator.rs @@ -1,10 +1,10 @@ -use std::ffi::CStr; +use std::ffi::{CStr, CString}; use bytes::{Bytes, BytesMut}; use crc32c::crc32c_append; use utils::lsn::Lsn; -use super::bindings::{XLogLongPageHeaderData, XLogPageHeaderData, XLOG_PAGE_MAGIC}; +use super::bindings::{RmgrId, XLogLongPageHeaderData, XLogPageHeaderData, XLOG_PAGE_MAGIC}; use super::xlog_utils::{ XlLogicalMessage, XLOG_RECORD_CRC_OFFS, XLOG_SIZE_OF_XLOG_RECORD, XLP_BKP_REMOVABLE, XLP_FIRST_IS_CONTRECORD, @@ -16,11 +16,65 @@ use crate::pg_constants::{ }; use crate::{WAL_SEGMENT_SIZE, XLOG_BLCKSZ}; -/// Generates binary WAL records for use in tests and benchmarks. Currently only generates logical -/// messages (effectively noops) with a fixed payload. It is used as an iterator which yields -/// encoded bytes for a single WAL record, including internal page headers if it spans pages. -/// Concatenating the bytes will yield a complete, well-formed WAL, which can be chunked at segment -/// boundaries if desired. Not optimized for performance. +/// A WAL record payload. Will be prefixed by an XLogRecord header when encoded. +pub struct Record { + pub rmid: RmgrId, + pub info: u8, + pub data: Bytes, +} + +impl Record { + /// Encodes the WAL record including an XLogRecord header. prev_lsn is the start position of + /// the previous record in the WAL -- this is ignored by the Safekeeper, but not Postgres. + pub fn encode(&self, prev_lsn: Lsn) -> Bytes { + // Prefix data with block ID and length. + let data_header = Bytes::from(match self.data.len() { + 0 => vec![], + 1..=255 => vec![XLR_BLOCK_ID_DATA_SHORT, self.data.len() as u8], + 256.. => { + let len_bytes = (self.data.len() as u32).to_le_bytes(); + [&[XLR_BLOCK_ID_DATA_LONG], len_bytes.as_slice()].concat() + } + }); + + // Construct the WAL record header. + let mut header = XLogRecord { + xl_tot_len: (XLOG_SIZE_OF_XLOG_RECORD + data_header.len() + self.data.len()) as u32, + xl_xid: 0, + xl_prev: prev_lsn.into(), + xl_info: self.info, + xl_rmid: self.rmid, + __bindgen_padding_0: [0; 2], + xl_crc: 0, // see below + }; + + // Compute the CRC checksum for the data, and the header up to the CRC field. + let mut crc = 0; + crc = crc32c_append(crc, &data_header); + crc = crc32c_append(crc, &self.data); + crc = crc32c_append(crc, &header.encode().unwrap()[0..XLOG_RECORD_CRC_OFFS]); + header.xl_crc = crc; + + // Encode the final header and record. + let header = header.encode().unwrap(); + + [header, data_header, self.data.clone()].concat().into() + } +} + +/// Generates WAL record payloads. +/// +/// TODO: currently only provides LogicalMessageGenerator for trivial noop messages. Add a generator +/// that creates a table and inserts rows. +pub trait RecordGenerator: Iterator {} + +impl> RecordGenerator for I {} + +/// Generates binary WAL for use in tests and benchmarks. The provided record generator constructs +/// the WAL records. It is used as an iterator which yields encoded bytes for a single WAL record, +/// including internal page headers if it spans pages. Concatenating the bytes will yield a +/// complete, well-formed WAL, which can be chunked at segment boundaries if desired. Not optimized +/// for performance. /// /// The WAL format is version-dependant (see e.g. `XLOG_PAGE_MAGIC`), so make sure to import this /// for the appropriate Postgres version (e.g. `postgres_ffi::v17::wal_generator::WalGenerator`). @@ -31,10 +85,10 @@ use crate::{WAL_SEGMENT_SIZE, XLOG_BLCKSZ}; /// | Segment 1 | Segment 2 | Segment 3 | /// | Page 1 | Page 2 | Page 3 | Page 4 | Page 5 | Page 6 | Page 7 | Page 8 | Page 9 | /// | R1 | R2 |R3| R4 | R5 | R6 | R7 | R8 | -/// -/// TODO: support generating actual tables and rows. #[derive(Default)] -pub struct WalGenerator { +pub struct WalGenerator { + /// Generates record payloads for the WAL. + pub record_generator: R, /// Current LSN to append the next record at. /// /// Callers can modify this (and prev_lsn) to restart generation at a different LSN, but should @@ -46,73 +100,35 @@ pub struct WalGenerator { pub prev_lsn: Lsn, } -impl WalGenerator { - // For now, hardcode the message payload. - // TODO: support specifying the payload size. - const PREFIX: &CStr = c"prefix"; - const MESSAGE: &[u8] = b"message"; - - // Hardcode the sys, timeline, and DB IDs. We can make them configurable if we care about them. +impl WalGenerator { + // Hardcode the sys and timeline ID. We can make them configurable if we care about them. const SYS_ID: u64 = 0; const TIMELINE_ID: u32 = 1; - const DB_ID: u32 = 0; - /// Creates a new WAL generator, which emits logical message records (noops). - pub fn new() -> Self { - Self::default() + /// Creates a new WAL generator with the given record generator. + pub fn new(record_generator: R) -> WalGenerator { + Self { + record_generator, + lsn: Lsn(0), + prev_lsn: Lsn(0), + } } - /// Encodes a logical message (basically a noop), with the given prefix and message. - pub(crate) fn encode_logical_message(prefix: &CStr, message: &[u8]) -> Bytes { - let prefix = prefix.to_bytes_with_nul(); - let header = XlLogicalMessage { - db_id: Self::DB_ID, - transactional: 0, - prefix_size: prefix.len() as u64, - message_size: message.len() as u64, - }; - [&header.encode(), prefix, message].concat().into() + /// Appends a record with an arbitrary payload at the current LSN, then increments the LSN. + /// Returns the WAL bytes for the record, including page headers and padding, and the start LSN. + fn append_record(&mut self, record: Record) -> (Lsn, Bytes) { + let record = record.encode(self.prev_lsn); + let record = Self::insert_pages(record, self.lsn); + let record = Self::pad_record(record, self.lsn); + let lsn = self.lsn; + self.prev_lsn = self.lsn; + self.lsn += record.len() as u64; + (lsn, record) } - /// Encode a WAL record with the given payload data (e.g. a logical message). - pub(crate) fn encode_record(data: Bytes, rmid: u8, info: u8, prev_lsn: Lsn) -> Bytes { - // Prefix data with block ID and length. - let data_header = Bytes::from(match data.len() { - 0 => vec![], - 1..=255 => vec![XLR_BLOCK_ID_DATA_SHORT, data.len() as u8], - 256.. => { - let len_bytes = (data.len() as u32).to_le_bytes(); - [&[XLR_BLOCK_ID_DATA_LONG], len_bytes.as_slice()].concat() - } - }); - - // Construct the WAL record header. - let mut header = XLogRecord { - xl_tot_len: (XLOG_SIZE_OF_XLOG_RECORD + data_header.len() + data.len()) as u32, - xl_xid: 0, - xl_prev: prev_lsn.into(), - xl_info: info, - xl_rmid: rmid, - __bindgen_padding_0: [0; 2], - xl_crc: 0, // see below - }; - - // Compute the CRC checksum for the data, and the header up to the CRC field. - let mut crc = 0; - crc = crc32c_append(crc, &data_header); - crc = crc32c_append(crc, &data); - crc = crc32c_append(crc, &header.encode().unwrap()[0..XLOG_RECORD_CRC_OFFS]); - header.xl_crc = crc; - - // Encode the final header and record. - let header = header.encode().unwrap(); - - [header, data_header, data].concat().into() - } - - /// Injects page headers on 8KB page boundaries. Takes the current LSN position where the record + /// Inserts page headers on 8KB page boundaries. Takes the current LSN position where the record /// is to be appended. - fn encode_pages(record: Bytes, mut lsn: Lsn) -> Bytes { + fn insert_pages(record: Bytes, mut lsn: Lsn) -> Bytes { // Fast path: record fits in current page, and the page already has a header. if lsn.remaining_in_block() as usize >= record.len() && lsn.block_offset() > 0 { return record; @@ -173,31 +189,71 @@ impl WalGenerator { } [record, Bytes::from(vec![0; padding])].concat().into() } - - /// Generates a record with an arbitrary payload at the current LSN, then increments the LSN. - pub fn generate_record(&mut self, data: Bytes, rmid: u8, info: u8) -> Bytes { - let record = Self::encode_record(data, rmid, info, self.prev_lsn); - let record = Self::encode_pages(record, self.lsn); - let record = Self::pad_record(record, self.lsn); - self.prev_lsn = self.lsn; - self.lsn += record.len() as u64; - record - } - - /// Generates a logical message at the current LSN. Can be used to construct arbitrary messages. - pub fn generate_logical_message(&mut self, prefix: &CStr, message: &[u8]) -> Bytes { - let data = Self::encode_logical_message(prefix, message); - self.generate_record(data, RM_LOGICALMSG_ID, XLOG_LOGICAL_MESSAGE) - } } -/// Generate WAL records as an iterator. -impl Iterator for WalGenerator { +/// Generates WAL records as an iterator. +impl Iterator for WalGenerator { type Item = (Lsn, Bytes); fn next(&mut self) -> Option { - let lsn = self.lsn; - let record = self.generate_logical_message(Self::PREFIX, Self::MESSAGE); - Some((lsn, record)) + let record = self.record_generator.next()?; + Some(self.append_record(record)) + } +} + +/// Generates logical message records (effectively noops) with a fixed message. +pub struct LogicalMessageGenerator { + prefix: CString, + message: Vec, +} + +impl LogicalMessageGenerator { + const DB_ID: u32 = 0; // hardcoded for now + const RM_ID: RmgrId = RM_LOGICALMSG_ID; + const INFO: u8 = XLOG_LOGICAL_MESSAGE; + + /// Creates a new LogicalMessageGenerator. + pub fn new(prefix: &CStr, message: &[u8]) -> Self { + Self { + prefix: prefix.to_owned(), + message: message.to_owned(), + } + } + + /// Encodes a logical message. + fn encode(prefix: &CStr, message: &[u8]) -> Bytes { + let prefix = prefix.to_bytes_with_nul(); + let header = XlLogicalMessage { + db_id: Self::DB_ID, + transactional: 0, + prefix_size: prefix.len() as u64, + message_size: message.len() as u64, + }; + [&header.encode(), prefix, message].concat().into() + } +} + +impl Iterator for LogicalMessageGenerator { + type Item = Record; + + fn next(&mut self) -> Option { + Some(Record { + rmid: Self::RM_ID, + info: Self::INFO, + data: Self::encode(&self.prefix, &self.message), + }) + } +} + +impl WalGenerator { + /// Convenience method for appending a WAL record with an arbitrary logical message at the + /// current WAL LSN position. Returns the start LSN and resulting WAL bytes. + pub fn append_logical_message(&mut self, prefix: &CStr, message: &[u8]) -> (Lsn, Bytes) { + let record = Record { + rmid: LogicalMessageGenerator::RM_ID, + info: LogicalMessageGenerator::INFO, + data: LogicalMessageGenerator::encode(prefix, message), + }; + self.append_record(record) } } diff --git a/libs/postgres_ffi/src/xlog_utils.rs b/libs/postgres_ffi/src/xlog_utils.rs index 78a965174f..852b20eace 100644 --- a/libs/postgres_ffi/src/xlog_utils.rs +++ b/libs/postgres_ffi/src/xlog_utils.rs @@ -12,9 +12,9 @@ use super::bindings::{ CheckPoint, ControlFileData, DBState_DB_SHUTDOWNED, FullTransactionId, TimeLineID, TimestampTz, XLogLongPageHeaderData, XLogPageHeaderData, XLogRecPtr, XLogRecord, XLogSegNo, XLOG_PAGE_MAGIC, }; -use super::wal_generator::WalGenerator; +use super::wal_generator::LogicalMessageGenerator; use super::PG_MAJORVERSION; -use crate::pg_constants::{self, RM_LOGICALMSG_ID, XLOG_LOGICAL_MESSAGE}; +use crate::pg_constants; use crate::PG_TLI; use crate::{uint32, uint64, Oid}; use crate::{WAL_SEGMENT_SIZE, XLOG_BLCKSZ}; @@ -493,12 +493,10 @@ pub fn encode_logical_message(prefix: &str, message: &str) -> Bytes { // This function can take untrusted input, so discard any NUL bytes in the prefix string. let prefix = CString::new(prefix.replace('\0', "")).expect("no NULs"); let message = message.as_bytes(); - WalGenerator::encode_record( - WalGenerator::encode_logical_message(&prefix, message), - RM_LOGICALMSG_ID, - XLOG_LOGICAL_MESSAGE, - Lsn(0), - ) + LogicalMessageGenerator::new(&prefix, message) + .next() + .unwrap() + .encode(Lsn(0)) } #[cfg(test)] diff --git a/safekeeper/tests/walproposer_sim/walproposer_disk.rs b/safekeeper/tests/walproposer_sim/walproposer_disk.rs index f70cd65dfc..aefb3919a1 100644 --- a/safekeeper/tests/walproposer_sim/walproposer_disk.rs +++ b/safekeeper/tests/walproposer_sim/walproposer_disk.rs @@ -1,7 +1,7 @@ use std::{ffi::CStr, sync::Arc}; use parking_lot::{Mutex, MutexGuard}; -use postgres_ffi::v16::wal_generator::WalGenerator; +use postgres_ffi::v16::wal_generator::{LogicalMessageGenerator, WalGenerator}; use utils::lsn::Lsn; use super::block_storage::BlockStorage; @@ -18,7 +18,7 @@ impl DiskWalProposer { internal_available_lsn: Lsn(0), prev_lsn: Lsn(0), disk: BlockStorage::new(), - wal_generator: WalGenerator::new(), + wal_generator: WalGenerator::new(LogicalMessageGenerator::new(c"", &[])), }), }) } @@ -36,7 +36,7 @@ pub struct State { // actual WAL storage disk: BlockStorage, // WAL record generator - wal_generator: WalGenerator, + wal_generator: WalGenerator, } impl State { @@ -64,7 +64,7 @@ impl State { /// Inserts a logical record in the WAL at the current LSN. pub fn insert_logical_message(&mut self, prefix: &CStr, msg: &[u8]) { - let record = self.wal_generator.generate_logical_message(prefix, msg); + let (_, record) = self.wal_generator.append_logical_message(prefix, msg); self.disk.write(self.internal_available_lsn.into(), &record); self.prev_lsn = self.internal_available_lsn; self.internal_available_lsn += record.len() as u64; From f54f0e8e2d4da4f2d297d73ffaceee37de4500d0 Mon Sep 17 00:00:00 2001 From: Arseny Sher Date: Thu, 7 Nov 2024 14:29:52 +0300 Subject: [PATCH 31/52] Fix direct reading from WAL buffers. (#9639) Fix direct reading from WAL buffers. Pointer wasn't advanced which resulted in sending corrupted WAL if part of read used WAL buffers and part read from the file. Also move it to neon_walreader so that e.g. replication could also make use of it. ref https://github.com/neondatabase/cloud/issues/19567 --- pgxn/neon/neon_walreader.c | 11 ++++++++ pgxn/neon/walproposer.c | 40 +++++++++++++++------------ pgxn/neon/walproposer_pg.c | 32 ++++----------------- test_runner/fixtures/neon_fixtures.py | 12 ++++++-- 4 files changed, 48 insertions(+), 47 deletions(-) diff --git a/pgxn/neon/neon_walreader.c b/pgxn/neon/neon_walreader.c index b575712dbe..5854a7ef0f 100644 --- a/pgxn/neon/neon_walreader.c +++ b/pgxn/neon/neon_walreader.c @@ -611,6 +611,17 @@ NeonWALReadLocal(NeonWALReader *state, char *buf, XLogRecPtr startptr, Size coun recptr = startptr; nbytes = count; +/* Try to read directly from WAL buffers first. */ +#if PG_MAJORVERSION_NUM >= 17 + { + Size rbytes; + rbytes = WALReadFromBuffers(p, recptr, nbytes, tli); + recptr += rbytes; + nbytes -= rbytes; + p += rbytes; + } +#endif + while (nbytes > 0) { uint32 startoff; diff --git a/pgxn/neon/walproposer.c b/pgxn/neon/walproposer.c index d2a6104c74..e89ffdb628 100644 --- a/pgxn/neon/walproposer.c +++ b/pgxn/neon/walproposer.c @@ -1361,29 +1361,35 @@ SendAppendRequests(Safekeeper *sk) if (sk->active_state == SS_ACTIVE_READ_WAL) { char *errmsg; + int req_len; req = &sk->appendRequest; + req_len = req->endLsn - req->beginLsn; - switch (wp->api.wal_read(sk, - &sk->outbuf.data[sk->outbuf.len], - req->beginLsn, - req->endLsn - req->beginLsn, - &errmsg)) + /* We send zero sized AppenRequests as heartbeats; don't wal_read for these. */ + if (req_len > 0) { - case NEON_WALREAD_SUCCESS: - break; - case NEON_WALREAD_WOULDBLOCK: - return true; - case NEON_WALREAD_ERROR: - wp_log(WARNING, "WAL reading for node %s:%s failed: %s", - sk->host, sk->port, errmsg); - ShutdownConnection(sk); - return false; - default: - Assert(false); + switch (wp->api.wal_read(sk, + &sk->outbuf.data[sk->outbuf.len], + req->beginLsn, + req_len, + &errmsg)) + { + case NEON_WALREAD_SUCCESS: + break; + case NEON_WALREAD_WOULDBLOCK: + return true; + case NEON_WALREAD_ERROR: + wp_log(WARNING, "WAL reading for node %s:%s failed: %s", + sk->host, sk->port, errmsg); + ShutdownConnection(sk); + return false; + default: + Assert(false); + } } - sk->outbuf.len += req->endLsn - req->beginLsn; + sk->outbuf.len += req_len; writeResult = wp->api.conn_async_write(sk, sk->outbuf.data, sk->outbuf.len); diff --git a/pgxn/neon/walproposer_pg.c b/pgxn/neon/walproposer_pg.c index 706941c3f0..86444084ff 100644 --- a/pgxn/neon/walproposer_pg.c +++ b/pgxn/neon/walproposer_pg.c @@ -1489,33 +1489,11 @@ walprop_pg_wal_read(Safekeeper *sk, char *buf, XLogRecPtr startptr, Size count, { NeonWALReadResult res; -#if PG_MAJORVERSION_NUM >= 17 - if (!sk->wp->config->syncSafekeepers) - { - Size rbytes; - rbytes = WALReadFromBuffers(buf, startptr, count, - walprop_pg_get_timeline_id()); - - startptr += rbytes; - count -= rbytes; - } -#endif - - if (count == 0) - { - res = NEON_WALREAD_SUCCESS; - } - else - { - Assert(count > 0); - - /* Now read the remaining WAL from the WAL file */ - res = NeonWALRead(sk->xlogreader, - buf, - startptr, - count, - walprop_pg_get_timeline_id()); - } + res = NeonWALRead(sk->xlogreader, + buf, + startptr, + count, + walprop_pg_get_timeline_id()); if (res == NEON_WALREAD_SUCCESS) { diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py index e4d6e6da5d..e23f46d1ca 100644 --- a/test_runner/fixtures/neon_fixtures.py +++ b/test_runner/fixtures/neon_fixtures.py @@ -4177,9 +4177,15 @@ class Safekeeper(LogUtils): return self def assert_no_errors(self): - assert not self.log_contains("manager task finished prematurely") - assert not self.log_contains("error while acquiring WalResidentTimeline guard") - assert not self.log_contains("timeout while acquiring WalResidentTimeline guard") + not_allowed = [ + "manager task finished prematurely", + "error while acquiring WalResidentTimeline guard", + "timeout while acquiring WalResidentTimeline guard", + "invalid xlog page header:", + "WAL record crc mismatch at", + ] + for na in not_allowed: + assert not self.log_contains(na) def append_logical_message( self, tenant_id: TenantId, timeline_id: TimelineId, request: dict[str, Any] From 01265b7bc66275a8f974fa42e8aeb8ad050404b3 Mon Sep 17 00:00:00 2001 From: Erik Grinaker Date: Thu, 7 Nov 2024 14:24:03 +0100 Subject: [PATCH 32/52] safekeeper: add basic WAL ingestion benchmarks (#9531) ## Problem We don't have any benchmarks for Safekeeper WAL ingestion. ## Summary of changes Add some basic benchmarks for WAL ingestion, specifically for `SafeKeeper::process_msg()` (single append) and `WalAcceptor` (pipelined batch ingestion). Also add some baseline file write benchmarks. --- Cargo.lock | 2 + safekeeper/Cargo.toml | 6 + safekeeper/benches/README.md | 22 ++ safekeeper/benches/benchutils.rs | 102 +++++++++ safekeeper/benches/receive_wal.rs | 341 ++++++++++++++++++++++++++++++ safekeeper/src/lib.rs | 4 +- safekeeper/src/state.rs | 1 - safekeeper/src/timeline.rs | 61 ++++-- 8 files changed, 516 insertions(+), 23 deletions(-) create mode 100644 safekeeper/benches/README.md create mode 100644 safekeeper/benches/benchutils.rs create mode 100644 safekeeper/benches/receive_wal.rs diff --git a/Cargo.lock b/Cargo.lock index 9c2a0b455e..7d18f44aec 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5146,6 +5146,7 @@ dependencies = [ "chrono", "clap", "crc32c", + "criterion", "desim", "fail", "futures", @@ -5153,6 +5154,7 @@ dependencies = [ "http 1.1.0", "humantime", "hyper 0.14.30", + "itertools 0.10.5", "metrics", "once_cell", "parking_lot 0.12.1", diff --git a/safekeeper/Cargo.toml b/safekeeper/Cargo.toml index ec08d02240..85561e4aff 100644 --- a/safekeeper/Cargo.toml +++ b/safekeeper/Cargo.toml @@ -61,8 +61,14 @@ utils.workspace = true workspace_hack.workspace = true [dev-dependencies] +criterion.workspace = true +itertools.workspace = true walproposer.workspace = true rand.workspace = true desim.workspace = true tracing.workspace = true tracing-subscriber = { workspace = true, features = ["json"] } + +[[bench]] +name = "receive_wal" +harness = false diff --git a/safekeeper/benches/README.md b/safekeeper/benches/README.md new file mode 100644 index 0000000000..4119cc8d6e --- /dev/null +++ b/safekeeper/benches/README.md @@ -0,0 +1,22 @@ +## Safekeeper Benchmarks + +To run benchmarks: + +```sh +# All benchmarks. +cargo bench --package safekeeper + +# Specific file. +cargo bench --package safekeeper --bench receive_wal + +# Specific benchmark. +cargo bench --package safekeeper --bench receive_wal process_msg/fsync=false + +# List available benchmarks. +cargo bench --package safekeeper --benches -- --list +``` + +Additional charts and statistics are available in `target/criterion/report/index.html`. + +Benchmarks are automatically compared against the previous run. To compare against other runs, see +`--baseline` and `--save-baseline`. \ No newline at end of file diff --git a/safekeeper/benches/benchutils.rs b/safekeeper/benches/benchutils.rs new file mode 100644 index 0000000000..4e8dc58c49 --- /dev/null +++ b/safekeeper/benches/benchutils.rs @@ -0,0 +1,102 @@ +use std::sync::Arc; + +use camino_tempfile::Utf8TempDir; +use safekeeper::rate_limit::RateLimiter; +use safekeeper::safekeeper::{ProposerAcceptorMessage, ProposerElected, SafeKeeper, TermHistory}; +use safekeeper::state::{TimelinePersistentState, TimelineState}; +use safekeeper::timeline::{get_timeline_dir, SharedState, StateSK, Timeline}; +use safekeeper::timelines_set::TimelinesSet; +use safekeeper::wal_backup::remote_timeline_path; +use safekeeper::{control_file, wal_storage, SafeKeeperConf}; +use tokio::fs::create_dir_all; +use utils::id::{NodeId, TenantTimelineId}; +use utils::lsn::Lsn; + +/// A Safekeeper benchmarking environment. Uses a tempdir for storage, removed on drop. +pub struct Env { + /// Whether to enable fsync. + pub fsync: bool, + /// Benchmark directory. Deleted when dropped. + pub tempdir: Utf8TempDir, +} + +impl Env { + /// Creates a new benchmarking environment in a temporary directory. fsync controls whether to + /// enable fsyncing. + pub fn new(fsync: bool) -> anyhow::Result { + let tempdir = camino_tempfile::tempdir()?; + Ok(Self { fsync, tempdir }) + } + + /// Constructs a Safekeeper config for the given node ID. + fn make_conf(&self, node_id: NodeId) -> SafeKeeperConf { + let mut conf = SafeKeeperConf::dummy(); + conf.my_id = node_id; + conf.no_sync = !self.fsync; + conf.workdir = self.tempdir.path().join(format!("safekeeper-{node_id}")); + conf + } + + /// Constructs a Safekeeper with the given node and tenant/timeline ID. + /// + /// TODO: we should support using in-memory storage, to measure non-IO costs. This would be + /// easier if SafeKeeper used trait objects for storage rather than generics. It's also not + /// currently possible to construct a timeline using non-file storage since StateSK only accepts + /// SafeKeeper. + pub async fn make_safekeeper( + &self, + node_id: NodeId, + ttid: TenantTimelineId, + ) -> anyhow::Result> { + let conf = self.make_conf(node_id); + + let timeline_dir = get_timeline_dir(&conf, &ttid); + create_dir_all(&timeline_dir).await?; + + let mut pstate = TimelinePersistentState::empty(); + pstate.tenant_id = ttid.tenant_id; + pstate.timeline_id = ttid.timeline_id; + + let wal = wal_storage::PhysicalStorage::new(&ttid, &timeline_dir, &pstate, conf.no_sync)?; + let ctrl = + control_file::FileStorage::create_new(&timeline_dir, pstate, conf.no_sync).await?; + let state = TimelineState::new(ctrl); + let mut safekeeper = SafeKeeper::new(state, wal, conf.my_id)?; + + // Emulate an initial election. + safekeeper + .process_msg(&ProposerAcceptorMessage::Elected(ProposerElected { + term: 1, + start_streaming_at: Lsn(0), + term_history: TermHistory(vec![(1, Lsn(0)).into()]), + timeline_start_lsn: Lsn(0), + })) + .await?; + + Ok(safekeeper) + } + + /// Constructs a timeline, including a new Safekeeper with the given node ID, and spawns its + /// manager task. + pub async fn make_timeline( + &self, + node_id: NodeId, + ttid: TenantTimelineId, + ) -> anyhow::Result> { + let conf = self.make_conf(node_id); + let timeline_dir = get_timeline_dir(&conf, &ttid); + let remote_path = remote_timeline_path(&ttid)?; + + let safekeeper = self.make_safekeeper(node_id, ttid).await?; + let shared_state = SharedState::new(StateSK::Loaded(safekeeper)); + + let timeline = Timeline::new(ttid, &timeline_dir, &remote_path, shared_state); + timeline.bootstrap( + &mut timeline.write_shared_state().await, + &conf, + Arc::new(TimelinesSet::default()), // ignored for now + RateLimiter::new(0, 0), + ); + Ok(timeline) + } +} diff --git a/safekeeper/benches/receive_wal.rs b/safekeeper/benches/receive_wal.rs new file mode 100644 index 0000000000..e32d7526ca --- /dev/null +++ b/safekeeper/benches/receive_wal.rs @@ -0,0 +1,341 @@ +//! WAL ingestion benchmarks. + +#[path = "benchutils.rs"] +mod benchutils; + +use std::io::Write as _; + +use benchutils::Env; +use camino_tempfile::tempfile; +use criterion::{criterion_group, criterion_main, BatchSize, Bencher, Criterion}; +use itertools::Itertools as _; +use postgres_ffi::v17::wal_generator::{LogicalMessageGenerator, WalGenerator}; +use safekeeper::receive_wal::{self, WalAcceptor}; +use safekeeper::safekeeper::{ + AcceptorProposerMessage, AppendRequest, AppendRequestHeader, ProposerAcceptorMessage, +}; +use tokio::io::AsyncWriteExt as _; +use utils::id::{NodeId, TenantTimelineId}; +use utils::lsn::Lsn; + +const KB: usize = 1024; +const MB: usize = 1024 * KB; +const GB: usize = 1024 * MB; + +// Register benchmarks with Criterion. +criterion_group!( + benches, + bench_process_msg, + bench_wal_acceptor, + bench_wal_acceptor_throughput, + bench_file_write +); +criterion_main!(benches); + +/// Benchmarks SafeKeeper::process_msg() as time per message and throughput. Each message is an +/// AppendRequest with a single WAL record containing an XlLogicalMessage of varying size. When +/// measuring throughput, only the logical message payload is considered, excluding +/// segment/page/record headers. +fn bench_process_msg(c: &mut Criterion) { + let mut g = c.benchmark_group("process_msg"); + for fsync in [false, true] { + for commit in [false, true] { + for size in [8, KB, 8 * KB, 128 * KB, MB] { + // Kind of weird to change the group throughput per benchmark, but it's the only way + // to vary it per benchmark. It works. + g.throughput(criterion::Throughput::Bytes(size as u64)); + g.bench_function(format!("fsync={fsync}/commit={commit}/size={size}"), |b| { + run_bench(b, size, fsync, commit).unwrap() + }); + } + } + } + + // The actual benchmark. If commit is true, advance the commit LSN on every message. + fn run_bench(b: &mut Bencher, size: usize, fsync: bool, commit: bool) -> anyhow::Result<()> { + let runtime = tokio::runtime::Builder::new_current_thread() // single is fine, sync IO only + .enable_all() + .build()?; + + // Construct the payload. The prefix counts towards the payload (including NUL terminator). + let prefix = c"p"; + let prefixlen = prefix.to_bytes_with_nul().len(); + assert!(size >= prefixlen); + let message = vec![0; size - prefixlen]; + + let walgen = &mut WalGenerator::new(LogicalMessageGenerator::new(prefix, &message)); + + // Set up the Safekeeper. + let env = Env::new(fsync)?; + let mut safekeeper = + runtime.block_on(env.make_safekeeper(NodeId(1), TenantTimelineId::generate()))?; + + b.iter_batched_ref( + // Pre-construct WAL records and requests. Criterion will batch them. + || { + let (lsn, record) = walgen.next().expect("endless WAL"); + ProposerAcceptorMessage::AppendRequest(AppendRequest { + h: AppendRequestHeader { + term: 1, + term_start_lsn: Lsn(0), + begin_lsn: lsn, + end_lsn: lsn + record.len() as u64, + commit_lsn: if commit { lsn } else { Lsn(0) }, // commit previous record + truncate_lsn: Lsn(0), + proposer_uuid: [0; 16], + }, + wal_data: record, + }) + }, + // Benchmark message processing (time per message). + |msg| { + runtime + .block_on(safekeeper.process_msg(msg)) + .expect("message failed") + }, + BatchSize::SmallInput, // automatically determine a batch size + ); + Ok(()) + } +} + +/// Benchmarks WalAcceptor message processing time by sending it a batch of WAL records and waiting +/// for it to confirm that the last LSN has been flushed to storage. We pipeline a bunch of messages +/// instead of measuring each individual message to amortize costs (e.g. fsync), which is more +/// realistic. Records are XlLogicalMessage with a tiny payload (~64 bytes per record including +/// headers). Records are pre-constructed to avoid skewing the benchmark. +/// +/// TODO: add benchmarks with in-memory storage, see comment on `Env::make_safekeeper()`: +fn bench_wal_acceptor(c: &mut Criterion) { + let mut g = c.benchmark_group("wal_acceptor"); + for fsync in [false, true] { + for n in [1, 100, 10000] { + g.bench_function(format!("fsync={fsync}/n={n}"), |b| { + run_bench(b, n, fsync).unwrap() + }); + } + } + + /// The actual benchmark. n is the number of WAL records to send in a pipelined batch. + fn run_bench(b: &mut Bencher, n: usize, fsync: bool) -> anyhow::Result<()> { + let runtime = tokio::runtime::Runtime::new()?; // needs multithreaded + + let env = Env::new(fsync)?; + let walgen = &mut WalGenerator::new(LogicalMessageGenerator::new(c"prefix", b"message")); + + // Create buffered channels that can fit all requests, to avoid blocking on channels. + let (msg_tx, msg_rx) = tokio::sync::mpsc::channel(n); + let (reply_tx, mut reply_rx) = tokio::sync::mpsc::channel(n); + + // Spawn the WalAcceptor task. + runtime.block_on(async { + // TODO: WalAcceptor doesn't actually need a full timeline, only + // Safekeeper::process_msg(). Consider decoupling them to simplify the setup. + let tli = env + .make_timeline(NodeId(1), TenantTimelineId::generate()) + .await? + .wal_residence_guard() + .await?; + WalAcceptor::spawn(tli, msg_rx, reply_tx, Some(0)); + anyhow::Ok(()) + })?; + + b.iter_batched( + // Pre-construct a batch of WAL records and requests. + || { + walgen + .take(n) + .map(|(lsn, record)| AppendRequest { + h: AppendRequestHeader { + term: 1, + term_start_lsn: Lsn(0), + begin_lsn: lsn, + end_lsn: lsn + record.len() as u64, + commit_lsn: Lsn(0), + truncate_lsn: Lsn(0), + proposer_uuid: [0; 16], + }, + wal_data: record, + }) + .collect_vec() + }, + // Benchmark batch ingestion (time per batch). + |reqs| { + runtime.block_on(async { + let final_lsn = reqs.last().unwrap().h.end_lsn; + // Stuff all the messages into the buffered channel to pipeline them. + for req in reqs { + let msg = ProposerAcceptorMessage::AppendRequest(req); + msg_tx.send(msg).await.expect("send failed"); + } + // Wait for the last message to get flushed. + while let Some(reply) = reply_rx.recv().await { + if let AcceptorProposerMessage::AppendResponse(resp) = reply { + if resp.flush_lsn >= final_lsn { + return; + } + } + } + panic!("disconnected") + }) + }, + BatchSize::PerIteration, // only run one request batch at a time + ); + Ok(()) + } +} + +/// Benchmarks WalAcceptor throughput by sending 1 GB of data with varying message sizes and waiting +/// for the last LSN to be flushed to storage. Only the actual message payload counts towards +/// throughput, headers are excluded and considered overhead. Records are XlLogicalMessage. +/// +/// To avoid running out of memory, messages are constructed during the benchmark. +fn bench_wal_acceptor_throughput(c: &mut Criterion) { + const VOLUME: usize = GB; // NB: excludes message/page/segment headers and padding + + let mut g = c.benchmark_group("wal_acceptor_throughput"); + g.sample_size(10); + g.throughput(criterion::Throughput::Bytes(VOLUME as u64)); + + for fsync in [false, true] { + for commit in [false, true] { + for size in [KB, 8 * KB, 128 * KB, MB] { + assert_eq!(VOLUME % size, 0, "volume must be divisible by size"); + let count = VOLUME / size; + g.bench_function(format!("fsync={fsync}/commit={commit}/size={size}"), |b| { + run_bench(b, count, size, fsync, commit).unwrap() + }); + } + } + } + + /// The actual benchmark. size is the payload size per message, count is the number of messages. + /// If commit is true, advance the commit LSN on each message. + fn run_bench( + b: &mut Bencher, + count: usize, + size: usize, + fsync: bool, + commit: bool, + ) -> anyhow::Result<()> { + let runtime = tokio::runtime::Runtime::new()?; // needs multithreaded + + // Construct the payload. The prefix counts towards the payload (including NUL terminator). + let prefix = c"p"; + let prefixlen = prefix.to_bytes_with_nul().len(); + assert!(size >= prefixlen); + let message = vec![0; size - prefixlen]; + + let walgen = &mut WalGenerator::new(LogicalMessageGenerator::new(prefix, &message)); + + // Construct and spawn the WalAcceptor task. + let env = Env::new(fsync)?; + + let (msg_tx, msg_rx) = tokio::sync::mpsc::channel(receive_wal::MSG_QUEUE_SIZE); + let (reply_tx, mut reply_rx) = tokio::sync::mpsc::channel(receive_wal::REPLY_QUEUE_SIZE); + + runtime.block_on(async { + let tli = env + .make_timeline(NodeId(1), TenantTimelineId::generate()) + .await? + .wal_residence_guard() + .await?; + WalAcceptor::spawn(tli, msg_rx, reply_tx, Some(0)); + anyhow::Ok(()) + })?; + + // Ingest the WAL. + b.iter(|| { + runtime.block_on(async { + let reqgen = walgen.take(count).map(|(lsn, record)| AppendRequest { + h: AppendRequestHeader { + term: 1, + term_start_lsn: Lsn(0), + begin_lsn: lsn, + end_lsn: lsn + record.len() as u64, + commit_lsn: if commit { lsn } else { Lsn(0) }, // commit previous record + truncate_lsn: Lsn(0), + proposer_uuid: [0; 16], + }, + wal_data: record, + }); + + // Send requests. + for req in reqgen { + _ = reply_rx.try_recv(); // discard any replies, to avoid blocking + let msg = ProposerAcceptorMessage::AppendRequest(req); + msg_tx.send(msg).await.expect("send failed"); + } + + // Wait for last message to get flushed. + while let Some(reply) = reply_rx.recv().await { + if let AcceptorProposerMessage::AppendResponse(resp) = reply { + if resp.flush_lsn >= walgen.lsn { + return; + } + } + } + panic!("disconnected") + }) + }); + Ok(()) + } +} + +/// Benchmarks OS write throughput by appending blocks of a given size to a file. This is intended +/// to compare Tokio and stdlib writes, and give a baseline for optimal WAL throughput. +fn bench_file_write(c: &mut Criterion) { + let mut g = c.benchmark_group("file_write"); + + for kind in ["stdlib", "tokio"] { + for fsync in [false, true] { + for size in [8, KB, 8 * KB, 128 * KB, MB] { + // Kind of weird to change the group throughput per benchmark, but it's the only way to + // vary it per benchmark. It works. + g.throughput(criterion::Throughput::Bytes(size as u64)); + g.bench_function( + format!("{kind}/fsync={fsync}/size={size}"), + |b| match kind { + "stdlib" => run_bench_stdlib(b, size, fsync).unwrap(), + "tokio" => run_bench_tokio(b, size, fsync).unwrap(), + name => panic!("unknown kind {name}"), + }, + ); + } + } + } + + fn run_bench_stdlib(b: &mut Bencher, size: usize, fsync: bool) -> anyhow::Result<()> { + let mut file = tempfile()?; + let buf = vec![0u8; size]; + + b.iter(|| { + file.write_all(&buf).unwrap(); + file.flush().unwrap(); + if fsync { + file.sync_data().unwrap(); + } + }); + + Ok(()) + } + + fn run_bench_tokio(b: &mut Bencher, size: usize, fsync: bool) -> anyhow::Result<()> { + let runtime = tokio::runtime::Runtime::new()?; // needs multithreaded + + let mut file = tokio::fs::File::from_std(tempfile()?); + let buf = vec![0u8; size]; + + b.iter(|| { + runtime.block_on(async { + file.write_all(&buf).await.unwrap(); + file.flush().await.unwrap(); + if fsync { + file.sync_data().await.unwrap(); + } + }) + }); + + Ok(()) + } +} diff --git a/safekeeper/src/lib.rs b/safekeeper/src/lib.rs index b1cddaf062..6d68b6b59b 100644 --- a/safekeeper/src/lib.rs +++ b/safekeeper/src/lib.rs @@ -112,9 +112,7 @@ impl SafeKeeperConf { } impl SafeKeeperConf { - #[cfg(test)] - #[allow(unused)] - fn dummy() -> Self { + pub fn dummy() -> Self { SafeKeeperConf { workdir: Utf8PathBuf::from("./"), no_sync: false, diff --git a/safekeeper/src/state.rs b/safekeeper/src/state.rs index 0826a148ec..b8925d785e 100644 --- a/safekeeper/src/state.rs +++ b/safekeeper/src/state.rs @@ -138,7 +138,6 @@ impl TimelinePersistentState { }) } - #[cfg(test)] pub fn empty() -> Self { TimelinePersistentState::new( &TenantTimelineId::empty(), diff --git a/safekeeper/src/timeline.rs b/safekeeper/src/timeline.rs index f0113978c4..fa91241177 100644 --- a/safekeeper/src/timeline.rs +++ b/safekeeper/src/timeline.rs @@ -2,7 +2,7 @@ //! to glue together SafeKeeper and all other background services. use anyhow::{anyhow, bail, Result}; -use camino::Utf8PathBuf; +use camino::{Utf8Path, Utf8PathBuf}; use remote_storage::RemotePath; use safekeeper_api::models::TimelineTermBumpResponse; use serde::{Deserialize, Serialize}; @@ -325,8 +325,17 @@ pub struct SharedState { } impl SharedState { + /// Creates a new SharedState. + pub fn new(sk: StateSK) -> Self { + Self { + sk, + peers_info: PeersInfo(vec![]), + wal_removal_on_hold: false, + } + } + /// Restore SharedState from control file. If file doesn't exist, bails out. - fn restore(conf: &SafeKeeperConf, ttid: &TenantTimelineId) -> Result { + pub fn restore(conf: &SafeKeeperConf, ttid: &TenantTimelineId) -> Result { let timeline_dir = get_timeline_dir(conf, ttid); let control_store = control_file::FileStorage::restore_new(&timeline_dir, conf.no_sync)?; if control_store.server.wal_seg_size == 0 { @@ -352,11 +361,7 @@ impl SharedState { } }; - Ok(Self { - sk, - peers_info: PeersInfo(vec![]), - wal_removal_on_hold: false, - }) + Ok(Self::new(sk)) } pub(crate) fn get_wal_seg_size(&self) -> usize { @@ -480,11 +485,13 @@ pub struct Timeline { } impl Timeline { - /// Load existing timeline from disk. - pub fn load_timeline(conf: &SafeKeeperConf, ttid: TenantTimelineId) -> Result> { - let _enter = info_span!("load_timeline", timeline = %ttid.timeline_id).entered(); - - let shared_state = SharedState::restore(conf, &ttid)?; + /// Constructs a new timeline. + pub fn new( + ttid: TenantTimelineId, + timeline_dir: &Utf8Path, + remote_path: &RemotePath, + shared_state: SharedState, + ) -> Arc { let (commit_lsn_watch_tx, commit_lsn_watch_rx) = watch::channel(shared_state.sk.state().commit_lsn); let (term_flush_lsn_watch_tx, term_flush_lsn_watch_rx) = watch::channel(TermLsn::from(( @@ -494,10 +501,11 @@ impl Timeline { let (shared_state_version_tx, shared_state_version_rx) = watch::channel(0); let walreceivers = WalReceivers::new(); - let remote_path = remote_timeline_path(&ttid)?; - Ok(Arc::new(Timeline { + + Arc::new(Self { ttid, - remote_path, + remote_path: remote_path.to_owned(), + timeline_dir: timeline_dir.to_owned(), commit_lsn_watch_tx, commit_lsn_watch_rx, term_flush_lsn_watch_tx, @@ -508,13 +516,28 @@ impl Timeline { walsenders: WalSenders::new(walreceivers.clone()), walreceivers, cancel: CancellationToken::default(), - timeline_dir: get_timeline_dir(conf, &ttid), manager_ctl: ManagerCtl::new(), broker_active: AtomicBool::new(false), wal_backup_active: AtomicBool::new(false), last_removed_segno: AtomicU64::new(0), mgr_status: AtomicStatus::new(), - })) + }) + } + + /// Load existing timeline from disk. + pub fn load_timeline(conf: &SafeKeeperConf, ttid: TenantTimelineId) -> Result> { + let _enter = info_span!("load_timeline", timeline = %ttid.timeline_id).entered(); + + let shared_state = SharedState::restore(conf, &ttid)?; + let timeline_dir = get_timeline_dir(conf, &ttid); + let remote_path = remote_timeline_path(&ttid)?; + + Ok(Timeline::new( + ttid, + &timeline_dir, + &remote_path, + shared_state, + )) } /// Initialize fresh timeline on disk and start background tasks. If init @@ -1128,13 +1151,13 @@ async fn delete_dir(path: &Utf8PathBuf) -> Result { /// Get a path to the tenant directory. If you just need to get a timeline directory, /// use WalResidentTimeline::get_timeline_dir instead. -pub(crate) fn get_tenant_dir(conf: &SafeKeeperConf, tenant_id: &TenantId) -> Utf8PathBuf { +pub fn get_tenant_dir(conf: &SafeKeeperConf, tenant_id: &TenantId) -> Utf8PathBuf { conf.workdir.join(tenant_id.to_string()) } /// Get a path to the timeline directory. If you need to read WAL files from disk, /// use WalResidentTimeline::get_timeline_dir instead. This function does not check /// timeline eviction status and WAL files might not be present on disk. -pub(crate) fn get_timeline_dir(conf: &SafeKeeperConf, ttid: &TenantTimelineId) -> Utf8PathBuf { +pub fn get_timeline_dir(conf: &SafeKeeperConf, ttid: &TenantTimelineId) -> Utf8PathBuf { get_tenant_dir(conf, &ttid.tenant_id).join(ttid.timeline_id.to_string()) } From f18aa04b902204f2356b7ff67dc04bb4d3176327 Mon Sep 17 00:00:00 2001 From: Erik Grinaker Date: Thu, 7 Nov 2024 16:09:57 +0100 Subject: [PATCH 33/52] safekeeper: use `set_len()` to zero out segments (#9665) ## Problem When we create a new segment, we zero it out in order to avoid changing the length and fsyncing metadata on every write. However, we zeroed it out by writing 8 KB zero-pages, and Tokio file writes have non-trivial overhead. ## Summary of changes Zero out the segment using [`File::set_len()`](https://docs.rs/tokio/latest/i686-unknown-linux-gnu/tokio/fs/struct.File.html#method.set_len) instead. This will typically (depending on the filesystem) just write a sparse file and omit the 16 MB of data entirely. This improves WAL append throughput for large messages by over 400% with fsync disabled, and 100% with fsync enabled. --- safekeeper/src/wal_storage.rs | 45 ++++++++----------- .../regress/test_wal_acceptor_async.py | 2 +- 2 files changed, 20 insertions(+), 27 deletions(-) diff --git a/safekeeper/src/wal_storage.rs b/safekeeper/src/wal_storage.rs index 33b8bfe28e..4e67940c51 100644 --- a/safekeeper/src/wal_storage.rs +++ b/safekeeper/src/wal_storage.rs @@ -31,7 +31,6 @@ use crate::state::TimelinePersistentState; use crate::wal_backup::{read_object, remote_timeline_path}; use postgres_ffi::waldecoder::WalStreamDecoder; use postgres_ffi::XLogFileName; -use postgres_ffi::XLOG_BLCKSZ; use pq_proto::SystemId; use utils::{id::TenantTimelineId, lsn::Lsn}; @@ -223,6 +222,15 @@ impl PhysicalStorage { ) } + /// Call fsync if config requires so. + async fn fsync_file(&mut self, file: &File) -> Result<()> { + if !self.no_sync { + self.metrics + .observe_flush_seconds(time_io_closure(file.sync_all()).await?); + } + Ok(()) + } + /// Call fdatasync if config requires so. async fn fdatasync_file(&mut self, file: &File) -> Result<()> { if !self.no_sync { @@ -256,11 +264,15 @@ impl PhysicalStorage { // half initialized segment, first bake it under tmp filename and // then rename. let tmp_path = self.timeline_dir.join("waltmp"); - let mut file = File::create(&tmp_path) + let file = File::create(&tmp_path) .await .with_context(|| format!("Failed to open tmp wal file {:?}", &tmp_path))?; - write_zeroes(&mut file, self.wal_seg_size).await?; + fail::fail_point!("sk-zero-segment", |_| { + info!("sk-zero-segment failpoint hit"); + Err(anyhow::anyhow!("failpoint: sk-zero-segment")) + }); + file.set_len(self.wal_seg_size as u64).await?; // Note: this doesn't get into observe_flush_seconds metric. But // segment init should be separate metric, if any. @@ -486,12 +498,12 @@ impl Storage for PhysicalStorage { // Remove all segments after the given LSN. remove_segments_from_disk(&self.timeline_dir, self.wal_seg_size, |x| x > segno).await?; - let (mut file, is_partial) = self.open_or_create(segno).await?; + let (file, is_partial) = self.open_or_create(segno).await?; // Fill end with zeroes - file.seek(SeekFrom::Start(xlogoff as u64)).await?; - write_zeroes(&mut file, self.wal_seg_size - xlogoff).await?; - self.fdatasync_file(&file).await?; + file.set_len(xlogoff as u64).await?; + file.set_len(self.wal_seg_size as u64).await?; + self.fsync_file(&file).await?; if !is_partial { // Make segment partial once again @@ -751,25 +763,6 @@ impl WalReader { } } -/// Zero block for filling created WAL segments. -const ZERO_BLOCK: &[u8] = &[0u8; XLOG_BLCKSZ]; - -/// Helper for filling file with zeroes. -async fn write_zeroes(file: &mut File, mut count: usize) -> Result<()> { - fail::fail_point!("sk-write-zeroes", |_| { - info!("write_zeroes hit failpoint"); - Err(anyhow::anyhow!("failpoint: sk-write-zeroes")) - }); - - while count >= XLOG_BLCKSZ { - file.write_all(ZERO_BLOCK).await?; - count -= XLOG_BLCKSZ; - } - file.write_all(&ZERO_BLOCK[0..count]).await?; - file.flush().await?; - Ok(()) -} - /// Helper function for opening WAL segment `segno` in `dir`. Returns file and /// whether it is .partial. pub(crate) async fn open_wal_file( diff --git a/test_runner/regress/test_wal_acceptor_async.py b/test_runner/regress/test_wal_acceptor_async.py index 92306469f8..f328974264 100644 --- a/test_runner/regress/test_wal_acceptor_async.py +++ b/test_runner/regress/test_wal_acceptor_async.py @@ -602,7 +602,7 @@ async def run_segment_init_failure(env: NeonEnv): sk = env.safekeepers[0] sk_http = sk.http_client() - sk_http.configure_failpoints([("sk-write-zeroes", "return")]) + sk_http.configure_failpoints([("sk-zero-segment", "return")]) conn = await ep.connect_async() ep.safe_psql("select pg_switch_wal()") # jump to the segment boundary # next insertion should hang until failpoint is disabled. From a8d9939ea959ad8c00b376fb8e189f630d3013ab Mon Sep 17 00:00:00 2001 From: "Alex Chi Z." <4198311+skyzh@users.noreply.github.com> Date: Thu, 7 Nov 2024 10:38:15 -0500 Subject: [PATCH 34/52] fix(pageserver): reduce aux compaction threshold (#9647) ref https://github.com/neondatabase/neon/issues/9441 The metrics from LR publisher testing project: ~300KB aux key deltas per 256MB files. Therefore, I think we can do compaction more aggressively as these deltas are small and compaction can reduce layer download latency. We also have a read path perf fix https://github.com/neondatabase/neon/pull/9631 but I'd still combine the read path fix with the reduce of the compaction threshold. ## Summary of changes * reduce metadata compaction threshold * use num of L1 delta layers as an indicator for metadata compaction * dump more logs Signed-off-by: Alex Chi Z --- pageserver/src/pgdatadir_mapping.rs | 2 +- pageserver/src/tenant/timeline.rs | 10 ++++------ 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/pageserver/src/pgdatadir_mapping.rs b/pageserver/src/pgdatadir_mapping.rs index 7b106569a4..7c1abbf3e2 100644 --- a/pageserver/src/pgdatadir_mapping.rs +++ b/pageserver/src/pgdatadir_mapping.rs @@ -45,7 +45,7 @@ use wal_decoder::serialized_batch::SerializedValueBatch; pub const MAX_AUX_FILE_DELTAS: usize = 1024; /// Max number of aux-file-related delta layers. The compaction will create a new image layer once this threshold is reached. -pub const MAX_AUX_FILE_V2_DELTAS: usize = 64; +pub const MAX_AUX_FILE_V2_DELTAS: usize = 16; #[derive(Debug)] pub enum LsnForTimestamp { diff --git a/pageserver/src/tenant/timeline.rs b/pageserver/src/tenant/timeline.rs index ee823beca8..6e082aecf5 100644 --- a/pageserver/src/tenant/timeline.rs +++ b/pageserver/src/tenant/timeline.rs @@ -4090,6 +4090,7 @@ impl Timeline { ) -> Result { // Metadata keys image layer creation. let mut reconstruct_state = ValuesReconstructState::default(); + let begin = Instant::now(); let data = self .get_vectored_impl(partition.clone(), lsn, &mut reconstruct_state, ctx) .await?; @@ -4106,14 +4107,11 @@ impl Timeline { (new_data, total_kb_retrieved / 1024, total_keys_retrieved) }; let delta_files_accessed = reconstruct_state.get_delta_layers_visited(); + let elapsed = begin.elapsed(); let trigger_generation = delta_files_accessed as usize >= MAX_AUX_FILE_V2_DELTAS; - debug!( - trigger_generation, - delta_files_accessed, - total_kb_retrieved, - total_keys_retrieved, - "generate metadata images" + info!( + "metadata key compaction: trigger_generation={trigger_generation}, delta_files_accessed={delta_files_accessed}, total_kb_retrieved={total_kb_retrieved}, total_keys_retrieved={total_keys_retrieved}, read_time={}s", elapsed.as_secs_f64() ); if !trigger_generation && mode == ImageLayerCreationMode::Try { From 75aa19aa2dcc3f8907acd0eeae9e7c866f226f23 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arpad=20M=C3=BCller?= Date: Thu, 7 Nov 2024 17:13:50 +0100 Subject: [PATCH 35/52] Don't attach is_archived to debug output (#9679) We are in branches where we know its value already. --- pageserver/src/tenant/timeline/offload.rs | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/pageserver/src/tenant/timeline/offload.rs b/pageserver/src/tenant/timeline/offload.rs index cccf24e303..2dc461c28d 100644 --- a/pageserver/src/tenant/timeline/offload.rs +++ b/pageserver/src/tenant/timeline/offload.rs @@ -47,15 +47,12 @@ pub(crate) async fn offload_timeline( match is_archived { Some(true) => (), Some(false) => { - tracing::warn!(?is_archived, "tried offloading a non-archived timeline"); + tracing::warn!("tried offloading a non-archived timeline"); return Err(OffloadError::NotArchived); } None => { // This is legal: calls to this function can race with the timeline shutting down - tracing::info!( - ?is_archived, - "tried offloading a timeline whose remote storage is not initialized" - ); + tracing::info!("tried offloading a timeline whose remote storage is not initialized"); return Err(OffloadError::Cancelled); } } From 82e3f0ecba8542cd5d1a95bf6d938aacbc073905 Mon Sep 17 00:00:00 2001 From: Conrad Ludgate Date: Thu, 7 Nov 2024 16:24:38 +0000 Subject: [PATCH 36/52] [proxy/authorize]: improve JWKS reliability (#9676) While setting up some tests, I noticed that we didn't support keycloak. They make use of encryption JWKs as well as signature ones. Our current jwks crate does not support parsing encryption keys which caused the entire jwk set to fail to parse. Switching to lazy parsing fixes this. Also while setting up tests, I couldn't use localhost jwks server as we require HTTPS and we were using webpki so it was impossible to add a custom CA. Enabling native roots addresses this possibility. I saw some of our current e2e tests against our custom JWKS in s3 were taking a while to fetch. I've added a timeout + retries to address this. --- Cargo.lock | 1 + proxy/Cargo.toml | 2 +- proxy/src/auth/backend/jwt.rs | 162 +++++++++++++++++++++++-- proxy/src/http/mod.rs | 22 ++-- proxy/src/serverless/conn_pool_lib.rs | 3 +- proxy/src/serverless/http_conn_pool.rs | 1 - workspace_hack/Cargo.toml | 2 +- 7 files changed, 168 insertions(+), 25 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7d18f44aec..00d58be2d5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4743,6 +4743,7 @@ dependencies = [ "percent-encoding", "pin-project-lite", "rustls 0.22.4", + "rustls-native-certs 0.7.0", "rustls-pemfile 2.1.1", "rustls-pki-types", "serde", diff --git a/proxy/Cargo.toml b/proxy/Cargo.toml index efd336dbea..1665d6361a 100644 --- a/proxy/Cargo.toml +++ b/proxy/Cargo.toml @@ -60,7 +60,7 @@ prometheus.workspace = true rand.workspace = true regex.workspace = true remote_storage = { version = "0.1", path = "../libs/remote_storage/" } -reqwest.workspace = true +reqwest = { workspace = true, features = ["rustls-tls-native-roots"] } reqwest-middleware = { workspace = true, features = ["json"] } reqwest-retry.workspace = true reqwest-tracing.workspace = true diff --git a/proxy/src/auth/backend/jwt.rs b/proxy/src/auth/backend/jwt.rs index 83c3617612..bfc674139b 100644 --- a/proxy/src/auth/backend/jwt.rs +++ b/proxy/src/auth/backend/jwt.rs @@ -7,8 +7,11 @@ use arc_swap::ArcSwapOption; use dashmap::DashMap; use jose_jwk::crypto::KeyInfo; use reqwest::{redirect, Client}; +use reqwest_retry::policies::ExponentialBackoff; +use reqwest_retry::RetryTransientMiddleware; use serde::de::Visitor; use serde::{Deserialize, Deserializer}; +use serde_json::value::RawValue; use signature::Verifier; use thiserror::Error; use tokio::time::Instant; @@ -16,7 +19,7 @@ use tokio::time::Instant; use crate::auth::backend::ComputeCredentialKeys; use crate::context::RequestMonitoring; use crate::control_plane::errors::GetEndpointJwksError; -use crate::http::parse_json_body_with_limit; +use crate::http::read_body_with_limit; use crate::intern::RoleNameInt; use crate::types::{EndpointId, RoleName}; @@ -28,6 +31,10 @@ const MAX_RENEW: Duration = Duration::from_secs(3600); const MAX_JWK_BODY_SIZE: usize = 64 * 1024; const JWKS_USER_AGENT: &str = "neon-proxy"; +const JWKS_CONNECT_TIMEOUT: Duration = Duration::from_secs(2); +const JWKS_FETCH_TIMEOUT: Duration = Duration::from_secs(5); +const JWKS_FETCH_RETRIES: u32 = 3; + /// How to get the JWT auth rules pub(crate) trait FetchAuthRules: Clone + Send + Sync + 'static { fn fetch_auth_rules( @@ -55,7 +62,7 @@ pub(crate) struct AuthRule { } pub struct JwkCache { - client: reqwest::Client, + client: reqwest_middleware::ClientWithMiddleware, map: DashMap<(EndpointId, RoleName), Arc>, } @@ -117,6 +124,14 @@ impl Default for JwkCacheEntryLock { } } +#[derive(Deserialize)] +struct JwkSet<'a> { + /// we parse into raw-value because not all keys in a JWKS are ones + /// we can parse directly, so we parse them lazily. + #[serde(borrow)] + keys: Vec<&'a RawValue>, +} + impl JwkCacheEntryLock { async fn acquire_permit<'a>(self: &'a Arc) -> JwkRenewalPermit<'a> { JwkRenewalPermit::acquire_permit(self).await @@ -130,7 +145,7 @@ impl JwkCacheEntryLock { &self, _permit: JwkRenewalPermit<'_>, ctx: &RequestMonitoring, - client: &reqwest::Client, + client: &reqwest_middleware::ClientWithMiddleware, endpoint: EndpointId, auth_rules: &F, ) -> Result, JwtError> { @@ -154,22 +169,73 @@ impl JwkCacheEntryLock { let req = client.get(rule.jwks_url.clone()); // TODO(conrad): eventually switch to using reqwest_middleware/`new_client_with_timeout`. // TODO(conrad): We need to filter out URLs that point to local resources. Public internet only. - match req.send().await.and_then(|r| r.error_for_status()) { + match req.send().await.and_then(|r| { + r.error_for_status() + .map_err(reqwest_middleware::Error::Reqwest) + }) { // todo: should we re-insert JWKs if we want to keep this JWKs URL? // I expect these failures would be quite sparse. Err(e) => tracing::warn!(url=?rule.jwks_url, error=?e, "could not fetch JWKs"), Ok(r) => { let resp: http::Response = r.into(); - match parse_json_body_with_limit::( - resp.into_body(), - MAX_JWK_BODY_SIZE, - ) - .await + + let bytes = match read_body_with_limit(resp.into_body(), MAX_JWK_BODY_SIZE) + .await { + Ok(bytes) => bytes, + Err(e) => { + tracing::warn!(url=?rule.jwks_url, error=?e, "could not decode JWKs"); + continue; + } + }; + + match serde_json::from_slice::(&bytes) { Err(e) => { tracing::warn!(url=?rule.jwks_url, error=?e, "could not decode JWKs"); } Ok(jwks) => { + // size_of::<&RawValue>() == 16 + // size_of::() == 288 + // better to not pre-allocate this as it might be pretty large - especially if it has many + // keys we don't want or need. + // trivial 'attack': `{"keys":[` + repeat(`0`).take(30000).join(`,`) + `]}` + // this would consume 8MiB just like that! + let mut keys = vec![]; + let mut failed = 0; + for key in jwks.keys { + match serde_json::from_str::(key.get()) { + Ok(key) => { + // if `use` (called `cls` in rust) is specified to be something other than signing, + // we can skip storing it. + if key + .prm + .cls + .as_ref() + .is_some_and(|c| *c != jose_jwk::Class::Signing) + { + continue; + } + + keys.push(key); + } + Err(e) => { + tracing::debug!(url=?rule.jwks_url, failed=?e, "could not decode JWK"); + failed += 1; + } + } + } + keys.shrink_to_fit(); + + if failed > 0 { + tracing::warn!(url=?rule.jwks_url, failed, "could not decode JWKs"); + } + + if keys.is_empty() { + tracing::warn!(url=?rule.jwks_url, "no valid JWKs found inside the response body"); + continue; + } + + let jwks = jose_jwk::JwkSet { keys }; key_sets.insert( rule.id, KeySet { @@ -179,7 +245,7 @@ impl JwkCacheEntryLock { }, ); } - } + }; } } } @@ -196,7 +262,7 @@ impl JwkCacheEntryLock { async fn get_or_update_jwk_cache( self: &Arc, ctx: &RequestMonitoring, - client: &reqwest::Client, + client: &reqwest_middleware::ClientWithMiddleware, endpoint: EndpointId, fetch: &F, ) -> Result, JwtError> { @@ -250,7 +316,7 @@ impl JwkCacheEntryLock { self: &Arc, ctx: &RequestMonitoring, jwt: &str, - client: &reqwest::Client, + client: &reqwest_middleware::ClientWithMiddleware, endpoint: EndpointId, role_name: &RoleName, fetch: &F, @@ -369,8 +435,19 @@ impl Default for JwkCache { let client = Client::builder() .user_agent(JWKS_USER_AGENT) .redirect(redirect::Policy::none()) + .tls_built_in_native_certs(true) + .connect_timeout(JWKS_CONNECT_TIMEOUT) + .timeout(JWKS_FETCH_TIMEOUT) .build() - .expect("using &str and standard redirect::Policy"); + .expect("client config should be valid"); + + // Retry up to 3 times with increasing intervals between attempts. + let retry_policy = ExponentialBackoff::builder().build_with_max_retries(JWKS_FETCH_RETRIES); + + let client = reqwest_middleware::ClientBuilder::new(client) + .with(RetryTransientMiddleware::new_with_policy(retry_policy)) + .build(); + JwkCache { client, map: DashMap::default(), @@ -1209,4 +1286,63 @@ X0n5X2/pBLJzxZc62ccvZYVnctBiFs6HbSnxpuMQCfkt/BcR/ttIepBQQIW86wHL } } } + + #[tokio::test] + async fn check_jwk_keycloak_regression() { + let (rs, valid_jwk) = new_rsa_jwk(RS1, "rs1".into()); + let valid_jwk = serde_json::to_value(valid_jwk).unwrap(); + + // This is valid, but we cannot parse it as we have no support for encryption JWKs, only signature based ones. + // This is taken directly from keycloak. + let invalid_jwk = serde_json::json! { + { + "kid": "U-Jc9xRli84eNqRpYQoIPF-GNuRWV3ZvAIhziRW2sbQ", + "kty": "RSA", + "alg": "RSA-OAEP", + "use": "enc", + "n": "yypYWsEKmM_wWdcPnSGLSm5ytw1WG7P7EVkKSulcDRlrM6HWj3PR68YS8LySYM2D9Z-79oAdZGKhIfzutqL8rK1vS14zDuPpAM-RWY3JuQfm1O_-1DZM8-07PmVRegP5KPxsKblLf_My8ByH6sUOIa1p2rbe2q_b0dSTXYu1t0dW-cGL5VShc400YymvTwpc-5uYNsaVxZajnB7JP1OunOiuCJ48AuVp3PqsLzgoXqlXEB1ZZdch3xT3bxaTtNruGvG4xmLZY68O_T3yrwTCNH2h_jFdGPyXdyZToCMSMK2qSbytlfwfN55pT9Vv42Lz1YmoB7XRjI9aExKPc5AxFw", + "e": "AQAB", + "x5c": [ + "MIICmzCCAYMCBgGS41E6azANBgkqhkiG9w0BAQsFADARMQ8wDQYDVQQDDAZtYXN0ZXIwHhcNMjQxMDMxMTYwMTQ0WhcNMzQxMDMxMTYwMzI0WjARMQ8wDQYDVQQDDAZtYXN0ZXIwggEiMA0GCSqGSIb3DQEBAQUAA4IBDwAwggEKAoIBAQDLKlhawQqYz/BZ1w+dIYtKbnK3DVYbs/sRWQpK6VwNGWszodaPc9HrxhLwvJJgzYP1n7v2gB1kYqEh/O62ovysrW9LXjMO4+kAz5FZjcm5B+bU7/7UNkzz7Ts+ZVF6A/ko/GwpuUt/8zLwHIfqxQ4hrWnatt7ar9vR1JNdi7W3R1b5wYvlVKFzjTRjKa9PClz7m5g2xpXFlqOcHsk/U66c6K4InjwC5Wnc+qwvOCheqVcQHVll1yHfFPdvFpO02u4a8bjGYtljrw79PfKvBMI0faH+MV0Y/Jd3JlOgIxIwrapJvK2V/B83nmlP1W/jYvPViagHtdGMj1oTEo9zkDEXAgMBAAEwDQYJKoZIhvcNAQELBQADggEBAECYX59+Q9v6c9sb6Q0/C6IgLWG2nVCgVE1YWwIzz+68WrhlmNCRuPjY94roB+tc2tdHbj+Nh3LMzJk7L1KCQoW1+LPK6A6E8W9ad0YPcuw8csV2pUA3+H56exQMH0fUAPQAU7tXWvnQ7otcpV1XA8afn/NTMTsnxi9mSkor8MLMYQ3aeRyh1+LAchHBthWiltqsSUqXrbJF59u5p0ghquuKcWR3TXsA7klGYBgGU5KAJifr9XT87rN0bOkGvbeWAgKvnQnjZwxdnLqTfp/pRY/PiJJHhgIBYPIA7STGnMPjmJ995i34zhnbnd8WHXJA3LxrIMqLW/l8eIdvtM1w8KI=" + ], + "x5t": "QhfzMMnuAfkReTgZ1HtrfyOeeZs", + "x5t#S256": "cmHDUdKgLiRCEN28D5FBy9IJLFmR7QWfm77SLhGTCTU" + } + }; + + let jwks = serde_json::json! {{ "keys": [invalid_jwk, valid_jwk ] }}; + let jwks_addr = jwks_server(move |path| match path { + "/" => Some(serde_json::to_vec(&jwks).unwrap()), + _ => None, + }) + .await; + + let role_name = RoleName::from("anonymous"); + let role = RoleNameInt::from(&role_name); + + let rules = vec![AuthRule { + id: "foo".to_owned(), + jwks_url: format!("http://{jwks_addr}/").parse().unwrap(), + audience: None, + role_names: vec![role], + }]; + + let fetch = Fetch(rules); + let jwk_cache = JwkCache::default(); + + let endpoint = EndpointId::from("ep"); + + let token = new_rsa_jwt("rs1".into(), rs); + + jwk_cache + .check_jwt( + &RequestMonitoring::test(), + endpoint.clone(), + &role_name, + &fetch, + &token, + ) + .await + .unwrap(); + } } diff --git a/proxy/src/http/mod.rs b/proxy/src/http/mod.rs index f1b632e704..b1642cedb3 100644 --- a/proxy/src/http/mod.rs +++ b/proxy/src/http/mod.rs @@ -6,7 +6,6 @@ pub mod health_server; use std::time::Duration; -use anyhow::bail; use bytes::Bytes; use http::Method; use http_body_util::BodyExt; @@ -16,7 +15,7 @@ use reqwest_middleware::RequestBuilder; pub(crate) use reqwest_middleware::{ClientWithMiddleware, Error}; pub(crate) use reqwest_retry::policies::ExponentialBackoff; pub(crate) use reqwest_retry::RetryTransientMiddleware; -use serde::de::DeserializeOwned; +use thiserror::Error; use crate::metrics::{ConsoleRequest, Metrics}; use crate::url::ApiUrl; @@ -122,10 +121,19 @@ impl Endpoint { } } -pub(crate) async fn parse_json_body_with_limit( +#[derive(Error, Debug)] +pub(crate) enum ReadBodyError { + #[error("Content length exceeds limit of {limit} bytes")] + BodyTooLarge { limit: usize }, + + #[error(transparent)] + Read(#[from] reqwest::Error), +} + +pub(crate) async fn read_body_with_limit( mut b: impl Body + Unpin, limit: usize, -) -> anyhow::Result { +) -> Result, ReadBodyError> { // We could use `b.limited().collect().await.to_bytes()` here // but this ends up being slightly more efficient as far as I can tell. @@ -133,20 +141,20 @@ pub(crate) async fn parse_json_body_with_limit( // in reqwest, this value is influenced by the Content-Length header. let lower_bound = match usize::try_from(b.size_hint().lower()) { Ok(bound) if bound <= limit => bound, - _ => bail!("Content length exceeds limit of {limit} bytes"), + _ => return Err(ReadBodyError::BodyTooLarge { limit }), }; let mut bytes = Vec::with_capacity(lower_bound); while let Some(frame) = b.frame().await.transpose()? { if let Ok(data) = frame.into_data() { if bytes.len() + data.len() > limit { - bail!("Content length exceeds limit of {limit} bytes") + return Err(ReadBodyError::BodyTooLarge { limit }); } bytes.extend_from_slice(&data); } } - Ok(serde_json::from_slice::(&bytes)?) + Ok(bytes) } #[cfg(test)] diff --git a/proxy/src/serverless/conn_pool_lib.rs b/proxy/src/serverless/conn_pool_lib.rs index 00a8ac4768..61c39c32c9 100644 --- a/proxy/src/serverless/conn_pool_lib.rs +++ b/proxy/src/serverless/conn_pool_lib.rs @@ -16,8 +16,7 @@ use super::http_conn_pool::ClientDataHttp; use super::local_conn_pool::ClientDataLocal; use crate::auth::backend::ComputeUserInfo; use crate::context::RequestMonitoring; -use crate::control_plane::messages::ColdStartInfo; -use crate::control_plane::messages::MetricsAuxInfo; +use crate::control_plane::messages::{ColdStartInfo, MetricsAuxInfo}; use crate::metrics::{HttpEndpointPoolsGuard, Metrics}; use crate::types::{DbName, EndpointCacheKey, RoleName}; use crate::usage_metrics::{Ids, MetricCounter, USAGE_METRICS}; diff --git a/proxy/src/serverless/http_conn_pool.rs b/proxy/src/serverless/http_conn_pool.rs index 56be70abec..a1d4473b01 100644 --- a/proxy/src/serverless/http_conn_pool.rs +++ b/proxy/src/serverless/http_conn_pool.rs @@ -7,7 +7,6 @@ use hyper::client::conn::http2; use hyper_util::rt::{TokioExecutor, TokioIo}; use parking_lot::RwLock; use rand::Rng; -use std::result::Result::Ok; use tokio::net::TcpStream; use tracing::{debug, error, info, info_span, Instrument}; diff --git a/workspace_hack/Cargo.toml b/workspace_hack/Cargo.toml index 02deecd385..ae4018a884 100644 --- a/workspace_hack/Cargo.toml +++ b/workspace_hack/Cargo.toml @@ -64,7 +64,7 @@ rand = { version = "0.8", features = ["small_rng"] } regex = { version = "1" } regex-automata = { version = "0.4", default-features = false, features = ["dfa-onepass", "hybrid", "meta", "nfa-backtrack", "perf-inline", "perf-literal", "unicode"] } regex-syntax = { version = "0.8" } -reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "stream"] } +reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "rustls-tls-native-roots", "stream"] } rustls = { version = "0.23", default-features = false, features = ["logging", "ring", "std", "tls12"] } scopeguard = { version = "1" } serde = { version = "1", features = ["alloc", "derive"] } From 9132d80aa3e20c58878b945b86c442bf7756a30e Mon Sep 17 00:00:00 2001 From: Peter Bendel Date: Thu, 7 Nov 2024 19:00:25 +0100 Subject: [PATCH 37/52] add pgcopydb tool to build tools image (#9658) ## Problem build-tools image does not provide superuser, so additional packages can not be installed during GitHub benchmarking workflows but need to be added to the image ## Summary of changes install pgcopydb version 0.17-1 or higher into build-tools bookworm image ```bash docker run -it neondatabase/build-tools:-bookworm-arm64 /bin/bash ... nonroot@c23c6f4901ce:~$ LD_LIBRARY_PATH=/pgcopydb/lib /pgcopydb/bin/pgcopydb --version; 13:58:19.768 8 INFO Running pgcopydb version 0.17 from "/pgcopydb/bin/pgcopydb" pgcopydb version 0.17 compiled with PostgreSQL 16.4 (Debian 16.4-1.pgdg120+2) on aarch64-unknown-linux-gnu, compiled by gcc (Debian 12.2.0-14) 12.2.0, 64-bit compatible with Postgres 11, 12, 13, 14, 15, and 16 ``` Example usage of that image in a workflow https://github.com/neondatabase/neon/actions/runs/11725718371/job/32662681172#step:7:14 --- build-tools.Dockerfile | 66 ++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 63 insertions(+), 3 deletions(-) diff --git a/build-tools.Dockerfile b/build-tools.Dockerfile index 93f1e48afa..c1190b13f4 100644 --- a/build-tools.Dockerfile +++ b/build-tools.Dockerfile @@ -1,12 +1,66 @@ ARG DEBIAN_VERSION=bullseye -FROM debian:${DEBIAN_VERSION}-slim +FROM debian:bookworm-slim AS pgcopydb_builder +ARG DEBIAN_VERSION + +RUN if [ "${DEBIAN_VERSION}" = "bookworm" ]; then \ + set -e && \ + apt update && \ + apt install -y --no-install-recommends \ + ca-certificates wget gpg && \ + wget -qO - https://www.postgresql.org/media/keys/ACCC4CF8.asc | gpg --dearmor -o /usr/share/keyrings/postgresql-keyring.gpg && \ + echo "deb [signed-by=/usr/share/keyrings/postgresql-keyring.gpg] http://apt.postgresql.org/pub/repos/apt bookworm-pgdg main" > /etc/apt/sources.list.d/pgdg.list && \ + apt-get update && \ + apt install -y --no-install-recommends \ + build-essential \ + autotools-dev \ + libedit-dev \ + libgc-dev \ + libpam0g-dev \ + libreadline-dev \ + libselinux1-dev \ + libxslt1-dev \ + libssl-dev \ + libkrb5-dev \ + zlib1g-dev \ + liblz4-dev \ + libpq5 \ + libpq-dev \ + libzstd-dev \ + postgresql-16 \ + postgresql-server-dev-16 \ + postgresql-common \ + python3-sphinx && \ + wget -O /tmp/pgcopydb.tar.gz https://github.com/dimitri/pgcopydb/archive/refs/tags/v0.17.tar.gz && \ + mkdir /tmp/pgcopydb && \ + tar -xzf /tmp/pgcopydb.tar.gz -C /tmp/pgcopydb --strip-components=1 && \ + cd /tmp/pgcopydb && \ + make -s clean && \ + make -s -j12 install && \ + libpq_path=$(find /lib /usr/lib -name "libpq.so.5" | head -n 1) && \ + mkdir -p /pgcopydb/lib && \ + cp "$libpq_path" /pgcopydb/lib/; \ + else \ + # copy command below will fail if we don't have dummy files, so we create them for other debian versions + mkdir -p /usr/lib/postgresql/16/bin && touch /usr/lib/postgresql/16/bin/pgcopydb && \ + mkdir -p mkdir -p /pgcopydb/lib && touch /pgcopydb/lib/libpq.so.5; \ + fi + +FROM debian:${DEBIAN_VERSION}-slim AS build_tools ARG DEBIAN_VERSION # Add nonroot user RUN useradd -ms /bin/bash nonroot -b /home SHELL ["/bin/bash", "-c"] +RUN mkdir -p /pgcopydb/bin && \ + mkdir -p /pgcopydb/lib && \ + chmod -R 755 /pgcopydb && \ + chown -R nonroot:nonroot /pgcopydb + +COPY --from=pgcopydb_builder /usr/lib/postgresql/16/bin/pgcopydb /pgcopydb/bin/pgcopydb +COPY --from=pgcopydb_builder /pgcopydb/lib/libpq.so.5 /pgcopydb/lib/libpq.so.5 + # System deps # # 'gdb' is included so that we get backtraces of core dumps produced in @@ -38,7 +92,7 @@ RUN set -e \ libseccomp-dev \ libsqlite3-dev \ libssl-dev \ - $([[ "${DEBIAN_VERSION}" = "bullseye" ]] && libstdc++-10-dev || libstdc++-11-dev) \ + $([[ "${DEBIAN_VERSION}" = "bullseye" ]] && echo libstdc++-10-dev || echo libstdc++-11-dev) \ libtool \ libxml2-dev \ libxmlsec1-dev \ @@ -235,7 +289,13 @@ RUN whoami \ && cargo --version --verbose \ && rustup --version --verbose \ && rustc --version --verbose \ - && clang --version + && clang --version + +RUN if [ "${DEBIAN_VERSION}" = "bookworm" ]; then \ + LD_LIBRARY_PATH=/pgcopydb/lib /pgcopydb/bin/pgcopydb --version; \ +else \ + echo "pgcopydb is not available for ${DEBIAN_VERSION}"; \ +fi # Set following flag to check in Makefile if its running in Docker RUN touch /home/nonroot/.docker_build From 79929bb1b654391e63d0f02f668993834806b837 Mon Sep 17 00:00:00 2001 From: Heikki Linnakangas Date: Fri, 8 Nov 2024 10:35:03 +0200 Subject: [PATCH 38/52] Disable `rust_2024_compatibility` lint option (#9615) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Compiling with nightly rust compiler, I'm getting a lot of errors like this: error: `if let` assigns a shorter lifetime since Edition 2024 --> proxy/src/auth/backend/jwt.rs:226:16 | 226 | if let Some(permit) = self.try_acquire_permit() { | ^^^^^^^^^^^^^^^^^^^------------------------- | | | this value has a significant drop implementation which may observe a major change in drop order and requires your discretion | = warning: this changes meaning in Rust 2024 = note: for more information, see issue #124085 help: the value is now dropped here in Edition 2024 --> proxy/src/auth/backend/jwt.rs:241:13 | 241 | } else { | ^ note: the lint level is defined here --> proxy/src/lib.rs:8:5 | 8 | rust_2024_compatibility | ^^^^^^^^^^^^^^^^^^^^^^^ = note: `#[deny(if_let_rescope)]` implied by `#[deny(rust_2024_compatibility)]` and this: error: these values and local bindings have significant drop implementation that will have a different drop order from that of Edition 2021 --> proxy/src/auth/backend/jwt.rs:376:18 | 369 | let client = Client::builder() | ------ these values have significant drop implementation and will observe changes in drop order under Edition 2024 ... 376 | map: DashMap::default(), | ^^^^^^^^^^^^^^^^^^ | = warning: this changes meaning in Rust 2024 = note: for more information, see issue #123739 = note: `#[deny(tail_expr_drop_order)]` implied by `#[deny(rust_2024_compatibility)]` They are caused by the `rust_2024_compatibility` lint option. When we actually switch to the 2024 edition, it makes sense to go through all these and check that the drop order changes don't break anything, but in the meanwhile, there's no easy way to avoid these errors. Disable it, to allow compiling with nightly again. Co-authored-by: Arpad Müller --- proxy/src/lib.rs | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/proxy/src/lib.rs b/proxy/src/lib.rs index f95d645c23..ad7e1d2771 100644 --- a/proxy/src/lib.rs +++ b/proxy/src/lib.rs @@ -1,12 +1,6 @@ // rustc lints/lint groups // https://doc.rust-lang.org/rustc/lints/groups.html -#![deny( - deprecated, - future_incompatible, - let_underscore, - nonstandard_style, - rust_2024_compatibility -)] +#![deny(deprecated, future_incompatible, let_underscore, nonstandard_style)] #![warn(clippy::all, clippy::pedantic, clippy::cargo)] // List of denied lints from the clippy::restriction group. // https://rust-lang.github.io/rust-clippy/master/index.html#?groups=restriction From 027889b06ca9324604575183d84aede5f0c4c906 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?JC=20Gr=C3=BCnhage?= Date: Fri, 8 Nov 2024 10:44:59 +0100 Subject: [PATCH 39/52] ci: use set-docker-config-dir from dev-actions (#9638) set-docker-config-dir was replicated over multiple repositories. The replica of this action was removed from this repository and it's using the version from github.com/neondatabase/dev-actions instead --- .../actions/set-docker-config-dir/action.yml | 36 ------------------- .github/workflows/build-build-tools-image.yml | 2 +- .github/workflows/build_and_test.yml | 8 ++--- 3 files changed, 5 insertions(+), 41 deletions(-) delete mode 100644 .github/actions/set-docker-config-dir/action.yml diff --git a/.github/actions/set-docker-config-dir/action.yml b/.github/actions/set-docker-config-dir/action.yml deleted file mode 100644 index 3ee8bec8c6..0000000000 --- a/.github/actions/set-docker-config-dir/action.yml +++ /dev/null @@ -1,36 +0,0 @@ -name: "Set custom docker config directory" -description: "Create a directory for docker config and set DOCKER_CONFIG" - -# Use custom DOCKER_CONFIG directory to avoid conflicts with default settings -runs: - using: "composite" - steps: - - name: Show warning on GitHub-hosted runners - if: runner.environment == 'github-hosted' - shell: bash -euo pipefail {0} - run: | - # Using the following environment variables to find a path to the workflow file - # ${GITHUB_WORKFLOW_REF} - octocat/hello-world/.github/workflows/my-workflow.yml@refs/heads/my_branch - # ${GITHUB_REPOSITORY} - octocat/hello-world - # ${GITHUB_REF} - refs/heads/my_branch - # From https://docs.github.com/en/actions/writing-workflows/choosing-what-your-workflow-does/variables - - filename_with_ref=${GITHUB_WORKFLOW_REF#"$GITHUB_REPOSITORY/"} - filename=${filename_with_ref%"@$GITHUB_REF"} - - # https://docs.github.com/en/actions/writing-workflows/choosing-what-your-workflow-does/workflow-commands-for-github-actions#setting-a-warning-message - title='Unnecessary usage of `.github/actions/set-docker-config-dir`' - message='No need to use `.github/actions/set-docker-config-dir` action on GitHub-hosted runners' - echo "::warning file=${filename},title=${title}::${message}" - - - uses: pyTooling/Actions/with-post-step@74afc5a42a17a046c90c68cb5cfa627e5c6c5b6b # v1.0.7 - env: - DOCKER_CONFIG: .docker-custom-${{ github.run_id }}-${{ github.run_attempt }} - with: - main: | - mkdir -p "${DOCKER_CONFIG}" - echo DOCKER_CONFIG=${DOCKER_CONFIG} | tee -a $GITHUB_ENV - post: | - if [ -d "${DOCKER_CONFIG}" ]; then - rm -r "${DOCKER_CONFIG}" - fi diff --git a/.github/workflows/build-build-tools-image.yml b/.github/workflows/build-build-tools-image.yml index 10750089b2..82b065c524 100644 --- a/.github/workflows/build-build-tools-image.yml +++ b/.github/workflows/build-build-tools-image.yml @@ -64,7 +64,7 @@ jobs: - uses: actions/checkout@v4 - - uses: ./.github/actions/set-docker-config-dir + - uses: neondatabase/dev-actions/set-docker-config-dir@6094485bf440001c94a94a3f9e221e81ff6b6193 - uses: docker/setup-buildx-action@v3 with: cache-binary: false diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index bba51ddc92..bcf021a9a1 100644 --- a/.github/workflows/build_and_test.yml +++ b/.github/workflows/build_and_test.yml @@ -552,7 +552,7 @@ jobs: with: submodules: true - - uses: ./.github/actions/set-docker-config-dir + - uses: neondatabase/dev-actions/set-docker-config-dir@6094485bf440001c94a94a3f9e221e81ff6b6193 - uses: docker/setup-buildx-action@v3 with: cache-binary: false @@ -643,7 +643,7 @@ jobs: with: submodules: true - - uses: ./.github/actions/set-docker-config-dir + - uses: neondatabase/dev-actions/set-docker-config-dir@6094485bf440001c94a94a3f9e221e81ff6b6193 - uses: docker/setup-buildx-action@v3 with: cache-binary: false @@ -824,7 +824,7 @@ jobs: curl -fL https://github.com/neondatabase/autoscaling/releases/download/$VM_BUILDER_VERSION/vm-builder -o vm-builder chmod +x vm-builder - - uses: ./.github/actions/set-docker-config-dir + - uses: neondatabase/dev-actions/set-docker-config-dir@6094485bf440001c94a94a3f9e221e81ff6b6193 - uses: docker/login-action@v3 with: username: ${{ secrets.NEON_DOCKERHUB_USERNAME }} @@ -860,7 +860,7 @@ jobs: steps: - uses: actions/checkout@v4 - - uses: ./.github/actions/set-docker-config-dir + - uses: neondatabase/dev-actions/set-docker-config-dir@6094485bf440001c94a94a3f9e221e81ff6b6193 - uses: docker/login-action@v3 with: username: ${{ secrets.NEON_DOCKERHUB_USERNAME }} From aa9112efce42869472fbee7bfa0048f12d3ff81a Mon Sep 17 00:00:00 2001 From: John Spray Date: Fri, 8 Nov 2024 10:16:04 +0000 Subject: [PATCH 40/52] pageserver: add `no_sync` for use in regression tests (1/2) (#9677) ## Problem In test environments, the `syncfs` that the pageserver does on startup can take a long time, as other tests running concurrently might have many gigabytes of dirty pages. ## Summary of changes - Add a `no_sync` option to the pageserver's config. - Skip syncfs on startup if this is set - A subsequent PR (https://github.com/neondatabase/neon/pull/9678) will enable this by default in tests. We need to wait until after the next release to avoid breaking compat tests, which would fail if we set no_sync & use an old pageserver binary. Q: Why is this a different mechanism than safekeeper, which as a --no-sync CLI? A: Because the way we manage pageservers in neon_local depends on the pageserver.toml containing the full configuration, whereas safekeepers have a config file which is neon-local-specific and can drive a CLI flag. Q: Why is the option no_sync rather than sync? A: For boolean configs with a dangerous value, it's preferable to make "false" the safe option, so that any downstream future config tooling that might have a "booleans are false by default" behavior (e.g. golang structs) is safe by default. Q: Why only skip the syncfs, and not all fsyncs? A: Skipping all fsyncs would require more code changes, and the most acute problem isn't fsyncs themselves (these just slow down a running test), it's the syncfs (which makes a pageserver startup slow as a result of _other_ tests) --- control_plane/src/bin/neon_local.rs | 3 +++ control_plane/src/local_env.rs | 10 ++++++++++ control_plane/src/pageserver.rs | 1 + libs/pageserver_api/src/config.rs | 3 +++ pageserver/src/bin/pageserver.rs | 18 +++++++++++------- pageserver/src/config.rs | 5 +++++ 6 files changed, 33 insertions(+), 7 deletions(-) diff --git a/control_plane/src/bin/neon_local.rs b/control_plane/src/bin/neon_local.rs index 48438adf43..c4063bbd1a 100644 --- a/control_plane/src/bin/neon_local.rs +++ b/control_plane/src/bin/neon_local.rs @@ -944,6 +944,9 @@ fn handle_init(args: &InitCmdArgs) -> anyhow::Result { pg_auth_type: AuthType::Trust, http_auth_type: AuthType::Trust, other: Default::default(), + // Typical developer machines use disks with slow fsync, and we don't care + // about data integrity: disable disk syncs. + no_sync: true, } }) .collect(), diff --git a/control_plane/src/local_env.rs b/control_plane/src/local_env.rs index 9dc2a0c36b..032c88a829 100644 --- a/control_plane/src/local_env.rs +++ b/control_plane/src/local_env.rs @@ -225,6 +225,7 @@ pub struct PageServerConf { pub listen_http_addr: String, pub pg_auth_type: AuthType, pub http_auth_type: AuthType, + pub no_sync: bool, } impl Default for PageServerConf { @@ -235,6 +236,7 @@ impl Default for PageServerConf { listen_http_addr: String::new(), pg_auth_type: AuthType::Trust, http_auth_type: AuthType::Trust, + no_sync: false, } } } @@ -249,6 +251,8 @@ pub struct NeonLocalInitPageserverConf { pub listen_http_addr: String, pub pg_auth_type: AuthType, pub http_auth_type: AuthType, + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub no_sync: bool, #[serde(flatten)] pub other: HashMap, } @@ -261,6 +265,7 @@ impl From<&NeonLocalInitPageserverConf> for PageServerConf { listen_http_addr, pg_auth_type, http_auth_type, + no_sync, other: _, } = conf; Self { @@ -269,6 +274,7 @@ impl From<&NeonLocalInitPageserverConf> for PageServerConf { listen_http_addr: listen_http_addr.clone(), pg_auth_type: *pg_auth_type, http_auth_type: *http_auth_type, + no_sync: *no_sync, } } } @@ -569,6 +575,8 @@ impl LocalEnv { listen_http_addr: String, pg_auth_type: AuthType, http_auth_type: AuthType, + #[serde(default)] + no_sync: bool, } let config_toml_path = dentry.path().join("pageserver.toml"); let config_toml: PageserverConfigTomlSubset = toml_edit::de::from_str( @@ -591,6 +599,7 @@ impl LocalEnv { listen_http_addr, pg_auth_type, http_auth_type, + no_sync, } = config_toml; let IdentityTomlSubset { id: identity_toml_id, @@ -607,6 +616,7 @@ impl LocalEnv { listen_http_addr, pg_auth_type, http_auth_type, + no_sync, }; pageservers.push(conf); } diff --git a/control_plane/src/pageserver.rs b/control_plane/src/pageserver.rs index eab76e14c3..ae5e22ddc6 100644 --- a/control_plane/src/pageserver.rs +++ b/control_plane/src/pageserver.rs @@ -273,6 +273,7 @@ impl PageServerNode { ) })?; let args = vec!["-D", datadir_path_str]; + background_process::start_process( "pageserver", &datadir, diff --git a/libs/pageserver_api/src/config.rs b/libs/pageserver_api/src/config.rs index 00cc426c3c..6de34fdd35 100644 --- a/libs/pageserver_api/src/config.rs +++ b/libs/pageserver_api/src/config.rs @@ -106,6 +106,8 @@ pub struct ConfigToml { pub ephemeral_bytes_per_memory_kb: usize, pub l0_flush: Option, pub virtual_file_io_mode: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub no_sync: Option, } #[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] @@ -389,6 +391,7 @@ impl Default for ConfigToml { l0_flush: None, virtual_file_io_mode: None, tenant_config: TenantConfigToml::default(), + no_sync: None, } } } diff --git a/pageserver/src/bin/pageserver.rs b/pageserver/src/bin/pageserver.rs index 782122139e..fe2a31167d 100644 --- a/pageserver/src/bin/pageserver.rs +++ b/pageserver/src/bin/pageserver.rs @@ -154,13 +154,17 @@ fn main() -> anyhow::Result<()> { }, }; - let started = Instant::now(); - syncfs(dirfd)?; - let elapsed = started.elapsed(); - info!( - elapsed_ms = elapsed.as_millis(), - "made tenant directory contents durable" - ); + if conf.no_sync { + info!("Skipping syncfs on startup"); + } else { + let started = Instant::now(); + syncfs(dirfd)?; + let elapsed = started.elapsed(); + info!( + elapsed_ms = elapsed.as_millis(), + "made tenant directory contents durable" + ); + } } // Initialize up failpoints support diff --git a/pageserver/src/config.rs b/pageserver/src/config.rs index 06d4326459..d62066ac22 100644 --- a/pageserver/src/config.rs +++ b/pageserver/src/config.rs @@ -178,6 +178,9 @@ pub struct PageServerConf { /// Direct IO settings pub virtual_file_io_mode: virtual_file::IoMode, + + /// Optionally disable disk syncs (unsafe!) + pub no_sync: bool, } /// Token for authentication to safekeepers @@ -332,6 +335,7 @@ impl PageServerConf { concurrent_tenant_size_logical_size_queries, virtual_file_io_engine, tenant_config, + no_sync, } = config_toml; let mut conf = PageServerConf { @@ -409,6 +413,7 @@ impl PageServerConf { .map(crate::l0_flush::L0FlushConfig::from) .unwrap_or_default(), virtual_file_io_mode: virtual_file_io_mode.unwrap_or(virtual_file::IoMode::preferred()), + no_sync: no_sync.unwrap_or(false), }; // ------------------------------------------------------------ From 17c002b660a173bb6cdec07ae77103cd8580ee98 Mon Sep 17 00:00:00 2001 From: Konstantin Knizhnik Date: Fri, 8 Nov 2024 14:54:58 +0200 Subject: [PATCH 41/52] Do not copy logical replicaiton slots to replica (#9458) ## Problem Replication slots are now persisted using AUX files mechanism and included in basebackup when replica is launched. This slots are not somehow used at replica but hold WAL, which may cause local disk space exhaustion. ## Summary of changes Add `--replica` parameter to basebackup request and do not include replication slot state files in basebackup for replica. ## Checklist before requesting a review - [ ] I have performed a self-review of my code. - [ ] If it is a core feature, I have added thorough tests. - [ ] Do we need to implement analytics? if so did you add the relevant metrics to the dashboard? - [ ] If this PR requires public announcement, mark it with /release-notes label and add several sentences in this section. ## Checklist before merging - [ ] Do not forget to reformat commit message to not include the above checklist --------- Co-authored-by: Konstantin Knizhnik --- compute_tools/src/compute.rs | 28 +++++++++-- .../test_physical_and_logical_replicaiton.py | 50 +++++++++++++++++++ 2 files changed, 73 insertions(+), 5 deletions(-) create mode 100644 test_runner/regress/test_physical_and_logical_replicaiton.py diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index d3e42fe618..0a8cb14058 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -364,11 +364,29 @@ impl ComputeNode { let pageserver_connect_micros = start_time.elapsed().as_micros() as u64; let basebackup_cmd = match lsn { - Lsn(0) => format!("basebackup {} {} --gzip", spec.tenant_id, spec.timeline_id), - _ => format!( - "basebackup {} {} {} --gzip", - spec.tenant_id, spec.timeline_id, lsn - ), + Lsn(0) => { + if spec.spec.mode != ComputeMode::Primary { + format!( + "basebackup {} {} --gzip --replica", + spec.tenant_id, spec.timeline_id + ) + } else { + format!("basebackup {} {} --gzip", spec.tenant_id, spec.timeline_id) + } + } + _ => { + if spec.spec.mode != ComputeMode::Primary { + format!( + "basebackup {} {} {} --gzip --replica", + spec.tenant_id, spec.timeline_id, lsn + ) + } else { + format!( + "basebackup {} {} {} --gzip", + spec.tenant_id, spec.timeline_id, lsn + ) + } + } }; let copyreader = client.copy_out(basebackup_cmd.as_str())?; diff --git a/test_runner/regress/test_physical_and_logical_replicaiton.py b/test_runner/regress/test_physical_and_logical_replicaiton.py new file mode 100644 index 0000000000..ec14e08a14 --- /dev/null +++ b/test_runner/regress/test_physical_and_logical_replicaiton.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +import time + +from fixtures.neon_fixtures import NeonEnv, logical_replication_sync + + +def test_physical_and_logical_replication(neon_simple_env: NeonEnv, vanilla_pg): + env = neon_simple_env + + n_records = 100000 + + primary = env.endpoints.create_start( + branch_name="main", + endpoint_id="primary", + config_lines=["min_wal_size=32MB", "max_wal_size=64MB"], + ) + p_con = primary.connect() + p_cur = p_con.cursor() + p_cur.execute("CREATE TABLE t(pk bigint primary key, payload text default repeat('?',200))") + p_cur.execute("create publication pub1 for table t") + + # start subscriber to primary + vanilla_pg.start() + vanilla_pg.safe_psql("CREATE TABLE t(pk bigint primary key, payload text)") + connstr = primary.connstr().replace("'", "''") + vanilla_pg.safe_psql(f"create subscription sub1 connection '{connstr}' publication pub1") + + time.sleep(1) + secondary = env.endpoints.new_replica_start( + origin=primary, + endpoint_id="secondary", + config_lines=["min_wal_size=32MB", "max_wal_size=64MB"], + ) + + s_con = secondary.connect() + s_cur = s_con.cursor() + + for pk in range(n_records): + p_cur.execute("insert into t (pk) values (%s)", (pk,)) + + s_cur.execute("select count(*) from t") + assert s_cur.fetchall()[0][0] == n_records + + logical_replication_sync(vanilla_pg, primary) + assert vanilla_pg.safe_psql("select count(*) from t")[0][0] == n_records + + # Check that LR slot is not copied to replica + s_cur.execute("select count(*) from pg_replication_slots") + assert s_cur.fetchall()[0][0] == 0 From 3525d2e381c008904d05347742771f021325c6f8 Mon Sep 17 00:00:00 2001 From: Tristan Partin Date: Fri, 8 Nov 2024 09:15:38 -0600 Subject: [PATCH 42/52] Update TimescaleDB to 2.17.1 for PG 17 Signed-off-by: Tristan Partin --- compute/compute-node.Dockerfile | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/compute/compute-node.Dockerfile b/compute/compute-node.Dockerfile index f070f66c0a..6efef9e969 100644 --- a/compute/compute-node.Dockerfile +++ b/compute/compute-node.Dockerfile @@ -559,8 +559,8 @@ RUN case "${PG_VERSION}" in \ export TIMESCALEDB_CHECKSUM=584a351c7775f0e067eaa0e7277ea88cab9077cc4c455cbbf09a5d9723dce95d \ ;; \ "v17") \ - export TIMESCALEDB_VERSION=2.17.0 \ - export TIMESCALEDB_CHECKSUM=155bf64391d3558c42f31ca0e523cfc6252921974f75298c9039ccad1c89811a \ + export TIMESCALEDB_VERSION=2.17.1 \ + export TIMESCALEDB_CHECKSUM=6277cf43f5695e23dae1c5cfeba00474d730b66ed53665a84b787a6bb1a57e28 \ ;; \ esac && \ wget https://github.com/timescale/timescaledb/archive/refs/tags/${TIMESCALEDB_VERSION}.tar.gz -O timescaledb.tar.gz && \ From f561cbe1c709f07c507ffe642e975838ee430ef6 Mon Sep 17 00:00:00 2001 From: "Alex Chi Z." <4198311+skyzh@users.noreply.github.com> Date: Fri, 8 Nov 2024 10:35:27 -0500 Subject: [PATCH 43/52] fix(pageserver): drain upload queue before detaching ancestor (#9651) In INC-317 https://neondb.slack.com/archives/C033RQ5SPDH/p1730815677932209, we saw an interesting series of operations that would remove valid layer files existing in the layer map. * Timeline A starts compaction and generates an image layer Z but not uploading it yet. * Timeline B/C starts ancestor detaching (which should not affect timeline A) * The tenant gets restarted as part of the ancestor detaching process, without increasing the generation number. * Timeline A reloads, discovering the layer Z is a future layer, and schedules a **deletion into the deletion queue**. This means that the file will be deleted any time in the future. * Timeline A starts compaction and generates layer Z again, adding it to the layer map. Note that because we don't bump generation number during ancestor detach, it has the same filename + generation number as the original Z. * Timeline A deletes layer Z from s3 + disk, and now we have a dangling reference in the layer map, blocking all compaction/logical_size_calculation process. ## Summary of changes * We wait until all layers to be uploaded before shutting down the tenants in `Flush` mode. * Ancestor detach restarts now use this mode. * Ancestor detach also waits for remote queue completion before starting the detaching process. * The patch ensures that we don't have any future image layer (or something similar) after restart, but not fixing the underlying problem around generation numbers. --------- Signed-off-by: Alex Chi Z --- pageserver/src/http/routes.rs | 15 +++++++++++ pageserver/src/tenant/mgr.rs | 2 +- .../src/tenant/remote_timeline_client.rs | 12 +++++++++ pageserver/src/tenant/timeline.rs | 25 ++++++++++++++----- 4 files changed, 47 insertions(+), 7 deletions(-) diff --git a/pageserver/src/http/routes.rs b/pageserver/src/http/routes.rs index 72eb3e7ade..d57bd98e95 100644 --- a/pageserver/src/http/routes.rs +++ b/pageserver/src/http/routes.rs @@ -2169,6 +2169,21 @@ async fn timeline_detach_ancestor_handler( let ctx = RequestContext::new(TaskKind::DetachAncestor, DownloadBehavior::Download); let ctx = &ctx; + // Flush the upload queues of all timelines before detaching ancestor. We do the same thing again + // during shutdown. This early upload ensures the pageserver does not need to upload too many + // things and creates downtime during timeline reloads. + for timeline in tenant.list_timelines() { + timeline + .remote_client + .wait_completion() + .await + .map_err(|e| { + ApiError::PreconditionFailed(format!("cannot drain upload queue: {e}").into()) + })?; + } + + tracing::info!("all timeline upload queues are drained"); + let timeline = tenant.get_timeline(timeline_id, true)?; let progress = timeline diff --git a/pageserver/src/tenant/mgr.rs b/pageserver/src/tenant/mgr.rs index a4c458b737..4fc9d740c8 100644 --- a/pageserver/src/tenant/mgr.rs +++ b/pageserver/src/tenant/mgr.rs @@ -1959,7 +1959,7 @@ impl TenantManager { attempt.before_reset_tenant(); let (_guard, progress) = utils::completion::channel(); - match tenant.shutdown(progress, ShutdownMode::Hard).await { + match tenant.shutdown(progress, ShutdownMode::Flush).await { Ok(()) => { slot_guard.drop_old_value().expect("it was just shutdown"); } diff --git a/pageserver/src/tenant/remote_timeline_client.rs b/pageserver/src/tenant/remote_timeline_client.rs index 0aa8d61036..b37c16e133 100644 --- a/pageserver/src/tenant/remote_timeline_client.rs +++ b/pageserver/src/tenant/remote_timeline_client.rs @@ -2201,6 +2201,18 @@ impl RemoteTimelineClient { inner.initialized_mut()?; Ok(UploadQueueAccessor { inner }) } + + pub(crate) fn no_pending_work(&self) -> bool { + let inner = self.upload_queue.lock().unwrap(); + match &*inner { + UploadQueue::Uninitialized + | UploadQueue::Stopped(UploadQueueStopped::Uninitialized) => true, + UploadQueue::Stopped(UploadQueueStopped::Deletable(x)) => { + x.upload_queue_for_deletion.no_pending_work() + } + UploadQueue::Initialized(x) => x.no_pending_work(), + } + } } pub(crate) struct UploadQueueAccessor<'a> { diff --git a/pageserver/src/tenant/timeline.rs b/pageserver/src/tenant/timeline.rs index 6e082aecf5..4d086df2d1 100644 --- a/pageserver/src/tenant/timeline.rs +++ b/pageserver/src/tenant/timeline.rs @@ -852,6 +852,10 @@ pub(crate) enum ShutdownMode { /// While we are flushing, we continue to accept read I/O for LSNs ingested before /// the call to [`Timeline::shutdown`]. FreezeAndFlush, + /// Only flush the layers to the remote storage without freezing any open layers. This is the + /// mode used by ancestor detach and any other operations that reloads a tenant but not increasing + /// the generation number. + Flush, /// Shut down immediately, without waiting for any open layers to flush. Hard, } @@ -1678,11 +1682,6 @@ impl Timeline { pub(crate) async fn shutdown(&self, mode: ShutdownMode) { debug_assert_current_span_has_tenant_and_timeline_id(); - let try_freeze_and_flush = match mode { - ShutdownMode::FreezeAndFlush => true, - ShutdownMode::Hard => false, - }; - // Regardless of whether we're going to try_freeze_and_flush // or not, stop ingesting any more data. Walreceiver only provides // cancellation but no "wait until gone", because it uses the Timeline::gate. @@ -1704,7 +1703,7 @@ impl Timeline { // ... and inform any waiters for newer LSNs that there won't be any. self.last_record_lsn.shutdown(); - if try_freeze_and_flush { + if let ShutdownMode::FreezeAndFlush = mode { if let Some((open, frozen)) = self .layers .read() @@ -1746,6 +1745,20 @@ impl Timeline { warn!("failed to freeze and flush: {e:#}"); } } + + // `self.remote_client.shutdown().await` above should have already flushed everything from the queue, but + // we also do a final check here to ensure that the queue is empty. + if !self.remote_client.no_pending_work() { + warn!("still have pending work in remote upload queue, but continuing shutting down anyways"); + } + } + + if let ShutdownMode::Flush = mode { + // drain the upload queue + self.remote_client.shutdown().await; + if !self.remote_client.no_pending_work() { + warn!("still have pending work in remote upload queue, but continuing shutting down anyways"); + } } // Signal any subscribers to our cancellation token to drop out From 30680d1f3289093b532ecf2a417b6fe3309ea57b Mon Sep 17 00:00:00 2001 From: Vlad Lazar Date: Fri, 8 Nov 2024 17:00:31 +0000 Subject: [PATCH 44/52] tests: use tigther storcon scopes (#9696) ## Problem https://github.com/neondatabase/neon/pull/9596 did not update tests because that would've broken the compat tests. ## Summary of Changes Use infra scope where possible. --- test_runner/fixtures/auth_tokens.py | 1 + test_runner/fixtures/neon_fixtures.py | 12 ++++++------ 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/test_runner/fixtures/auth_tokens.py b/test_runner/fixtures/auth_tokens.py index 8ebaf61e5e..be16be81de 100644 --- a/test_runner/fixtures/auth_tokens.py +++ b/test_runner/fixtures/auth_tokens.py @@ -45,3 +45,4 @@ class TokenScope(str, Enum): SAFEKEEPER_DATA = "safekeeperdata" TENANT = "tenant" SCRUBBER = "scrubber" + INFRA = "infra" diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py index e23f46d1ca..83c68794c3 100644 --- a/test_runner/fixtures/neon_fixtures.py +++ b/test_runner/fixtures/neon_fixtures.py @@ -1782,7 +1782,7 @@ class NeonStorageController(MetricsGetter, LogUtils): self.request( "PUT", f"{self.api}/control/v1/node/{node_id}/drain", - headers=self.headers(TokenScope.ADMIN), + headers=self.headers(TokenScope.INFRA), ) def cancel_node_drain(self, node_id): @@ -1790,7 +1790,7 @@ class NeonStorageController(MetricsGetter, LogUtils): self.request( "DELETE", f"{self.api}/control/v1/node/{node_id}/drain", - headers=self.headers(TokenScope.ADMIN), + headers=self.headers(TokenScope.INFRA), ) def node_fill(self, node_id): @@ -1798,7 +1798,7 @@ class NeonStorageController(MetricsGetter, LogUtils): self.request( "PUT", f"{self.api}/control/v1/node/{node_id}/fill", - headers=self.headers(TokenScope.ADMIN), + headers=self.headers(TokenScope.INFRA), ) def cancel_node_fill(self, node_id): @@ -1806,14 +1806,14 @@ class NeonStorageController(MetricsGetter, LogUtils): self.request( "DELETE", f"{self.api}/control/v1/node/{node_id}/fill", - headers=self.headers(TokenScope.ADMIN), + headers=self.headers(TokenScope.INFRA), ) def node_status(self, node_id): response = self.request( "GET", f"{self.api}/control/v1/node/{node_id}", - headers=self.headers(TokenScope.ADMIN), + headers=self.headers(TokenScope.INFRA), ) return response.json() @@ -1829,7 +1829,7 @@ class NeonStorageController(MetricsGetter, LogUtils): response = self.request( "GET", f"{self.api}/control/v1/node", - headers=self.headers(TokenScope.ADMIN), + headers=self.headers(TokenScope.INFRA), ) return response.json() From b6bc954c5d3846214ee0a38010dd0228a7c2d7f5 Mon Sep 17 00:00:00 2001 From: Alexander Bayandin Date: Fri, 8 Nov 2024 17:32:56 +0000 Subject: [PATCH 45/52] CI: move check codestyle python to reusable workflow and run on a merge_group (#9683) ## Problem To prevent breaking main after Python 3.11 PR get merged we need to enable merge queue and run `check-codestyle-python` job on it ## Summary of changes - Move `check-codestyle-python` to a reusable workflow - Run this workflow on `merge_group` event --- .github/workflows/_check-codestyle-python.yml | 37 +++++++++++++++ .github/workflows/build_and_test.yml | 34 ++------------ .github/workflows/pre-merge-checks.yml | 47 +++++++++++++++++++ 3 files changed, 89 insertions(+), 29 deletions(-) create mode 100644 .github/workflows/_check-codestyle-python.yml create mode 100644 .github/workflows/pre-merge-checks.yml diff --git a/.github/workflows/_check-codestyle-python.yml b/.github/workflows/_check-codestyle-python.yml new file mode 100644 index 0000000000..9ae28a1379 --- /dev/null +++ b/.github/workflows/_check-codestyle-python.yml @@ -0,0 +1,37 @@ +name: Check Codestyle Python + +on: + workflow_call: + inputs: + build-tools-image: + description: 'build-tools image' + required: true + type: string + +defaults: + run: + shell: bash -euxo pipefail {0} + +jobs: + check-codestyle-python: + runs-on: [ self-hosted, small ] + container: + image: ${{ inputs.build-tools-image }} + credentials: + username: ${{ secrets.NEON_DOCKERHUB_USERNAME }} + password: ${{ secrets.NEON_DOCKERHUB_PASSWORD }} + options: --init + + steps: + - uses: actions/checkout@v4 + + - uses: actions/cache@v4 + with: + path: ~/.cache/pypoetry/virtualenvs + key: v2-${{ runner.os }}-${{ runner.arch }}-python-deps-bookworm-${{ hashFiles('poetry.lock') }} + + - run: ./scripts/pysync + + - run: poetry run ruff check . + - run: poetry run ruff format --check . + - run: poetry run mypy . diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index bcf021a9a1..d415e20db8 100644 --- a/.github/workflows/build_and_test.yml +++ b/.github/workflows/build_and_test.yml @@ -90,35 +90,10 @@ jobs: check-codestyle-python: needs: [ check-permissions, build-build-tools-image ] - runs-on: [ self-hosted, small ] - container: - image: ${{ needs.build-build-tools-image.outputs.image }}-bookworm - credentials: - username: ${{ secrets.NEON_DOCKERHUB_USERNAME }} - password: ${{ secrets.NEON_DOCKERHUB_PASSWORD }} - options: --init - - steps: - - name: Checkout - uses: actions/checkout@v4 - - - name: Cache poetry deps - uses: actions/cache@v4 - with: - path: ~/.cache/pypoetry/virtualenvs - key: v2-${{ runner.os }}-${{ runner.arch }}-python-deps-bookworm-${{ hashFiles('poetry.lock') }} - - - name: Install Python deps - run: ./scripts/pysync - - - name: Run `ruff check` to ensure code format - run: poetry run ruff check . - - - name: Run `ruff format` to ensure code format - run: poetry run ruff format --check . - - - name: Run mypy to check types - run: poetry run mypy . + uses: ./.github/workflows/_check-codestyle-python.yml + with: + build-tools-image: ${{ needs.build-build-tools-image.outputs.image }}-bookworm + secrets: inherit check-codestyle-jsonnet: needs: [ check-permissions, build-build-tools-image ] @@ -141,6 +116,7 @@ jobs: # Check that the vendor/postgres-* submodules point to the # corresponding REL_*_STABLE_neon branches. check-submodules: + needs: [ check-permissions ] runs-on: ubuntu-22.04 steps: - name: Checkout diff --git a/.github/workflows/pre-merge-checks.yml b/.github/workflows/pre-merge-checks.yml new file mode 100644 index 0000000000..40ce644eb6 --- /dev/null +++ b/.github/workflows/pre-merge-checks.yml @@ -0,0 +1,47 @@ +name: + +on: + merge_group: + branches: + - main + +# No permission for GITHUB_TOKEN by default; the **minimal required** set of permissions should be granted in each job. +permissions: {} + +jobs: + get-changed-files: + runs-on: ubuntu-22.04 + outputs: + any_changed: ${{ steps.src.outputs.any_changed }} + steps: + - uses: actions/checkout@v4 + - uses: tj-actions/changed-files@c3a1bb2c992d77180ae65be6ae6c166cf40f857c # v45.0.3 + id: src + with: + files: | + .github/workflows/pre-merge-checks.yml + **/**.py + poetry.lock + pyproject.toml + + - name: PRINT ALL CHANGED FILES FOR DEBUG PURPOSES + env: + ALL_CHANGED_FILES: ${{ steps.src.outputs.all_changed_files }} + run: echo "${ALL_CHANGED_FILES}" + + check-build-tools-image: + uses: ./.github/workflows/check-build-tools-image.yml + + build-build-tools-image: + needs: [ check-build-tools-image ] + uses: ./.github/workflows/build-build-tools-image.yml + with: + image-tag: ${{ needs.check-build-tools-image.outputs.image-tag }} + secrets: inherit + + check-codestyle-python: + needs: [ build-build-tools-image ] + uses: ./.github/workflows/_check-codestyle-python.yml + with: + build-tools-image: ${{ needs.build-build-tools-image.outputs.image }}-bookworm + secrets: inherit From 34a4eb6f2a7ddb9bd98b1b1f7b8959fa57b3007a Mon Sep 17 00:00:00 2001 From: Tristan Partin Date: Fri, 8 Nov 2024 12:19:18 -0600 Subject: [PATCH 46/52] Switch compute-related locales to C.UTF-8 by default Right now, our environments create databases with the C locale, which is really unfortunate for users who have data stored in other languages that they want to analyze. For instance, show_trgm on Hebrew text currently doesn't work in staging or production. I don't envision this being the final solution. I think this is just a way to set a known value so the pageserver doesn't use its parent environment. The final solution to me is exposing initdb parameters to users in the console. Then they could use a different locale or encoding if they so chose. Signed-off-by: Tristan Partin --- compute_tools/src/config.rs | 6 ++ libs/pageserver_api/src/config.rs | 3 + libs/utils/scripts/restore_from_wal.sh | 39 ++++++++++++- pageserver/src/config.rs | 3 + pageserver/src/tenant.rs | 28 ++++++++-- test_runner/regress/test_compute_locales.py | 61 +++++++++++++++++++++ test_runner/regress/test_wal_restore.py | 2 + 7 files changed, 136 insertions(+), 6 deletions(-) create mode 100644 test_runner/regress/test_compute_locales.py diff --git a/compute_tools/src/config.rs b/compute_tools/src/config.rs index 479100eb89..50e2a95e9d 100644 --- a/compute_tools/src/config.rs +++ b/compute_tools/src/config.rs @@ -73,6 +73,12 @@ pub fn write_postgres_conf( )?; } + // Locales + writeln!(file, "lc_messages='C.UTF-8'")?; + writeln!(file, "lc_monetary='C.UTF-8'")?; + writeln!(file, "lc_time='C.UTF-8'")?; + writeln!(file, "lc_numeric='C.UTF-8'")?; + match spec.mode { ComputeMode::Primary => {} ComputeMode::Static(lsn) => { diff --git a/libs/pageserver_api/src/config.rs b/libs/pageserver_api/src/config.rs index 6de34fdd35..4272181954 100644 --- a/libs/pageserver_api/src/config.rs +++ b/libs/pageserver_api/src/config.rs @@ -64,6 +64,7 @@ pub struct ConfigToml { #[serde(with = "humantime_serde")] pub wal_redo_timeout: Duration, pub superuser: String, + pub locale: String, pub page_cache_size: usize, pub max_file_descriptors: usize, pub pg_distrib_dir: Option, @@ -276,6 +277,7 @@ pub mod defaults { pub const DEFAULT_WAL_REDO_TIMEOUT: &str = "60 s"; pub const DEFAULT_SUPERUSER: &str = "cloud_admin"; + pub const DEFAULT_LOCALE: &str = "C.UTF-8"; pub const DEFAULT_PAGE_CACHE_SIZE: usize = 8192; pub const DEFAULT_MAX_FILE_DESCRIPTORS: usize = 100; @@ -326,6 +328,7 @@ impl Default for ConfigToml { wal_redo_timeout: (humantime::parse_duration(DEFAULT_WAL_REDO_TIMEOUT) .expect("cannot parse default wal redo timeout")), superuser: (DEFAULT_SUPERUSER.to_string()), + locale: DEFAULT_LOCALE.to_string(), page_cache_size: (DEFAULT_PAGE_CACHE_SIZE), max_file_descriptors: (DEFAULT_MAX_FILE_DESCRIPTORS), pg_distrib_dir: None, // Utf8PathBuf::from("./pg_install"), // TODO: formely, this was std::env::current_dir() diff --git a/libs/utils/scripts/restore_from_wal.sh b/libs/utils/scripts/restore_from_wal.sh index 316ec8ed0d..93448369a0 100755 --- a/libs/utils/scripts/restore_from_wal.sh +++ b/libs/utils/scripts/restore_from_wal.sh @@ -1,4 +1,4 @@ -#!/bin/bash +#!/usr/bin/env bash set -euxo pipefail @@ -6,9 +6,44 @@ PG_BIN=$1 WAL_PATH=$2 DATA_DIR=$3 PORT=$4 +PG_VERSION=$5 SYSID=$(od -A n -j 24 -N 8 -t d8 "$WAL_PATH"/000000010000000000000002* | cut -c 3-) + +# The way that initdb is invoked must match how the pageserver runs initdb. +function initdb_with_args { + local cmd=( + "$PG_BIN"/initdb + -E utf8 + -U cloud_admin + -D "$DATA_DIR" + --locale 'C.UTF-8' + --lc-collate 'C.UTF-8' + --lc-ctype 'C.UTF-8' + --lc-messages 'C.UTF-8' + --lc-monetary 'C.UTF-8' + --lc-numeric 'C.UTF-8' + --lc-time 'C.UTF-8' + --sysid="$SYSID" + ) + + case "$PG_VERSION" in + 14) + # Postgres 14 and below didn't support --locale-provider + ;; + 15 | 16) + cmd+=(--locale-provider 'libc') + ;; + *) + # Postgres 17 added the builtin provider + cmd+=(--locale-provider 'builtin') + ;; + esac + + eval env -i LD_LIBRARY_PATH="$PG_BIN"/../lib "${cmd[*]}" +} + rm -fr "$DATA_DIR" -env -i LD_LIBRARY_PATH="$PG_BIN"/../lib "$PG_BIN"/initdb -E utf8 -U cloud_admin -D "$DATA_DIR" --sysid="$SYSID" +initdb_with_args echo "port=$PORT" >> "$DATA_DIR"/postgresql.conf echo "shared_preload_libraries='\$libdir/neon_rmgr.so'" >> "$DATA_DIR"/postgresql.conf REDO_POS=0x$("$PG_BIN"/pg_controldata -D "$DATA_DIR" | grep -F "REDO location"| cut -c 42-) diff --git a/pageserver/src/config.rs b/pageserver/src/config.rs index d62066ac22..b694a43599 100644 --- a/pageserver/src/config.rs +++ b/pageserver/src/config.rs @@ -69,6 +69,7 @@ pub struct PageServerConf { pub wal_redo_timeout: Duration, pub superuser: String, + pub locale: String, pub page_cache_size: usize, pub max_file_descriptors: usize, @@ -301,6 +302,7 @@ impl PageServerConf { wait_lsn_timeout, wal_redo_timeout, superuser, + locale, page_cache_size, max_file_descriptors, pg_distrib_dir, @@ -348,6 +350,7 @@ impl PageServerConf { wait_lsn_timeout, wal_redo_timeout, superuser, + locale, page_cache_size, max_file_descriptors, http_auth_type, diff --git a/pageserver/src/tenant.rs b/pageserver/src/tenant.rs index d45c99a41b..34ea6dae1f 100644 --- a/pageserver/src/tenant.rs +++ b/pageserver/src/tenant.rs @@ -4779,10 +4779,18 @@ async fn run_initdb( let _permit = INIT_DB_SEMAPHORE.acquire().await; - let initdb_command = tokio::process::Command::new(&initdb_bin_path) + let mut initdb_command = tokio::process::Command::new(&initdb_bin_path); + initdb_command .args(["--pgdata", initdb_target_dir.as_ref()]) .args(["--username", &conf.superuser]) .args(["--encoding", "utf8"]) + .args(["--locale", &conf.locale]) + .args(["--lc-collate", &conf.locale]) + .args(["--lc-ctype", &conf.locale]) + .args(["--lc-messages", &conf.locale]) + .args(["--lc-monetary", &conf.locale]) + .args(["--lc-numeric", &conf.locale]) + .args(["--lc-time", &conf.locale]) .arg("--no-instructions") .arg("--no-sync") .env_clear() @@ -4792,15 +4800,27 @@ async fn run_initdb( // stdout invocation produces the same output every time, we don't need it .stdout(std::process::Stdio::null()) // we would be interested in the stderr output, if there was any - .stderr(std::process::Stdio::piped()) - .spawn()?; + .stderr(std::process::Stdio::piped()); + + // Before version 14, only the libc provide was available. + if pg_version > 14 { + // Version 17 brought with it a builtin locale provider which only provides + // C and C.UTF-8. While being safer for collation purposes since it is + // guaranteed to be consistent throughout a major release, it is also more + // performant. + let locale_provider = if pg_version >= 17 { "builtin" } else { "libc" }; + + initdb_command.args(["--locale-provider", locale_provider]); + } + + let initdb_proc = initdb_command.spawn()?; // Ideally we'd select here with the cancellation token, but the problem is that // we can't safely terminate initdb: it launches processes of its own, and killing // initdb doesn't kill them. After we return from this function, we want the target // directory to be able to be cleaned up. // See https://github.com/neondatabase/neon/issues/6385 - let initdb_output = initdb_command.wait_with_output().await?; + let initdb_output = initdb_proc.wait_with_output().await?; if !initdb_output.status.success() { return Err(InitdbError::Failed( initdb_output.status, diff --git a/test_runner/regress/test_compute_locales.py b/test_runner/regress/test_compute_locales.py new file mode 100644 index 0000000000..00ef32fb5e --- /dev/null +++ b/test_runner/regress/test_compute_locales.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, cast + +from fixtures.pg_version import PgVersion + +if TYPE_CHECKING: + from collections.abc import Sequence + + from fixtures.neon_fixtures import NeonEnv + + +def test_default_locales(neon_simple_env: NeonEnv): + """ + Test that the default locales for compute databases is C.UTF-8. + """ + env = neon_simple_env + + endpoint = env.endpoints.create_start("main") + + domain_locales = cast( + "Sequence[str]", + endpoint.safe_psql( + "SELECT current_setting('lc_messages') AS lc_messages," + + "current_setting('lc_monetary') AS lc_monetary," + + "current_setting('lc_numeric') AS lc_numeric," + + "current_setting('lc_time') AS lc_time" + )[0], + ) + for dl in domain_locales: + assert dl == "C.UTF-8" + + # Postgres 15 added the locale providers + if env.pg_version < PgVersion.V15: + results = cast( + "Sequence[str]", + endpoint.safe_psql( + "SELECT datcollate, datctype FROM pg_database WHERE datname = current_database()" + )[0], + ) + + datcollate = results[0] + datctype = results[1] + else: + results = cast( + "Sequence[str]", + endpoint.safe_psql( + "SELECT datlocprovider, datcollate, datctype FROM pg_database WHERE datname = current_database()" + )[0], + ) + datlocprovider = results[0] + datcollate = results[1] + datctype = results[2] + + if env.pg_version >= PgVersion.V17: + assert datlocprovider == "b", "The locale provider is not builtin" + else: + assert datlocprovider == "c", "The locale provider is not libc" + + assert datcollate == "C.UTF-8" + assert datctype == "C.UTF-8" diff --git a/test_runner/regress/test_wal_restore.py b/test_runner/regress/test_wal_restore.py index 05b6ad8a9b..c8e51fde13 100644 --- a/test_runner/regress/test_wal_restore.py +++ b/test_runner/regress/test_wal_restore.py @@ -64,6 +64,7 @@ def test_wal_restore( ), str(data_dir), str(port), + env.pg_version, ] ) restored.start() @@ -127,6 +128,7 @@ def test_wal_restore_initdb( ), str(data_dir), str(port), + env.pg_version, ] ) restored.start() From ecca62a45dfa0be134c41c113adc8e2519b827af Mon Sep 17 00:00:00 2001 From: "Alex Chi Z." <4198311+skyzh@users.noreply.github.com> Date: Fri, 8 Nov 2024 13:44:00 -0500 Subject: [PATCH 47/52] feat(pageserver): more log lines around frozen layers (#9697) We saw pageserver OOMs https://github.com/neondatabase/cloud/issues/19715 for tenants doing large writes. Add log lines around in-memory layers to hopefully collect some info during my on-call shift next week. ## Summary of changes * Estimate in-memory size of an in-mem layer. * Print frozen layer number if there are too many layers accumulated in memory. --------- Signed-off-by: Alex Chi Z --- .../tenant/storage_layer/inmemory_layer.rs | 13 ++++++++++++ pageserver/src/tenant/timeline.rs | 20 +++++++++++++++++++ 2 files changed, 33 insertions(+) diff --git a/pageserver/src/tenant/storage_layer/inmemory_layer.rs b/pageserver/src/tenant/storage_layer/inmemory_layer.rs index 2ce26ed2eb..af6112d535 100644 --- a/pageserver/src/tenant/storage_layer/inmemory_layer.rs +++ b/pageserver/src/tenant/storage_layer/inmemory_layer.rs @@ -67,6 +67,8 @@ pub struct InMemoryLayer { /// The above fields never change, except for `end_lsn`, which is only set once. /// All other changing parts are in `inner`, and protected by a mutex. inner: RwLock, + + estimated_in_mem_size: AtomicU64, } impl std::fmt::Debug for InMemoryLayer { @@ -543,6 +545,10 @@ impl InMemoryLayer { Ok(inner.file.len()) } + pub fn estimated_in_mem_size(&self) -> u64 { + self.estimated_in_mem_size.load(AtomicOrdering::Relaxed) + } + /// Create a new, empty, in-memory layer pub async fn create( conf: &'static PageServerConf, @@ -572,6 +578,7 @@ impl InMemoryLayer { file, resource_units: GlobalResourceUnits::new(), }), + estimated_in_mem_size: AtomicU64::new(0), }) } @@ -642,6 +649,12 @@ impl InMemoryLayer { // because this case is unexpected, and we would like tests to fail if this happens. warn!("Key {} at {} written twice at same LSN", key, lsn); } + self.estimated_in_mem_size.fetch_add( + (std::mem::size_of::() + + std::mem::size_of::() + + std::mem::size_of::()) as u64, + AtomicOrdering::Relaxed, + ); } inner.resource_units.maybe_publish_size(new_size); diff --git a/pageserver/src/tenant/timeline.rs b/pageserver/src/tenant/timeline.rs index 4d086df2d1..60cc689c5e 100644 --- a/pageserver/src/tenant/timeline.rs +++ b/pageserver/src/tenant/timeline.rs @@ -23,6 +23,7 @@ use handle::ShardTimelineId; use offload::OffloadError; use once_cell::sync::Lazy; use pageserver_api::{ + config::tenant_conf_defaults::DEFAULT_COMPACTION_THRESHOLD, key::{ KEY_SIZE, METADATA_KEY_BEGIN_PREFIX, METADATA_KEY_END_PREFIX, NON_INHERITED_RANGE, NON_INHERITED_SPARSE_RANGE, @@ -3501,18 +3502,37 @@ impl Timeline { let timer = self.metrics.flush_time_histo.start_timer(); + let num_frozen_layers; + let frozen_layer_total_size; let layer_to_flush = { let guard = self.layers.read().await; let Ok(lm) = guard.layer_map() else { info!("dropping out of flush loop for timeline shutdown"); return; }; + num_frozen_layers = lm.frozen_layers.len(); + frozen_layer_total_size = lm + .frozen_layers + .iter() + .map(|l| l.estimated_in_mem_size()) + .sum::(); lm.frozen_layers.front().cloned() // drop 'layers' lock to allow concurrent reads and writes }; let Some(layer_to_flush) = layer_to_flush else { break Ok(()); }; + if num_frozen_layers + > std::cmp::max( + self.get_compaction_threshold(), + DEFAULT_COMPACTION_THRESHOLD, + ) + && frozen_layer_total_size >= /* 64 MB */ 64000000 + { + tracing::warn!( + "too many frozen layers: {num_frozen_layers} layers with estimated in-mem size of {frozen_layer_total_size} bytes", + ); + } match self.flush_frozen_layer(layer_to_flush, ctx).await { Ok(this_layer_to_lsn) => { flushed_to_lsn = std::cmp::max(flushed_to_lsn, this_layer_to_lsn); From ab47804d000addd668e2583275bfdeb8209502e4 Mon Sep 17 00:00:00 2001 From: Erik Grinaker Date: Fri, 8 Nov 2024 20:25:31 +0100 Subject: [PATCH 48/52] safekeeper: remove unused `WriteGuardSharedState::skip_update` (#9699) --- safekeeper/src/timeline.rs | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/safekeeper/src/timeline.rs b/safekeeper/src/timeline.rs index fa91241177..85add6bfea 100644 --- a/safekeeper/src/timeline.rs +++ b/safekeeper/src/timeline.rs @@ -108,16 +108,11 @@ pub type ReadGuardSharedState<'a> = RwLockReadGuard<'a, SharedState>; pub struct WriteGuardSharedState<'a> { tli: Arc, guard: RwLockWriteGuard<'a, SharedState>, - skip_update: bool, } impl<'a> WriteGuardSharedState<'a> { fn new(tli: Arc, guard: RwLockWriteGuard<'a, SharedState>) -> Self { - WriteGuardSharedState { - tli, - guard, - skip_update: false, - } + WriteGuardSharedState { tli, guard } } } @@ -159,12 +154,10 @@ impl Drop for WriteGuardSharedState<'_> { } }); - if !self.skip_update { - // send notification about shared state update - self.tli.shared_state_version_tx.send_modify(|old| { - *old += 1; - }); - } + // send notification about shared state update + self.tli.shared_state_version_tx.send_modify(|old| { + *old += 1; + }); } } From af8238ae52aaf81cb02fdc246f9a7914538ded7d Mon Sep 17 00:00:00 2001 From: "Alex Chi Z." <4198311+skyzh@users.noreply.github.com> Date: Fri, 8 Nov 2024 14:28:55 -0500 Subject: [PATCH 49/52] fix(pageserver): drain upload queue before offloading timeline (#9682) It is possible at the point we shutdown the timeline, there are still layer files we did not upload. ## Summary of changes * If the queue is not empty, avoid offloading. * Shutdown the timeline gracefully using the flush mode to ensure all local files are uploaded before deleting the timeline directory. --------- Signed-off-by: Alex Chi Z --- pageserver/src/http/routes.rs | 4 ++-- pageserver/src/tenant.rs | 3 ++- pageserver/src/tenant/timeline.rs | 10 +++++++--- pageserver/src/tenant/timeline/offload.rs | 2 +- 4 files changed, 12 insertions(+), 7 deletions(-) diff --git a/pageserver/src/http/routes.rs b/pageserver/src/http/routes.rs index d57bd98e95..dde9c5dd0b 100644 --- a/pageserver/src/http/routes.rs +++ b/pageserver/src/http/routes.rs @@ -2002,9 +2002,9 @@ async fn timeline_offload_handler( "timeline has attached children".into(), )); } - if !timeline.can_offload() { + if let (false, reason) = timeline.can_offload() { return Err(ApiError::PreconditionFailed( - "Timeline::can_offload() returned false".into(), + format!("Timeline::can_offload() check failed: {}", reason) .into(), )); } offload_timeline(&tenant, &timeline) diff --git a/pageserver/src/tenant.rs b/pageserver/src/tenant.rs index 34ea6dae1f..903174680e 100644 --- a/pageserver/src/tenant.rs +++ b/pageserver/src/tenant.rs @@ -2493,7 +2493,8 @@ impl Tenant { timelines_to_compact_or_offload = timelines .iter() .filter_map(|(timeline_id, timeline)| { - let (is_active, can_offload) = (timeline.is_active(), timeline.can_offload()); + let (is_active, (can_offload, _)) = + (timeline.is_active(), timeline.can_offload()); let has_no_unoffloaded_children = { !timelines .iter() diff --git a/pageserver/src/tenant/timeline.rs b/pageserver/src/tenant/timeline.rs index 60cc689c5e..56faacbaee 100644 --- a/pageserver/src/tenant/timeline.rs +++ b/pageserver/src/tenant/timeline.rs @@ -1570,12 +1570,16 @@ impl Timeline { /// /// This is neccessary but not sufficient for offloading of the timeline as it might have /// child timelines that are not offloaded yet. - pub(crate) fn can_offload(&self) -> bool { + pub(crate) fn can_offload(&self) -> (bool, &'static str) { if self.remote_client.is_archived() != Some(true) { - return false; + return (false, "the timeline is not archived"); + } + if !self.remote_client.no_pending_work() { + // if the remote client is still processing some work, we can't offload + return (false, "the upload queue is not drained yet"); } - true + (true, "ok") } /// Outermost timeline compaction operation; downloads needed layers. Returns whether we have pending diff --git a/pageserver/src/tenant/timeline/offload.rs b/pageserver/src/tenant/timeline/offload.rs index 2dc461c28d..1394843467 100644 --- a/pageserver/src/tenant/timeline/offload.rs +++ b/pageserver/src/tenant/timeline/offload.rs @@ -58,7 +58,7 @@ pub(crate) async fn offload_timeline( } // Now that the Timeline is in Stopping state, request all the related tasks to shut down. - timeline.shutdown(super::ShutdownMode::Hard).await; + timeline.shutdown(super::ShutdownMode::Flush).await; // TODO extend guard mechanism above with method // to make deletions possible while offloading is in progress From ecde8d763257703f143e3fd74c024fc73ff9f13f Mon Sep 17 00:00:00 2001 From: Tristan Partin Date: Fri, 8 Nov 2024 14:43:15 -0600 Subject: [PATCH 50/52] Improve type safety according to pyright Pyright found many issues that mypy doesn't seem to want to catch or mypy isn't configured to catch. Signed-off-by: Tristan Partin --- test_runner/fixtures/benchmark_fixture.py | 6 ++++ test_runner/fixtures/compare_fixtures.py | 9 +++-- test_runner/fixtures/h2server.py | 36 +++++++++++++------ test_runner/fixtures/neon_fixtures.py | 2 +- test_runner/fixtures/pageserver/http.py | 2 +- test_runner/fixtures/pageserver/utils.py | 4 +++ test_runner/fixtures/paths.py | 4 +-- test_runner/performance/test_copy.py | 5 ++- .../regress/test_pageserver_generations.py | 1 + test_runner/regress/test_proxy_websockets.py | 2 +- test_runner/regress/test_sharding.py | 1 + .../regress/test_storage_controller.py | 2 ++ test_runner/regress/test_storage_scrubber.py | 7 +++- test_runner/regress/test_tenant_size.py | 3 +- .../regress/test_threshold_based_eviction.py | 1 + test_runner/regress/test_wal_acceptor.py | 13 +++---- 16 files changed, 67 insertions(+), 31 deletions(-) diff --git a/test_runner/fixtures/benchmark_fixture.py b/test_runner/fixtures/benchmark_fixture.py index 74fe39ef53..d3419bd8b1 100644 --- a/test_runner/fixtures/benchmark_fixture.py +++ b/test_runner/fixtures/benchmark_fixture.py @@ -80,7 +80,13 @@ class PgBenchRunResult: ): stdout_lines = stdout.splitlines() + number_of_clients = 0 + number_of_threads = 0 + number_of_transactions_actually_processed = 0 + latency_average = 0.0 latency_stddev = None + tps = 0.0 + scale = 0 # we know significant parts of these values from test input # but to be precise take them from output diff --git a/test_runner/fixtures/compare_fixtures.py b/test_runner/fixtures/compare_fixtures.py index 2195ae8225..85b6e7a3b8 100644 --- a/test_runner/fixtures/compare_fixtures.py +++ b/test_runner/fixtures/compare_fixtures.py @@ -8,7 +8,7 @@ from contextlib import _GeneratorContextManager, contextmanager # Type-related stuff from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, final import pytest from _pytest.fixtures import FixtureRequest @@ -70,12 +70,12 @@ class PgCompare(ABC): @contextmanager @abstractmethod - def record_pageserver_writes(self, out_name: str): + def record_pageserver_writes(self, out_name: str) -> Iterator[None]: pass @contextmanager @abstractmethod - def record_duration(self, out_name: str): + def record_duration(self, out_name: str) -> Iterator[None]: pass @contextmanager @@ -105,6 +105,7 @@ class PgCompare(ABC): return results +@final class NeonCompare(PgCompare): """PgCompare interface for the neon stack.""" @@ -206,6 +207,7 @@ class NeonCompare(PgCompare): return self.zenbenchmark.record_duration(out_name) +@final class VanillaCompare(PgCompare): """PgCompare interface for vanilla postgres.""" @@ -271,6 +273,7 @@ class VanillaCompare(PgCompare): return self.zenbenchmark.record_duration(out_name) +@final class RemoteCompare(PgCompare): """PgCompare interface for a remote postgres instance.""" diff --git a/test_runner/fixtures/h2server.py b/test_runner/fixtures/h2server.py index 92783e1fb2..e890b2bcf1 100644 --- a/test_runner/fixtures/h2server.py +++ b/test_runner/fixtures/h2server.py @@ -4,11 +4,14 @@ https://python-hyper.org/projects/hyper-h2/en/stable/asyncio-example.html auth-broker -> local-proxy needs a h2 connection, so we need a h2 server :) """ +from __future__ import annotations + import asyncio import collections import io import json from collections.abc import AsyncIterable +from typing import TYPE_CHECKING, final import pytest_asyncio from h2.config import H2Configuration @@ -25,34 +28,45 @@ from h2.events import ( ) from h2.exceptions import ProtocolError, StreamClosedError from h2.settings import SettingCodes +from typing_extensions import override + +if TYPE_CHECKING: + from typing import Any, Optional + RequestData = collections.namedtuple("RequestData", ["headers", "data"]) +@final class H2Server: - def __init__(self, host, port) -> None: + def __init__(self, host: str, port: int) -> None: self.host = host self.port = port +@final class H2Protocol(asyncio.Protocol): def __init__(self): config = H2Configuration(client_side=False, header_encoding="utf-8") self.conn = H2Connection(config=config) - self.transport = None - self.stream_data = {} - self.flow_control_futures = {} + self.transport: Optional[asyncio.Transport] = None + self.stream_data: dict[int, RequestData] = {} + self.flow_control_futures: dict[int, asyncio.Future[Any]] = {} - def connection_made(self, transport: asyncio.Transport): # type: ignore[override] + @override + def connection_made(self, transport: asyncio.BaseTransport): + assert isinstance(transport, asyncio.Transport) self.transport = transport self.conn.initiate_connection() self.transport.write(self.conn.data_to_send()) - def connection_lost(self, _exc): + @override + def connection_lost(self, exc: Optional[Exception]): for future in self.flow_control_futures.values(): future.cancel() self.flow_control_futures = {} + @override def data_received(self, data: bytes): assert self.transport is not None try: @@ -77,7 +91,7 @@ class H2Protocol(asyncio.Protocol): self.window_updated(event.stream_id, event.delta) elif isinstance(event, RemoteSettingsChanged): if SettingCodes.INITIAL_WINDOW_SIZE in event.changed_settings: - self.window_updated(None, 0) + self.window_updated(0, 0) self.transport.write(self.conn.data_to_send()) @@ -123,7 +137,7 @@ class H2Protocol(asyncio.Protocol): else: stream_data.data.write(data) - def stream_reset(self, stream_id): + def stream_reset(self, stream_id: int): """ A stream reset was sent. Stop sending data. """ @@ -131,7 +145,7 @@ class H2Protocol(asyncio.Protocol): future = self.flow_control_futures.pop(stream_id) future.cancel() - async def send_data(self, data, stream_id): + async def send_data(self, data: bytes, stream_id: int): """ Send data according to the flow control rules. """ @@ -161,7 +175,7 @@ class H2Protocol(asyncio.Protocol): self.transport.write(self.conn.data_to_send()) data = data[chunk_size:] - async def wait_for_flow_control(self, stream_id): + async def wait_for_flow_control(self, stream_id: int): """ Waits for a Future that fires when the flow control window is opened. """ @@ -169,7 +183,7 @@ class H2Protocol(asyncio.Protocol): self.flow_control_futures[stream_id] = f await f - def window_updated(self, stream_id, delta): + def window_updated(self, stream_id: int, delta): """ A window update frame was received. Unblock some number of flow control Futures. diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py index 83c68794c3..79baa8a32d 100644 --- a/test_runner/fixtures/neon_fixtures.py +++ b/test_runner/fixtures/neon_fixtures.py @@ -1857,7 +1857,7 @@ class NeonStorageController(MetricsGetter, LogUtils): shard_count: Optional[int] = None, shard_stripe_size: Optional[int] = None, tenant_config: Optional[dict[Any, Any]] = None, - placement_policy: Optional[Union[dict[Any, Any] | str]] = None, + placement_policy: Optional[Union[dict[Any, Any], str]] = None, ): """ Use this rather than pageserver_api() when you need to include shard parameters diff --git a/test_runner/fixtures/pageserver/http.py b/test_runner/fixtures/pageserver/http.py index 57a5d6875e..d1a9b5921a 100644 --- a/test_runner/fixtures/pageserver/http.py +++ b/test_runner/fixtures/pageserver/http.py @@ -316,7 +316,7 @@ class PageserverHttpClient(requests.Session, MetricsGetter): def tenant_location_conf( self, tenant_id: Union[TenantId, TenantShardId], - location_conf=dict[str, Any], + location_conf: dict[str, Any], flush_ms=None, lazy: Optional[bool] = None, ): diff --git a/test_runner/fixtures/pageserver/utils.py b/test_runner/fixtures/pageserver/utils.py index 4c4306be9e..ac7497ee6c 100644 --- a/test_runner/fixtures/pageserver/utils.py +++ b/test_runner/fixtures/pageserver/utils.py @@ -56,6 +56,8 @@ def wait_for_upload( lsn: Lsn, ): """waits for local timeline upload up to specified lsn""" + + current_lsn = Lsn(0) for i in range(20): current_lsn = remote_consistent_lsn(pageserver_http, tenant, timeline) if current_lsn >= lsn: @@ -203,6 +205,8 @@ def wait_for_last_record_lsn( lsn: Lsn, ) -> Lsn: """waits for pageserver to catch up to a certain lsn, returns the last observed lsn.""" + + current_lsn = Lsn(0) for i in range(1000): current_lsn = last_record_lsn(pageserver_http, tenant, timeline) if current_lsn >= lsn: diff --git a/test_runner/fixtures/paths.py b/test_runner/fixtures/paths.py index d950f2356d..60221573eb 100644 --- a/test_runner/fixtures/paths.py +++ b/test_runner/fixtures/paths.py @@ -112,7 +112,7 @@ def compatibility_snapshot_dir() -> Iterator[Path]: @pytest.fixture(scope="session") -def compatibility_neon_binpath() -> Optional[Iterator[Path]]: +def compatibility_neon_binpath() -> Iterator[Optional[Path]]: if os.getenv("REMOTE_ENV"): return comp_binpath = None @@ -133,7 +133,7 @@ def pg_distrib_dir(base_dir: Path) -> Iterator[Path]: @pytest.fixture(scope="session") -def compatibility_pg_distrib_dir() -> Optional[Iterator[Path]]: +def compatibility_pg_distrib_dir() -> Iterator[Optional[Path]]: compat_distrib_dir = None if env_compat_postgres_bin := os.environ.get("COMPATIBILITY_POSTGRES_DISTRIB_DIR"): compat_distrib_dir = Path(env_compat_postgres_bin).resolve() diff --git a/test_runner/performance/test_copy.py b/test_runner/performance/test_copy.py index 743604a381..d571fab6b5 100644 --- a/test_runner/performance/test_copy.py +++ b/test_runner/performance/test_copy.py @@ -2,11 +2,13 @@ from __future__ import annotations from contextlib import closing from io import BufferedReader, RawIOBase -from typing import Optional +from typing import Optional, final from fixtures.compare_fixtures import PgCompare +from typing_extensions import override +@final class CopyTestData(RawIOBase): def __init__(self, rows: int): self.rows = rows @@ -14,6 +16,7 @@ class CopyTestData(RawIOBase): self.linebuf: Optional[bytes] = None self.ptr = 0 + @override def readable(self): return True diff --git a/test_runner/regress/test_pageserver_generations.py b/test_runner/regress/test_pageserver_generations.py index 11ebb81023..8f6c9f16fd 100644 --- a/test_runner/regress/test_pageserver_generations.py +++ b/test_runner/regress/test_pageserver_generations.py @@ -656,6 +656,7 @@ def test_upgrade_generationless_local_file_paths( workload.write_rows(1000) attached_pageserver = env.get_tenant_pageserver(tenant_id) + assert attached_pageserver is not None secondary_pageserver = list([ps for ps in env.pageservers if ps.id != attached_pageserver.id])[ 0 ] diff --git a/test_runner/regress/test_proxy_websockets.py b/test_runner/regress/test_proxy_websockets.py index 071ca7c54e..ea01252ce4 100644 --- a/test_runner/regress/test_proxy_websockets.py +++ b/test_runner/regress/test_proxy_websockets.py @@ -37,7 +37,7 @@ async def test_websockets(static_proxy: NeonProxy): startup_message.extend(b"\0") length = (4 + len(startup_message)).to_bytes(4, byteorder="big") - await websocket.send([length, startup_message]) + await websocket.send([length, bytes(startup_message)]) startup_response = await websocket.recv() assert isinstance(startup_response, bytes) diff --git a/test_runner/regress/test_sharding.py b/test_runner/regress/test_sharding.py index 3a249bbdb4..ec633e352c 100644 --- a/test_runner/regress/test_sharding.py +++ b/test_runner/regress/test_sharding.py @@ -256,6 +256,7 @@ def test_sharding_split_compaction( # Cleanup part 1: while layers are still in PITR window, we should only drop layers that are fully redundant for shard in shards: ps = env.get_tenant_pageserver(shard) + assert ps is not None # Invoke compaction: this should drop any layers that don't overlap with the shard's key stripes detail_before = ps.http_client().timeline_detail(shard, timeline_id) diff --git a/test_runner/regress/test_storage_controller.py b/test_runner/regress/test_storage_controller.py index c8de292588..a069e0d01c 100644 --- a/test_runner/regress/test_storage_controller.py +++ b/test_runner/regress/test_storage_controller.py @@ -1237,6 +1237,7 @@ def test_storage_controller_tenant_deletion( # Assert attachments all have local content for shard_id in shard_ids: pageserver = env.get_tenant_pageserver(shard_id) + assert pageserver is not None assert pageserver.tenant_dir(shard_id).exists() # Assert all shards have some content in remote storage @@ -2745,6 +2746,7 @@ def test_storage_controller_validate_during_migration(neon_env_builder: NeonEnvB # Upload but don't compact origin_pageserver = env.get_tenant_pageserver(tenant_id) + assert origin_pageserver is not None dest_ps_id = [p.id for p in env.pageservers if p.id != origin_pageserver.id][0] origin_pageserver.http_client().timeline_checkpoint( tenant_id, timeline_id, wait_until_uploaded=True, compact=False diff --git a/test_runner/regress/test_storage_scrubber.py b/test_runner/regress/test_storage_scrubber.py index 05db0fe977..11ad2173ae 100644 --- a/test_runner/regress/test_storage_scrubber.py +++ b/test_runner/regress/test_storage_scrubber.py @@ -245,6 +245,7 @@ def test_scrubber_physical_gc_ancestors( workload.write_rows(100, upload=False) for shard in shards: ps = env.get_tenant_pageserver(shard) + assert ps is not None log.info(f"Waiting for shard {shard} on pageserver {ps.id}") ps.http_client().timeline_checkpoint( shard, timeline_id, compact=False, wait_until_uploaded=True @@ -270,6 +271,7 @@ def test_scrubber_physical_gc_ancestors( workload.churn_rows(100) for shard in shards: ps = env.get_tenant_pageserver(shard) + assert ps is not None ps.http_client().timeline_compact(shard, timeline_id, force_image_layer_creation=True) ps.http_client().timeline_gc(shard, timeline_id, 0) @@ -336,12 +338,15 @@ def test_scrubber_physical_gc_timeline_deletion(neon_env_builder: NeonEnvBuilder # Issue a deletion queue flush so that the parent shard can't leave behind layers # that will look like unexpected garbage to the scrubber - env.get_tenant_pageserver(tenant_id).http_client().deletion_queue_flush(execute=True) + ps = env.get_tenant_pageserver(tenant_id) + assert ps is not None + ps.http_client().deletion_queue_flush(execute=True) new_shard_count = 4 shards = env.storage_controller.tenant_shard_split(tenant_id, shard_count=new_shard_count) for shard in shards: ps = env.get_tenant_pageserver(shard) + assert ps is not None log.info(f"Waiting for shard {shard} on pageserver {ps.id}") ps.http_client().timeline_checkpoint( shard, timeline_id, compact=False, wait_until_uploaded=True diff --git a/test_runner/regress/test_tenant_size.py b/test_runner/regress/test_tenant_size.py index b41f1709bd..0c431fa453 100644 --- a/test_runner/regress/test_tenant_size.py +++ b/test_runner/regress/test_tenant_size.py @@ -315,6 +315,7 @@ def test_single_branch_get_tenant_size_grows( tenant_id: TenantId, timeline_id: TimelineId, ) -> tuple[Lsn, int]: + size = 0 consistent = False size_debug = None @@ -360,7 +361,7 @@ def test_single_branch_get_tenant_size_grows( collected_responses.append(("CREATE", current_lsn, size)) batch_size = 100 - + prev_size = 0 for i in range(3): with endpoint.cursor() as cur: cur.execute( diff --git a/test_runner/regress/test_threshold_based_eviction.py b/test_runner/regress/test_threshold_based_eviction.py index 5f211ec4d4..68e9385035 100644 --- a/test_runner/regress/test_threshold_based_eviction.py +++ b/test_runner/regress/test_threshold_based_eviction.py @@ -146,6 +146,7 @@ def test_threshold_based_eviction( out += [f" {remote} {layer.layer_file_name}"] return "\n".join(out) + stable_for: float = 0 observation_window = 8 * eviction_threshold consider_stable_when_no_change_for_seconds = 3 * eviction_threshold poll_interval = eviction_threshold / 3 diff --git a/test_runner/regress/test_wal_acceptor.py b/test_runner/regress/test_wal_acceptor.py index 157390c01c..e224d5eb01 100644 --- a/test_runner/regress/test_wal_acceptor.py +++ b/test_runner/regress/test_wal_acceptor.py @@ -1506,15 +1506,10 @@ class SafekeeperEnv: port=port.http, auth_token=None, ) - try: - safekeeper_process = start_in_background( - cmd, safekeeper_dir, "safekeeper.log", safekeeper_client.check_status - ) - return safekeeper_process - except Exception as e: - log.error(e) - safekeeper_process.kill() - raise Exception(f"Failed to start safekepeer as {cmd}, reason: {e}") from e + safekeeper_process = start_in_background( + cmd, safekeeper_dir, "safekeeper.log", safekeeper_client.check_status + ) + return safekeeper_process def get_safekeeper_connstrs(self): assert self.safekeepers is not None, "safekeepers are not initialized" From 2fcac0e66b1e4c5b6fb7adb01793c49850191f93 Mon Sep 17 00:00:00 2001 From: Alexander Bayandin Date: Sat, 9 Nov 2024 01:02:54 +0000 Subject: [PATCH 51/52] CI(pre-merge-checks): add required checks (#9700) ## Problem The Merge queue doesn't work because it expects certain jobs, which we don't have in the `pre-merge-checks` workflow. But it turns out we can just create jobs/checks with the same names in any workflow that we run. ## Summary of changes - Add `conclusion` jobs - Create `neon-cloud-e2e` status check - Add a bunch of `if`s to handle cases with no relevant changes found and prepare the workflow to run rust checks in the future - List the workflow in `report-workflow-stats` to collect stats about it --- .github/workflows/pre-merge-checks.yml | 61 ++++++++++++++++++--- .github/workflows/report-workflow-stats.yml | 1 + 2 files changed, 55 insertions(+), 7 deletions(-) diff --git a/.github/workflows/pre-merge-checks.yml b/.github/workflows/pre-merge-checks.yml index 40ce644eb6..137faa7abc 100644 --- a/.github/workflows/pre-merge-checks.yml +++ b/.github/workflows/pre-merge-checks.yml @@ -1,10 +1,14 @@ -name: +name: Pre-merge checks on: merge_group: branches: - main +defaults: + run: + shell: bash -euxo pipefail {0} + # No permission for GITHUB_TOKEN by default; the **minimal required** set of permissions should be granted in each job. permissions: {} @@ -12,11 +16,11 @@ jobs: get-changed-files: runs-on: ubuntu-22.04 outputs: - any_changed: ${{ steps.src.outputs.any_changed }} + python-changed: ${{ steps.python-src.outputs.any_changed }} steps: - uses: actions/checkout@v4 - - uses: tj-actions/changed-files@c3a1bb2c992d77180ae65be6ae6c166cf40f857c # v45.0.3 - id: src + - uses: tj-actions/changed-files@4edd678ac3f81e2dc578756871e4d00c19191daf # v45.0.4 + id: python-src with: files: | .github/workflows/pre-merge-checks.yml @@ -26,10 +30,13 @@ jobs: - name: PRINT ALL CHANGED FILES FOR DEBUG PURPOSES env: - ALL_CHANGED_FILES: ${{ steps.src.outputs.all_changed_files }} - run: echo "${ALL_CHANGED_FILES}" + PYTHON_CHANGED_FILES: ${{ steps.python-src.outputs.all_changed_files }} + run: | + echo "${PYTHON_CHANGED_FILES}" check-build-tools-image: + if: needs.get-changed-files.outputs.python-changed == 'true' + needs: [ get-changed-files ] uses: ./.github/workflows/check-build-tools-image.yml build-build-tools-image: @@ -40,8 +47,48 @@ jobs: secrets: inherit check-codestyle-python: - needs: [ build-build-tools-image ] + if: needs.get-changed-files.outputs.python-changed == 'true' + needs: [ get-changed-files, build-build-tools-image ] uses: ./.github/workflows/_check-codestyle-python.yml with: build-tools-image: ${{ needs.build-build-tools-image.outputs.image }}-bookworm secrets: inherit + + # To get items from the merge queue merged into main we need to satisfy "Status checks that are required". + # Currently we require 2 jobs (checks with exact name): + # - conclusion + # - neon-cloud-e2e + conclusion: + if: always() + permissions: + statuses: write # for `github.repos.createCommitStatus(...)` + needs: + - get-changed-files + - check-codestyle-python + runs-on: ubuntu-22.04 + steps: + - name: Create fake `neon-cloud-e2e` check + uses: actions/github-script@v7 + with: + # Retry script for 5XX server errors: https://github.com/actions/github-script#retries + retries: 5 + script: | + const { repo, owner } = context.repo; + const targetUrl = `${context.serverUrl}/${owner}/${repo}/actions/runs/${context.runId}`; + + await github.rest.repos.createCommitStatus({ + owner: owner, + repo: repo, + sha: context.sha, + context: `neon-cloud-e2e`, + state: `success`, + target_url: targetUrl, + description: `fake check for merge queue`, + }); + + - name: Fail the job if any of the dependencies do not succeed or skipped + run: exit 1 + if: | + (contains(needs.check-codestyle-python.result, 'skipped') && needs.get-changed-files.outputs.python-changed == 'true') + || contains(needs.*.result, 'failure') + || contains(needs.*.result, 'cancelled') diff --git a/.github/workflows/report-workflow-stats.yml b/.github/workflows/report-workflow-stats.yml index 6abeff7695..0d135a257c 100644 --- a/.github/workflows/report-workflow-stats.yml +++ b/.github/workflows/report-workflow-stats.yml @@ -23,6 +23,7 @@ on: - Test Postgres client libraries - Trigger E2E Tests - cleanup caches by a branch + - Pre-merge checks types: [completed] jobs: From ceaa80ffebca3050e06c6a5d75f184c6e637ef50 Mon Sep 17 00:00:00 2001 From: Vlad Lazar Date: Mon, 11 Nov 2024 09:58:41 +0000 Subject: [PATCH 52/52] storcon: add peer token for peer to peer communication (#9695) ## Problem We wish to stop using admin tokens in the infra repo, but step down requests use the admin token. ## Summary of Changes Introduce a new "ControllerPeer" scope and use it for step-down requests. --- libs/utils/src/auth.rs | 5 +++++ pageserver/src/auth.rs | 3 ++- safekeeper/src/auth.rs | 3 ++- storage_controller/src/http.rs | 2 +- 4 files changed, 10 insertions(+), 3 deletions(-) diff --git a/libs/utils/src/auth.rs b/libs/utils/src/auth.rs index 5bd6f4bedc..f7acc61ac1 100644 --- a/libs/utils/src/auth.rs +++ b/libs/utils/src/auth.rs @@ -40,6 +40,11 @@ pub enum Scope { /// Allows access to storage controller APIs used by the scrubber, to interrogate the state /// of a tenant & post scrub results. Scrubber, + + /// This scope is used for communication with other storage controller instances. + /// At the time of writing, this is only used for the step down request. + #[serde(rename = "controller_peer")] + ControllerPeer, } /// JWT payload. See docs/authentication.md for the format diff --git a/pageserver/src/auth.rs b/pageserver/src/auth.rs index 5c931fcfdb..4075427ab4 100644 --- a/pageserver/src/auth.rs +++ b/pageserver/src/auth.rs @@ -19,7 +19,8 @@ pub fn check_permission(claims: &Claims, tenant_id: Option) -> Result< | Scope::SafekeeperData | Scope::GenerationsApi | Scope::Infra - | Scope::Scrubber, + | Scope::Scrubber + | Scope::ControllerPeer, _, ) => Err(AuthError( format!( diff --git a/safekeeper/src/auth.rs b/safekeeper/src/auth.rs index fdd0830b02..81c79fae30 100644 --- a/safekeeper/src/auth.rs +++ b/safekeeper/src/auth.rs @@ -20,7 +20,8 @@ pub fn check_permission(claims: &Claims, tenant_id: Option) -> Result< | Scope::PageServerApi | Scope::GenerationsApi | Scope::Infra - | Scope::Scrubber, + | Scope::Scrubber + | Scope::ControllerPeer, _, ) => Err(AuthError( format!( diff --git a/storage_controller/src/http.rs b/storage_controller/src/http.rs index f6ea1aedc6..9b5d4caf31 100644 --- a/storage_controller/src/http.rs +++ b/storage_controller/src/http.rs @@ -1033,7 +1033,7 @@ async fn handle_update_preferred_azs(req: Request) -> Result) -> Result, ApiError> { - check_permissions(&req, Scope::Admin)?; + check_permissions(&req, Scope::ControllerPeer)?; let req = match maybe_forward(req).await { ForwardOutcome::Forwarded(res) => {