Cap WAL download for LR on compute start by max_slot_wal_keep_size.

Support custom types in broker (#5761 )
Old methods are unchanged for backwards compatibility. Added `SafekeeperDiscoveryRequest` and `SafekeeperDiscoveryResponse` types to serve as example, and also as a prerequisite for https://github.com/neondatabase/neon/issues/5471
2026-07-08 14:40:37 +00:00 · 2023-12-22 15:55:45 +03:00 · 2023-12-19 17:06:43 +00:00 · 2023-12-19 17:44:49 +01:00 · 2023-12-19 16:27:47 +00:00 · 2023-12-19 18:17:11 +02:00
18 changed files with 560 additions and 368 deletions
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -572,7 +572,7 @@ dependencies = [
 "once_cell",
 "pin-project-lite",
 "pin-utils",
- "rustls 0.21.9",
+ "rustls",
 "tokio",
 "tracing",
 ]
@@ -2278,10 +2278,10 @@ dependencies = [
 "http",
 "hyper",
 "log",
- "rustls 0.21.9",
+ "rustls",
 "rustls-native-certs",
 "tokio",
- "tokio-rustls 0.24.0",
+ "tokio-rustls",
 ]

 [[package]]
@@ -2493,7 +2493,7 @@ checksum = "5c7ea04a7c5c055c175f189b6dc6ba036fd62306b58c66c9f6389036c503a3f4"
 dependencies = [
 "base64 0.21.1",
 "js-sys",
- "pem",
+ "pem 3.0.3",
 "ring 0.17.6",
 "serde",
 "serde_json",
@@ -3290,6 +3290,16 @@ version = "0.1.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "19b17cddbe7ec3f8bc800887bab5e717348c95ea2ca0b1bf0837fb964dc67099"

+[[package]]
+name = "pem"
+version = "2.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6b13fe415cdf3c8e44518e18a7c95a13431d9bdf6d15367d82b23c377fdd441a"
+dependencies = [
+ "base64 0.21.1",
+ "serde",
+]
+
 [[package]]
 name = "pem"
 version = "3.0.3"
@@ -3472,14 +3482,14 @@ dependencies = [
 "futures",
 "once_cell",
 "pq_proto",
- "ring 0.17.6",
- "rustls 0.22.1",
- "rustls-pemfile 2.0.0",
+ "rustls",
+ "rustls-pemfile",
 "serde",
 "thiserror",
 "tokio",
 "tokio-postgres",
- "tokio-rustls 0.25.0",
+ "tokio-postgres-rustls",
+ "tokio-rustls",
 "tracing",
 "workspace_hack",
 ]
@@ -3707,8 +3717,8 @@ dependencies = [
 "routerify",
 "rstest",
 "rustc-hash",
- "rustls 0.22.1",
- "rustls-pemfile 2.0.0",
+ "rustls",
+ "rustls-pemfile",
 "scopeguard",
 "serde",
 "serde_json",
@@ -3722,7 +3732,7 @@ dependencies = [
 "tokio",
 "tokio-postgres",
 "tokio-postgres-rustls",
- "tokio-rustls 0.25.0",
+ "tokio-rustls",
 "tokio-util",
 "tracing",
 "tracing-opentelemetry",
@@ -3731,7 +3741,7 @@ dependencies = [
 "url",
 "utils",
 "uuid",
- "webpki-roots",
+ "webpki-roots 0.25.2",
 "workspace_hack",
 "x509-parser",
 ]
@@ -3850,12 +3860,12 @@ dependencies = [

 [[package]]
 name = "rcgen"
-version = "0.12.0"
+version = "0.11.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5d918c80c5a4c7560db726763020bd16db179e4d5b828078842274a443addb5d"
+checksum = "4954fbc00dcd4d8282c987710e50ba513d351400dbdd00e803a05172a90d8976"
 dependencies = [
- "pem",
- "ring 0.17.6",
+ "pem 2.0.1",
+ "ring 0.16.20",
 "time",
 "yasna",
 ]
@@ -3993,14 +4003,14 @@ dependencies = [
 "once_cell",
 "percent-encoding",
 "pin-project-lite",
- "rustls 0.21.9",
- "rustls-pemfile 1.0.2",
+ "rustls",
+ "rustls-pemfile",
 "serde",
 "serde_json",
 "serde_urlencoded",
 "tokio",
 "tokio-native-tls",
- "tokio-rustls 0.24.0",
+ "tokio-rustls",
 "tokio-util",
 "tower-service",
 "url",
@@ -4008,7 +4018,7 @@ dependencies = [
 "wasm-bindgen-futures",
 "wasm-streams",
 "web-sys",
- "webpki-roots",
+ "webpki-roots 0.25.2",
 "winreg",
 ]

@@ -4240,20 +4250,6 @@ dependencies = [
 "sct",
 ]

-[[package]]
-name = "rustls"
-version = "0.22.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fe6b63262c9fcac8659abfaa96cac103d28166d3ff3eaf8f412e19f3ae9e5a48"
-dependencies = [
- "log",
- "ring 0.17.6",
- "rustls-pki-types",
- "rustls-webpki 0.102.0",
- "subtle",
- "zeroize",
-]
-
 [[package]]
 name = "rustls-native-certs"
 version = "0.6.2"
@@ -4261,7 +4257,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0167bac7a9f490495f3c33013e7722b53cb087ecbe082fb0c6387c96f634ea50"
 dependencies = [
 "openssl-probe",
- "rustls-pemfile 1.0.2",
+ "rustls-pemfile",
 "schannel",
 "security-framework",
 ]
@@ -4276,21 +4272,15 @@ dependencies = [
 ]

 [[package]]
-name = "rustls-pemfile"
-version = "2.0.0"
+name = "rustls-webpki"
+version = "0.100.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "35e4980fa29e4c4b212ffb3db068a564cbf560e51d3944b7c88bd8bf5bec64f4"
+checksum = "e98ff011474fa39949b7e5c0428f9b4937eda7da7848bbb947786b7be0b27dab"
 dependencies = [
- "base64 0.21.1",
- "rustls-pki-types",
+ "ring 0.16.20",
+ "untrusted 0.7.1",
 ]

-[[package]]
-name = "rustls-pki-types"
-version = "1.0.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e7673e0aa20ee4937c6aacfc12bb8341cfbf054cdd21df6bec5fd0629fe9339b"
-
 [[package]]
 name = "rustls-webpki"
 version = "0.101.7"
@@ -4301,17 +4291,6 @@ dependencies = [
 "untrusted 0.9.0",
 ]

-[[package]]
-name = "rustls-webpki"
-version = "0.102.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "de2635c8bc2b88d367767c5de8ea1d8db9af3f6219eba28442242d9ab81d1b89"
-dependencies = [
- "ring 0.17.6",
- "rustls-pki-types",
- "untrusted 0.9.0",
-]
-
 [[package]]
 name = "rustversion"
 version = "1.0.12"
@@ -4352,7 +4331,7 @@ dependencies = [
 "serde_with",
 "thiserror",
 "tokio",
- "tokio-rustls 0.25.0",
+ "tokio-rustls",
 "tokio-stream",
 "tracing",
 "tracing-appender",
@@ -4516,7 +4495,7 @@ checksum = "2e95efd0cefa32028cdb9766c96de71d96671072f9fb494dc9fb84c0ef93e52b"
 dependencies = [
 "httpdate",
 "reqwest",
- "rustls 0.21.9",
+ "rustls",
 "sentry-backtrace",
 "sentry-contexts",
 "sentry-core",
@@ -4524,7 +4503,7 @@ dependencies = [
 "sentry-tracing",
 "tokio",
 "ureq",
- "webpki-roots",
+ "webpki-roots 0.25.2",
 ]

 [[package]]
@@ -5182,14 +5161,16 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"

 [[package]]
 name = "tls-listener"
-version = "0.9.0"
-source = "git+https://github.com/conradludgate/tls-listener?branch=main#4801141b5660613e77816044da6540aa64f388ec"
+version = "0.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "81294c017957a1a69794f506723519255879e15a870507faf45dfed288b763dd"
 dependencies = [
 "futures-util",
+ "hyper",
 "pin-project-lite",
 "thiserror",
 "tokio",
- "tokio-rustls 0.25.0",
+ "tokio-rustls",
 ]

 [[package]]
@@ -5272,10 +5253,10 @@ checksum = "dd5831152cb0d3f79ef5523b357319ba154795d64c7078b2daa95a803b54057f"
 dependencies = [
 "futures",
 "ring 0.16.20",
- "rustls 0.21.9",
+ "rustls",
 "tokio",
 "tokio-postgres",
- "tokio-rustls 0.24.0",
+ "tokio-rustls",
 ]

 [[package]]
@@ -5284,18 +5265,7 @@ version = "0.24.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e0d409377ff5b1e3ca6437aa86c1eb7d40c134bfec254e44c830defa92669db5"
 dependencies = [
- "rustls 0.21.9",
- "tokio",
-]
-
-[[package]]
-name = "tokio-rustls"
-version = "0.25.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "775e0c0f0adb3a2f22a00c4745d728b479985fc15ee7ca6a2608388c5569860f"
-dependencies = [
- "rustls 0.22.1",
- "rustls-pki-types",
+ "rustls",
 "tokio",
 ]

@@ -5442,9 +5412,9 @@ dependencies = [
 "pin-project",
 "prost",
 "rustls-native-certs",
- "rustls-pemfile 1.0.2",
+ "rustls-pemfile",
 "tokio",
- "tokio-rustls 0.24.0",
+ "tokio-rustls",
 "tokio-stream",
 "tower",
 "tower-layer",
@@ -5749,17 +5719,17 @@ checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1"

 [[package]]
 name = "ureq"
-version = "2.9.1"
+version = "2.7.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f8cdd25c339e200129fe4de81451814e5228c9b771d57378817d6117cc2b3f97"
+checksum = "0b11c96ac7ee530603dcdf68ed1557050f374ce55a5a07193ebf8cbc9f8927e9"
 dependencies = [
 "base64 0.21.1",
 "log",
 "once_cell",
- "rustls 0.21.9",
- "rustls-webpki 0.101.7",
+ "rustls",
+ "rustls-webpki 0.100.2",
 "url",
- "webpki-roots",
+ "webpki-roots 0.23.1",
 ]

 [[package]]
@@ -6068,6 +6038,15 @@ dependencies = [
 "wasm-bindgen",
 ]

+[[package]]
+name = "webpki-roots"
+version = "0.23.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b03058f88386e5ff5310d9111d53f48b17d732b401aeb83a8d5190f2ac459338"
+dependencies = [
+ "rustls-webpki 0.100.2",
+]
+
 [[package]]
 name = "webpki-roots"
 version = "0.25.2"
@@ -6316,8 +6295,11 @@ dependencies = [
 "either",
 "fail",
 "futures",
+ "futures-channel",
+ "futures-core",
 "futures-executor",
 "futures-io",
+ "futures-sink",
 "futures-util",
 "hex",
 "hmac",
@@ -6336,7 +6318,8 @@ dependencies = [
 "regex-automata 0.4.3",
 "regex-syntax 0.8.2",
 "reqwest",
- "rustls 0.21.9",
+ "ring 0.16.20",
+ "rustls",
 "scopeguard",
 "serde",
 "serde_json",
@@ -6347,7 +6330,7 @@ dependencies = [
 "time",
 "time-macros",
 "tokio",
- "tokio-rustls 0.24.0",
+ "tokio-rustls",
 "tokio-util",
 "toml_datetime",
 "toml_edit",
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -115,12 +115,11 @@ reqwest = { version = "0.11", default-features = false, features = ["rustls-tls"
 reqwest-tracing = { version = "0.4.0", features = ["opentelemetry_0_19"] }
 reqwest-middleware = "0.2.0"
 reqwest-retry = "0.2.2"
-ring = "0.17"
 routerify = "3"
 rpds = "0.13"
 rustc-hash = "1.1.0"
-rustls = "0.22.1"
-rustls-pemfile = "2.0.0"
+rustls = "0.21"
+rustls-pemfile = "1"
 rustls-split = "0.3"
 scopeguard = "1.1"
 sysinfo = "0.29.2"
@@ -144,11 +143,11 @@ tar = "0.4"
 task-local-extensions = "0.1.4"
 test-context = "0.1"
 thiserror = "1.0"
-tls-listener = { version = "0.9.0", features = ["rustls"] }
+tls-listener = { version = "0.7", features = ["rustls", "hyper-h1"] }
 tokio = { version = "1.17", features = ["macros"] }
 tokio-io-timeout = "1.2.0"
 tokio-postgres-rustls = "0.10.0"
-tokio-rustls = "0.25.0"
+tokio-rustls = "0.24"
 tokio-stream = "0.1"
 tokio-tar = "0.3"
 tokio-util = { version = "0.7.10", features = ["io", "rt"] }
@@ -203,7 +202,7 @@ workspace_hack = { version = "0.1", path = "./workspace_hack/" }

 ## Build dependencies
 criterion = "0.5.1"
-rcgen = "0.12"
+rcgen = "0.11"
 rstest = "0.18"
 camino-tempfile = "1.0.2"
 tonic-build = "0.9"
@@ -214,8 +213,6 @@ tonic-build = "0.9"
 # TODO: we should probably fork `tokio-postgres-rustls` instead.
 tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", branch="neon" }

-tls-listener = { git = "https://github.com/conradludgate/tls-listener", branch="main" }
-
 ################# Binary contents sections

 [profile.release]
--- a/compute_tools/src/spec.rs
+++ b/compute_tools/src/spec.rs
@@ -370,33 +370,49 @@ pub fn handle_role_deletions(spec: &ComputeSpec, connstr: &str, client: &mut Cli
    Ok(())
 }

+fn reassign_owned_objects_in_one_db(
+    conf: Config,
+    role_name: &PgIdent,
+    db_owner: &PgIdent,
+) -> Result<()> {
+    let mut client = conf.connect(NoTls)?;
+
+    // This will reassign all dependent objects to the db owner
+    let reassign_query = format!(
+        "REASSIGN OWNED BY {} TO {}",
+        role_name.pg_quote(),
+        db_owner.pg_quote()
+    );
+    info!(
+        "reassigning objects owned by '{}' in db '{}' to '{}'",
+        role_name,
+        conf.get_dbname().unwrap_or(""),
+        db_owner
+    );
+    client.simple_query(&reassign_query)?;
+
+    // This now will only drop privileges of the role
+    let drop_query = format!("DROP OWNED BY {}", role_name.pg_quote());
+    client.simple_query(&drop_query)?;
+    Ok(())
+}
+
 // Reassign all owned objects in all databases to the owner of the database.
 fn reassign_owned_objects(spec: &ComputeSpec, connstr: &str, role_name: &PgIdent) -> Result<()> {
    for db in &spec.cluster.databases {
        if db.owner != *role_name {
            let mut conf = Config::from_str(connstr)?;
            conf.dbname(&db.name);
-
-            let mut client = conf.connect(NoTls)?;
-
-            // This will reassign all dependent objects to the db owner
-            let reassign_query = format!(
-                "REASSIGN OWNED BY {} TO {}",
-                role_name.pg_quote(),
-                db.owner.pg_quote()
-            );
-            info!(
-                "reassigning objects owned by '{}' in db '{}' to '{}'",
-                role_name, &db.name, &db.owner
-            );
-            client.simple_query(&reassign_query)?;
-
-            // This now will only drop privileges of the role
-            let drop_query = format!("DROP OWNED BY {}", role_name.pg_quote());
-            client.simple_query(&drop_query)?;
+            reassign_owned_objects_in_one_db(conf, role_name, &db.owner)?;
        }
    }

+    // Also handle case when there are no databases in the spec.
+    // In this case we need to reassign objects in the default database.
+    let conf = Config::from_str(connstr)?;
+    let db_owner = PgIdent::from_str("cloud_admin")?;
+    reassign_owned_objects_in_one_db(conf, role_name, &db_owner)?;
+
    Ok(())
 }

--- a/libs/postgres_backend/Cargo.toml
+++ b/libs/postgres_backend/Cargo.toml
@@ -9,12 +9,10 @@ async-trait.workspace = true
 anyhow.workspace = true
 bytes.workspace = true
 futures.workspace = true
-ring.workspace = true
 rustls.workspace = true
 serde.workspace = true
 thiserror.workspace = true
 tokio.workspace = true
-tokio-postgres.workspace = true
 tokio-rustls.workspace = true
 tracing.workspace = true

@@ -24,4 +22,5 @@ workspace_hack.workspace = true
 [dev-dependencies]
 once_cell.workspace = true
 rustls-pemfile.workspace = true
-# tokio-postgres-rustls.workspace = true
+tokio-postgres.workspace = true
+tokio-postgres-rustls.workspace = true
--- a/libs/postgres_backend/src/lib.rs
+++ b/libs/postgres_backend/src/lib.rs
@@ -6,7 +6,7 @@
 #![deny(clippy::undocumented_unsafe_blocks)]
 use anyhow::Context;
 use bytes::Bytes;
-use futures::{pin_mut, TryFutureExt, FutureExt};
+use futures::pin_mut;
 use serde::{Deserialize, Serialize};
 use std::io::ErrorKind;
 use std::net::SocketAddr;
@@ -1030,115 +1030,3 @@ pub enum CopyStreamHandlerEnd {
    #[error(transparent)]
    Other(#[from] anyhow::Error),
 }
-
-#[derive(Clone)]
-pub struct MakeRustlsConnect {
-    config: Arc<rustls::ClientConfig>,
-}
-
-impl MakeRustlsConnect {
-    pub fn new(config: rustls::ClientConfig) -> Self {
-        Self {
-            config: Arc::new(config),
-        }
-    }
-}
-
-impl<S> tokio_postgres::tls::MakeTlsConnect<S> for MakeRustlsConnect
-where
-    S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
-{
-    type Stream = RustlsStream<S>;
-    type TlsConnect = RustlsConnect;
-    type Error = io::Error;
-
-    fn make_tls_connect(&mut self, hostname: &str) -> io::Result<RustlsConnect> {
-        rustls::pki_types::ServerName::try_from(hostname)
-            .map(|dns_name| {
-                RustlsConnect(Some(RustlsConnectData {
-                    hostname: dns_name.to_owned(),
-                    connector: Arc::clone(&self.config).into(),
-                }))
-            })
-            .or(Ok(RustlsConnect(None)))
-    }
-}
-
-pub struct RustlsConnect(Option<RustlsConnectData>);
-
-struct RustlsConnectData {
-    hostname: rustls::pki_types::ServerName<'static>,
-    connector: tokio_rustls::TlsConnector,
-}
-
-impl<S> tokio_postgres::tls::TlsConnect<S> for RustlsConnect
-where
-    S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
-{
-    type Stream = RustlsStream<S>;
-    type Error = io::Error;
-    type Future = Pin<Box<dyn Future<Output = io::Result<RustlsStream<S>>> + Send>>;
-
-    fn connect(self, stream: S) -> Self::Future {
-        match self.0 {
-            None => Box::pin(core::future::ready(Err(io::ErrorKind::InvalidInput.into()))),
-            Some(c) => c
-                .connector
-                .connect(c.hostname, stream)
-                .map_ok(|s| RustlsStream(Box::pin(s)))
-                .boxed(),
-        }
-    }
-}
-
-pub struct RustlsStream<S>(Pin<Box<tokio_rustls::client:: TlsStream<S>>>);
-
-impl<S> tokio_postgres::tls::TlsStream for RustlsStream<S>
-where
-    S: AsyncRead + AsyncWrite + Unpin,
-{
-    fn channel_binding(&self) -> tokio_postgres::tls::ChannelBinding {
-        let (_, session) = self.0.get_ref();
-        match session.peer_certificates() {
-            Some(certs) if !certs.is_empty() => {
-                let sha256 = ring::digest::digest(&ring::digest::SHA256, certs[0].as_ref());
-                tokio_postgres::tls::ChannelBinding::tls_server_end_point(sha256.as_ref().into())
-            }
-            _ => tokio_postgres::tls::ChannelBinding::none(),
-        }
-    }
-}
-
-impl<S> AsyncRead for RustlsStream<S>
-where
-    S: AsyncRead + AsyncWrite + Unpin,
-{
-    fn poll_read(
-        mut self: Pin<&mut Self>,
-        cx: &mut std::task:: Context,
-        buf: &mut tokio::io::ReadBuf<'_>,
-    ) -> Poll<tokio::io::Result<()>> {
-        self.0.as_mut().poll_read(cx, buf)
-    }
-}
-
-impl<S> AsyncWrite for RustlsStream<S>
-where
-    S: AsyncRead + AsyncWrite + Unpin,
-{
-    fn poll_write(
-        mut self: Pin<&mut Self>,
-        cx: &mut std::task:: Context,
-        buf: &[u8],
-    ) -> Poll<tokio::io::Result<usize>> {
-        self.0.as_mut().poll_write(cx, buf)
-    }
-
-    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut std::task:: Context) -> Poll<tokio::io::Result<()>> {
-        self.0.as_mut().poll_flush(cx)
-    }
-
-    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut std::task:: Context) -> Poll<tokio::io::Result<()>> {
-        self.0.as_mut().poll_shutdown(cx)
-    }
-}
--- a/libs/postgres_backend/tests/simple_select.rs
+++ b/libs/postgres_backend/tests/simple_select.rs
@@ -1,6 +1,5 @@
 /// Test postgres_backend_async with tokio_postgres
 use once_cell::sync::Lazy;
-use postgres_backend::MakeRustlsConnect;
 use postgres_backend::{AuthType, Handler, PostgresBackend, QueryError};
 use pq_proto::{BeMessage, RowDescriptor};
 use std::io::Cursor;
@@ -10,6 +9,7 @@ use tokio::net::{TcpListener, TcpStream};
 use tokio_postgres::config::SslMode;
 use tokio_postgres::tls::MakeTlsConnect;
 use tokio_postgres::{Config, NoTls, SimpleQueryMessage};
+use tokio_postgres_rustls::MakeRustlsConnect;

 // generate client, server test streams
 async fn make_tcp_pair() -> (TcpStream, TcpStream) {
@@ -72,21 +72,14 @@ async fn simple_select() {
    }
 }

-static KEY: Lazy<rustls::pki_types::PrivatePkcs1KeyDer<'static>> = Lazy::new(|| {
+static KEY: Lazy<rustls::PrivateKey> = Lazy::new(|| {
    let mut cursor = Cursor::new(include_bytes!("key.pem"));
-
-    let key = rustls_pemfile::rsa_private_keys(&mut cursor)
-        .next()
-        .unwrap()
-        .unwrap();
-    key.secret_pkcs1_der().to_owned().into()
+    rustls::PrivateKey(rustls_pemfile::rsa_private_keys(&mut cursor).unwrap()[0].clone())
 });

-static CERT: Lazy<rustls::pki_types::CertificateDer<'static>> = Lazy::new(|| {
+static CERT: Lazy<rustls::Certificate> = Lazy::new(|| {
    let mut cursor = Cursor::new(include_bytes!("cert.pem"));
-    let cert = rustls_pemfile::certs(&mut cursor).next().unwrap().unwrap();
-
-    cert.into_owned()
+    rustls::Certificate(rustls_pemfile::certs(&mut cursor).unwrap()[0].clone())
 });

 // test that basic select with ssl works
@@ -94,10 +87,10 @@ static CERT: Lazy<rustls::pki_types::CertificateDer<'static>> = Lazy::new(|| {
 async fn simple_select_ssl() {
    let (client_sock, server_sock) = make_tcp_pair().await;

-    let key = rustls::pki_types::PrivateKeyDer::Pkcs1(KEY.secret_pkcs1_der().to_owned().into());
    let server_cfg = rustls::ServerConfig::builder()
+        .with_safe_defaults()
        .with_no_client_auth()
-        .with_single_cert(vec![CERT.clone()], key)
+        .with_single_cert(vec![CERT.clone()], KEY.clone())
        .unwrap();
    let tls_config = Some(Arc::new(server_cfg));
    let pgbackend =
@@ -109,13 +102,14 @@ async fn simple_select_ssl() {
    });

    let client_cfg = rustls::ClientConfig::builder()
+        .with_safe_defaults()
        .with_root_certificates({
            let mut store = rustls::RootCertStore::empty();
-            store.add(CERT.clone()).unwrap();
+            store.add(&CERT).unwrap();
            store
        })
        .with_no_client_auth();
-    let mut make_tls_connect = MakeRustlsConnect::new(client_cfg);
+    let mut make_tls_connect = tokio_postgres_rustls::MakeRustlsConnect::new(client_cfg);
    let tls_connect = <MakeRustlsConnect as MakeTlsConnect<TcpStream>>::make_tls_connect(
        &mut make_tls_connect,
        "localhost",
--- a/pageserver/src/metrics.rs
+++ b/pageserver/src/metrics.rs
@@ -522,14 +522,18 @@ pub(crate) mod initial_logical_size {
    impl StartCalculation {
        pub(crate) fn first(&self, circumstances: StartCircumstances) -> OngoingCalculationGuard {
            let circumstances_label: &'static str = circumstances.into();
-            self.0.with_label_values(&["first", circumstances_label]);
+            self.0
+                .with_label_values(&["first", circumstances_label])
+                .inc();
            OngoingCalculationGuard {
                inc_drop_calculation: Some(DROP_CALCULATION.first.clone()),
            }
        }
        pub(crate) fn retry(&self, circumstances: StartCircumstances) -> OngoingCalculationGuard {
            let circumstances_label: &'static str = circumstances.into();
-            self.0.with_label_values(&["retry", circumstances_label]);
+            self.0
+                .with_label_values(&["retry", circumstances_label])
+                .inc();
            OngoingCalculationGuard {
                inc_drop_calculation: Some(DROP_CALCULATION.retry.clone()),
            }
--- a/pgxn/neon/walproposer_pg.c
+++ b/pgxn/neon/walproposer_pg.c
@@ -1712,11 +1712,29 @@ walprop_pg_after_election(WalProposer *wp)
 	f = fopen("restart.lsn", "rb");
 	if (f != NULL && !wp->config->syncSafekeepers)
 	{
-		fread(&lrRestartLsn, sizeof(lrRestartLsn), 1, f);
+		size_t rc = fread(&lrRestartLsn, sizeof(lrRestartLsn), 1, f);
 		fclose(f);
-		if (lrRestartLsn != InvalidXLogRecPtr)
+		if (rc == 1 && lrRestartLsn != InvalidXLogRecPtr)
 		{
-			elog(LOG, "Logical replication restart LSN %X/%X", LSN_FORMAT_ARGS(lrRestartLsn));
+			uint64		download_range_mb;
+
+			elog(LOG, "Logical replication restart LSN %X/%X, epochStartLsn %X/%X, max_slot_wal_keep_size_mb=%d",
+				LSN_FORMAT_ARGS(lrRestartLsn), LSN_FORMAT_ARGS(wp->propEpochStartLsn), max_slot_wal_keep_size_mb);
+
+			/*
+			* If we need to download more than a max_slot_wal_keep_size, cap to it to
+			* avoid risk of exploding pg_wal. Logical replication won't work until
+			* recreated, but at least compute would start; this also follows
+			* max_slot_wal_keep_size semantics.
+			*/
+			download_range_mb = (wp->propEpochStartLsn - lrRestartLsn) / 1024 / 1024;
+			if (max_slot_wal_keep_size_mb > 0 && download_range_mb >= max_slot_wal_keep_size_mb)
+			{
+				lrRestartLsn = wp->propEpochStartLsn - max_slot_wal_keep_size_mb * 1024 * 1024;
+				elog(WARNING, "capped WAL download for logical replication to %X/%X as max_slot_wal_keep_size=%dMB",
+							LSN_FORMAT_ARGS(lrRestartLsn), max_slot_wal_keep_size_mb);
+			}
+

 			/*
 			 * start from the beginning of the segment to fetch page headers
--- a/proxy/src/bin/pg_sni_router.rs
+++ b/proxy/src/bin/pg_sni_router.rs
@@ -6,6 +6,7 @@
 use std::{net::SocketAddr, sync::Arc};

 use futures::future::Either;
+use itertools::Itertools;
 use proxy::config::TlsServerEndPoint;
 use proxy::proxy::run_until_cancelled;
 use tokio::net::TcpListener;
@@ -75,12 +76,10 @@ async fn main() -> anyhow::Result<()> {
            let key = {
                let key_bytes = std::fs::read(key_path).context("TLS key file")?;
                let mut keys = rustls_pemfile::pkcs8_private_keys(&mut &key_bytes[..])
-                    .collect::<Result<Vec<_>, _>>()
                    .context(format!("Failed to read TLS keys at '{key_path}'"))?;

                ensure!(keys.len() == 1, "keys.len() = {} (should be 1)", keys.len());
-                let bytes = keys.pop().unwrap().secret_pkcs8_der().to_owned();
-                rustls::pki_types::PrivateKeyDer::Pkcs1(bytes.into())
+                keys.pop().map(rustls::PrivateKey).unwrap()
            };

            let cert_chain_bytes = std::fs::read(cert_path)
@@ -88,23 +87,25 @@ async fn main() -> anyhow::Result<()> {

            let cert_chain = {
                rustls_pemfile::certs(&mut &cert_chain_bytes[..])
-                .collect::<Result<Vec<_>,_>>()
                    .context(format!(
                        "Failed to read TLS certificate chain from bytes from file at '{cert_path}'."
                    ))?
+                    .into_iter()
+                    .map(rustls::Certificate)
+                    .collect_vec()
            };

            // needed for channel bindings
            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_protocol_versions(&[
-                &rustls::version::TLS13,
-                &rustls::version::TLS12,
-            ])
-            .with_no_client_auth()
-            .with_single_cert(cert_chain, key)?
-            .into();
+            let tls_config = rustls::ServerConfig::builder()
+                .with_safe_default_cipher_suites()
+                .with_safe_default_kx_groups()
+                .with_protocol_versions(&[&rustls::version::TLS13, &rustls::version::TLS12])?
+                .with_no_client_auth()
+                .with_single_cert(cert_chain, key)?
+                .into();

            (tls_config, tls_server_end_point)
        }
--- a/proxy/src/config.rs
+++ b/proxy/src/config.rs
@@ -1,9 +1,6 @@
 use crate::{auth, rate_limiter::RateBucketInfo};
 use anyhow::{bail, ensure, Context, Ok};
-use rustls::{
-    crypto::ring::sign,
-    pki_types::{CertificateDer, PrivateKeyDer},
-};
+use rustls::{sign, Certificate, PrivateKey};
 use sha2::{Digest, Sha256};
 use std::{
    collections::{HashMap, HashSet},
@@ -88,14 +85,14 @@ pub fn configure_tls(

    let cert_resolver = Arc::new(cert_resolver);

-    // allow TLS 1.2 to be compatible with older client libraries
-    let config = rustls::ServerConfig::builder_with_protocol_versions(&[
-        &rustls::version::TLS13,
-        &rustls::version::TLS12,
-    ])
-    .with_no_client_auth()
-    .with_cert_resolver(cert_resolver.clone())
-    .into();
+    let config = rustls::ServerConfig::builder()
+        .with_safe_default_cipher_suites()
+        .with_safe_default_kx_groups()
+        // allow TLS 1.2 to be compatible with older client libraries
+        .with_protocol_versions(&[&rustls::version::TLS13, &rustls::version::TLS12])?
+        .with_no_client_auth()
+        .with_cert_resolver(cert_resolver.clone())
+        .into();

    Ok(TlsConfig {
        config,
@@ -133,14 +130,14 @@ pub enum TlsServerEndPoint {
 }

 impl TlsServerEndPoint {
-    pub fn new(cert: &CertificateDer) -> anyhow::Result<Self> {
+    pub fn new(cert: &Certificate) -> anyhow::Result<Self> {
        let sha256_oids = [
            // I'm explicitly not adding MD5 or SHA1 here... They're bad.
            oid_registry::OID_SIG_ECDSA_WITH_SHA256,
            oid_registry::OID_PKCS1_SHA256WITHRSA,
        ];

-        let pem = x509_parser::parse_x509_certificate(cert)
+        let pem = x509_parser::parse_x509_certificate(&cert.0)
            .context("Failed to parse PEM object from cerficiate")?
            .1;

@@ -150,7 +147,8 @@ impl TlsServerEndPoint {
        let oid = pem.signature_algorithm.oid();
        let alg = reg.get(oid);
        if sha256_oids.contains(oid) {
-            let tls_server_end_point: [u8; 32] = Sha256::new().chain_update(cert).finalize().into();
+            let tls_server_end_point: [u8; 32] =
+                Sha256::new().chain_update(&cert.0).finalize().into();
            info!(subject = %pem.subject, signature_algorithm = alg.map(|a| a.description()), tls_server_end_point = %base64::encode(tls_server_end_point), "determined channel binding");
            Ok(Self::Sha256(tls_server_end_point))
        } else {
@@ -164,7 +162,7 @@ impl TlsServerEndPoint {
    }
 }

-#[derive(Default, Debug)]
+#[derive(Default)]
 pub struct CertResolver {
    certs: HashMap<String, (Arc<rustls::sign::CertifiedKey>, TlsServerEndPoint)>,
    default: Option<(Arc<rustls::sign::CertifiedKey>, TlsServerEndPoint)>,
@@ -184,12 +182,11 @@ impl CertResolver {
        let priv_key = {
            let key_bytes = std::fs::read(key_path)
                .context(format!("Failed to read TLS keys at '{key_path}'"))?;
-            let keys: Result<Vec<_>, _> =
-                rustls_pemfile::pkcs8_private_keys(&mut &key_bytes[..]).collect();
-            let mut keys = keys.context(format!("Failed to parse TLS keys at '{key_path}'"))?;
+            let mut keys = rustls_pemfile::pkcs8_private_keys(&mut &key_bytes[..])
+                .context(format!("Failed to parse TLS keys at '{key_path}'"))?;

            ensure!(keys.len() == 1, "keys.len() = {} (should be 1)", keys.len());
-            keys.pop().unwrap()
+            keys.pop().map(rustls::PrivateKey).unwrap()
        };

        let cert_chain_bytes = std::fs::read(cert_path)
@@ -197,28 +194,30 @@ impl CertResolver {

        let cert_chain = {
            rustls_pemfile::certs(&mut &cert_chain_bytes[..])
-                .collect::<Result<Vec<_>, _>>()
                .with_context(|| {
                    format!(
                    "Failed to read TLS certificate chain from bytes from file at '{cert_path}'."
                )
                })?
+                .into_iter()
+                .map(rustls::Certificate)
+                .collect()
        };

-        self.add_cert(PrivateKeyDer::Pkcs8(priv_key), cert_chain, is_default)
+        self.add_cert(priv_key, cert_chain, is_default)
    }

    pub fn add_cert(
        &mut self,
-        priv_key: PrivateKeyDer,
-        cert_chain: Vec<CertificateDer<'static>>,
+        priv_key: PrivateKey,
+        cert_chain: Vec<Certificate>,
        is_default: bool,
    ) -> anyhow::Result<()> {
        let key = sign::any_supported_type(&priv_key).context("invalid private key")?;

        let first_cert = &cert_chain[0];
        let tls_server_end_point = TlsServerEndPoint::new(first_cert)?;
-        let pem = x509_parser::parse_x509_certificate(first_cert)
+        let pem = x509_parser::parse_x509_certificate(&first_cert.0)
            .context("Failed to parse PEM object from cerficiate")?
            .1;

--- a/proxy/src/protocol2.rs
+++ b/proxy/src/protocol2.rs
@@ -328,23 +328,19 @@ impl<T: AsyncRead> AsyncRead for WithClientIp<T> {

 impl AsyncAccept for ProxyProtocolAccept {
    type Connection = WithClientIp<AddrStream>;
-    type Address = std::net::SocketAddr;
+
    type Error = io::Error;

    fn poll_accept(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
-    ) -> Poll<Result<(Self::Connection, Self::Address), Self::Error>> {
-        use hyper::server::accept::Accept;
+    ) -> Poll<Option<Result<Self::Connection, Self::Error>>> {
        let conn = ready!(Pin::new(&mut self.incoming).poll_accept(cx)?);
        let Some(conn) = conn else {
-            return Poll::Ready(Err(io::Error::new(
-                io::ErrorKind::NotConnected,
-                "no incoming connection?",
-            )));
+            return Poll::Ready(None);
        };
-        let addr = conn.remote_addr();
-        Poll::Ready(Ok((WithClientIp::new(conn), addr)))
+
+        Poll::Ready(Some(Ok(WithClientIp::new(conn))))
    }
 }

--- a/proxy/src/proxy/tests.rs
+++ b/proxy/src/proxy/tests.rs
@@ -11,21 +11,16 @@ use crate::console::{CachedNodeInfo, NodeInfo};
 use crate::proxy::retry::{retry_after, NUM_RETRIES_CONNECT};
 use crate::{auth, http, sasl, scram};
 use async_trait::async_trait;
-use postgres_backend::{MakeRustlsConnect, RustlsStream};
 use rstest::rstest;
-use rustls::pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs1KeyDer};
 use tokio_postgres::config::SslMode;
 use tokio_postgres::tls::{MakeTlsConnect, NoTls};
+use tokio_postgres_rustls::{MakeRustlsConnect, RustlsStream};

 /// Generate a set of TLS certificates: CA + server.
 fn generate_certs(
    hostname: &str,
    common_name: &str,
-) -> anyhow::Result<(
-    CertificateDer<'static>,
-    CertificateDer<'static>,
-    PrivateKeyDer<'static>,
-)> {
+) -> anyhow::Result<(rustls::Certificate, rustls::Certificate, rustls::PrivateKey)> {
    let ca = rcgen::Certificate::from_params({
        let mut params = rcgen::CertificateParams::default();
        params.is_ca = rcgen::IsCa::Ca(rcgen::BasicConstraints::Unconstrained);
@@ -42,9 +37,9 @@ fn generate_certs(
    })?;

    Ok((
-        CertificateDer::from(ca.serialize_der()?),
-        CertificateDer::from(cert.serialize_der_with_signer(&ca)?),
-        PrivateKeyDer::Pkcs1(PrivatePkcs1KeyDer::from(cert.serialize_private_key_der())),
+        rustls::Certificate(ca.serialize_der()?),
+        rustls::Certificate(cert.serialize_der_with_signer(&ca)?),
+        rustls::PrivateKey(cert.serialize_private_key_der()),
    ))
 }

@@ -78,10 +73,10 @@ fn generate_tls_config<'a>(
    let (ca, cert, key) = generate_certs(hostname, common_name)?;

    let tls_config = {
-        let key_clone = rustls::pki_types::PrivateKeyDer::Pkcs1(key.secret_der().to_owned().into());
        let config = rustls::ServerConfig::builder()
+            .with_safe_defaults()
            .with_no_client_auth()
-            .with_single_cert(vec![cert.clone()], key_clone)?
+            .with_single_cert(vec![cert.clone()], key.clone())?
            .into();

        let mut cert_resolver = CertResolver::new();
@@ -98,9 +93,10 @@ fn generate_tls_config<'a>(

    let client_config = {
        let config = rustls::ClientConfig::builder()
+            .with_safe_defaults()
            .with_root_certificates({
                let mut store = rustls::RootCertStore::empty();
-                store.add(ca)?;
+                store.add(&ca)?;
                store
            })
            .with_no_client_auth();
--- a/proxy/src/serverless.rs
+++ b/proxy/src/serverless.rs
@@ -77,19 +77,14 @@ pub async fn task_main(
    let ws_connections = tokio_util::task::task_tracker::TaskTracker::new();
    ws_connections.close(); // allows `ws_connections.wait to complete`

-    let tls_listener = TlsListener::new(tls_acceptor, addr_incoming)
-        .map(|x| match x {
-            Ok((conn, _)) => Ok(conn),
-            Err(e) => Err(e),
-        })
-        .filter(|conn| {
-            if let Err(err) = conn {
-                error!("failed to accept TLS connection for websockets: {err:?}");
-                ready(false)
-            } else {
-                ready(true)
-            }
-        });
+    let tls_listener = TlsListener::new(tls_acceptor, addr_incoming).filter(|conn| {
+        if let Err(err) = conn {
+            error!("failed to accept TLS connection for websockets: {err:?}");
+            ready(false)
+        } else {
+            ready(true)
+        }
+    });

    let make_svc = hyper::service::make_service_fn(
        |stream: &tokio_rustls::server::TlsStream<WithClientIp<AddrStream>>| {
--- a/storage_broker/benches/rps.rs
+++ b/storage_broker/benches/rps.rs
@@ -3,9 +3,12 @@ use std::sync::Arc;
 use std::time::{Duration, Instant};

 use clap::Parser;
-use storage_broker::proto::subscribe_safekeeper_info_request::SubscriptionKey;
-use storage_broker::proto::TenantTimelineId as ProtoTenantTimelineId;
-use storage_broker::proto::{SafekeeperTimelineInfo, SubscribeSafekeeperInfoRequest};
+
+use storage_broker::proto::SafekeeperTimelineInfo;
+use storage_broker::proto::{
+    FilterTenantTimelineId, MessageType, SubscribeByFilterRequest,
+    TenantTimelineId as ProtoTenantTimelineId, TypeSubscription, TypedMessage,
+};

 use storage_broker::{BrokerClientChannel, DEFAULT_ENDPOINT};
 use tokio::time;
@@ -91,15 +94,23 @@ async fn subscribe(client: Option<BrokerClientChannel>, counter: Arc<AtomicU64>,
        None => storage_broker::connect(DEFAULT_ENDPOINT, Duration::from_secs(5)).unwrap(),
    };

-    let key = SubscriptionKey::TenantTimelineId(ProtoTenantTimelineId {
+    let ttid = ProtoTenantTimelineId {
        tenant_id: vec![0xFF; 16],
        timeline_id: tli_from_u64(i),
-    });
-    let request = SubscribeSafekeeperInfoRequest {
-        subscription_key: Some(key),
    };
-    let mut stream = client
-        .subscribe_safekeeper_info(request)
+
+    let request = SubscribeByFilterRequest {
+        types: vec![TypeSubscription {
+            r#type: MessageType::SafekeeperTimelineInfo.into(),
+        }],
+        tenant_timeline_id: Some(FilterTenantTimelineId {
+            enabled: true,
+            tenant_timeline_id: Some(ttid),
+        }),
+    };
+
+    let mut stream: tonic::Streaming<TypedMessage> = client
+        .subscribe_by_filter(request)
        .await
        .unwrap()
        .into_inner();
--- a/storage_broker/proto/broker.proto
+++ b/storage_broker/proto/broker.proto
@@ -10,6 +10,12 @@ service BrokerService {

    // Publish safekeeper updates.
    rpc PublishSafekeeperInfo(stream SafekeeperTimelineInfo) returns (google.protobuf.Empty) {};
+
+    // Subscribe to all messages, limited by a filter.
+    rpc SubscribeByFilter(SubscribeByFilterRequest) returns (stream TypedMessage) {};
+
+    // Publish one message.
+    rpc PublishOne(TypedMessage) returns (google.protobuf.Empty) {};
 }

 message SubscribeSafekeeperInfoRequest {
@@ -48,3 +54,55 @@ message TenantTimelineId {
    bytes tenant_id = 1;
    bytes timeline_id = 2;
 }
+
+message FilterTenantTimelineId {
+    // If true, only messages related to `tenant_timeline_id` will be emitted.
+    // Otherwise, messages for all timelines will be emitted.
+    bool enabled = 1;
+    TenantTimelineId tenant_timeline_id = 2;
+}
+
+message TypeSubscription {
+    MessageType type = 1;
+}
+
+message SubscribeByFilterRequest {
+    // Subscription will emit messages only of the specified types. You need to specify
+    // at least one type to receive any messages.
+    repeated TypeSubscription types = 1;
+
+    // If set and enabled, subscription will emit messages only for the specified tenant/timeline.
+    optional FilterTenantTimelineId tenant_timeline_id = 2;
+}
+
+enum MessageType {
+    UNKNOWN = 0;
+    SAFEKEEPER_TIMELINE_INFO = 2;
+    SAFEKEEPER_DISCOVERY_REQUEST = 3;
+    SAFEKEEPER_DISCOVERY_RESPONSE = 4;
+}
+
+// A message with a type.
+message TypedMessage {
+    MessageType type = 1;
+
+    optional SafekeeperTimelineInfo safekeeper_timeline_info = 2;
+    optional SafekeeperDiscoveryRequest safekeeper_discovery_request = 3;
+    optional SafekeeperDiscoveryResponse safekeeper_discovery_response = 4;
+}
+
+message SafekeeperDiscoveryRequest {
+    TenantTimelineId tenant_timeline_id = 1;
+}
+
+// Shorter version of SafekeeperTimelineInfo, contains only necessary fields.
+message SafekeeperDiscoveryResponse {
+    uint64 safekeeper_id = 1;
+    TenantTimelineId tenant_timeline_id = 2;
+    // WAL available to download.
+    uint64 commit_lsn = 3;
+    // A connection string to use for WAL downloading.
+    string safekeeper_connstr = 4;
+    // Availability zone of a safekeeper.
+    optional string availability_zone = 5;
+}
--- a/storage_broker/src/bin/storage_broker.rs
+++ b/storage_broker/src/bin/storage_broker.rs
@@ -35,10 +35,16 @@ use tracing::*;
 use utils::signals::ShutdownSignals;

 use metrics::{Encoder, TextEncoder};
-use storage_broker::metrics::{NUM_PUBS, NUM_SUBS_ALL, NUM_SUBS_TIMELINE};
+use storage_broker::metrics::{
+    BROADCASTED_MESSAGES_TOTAL, BROADCAST_DROPPED_MESSAGES_TOTAL, NUM_PUBS, NUM_SUBS_ALL,
+    NUM_SUBS_TIMELINE, PROCESSED_MESSAGES_TOTAL, PUBLISHED_ONEOFF_MESSAGES_TOTAL,
+};
 use storage_broker::proto::broker_service_server::{BrokerService, BrokerServiceServer};
 use storage_broker::proto::subscribe_safekeeper_info_request::SubscriptionKey as ProtoSubscriptionKey;
-use storage_broker::proto::{SafekeeperTimelineInfo, SubscribeSafekeeperInfoRequest};
+use storage_broker::proto::{
+    FilterTenantTimelineId, MessageType, SafekeeperDiscoveryRequest, SafekeeperDiscoveryResponse,
+    SafekeeperTimelineInfo, SubscribeByFilterRequest, SubscribeSafekeeperInfoRequest, TypedMessage,
+};
 use storage_broker::{
    parse_proto_ttid, EitherBody, DEFAULT_KEEPALIVE_INTERVAL, DEFAULT_LISTEN_ADDR,
 };
@@ -73,8 +79,103 @@ struct Args {
    log_format: String,
 }

-type PubId = u64; // id of publisher for registering in maps
-type SubId = u64; // id of subscriber for registering in maps
+/// Id of publisher for registering in maps
+type PubId = u64;
+
+/// Id of subscriber for registering in maps
+type SubId = u64;
+
+/// Single enum type for all messages.
+#[derive(Clone, Debug, PartialEq)]
+#[allow(clippy::enum_variant_names)]
+enum Message {
+    SafekeeperTimelineInfo(SafekeeperTimelineInfo),
+    SafekeeperDiscoveryRequest(SafekeeperDiscoveryRequest),
+    SafekeeperDiscoveryResponse(SafekeeperDiscoveryResponse),
+}
+
+impl Message {
+    /// Convert proto message to internal message.
+    pub fn from(proto_msg: TypedMessage) -> Result<Self, Status> {
+        match proto_msg.r#type() {
+            MessageType::SafekeeperTimelineInfo => Ok(Message::SafekeeperTimelineInfo(
+                proto_msg.safekeeper_timeline_info.ok_or_else(|| {
+                    Status::new(Code::InvalidArgument, "missing safekeeper_timeline_info")
+                })?,
+            )),
+            MessageType::SafekeeperDiscoveryRequest => Ok(Message::SafekeeperDiscoveryRequest(
+                proto_msg.safekeeper_discovery_request.ok_or_else(|| {
+                    Status::new(
+                        Code::InvalidArgument,
+                        "missing safekeeper_discovery_request",
+                    )
+                })?,
+            )),
+            MessageType::SafekeeperDiscoveryResponse => Ok(Message::SafekeeperDiscoveryResponse(
+                proto_msg.safekeeper_discovery_response.ok_or_else(|| {
+                    Status::new(
+                        Code::InvalidArgument,
+                        "missing safekeeper_discovery_response",
+                    )
+                })?,
+            )),
+            MessageType::Unknown => Err(Status::new(
+                Code::InvalidArgument,
+                format!("invalid message type: {:?}", proto_msg.r#type),
+            )),
+        }
+    }
+
+    /// Get the tenant_timeline_id from the message.
+    pub fn tenant_timeline_id(&self) -> Result<Option<TenantTimelineId>, Status> {
+        match self {
+            Message::SafekeeperTimelineInfo(msg) => Ok(msg
+                .tenant_timeline_id
+                .as_ref()
+                .map(parse_proto_ttid)
+                .transpose()?),
+            Message::SafekeeperDiscoveryRequest(msg) => Ok(msg
+                .tenant_timeline_id
+                .as_ref()
+                .map(parse_proto_ttid)
+                .transpose()?),
+            Message::SafekeeperDiscoveryResponse(msg) => Ok(msg
+                .tenant_timeline_id
+                .as_ref()
+                .map(parse_proto_ttid)
+                .transpose()?),
+        }
+    }
+
+    /// Convert internal message to the protobuf struct.
+    pub fn as_typed_message(&self) -> TypedMessage {
+        let mut res = TypedMessage {
+            r#type: self.message_type() as i32,
+            ..Default::default()
+        };
+        match self {
+            Message::SafekeeperTimelineInfo(msg) => {
+                res.safekeeper_timeline_info = Some(msg.clone())
+            }
+            Message::SafekeeperDiscoveryRequest(msg) => {
+                res.safekeeper_discovery_request = Some(msg.clone())
+            }
+            Message::SafekeeperDiscoveryResponse(msg) => {
+                res.safekeeper_discovery_response = Some(msg.clone())
+            }
+        }
+        res
+    }
+
+    /// Get the message type.
+    pub fn message_type(&self) -> MessageType {
+        match self {
+            Message::SafekeeperTimelineInfo(_) => MessageType::SafekeeperTimelineInfo,
+            Message::SafekeeperDiscoveryRequest(_) => MessageType::SafekeeperDiscoveryRequest,
+            Message::SafekeeperDiscoveryResponse(_) => MessageType::SafekeeperDiscoveryResponse,
+        }
+    }
+}

 #[derive(Copy, Clone, Debug)]
 enum SubscriptionKey {
@@ -83,7 +184,7 @@ enum SubscriptionKey {
 }

 impl SubscriptionKey {
-    // Parse protobuf subkey (protobuf doesn't have fixed size bytes, we get vectors).
+    /// Parse protobuf subkey (protobuf doesn't have fixed size bytes, we get vectors).
    pub fn from_proto_subscription_key(key: ProtoSubscriptionKey) -> Result<Self, Status> {
        match key {
            ProtoSubscriptionKey::All(_) => Ok(SubscriptionKey::All),
@@ -92,14 +193,29 @@ impl SubscriptionKey {
            }
        }
    }
+
+    /// Parse from FilterTenantTimelineId
+    pub fn from_proto_filter_tenant_timeline_id(
+        f: &FilterTenantTimelineId,
+    ) -> Result<Self, Status> {
+        if !f.enabled {
+            return Ok(SubscriptionKey::All);
+        }
+
+        let ttid =
+            parse_proto_ttid(f.tenant_timeline_id.as_ref().ok_or_else(|| {
+                Status::new(Code::InvalidArgument, "missing tenant_timeline_id")
+            })?)?;
+        Ok(SubscriptionKey::Timeline(ttid))
+    }
 }

-// Channel to timeline subscribers.
+/// Channel to timeline subscribers.
 struct ChanToTimelineSub {
-    chan: broadcast::Sender<SafekeeperTimelineInfo>,
-    // Tracked separately to know when delete the shmem entry. receiver_count()
-    // is unhandy for that as unregistering and dropping the receiver side
-    // happens at different moments.
+    chan: broadcast::Sender<Message>,
+    /// Tracked separately to know when delete the shmem entry. receiver_count()
+    /// is unhandy for that as unregistering and dropping the receiver side
+    /// happens at different moments.
    num_subscribers: u64,
 }

@@ -110,7 +226,7 @@ struct SharedState {
    num_subs_to_timelines: i64,
    chans_to_timeline_subs: HashMap<TenantTimelineId, ChanToTimelineSub>,
    num_subs_to_all: i64,
-    chan_to_all_subs: broadcast::Sender<SafekeeperTimelineInfo>,
+    chan_to_all_subs: broadcast::Sender<Message>,
 }

 impl SharedState {
@@ -146,7 +262,7 @@ impl SharedState {
        &mut self,
        sub_key: SubscriptionKey,
        timeline_chan_size: usize,
-    ) -> (SubId, broadcast::Receiver<SafekeeperTimelineInfo>) {
+    ) -> (SubId, broadcast::Receiver<Message>) {
        let sub_id = self.next_sub_id;
        self.next_sub_id += 1;
        let sub_rx = match sub_key {
@@ -262,6 +378,29 @@ impl Registry {
            subscriber.id, subscriber.key, subscriber.remote_addr
        );
    }
+
+    /// Send msg to relevant subscribers.
+    pub fn send_msg(&self, msg: &Message) -> Result<(), Status> {
+        PROCESSED_MESSAGES_TOTAL.inc();
+
+        // send message to subscribers for everything
+        let shared_state = self.shared_state.read();
+        // Err means there is no subscribers, it is fine.
+        shared_state.chan_to_all_subs.send(msg.clone()).ok();
+
+        // send message to per timeline subscribers, if there is ttid
+        let ttid = msg.tenant_timeline_id()?;
+        if let Some(ttid) = ttid {
+            if let Some(subs) = shared_state.chans_to_timeline_subs.get(&ttid) {
+                // Err can't happen here, as tx is destroyed only after removing
+                // from the map the last subscriber along with tx.
+                subs.chan
+                    .send(msg.clone())
+                    .expect("rx is still in the map with zero subscribers");
+            }
+        }
+        Ok(())
+    }
 }

 // Private subscriber state.
@@ -269,7 +408,7 @@ struct Subscriber {
    id: SubId,
    key: SubscriptionKey,
    // Subscriber receives messages from publishers here.
-    sub_rx: broadcast::Receiver<SafekeeperTimelineInfo>,
+    sub_rx: broadcast::Receiver<Message>,
    // to unregister itself from shared state in Drop
    registry: Registry,
    // for logging
@@ -291,26 +430,9 @@ struct Publisher {
 }

 impl Publisher {
-    // Send msg to relevant subscribers.
-    pub fn send_msg(&mut self, msg: &SafekeeperTimelineInfo) -> Result<(), Status> {
-        // send message to subscribers for everything
-        let shared_state = self.registry.shared_state.read();
-        // Err means there is no subscribers, it is fine.
-        shared_state.chan_to_all_subs.send(msg.clone()).ok();
-
-        // send message to per timeline subscribers
-        let ttid =
-            parse_proto_ttid(msg.tenant_timeline_id.as_ref().ok_or_else(|| {
-                Status::new(Code::InvalidArgument, "missing tenant_timeline_id")
-            })?)?;
-        if let Some(subs) = shared_state.chans_to_timeline_subs.get(&ttid) {
-            // Err can't happen here, as tx is destroyed only after removing
-            // from the map the last subscriber along with tx.
-            subs.chan
-                .send(msg.clone())
-                .expect("rx is still in the map with zero subscribers");
-        }
-        Ok(())
+    /// Send msg to relevant subscribers.
+    pub fn send_msg(&mut self, msg: &Message) -> Result<(), Status> {
+        self.registry.send_msg(msg)
    }
 }

@@ -339,7 +461,7 @@ impl BrokerService for Broker {

        loop {
            match stream.next().await {
-                Some(Ok(msg)) => publisher.send_msg(&msg)?,
+                Some(Ok(msg)) => publisher.send_msg(&Message::SafekeeperTimelineInfo(msg))?,
                Some(Err(e)) => return Err(e), // grpc error from the stream
                None => break,                 // closed stream
            }
@@ -371,8 +493,15 @@ impl BrokerService for Broker {
            let mut missed_msgs: u64 = 0;
            loop {
                match subscriber.sub_rx.recv().await {
-                    Ok(info) => yield info,
+                    Ok(info) => {
+                        match info {
+                            Message::SafekeeperTimelineInfo(info) => yield info,
+                            _ => {},
+                        }
+                        BROADCASTED_MESSAGES_TOTAL.inc();
+                    },
                    Err(RecvError::Lagged(skipped_msg)) => {
+                        BROADCAST_DROPPED_MESSAGES_TOTAL.inc_by(skipped_msg);
                        missed_msgs += skipped_msg;
                        if (futures::poll!(Box::pin(warn_interval.tick()))).is_ready() {
                            warn!("subscription id={}, key={:?} addr={:?} dropped {} messages, channel is full",
@@ -392,6 +521,78 @@ impl BrokerService for Broker {
            Box::pin(output) as Self::SubscribeSafekeeperInfoStream
        ))
    }
+
+    type SubscribeByFilterStream =
+        Pin<Box<dyn Stream<Item = Result<TypedMessage, Status>> + Send + 'static>>;
+
+    /// Subscribe to all messages, limited by a filter.
+    async fn subscribe_by_filter(
+        &self,
+        request: Request<SubscribeByFilterRequest>,
+    ) -> std::result::Result<Response<Self::SubscribeByFilterStream>, Status> {
+        let remote_addr = request
+            .remote_addr()
+            .expect("TCPConnectInfo inserted by handler");
+        let proto_filter = request.into_inner();
+        let ttid_filter = proto_filter
+            .tenant_timeline_id
+            .as_ref()
+            .ok_or_else(|| Status::new(Code::InvalidArgument, "missing tenant_timeline_id"))?;
+
+        let sub_key = SubscriptionKey::from_proto_filter_tenant_timeline_id(ttid_filter)?;
+        let types_set = proto_filter
+            .types
+            .iter()
+            .map(|t| t.r#type)
+            .collect::<std::collections::HashSet<_>>();
+
+        let mut subscriber = self.registry.register_subscriber(sub_key, remote_addr);
+
+        // transform rx into stream with item = Result, as method result demands
+        let output = async_stream::try_stream! {
+            let mut warn_interval = time::interval(Duration::from_millis(1000));
+            let mut missed_msgs: u64 = 0;
+            loop {
+                match subscriber.sub_rx.recv().await {
+                    Ok(msg) => {
+                        let msg_type = msg.message_type() as i32;
+                        if types_set.contains(&msg_type) {
+                            yield msg.as_typed_message();
+                            BROADCASTED_MESSAGES_TOTAL.inc();
+                        }
+                    },
+                    Err(RecvError::Lagged(skipped_msg)) => {
+                        BROADCAST_DROPPED_MESSAGES_TOTAL.inc_by(skipped_msg);
+                        missed_msgs += skipped_msg;
+                        if (futures::poll!(Box::pin(warn_interval.tick()))).is_ready() {
+                            warn!("subscription id={}, key={:?} addr={:?} dropped {} messages, channel is full",
+                                subscriber.id, subscriber.key, subscriber.remote_addr, missed_msgs);
+                            missed_msgs = 0;
+                        }
+                    }
+                    Err(RecvError::Closed) => {
+                        // can't happen, we never drop the channel while there is a subscriber
+                        Err(Status::new(Code::Internal, "channel unexpectantly closed"))?;
+                    }
+                }
+            }
+        };
+
+        Ok(Response::new(
+            Box::pin(output) as Self::SubscribeByFilterStream
+        ))
+    }
+
+    /// Publish one message.
+    async fn publish_one(
+        &self,
+        request: Request<TypedMessage>,
+    ) -> std::result::Result<Response<()>, Status> {
+        let msg = Message::from(request.into_inner())?;
+        PUBLISHED_ONEOFF_MESSAGES_TOTAL.inc();
+        self.registry.send_msg(&msg)?;
+        Ok(Response::new(()))
+    }
 }

 // We serve only metrics and healthcheck through http1.
@@ -515,8 +716,8 @@ mod tests {
    use tokio::sync::broadcast::error::TryRecvError;
    use utils::id::{TenantId, TimelineId};

-    fn msg(timeline_id: Vec<u8>) -> SafekeeperTimelineInfo {
-        SafekeeperTimelineInfo {
+    fn msg(timeline_id: Vec<u8>) -> Message {
+        Message::SafekeeperTimelineInfo(SafekeeperTimelineInfo {
            safekeeper_id: 1,
            tenant_timeline_id: Some(ProtoTenantTimelineId {
                tenant_id: vec![0x00; 16],
@@ -533,7 +734,7 @@ mod tests {
            http_connstr: "neon-1-sk-1.local:7677".to_owned(),
            local_start_lsn: 0,
            availability_zone: None,
-        }
+        })
    }

    fn tli_from_u64(i: u64) -> Vec<u8> {
--- a/storage_broker/src/metrics.rs
+++ b/storage_broker/src/metrics.rs
@@ -1,6 +1,6 @@
 //! Broker metrics.

-use metrics::{register_int_gauge, IntGauge};
+use metrics::{register_int_counter, register_int_gauge, IntCounter, IntGauge};
 use once_cell::sync::Lazy;

 pub static NUM_PUBS: Lazy<IntGauge> = Lazy::new(|| {
@@ -23,3 +23,35 @@ pub static NUM_SUBS_ALL: Lazy<IntGauge> = Lazy::new(|| {
    )
    .expect("Failed to register metric")
 });
+
+pub static PROCESSED_MESSAGES_TOTAL: Lazy<IntCounter> = Lazy::new(|| {
+    register_int_counter!(
+        "storage_broker_processed_messages_total",
+        "Number of messages received by storage broker, before routing and broadcasting"
+    )
+    .expect("Failed to register metric")
+});
+
+pub static BROADCASTED_MESSAGES_TOTAL: Lazy<IntCounter> = Lazy::new(|| {
+    register_int_counter!(
+        "storage_broker_broadcasted_messages_total",
+        "Number of messages broadcasted (sent over network) to subscribers"
+    )
+    .expect("Failed to register metric")
+});
+
+pub static BROADCAST_DROPPED_MESSAGES_TOTAL: Lazy<IntCounter> = Lazy::new(|| {
+    register_int_counter!(
+        "storage_broker_broadcast_dropped_messages_total",
+        "Number of messages dropped due to channel capacity overflow"
+    )
+    .expect("Failed to register metric")
+});
+
+pub static PUBLISHED_ONEOFF_MESSAGES_TOTAL: Lazy<IntCounter> = Lazy::new(|| {
+    register_int_counter!(
+        "storage_broker_published_oneoff_messages_total",
+        "Number of one-off messages sent via PublishOne method"
+    )
+    .expect("Failed to register metric")
+});
--- a/workspace_hack/Cargo.toml
+++ b/workspace_hack/Cargo.toml
@@ -33,8 +33,11 @@ dashmap = { version = "5", default-features = false, features = ["raw-api"] }
 either = { version = "1" }
 fail = { version = "0.5", default-features = false, features = ["failpoints"] }
 futures = { version = "0.3" }
+futures-channel = { version = "0.3", features = ["sink"] }
+futures-core = { version = "0.3" }
 futures-executor = { version = "0.3" }
 futures-io = { version = "0.3" }
+futures-sink = { version = "0.3" }
 futures-util = { version = "0.3", features = ["channel", "io", "sink"] }
 hex = { version = "0.4", features = ["serde"] }
 hmac = { version = "0.12", default-features = false, features = ["reset"] }
@@ -53,6 +56,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.11", default-features = false, features = ["blocking", "default-tls", "json", "multipart", "rustls-tls", "stream"] }
+ring = { version = "0.16" }
 rustls = { version = "0.21", features = ["dangerous_configuration"] }
 scopeguard = { version = "1" }
 serde = { version = "1", features = ["alloc", "derive"] }
Author	SHA1	Message	Date
Arseny Sher	b8fad41fbe	Cap WAL download for LR on compute start by max_slot_wal_keep_size.	2023-12-22 15:55:45 +03:00
Arthur Petukhovsky	613906acea	Support custom types in broker (#5761 ) Old methods are unchanged for backwards compatibility. Added `SafekeeperDiscoveryRequest` and `SafekeeperDiscoveryResponse` types to serve as example, and also as a prerequisite for https://github.com/neondatabase/neon/issues/5471	2023-12-19 17:06:43 +00:00
Christian Schwarz	82809d2ec2	fix metric `pageserver_initial_logical_size_start_calculation` (#6191 ) It wasn't being incremented. Fixup of commit `1c88824ed0` Author: Christian Schwarz <christian@neon.tech> Date: Fri Dec 1 12:52:59 2023 +0100 initial logical size calculation: add a bunch of metrics (#5995)	2023-12-19 17:44:49 +01:00
Anastasia Lubennikova	0bd79eb063	Handle role deletion when project has no databases. (#6170 ) There is still default 'postgres' database, that may contain objects owned by the role or some ACLs. We need to reassign objects in this database too. ## Problem If customer deleted all databases and then tries to delete role, that has some non-standard ACLs, `apply_config` operation will stuck because of failing role deletion.	2023-12-19 16:27:47 +00:00
Konstantin Knizhnik	8ff5387da1	eliminate GCC warning for unchecked result of fread (#6167 ) ## Problem GCCproduce warning that bread result is not checked. It doesn't affect program logic, but better live without warnings. ## Summary of changes Check read result. ## 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	2023-12-19 18:17:11 +02:00