use smallvec and add comment about memory usage

## Problem

Gc currently doesn't work properly.

## Summary of changes

Change statement on running gc.

.config/nextest.toml

@@ -1,2 +1,2 @@
 [profile.default]
-slow-timeout = "1m"
+slow-timeout = { period = "20s", terminate-after = 3 }

Cargo.lock

@@ -3973,6 +3973,7 @@ dependencies = [
  "serde",
  "serde_json",
  "sha2",
+ "smallvec",
  "smol_str",
  "socket2 0.5.5",
  "sync_wrapper",
@@ -5031,9 +5032,9 @@ dependencies = [
 
 [[package]]
 name = "shlex"
-version = "1.1.0"
+version = "1.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "43b2853a4d09f215c24cc5489c992ce46052d359b5109343cbafbf26bc62f8a3"
+checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64"
 
 [[package]]
 name = "signal-hook"
@@ -5107,6 +5108,9 @@ name = "smallvec"
 version = "1.11.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "62bb4feee49fdd9f707ef802e22365a35de4b7b299de4763d44bfea899442ff9"
+dependencies = [
+ "serde",
+]
 
 [[package]]
 name = "smol_str"

Dockerfile.compute-node

@@ -143,6 +143,8 @@ RUN wget https://github.com/pgRouting/pgrouting/archive/v3.4.2.tar.gz -O pgrouti
 #########################################################################################
 FROM build-deps AS plv8-build
 COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
+
+ARG PG_VERSION
 RUN apt update && \
     apt install -y ninja-build python3-dev libncurses5 binutils clang
 
@@ -617,6 +619,7 @@ RUN wget https://github.com/theory/pg-semver/archive/refs/tags/v0.32.1.tar.gz -O
 FROM build-deps AS pg-embedding-pg-build
 COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
 
+ARG PG_VERSION
 ENV PATH "/usr/local/pgsql/bin/:$PATH"
 RUN case "${PG_VERSION}" in \
       "v14" | "v15") \
@@ -779,6 +782,8 @@ RUN wget https://github.com/eulerto/wal2json/archive/refs/tags/wal2json_2_5.tar.
 #
 #########################################################################################
 FROM build-deps AS neon-pg-ext-build
+ARG PG_VERSION
+
 # Public extensions
 COPY --from=postgis-build /usr/local/pgsql/ /usr/local/pgsql/
 COPY --from=postgis-build /sfcgal/* /

libs/postgres_ffi/src/xlog_utils.rs

@@ -329,8 +329,8 @@ impl CheckPoint {
     ///
     /// Returns 'true' if the XID was updated.
     pub fn update_next_xid(&mut self, xid: u32) -> bool {
-        // nextXid should nw greater than any XID in WAL, so increment provided XID and check for wraparround.
-        let mut new_xid = std::cmp::max(xid + 1, pg_constants::FIRST_NORMAL_TRANSACTION_ID);
+        // nextXid should be greater than any XID in WAL, so increment provided XID and check for wraparround.
+        let mut new_xid = std::cmp::max(xid.wrapping_add(1), pg_constants::FIRST_NORMAL_TRANSACTION_ID);
         // To reduce number of metadata checkpoints, we forward align XID on XID_CHECKPOINT_INTERVAL.
         // XID_CHECKPOINT_INTERVAL should not be larger than BLCKSZ*CLOG_XACTS_PER_BYTE
         new_xid =

pageserver/src/walingest.rs

@@ -102,7 +102,9 @@ impl WalIngest {
         buf.advance(decoded.main_data_offset);
 
         assert!(!self.checkpoint_modified);
-        if self.checkpoint.update_next_xid(decoded.xl_xid) {
+        if decoded.xl_xid != pg_constants::INVALID_TRANSACTION_ID
+            && self.checkpoint.update_next_xid(decoded.xl_xid)
+        {
             self.checkpoint_modified = true;
         }
 
@@ -330,8 +332,13 @@ impl WalIngest {
                         < 0
                     {
                         self.checkpoint.oldestXid = xlog_checkpoint.oldestXid;
-                        self.checkpoint_modified = true;
                     }
+
+                    // Write a new checkpoint key-value pair on every checkpoint record, even
+                    // if nothing really changed. Not strictly required, but it seems nice to
+                    // have some trace of the checkpoint records in the layer files at the same
+                    // LSNs.
+                    self.checkpoint_modified = true;
                 }
             }
             pg_constants::RM_LOGICALMSG_ID => {

proxy/Cargo.toml

@@ -5,7 +5,7 @@ edition.workspace = true
 license.workspace = true
 
 [features]
-default = ["testing"]
+default = []
 testing = []
 
 [dependencies]
@@ -81,6 +81,7 @@ postgres-native-tls.workspace = true
 postgres-protocol.workspace = true
 redis.workspace = true
 smol_str.workspace = true
+smallvec = { workspace = true, features = ["serde"] }
 
 workspace_hack.workspace = true
 

proxy/src/auth.rs

@@ -4,7 +4,9 @@ pub mod backend;
 pub use backend::BackendType;
 
 mod credentials;
-pub use credentials::{check_peer_addr_is_in_list, endpoint_sni, ComputeUserInfoMaybeEndpoint};
+pub use credentials::{
+    check_peer_addr_is_in_list, endpoint_sni, ComputeUserInfoMaybeEndpoint, IpPattern,
+};
 
 mod password_hack;
 pub use password_hack::parse_endpoint_param;

proxy/src/auth/backend.rs

@@ -10,6 +10,7 @@ use crate::auth::credentials::check_peer_addr_is_in_list;
 use crate::auth::validate_password_and_exchange;
 use crate::cache::Cached;
 use crate::console::errors::GetAuthInfoError;
+use crate::console::provider::ConsoleBackend;
 use crate::console::AuthSecret;
 use crate::context::RequestMonitoring;
 use crate::proxy::connect_compute::handle_try_wake;
@@ -34,6 +35,8 @@ use std::sync::Arc;
 use tokio::io::{AsyncRead, AsyncWrite};
 use tracing::{error, info, warn};
 
+use super::IpPattern;
+
 /// This type serves two purposes:
 ///
 /// * When `T` is `()`, it's just a regular auth backend selector
@@ -43,11 +46,8 @@ use tracing::{error, info, warn};
 ///   this helps us provide the credentials only to those auth
 ///   backends which require them for the authentication process.
 pub enum BackendType<'a, T> {
-    /// Current Cloud API (V2).
-    Console(Cow<'a, console::provider::neon::Api>, T),
-    /// Local mock of Cloud API (V2).
-    #[cfg(feature = "testing")]
-    Postgres(Cow<'a, console::provider::mock::Api>, T),
+    /// Cloud API (V2).
+    Console(Cow<'a, ConsoleBackend>, T),
     /// Authentication via a web browser.
     Link(Cow<'a, url::ApiUrl>),
     #[cfg(test)]
@@ -57,16 +57,22 @@ pub enum BackendType<'a, T> {
 
 pub trait TestBackend: Send + Sync + 'static {
     fn wake_compute(&self) -> Result<CachedNodeInfo, console::errors::WakeComputeError>;
-    fn get_allowed_ips(&self) -> Result<Vec<SmolStr>, console::errors::GetAuthInfoError>;
+    fn get_allowed_ips(&self) -> Result<Vec<IpPattern>, console::errors::GetAuthInfoError>;
 }
 
 impl std::fmt::Display for BackendType<'_, ()> {
     fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         use BackendType::*;
         match self {
-            Console(endpoint, _) => fmt.debug_tuple("Console").field(&endpoint.url()).finish(),
-            #[cfg(feature = "testing")]
-            Postgres(endpoint, _) => fmt.debug_tuple("Postgres").field(&endpoint.url()).finish(),
+            Console(api, _) => match &**api {
+                ConsoleBackend::Console(endpoint) => {
+                    fmt.debug_tuple("Console").field(&endpoint.url()).finish()
+                }
+                #[cfg(feature = "testing")]
+                ConsoleBackend::Postgres(endpoint) => {
+                    fmt.debug_tuple("Postgres").field(&endpoint.url()).finish()
+                }
+            },
             Link(url) => fmt.debug_tuple("Link").field(&url.as_str()).finish(),
             #[cfg(test)]
             Test(_) => fmt.debug_tuple("Test").finish(),
@@ -81,8 +87,6 @@ impl<T> BackendType<'_, T> {
         use BackendType::*;
         match self {
             Console(c, x) => Console(Cow::Borrowed(c), x),
-            #[cfg(feature = "testing")]
-            Postgres(c, x) => Postgres(Cow::Borrowed(c), x),
             Link(c) => Link(Cow::Borrowed(c)),
             #[cfg(test)]
             Test(x) => Test(*x),
@@ -98,8 +102,6 @@ impl<'a, T> BackendType<'a, T> {
         use BackendType::*;
         match self {
             Console(c, x) => Console(c, f(x)),
-            #[cfg(feature = "testing")]
-            Postgres(c, x) => Postgres(c, f(x)),
             Link(c) => Link(c),
             #[cfg(test)]
             Test(x) => Test(x),
@@ -114,8 +116,6 @@ impl<'a, T, E> BackendType<'a, Result<T, E>> {
         use BackendType::*;
         match self {
             Console(c, x) => x.map(|x| Console(c, x)),
-            #[cfg(feature = "testing")]
-            Postgres(c, x) => x.map(|x| Postgres(c, x)),
             Link(c) => Ok(Link(c)),
             #[cfg(test)]
             Test(x) => Ok(Test(x)),
@@ -325,8 +325,6 @@ impl<'a> BackendType<'a, ComputeUserInfoMaybeEndpoint> {
 
         match self {
             Console(_, user_info) => user_info.project.clone(),
-            #[cfg(feature = "testing")]
-            Postgres(_, user_info) => user_info.project.clone(),
             Link(_) => Some("link".into()),
             #[cfg(test)]
             Test(_) => Some("test".into()),
@@ -339,8 +337,6 @@ impl<'a> BackendType<'a, ComputeUserInfoMaybeEndpoint> {
 
         match self {
             Console(_, user_info) => &user_info.user,
-            #[cfg(feature = "testing")]
-            Postgres(_, user_info) => &user_info.user,
             Link(_) => "link",
             #[cfg(test)]
             Test(_) => "test",
@@ -371,19 +367,6 @@ impl<'a> BackendType<'a, ComputeUserInfoMaybeEndpoint> {
                         .await?;
                 (cache_info, BackendType::Console(api, user_info))
             }
-            #[cfg(feature = "testing")]
-            Postgres(api, user_info) => {
-                info!(
-                    user = &*user_info.user,
-                    project = user_info.project(),
-                    "performing authentication using a local postgres instance"
-                );
-
-                let (cache_info, user_info) =
-                    auth_and_wake_compute(ctx, &*api, user_info, client, allow_cleartext, config)
-                        .await?;
-                (cache_info, BackendType::Postgres(api, user_info))
-            }
             // NOTE: this auth backend doesn't use client credentials.
             Link(url) => {
                 info!("performing link authentication");
@@ -414,8 +397,6 @@ impl BackendType<'_, ComputeUserInfo> {
         use BackendType::*;
         match self {
             Console(api, user_info) => api.get_allowed_ips(ctx, user_info).await,
-            #[cfg(feature = "testing")]
-            Postgres(api, user_info) => api.get_allowed_ips(ctx, user_info).await,
             Link(_) => Ok(Cached::new_uncached(Arc::new(vec![]))),
             #[cfg(test)]
             Test(x) => Ok(Cached::new_uncached(Arc::new(x.get_allowed_ips()?))),
@@ -432,8 +413,6 @@ impl BackendType<'_, ComputeUserInfo> {
 
         match self {
             Console(api, user_info) => api.wake_compute(ctx, user_info).map_ok(Some).await,
-            #[cfg(feature = "testing")]
-            Postgres(api, user_info) => api.wake_compute(ctx, user_info).map_ok(Some).await,
             Link(_) => Ok(None),
             #[cfg(test)]
             Test(x) => x.wake_compute().map(Some),

proxy/src/auth/backend/link.rs

@@ -57,24 +57,31 @@ pub(super) async fn authenticate(
     link_uri: &reqwest::Url,
     client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
 ) -> auth::Result<NodeInfo> {
-    let psql_session_id = new_psql_session_id();
+    // registering waiter can fail if we get unlucky with rng.
+    // just try again.
+    let (psql_session_id, waiter) = loop {
+        let psql_session_id = new_psql_session_id();
+
+        match console::mgmt::get_waiter(&psql_session_id) {
+            Ok(waiter) => break (psql_session_id, waiter),
+            Err(_e) => continue,
+        }
+    };
+
     let span = info_span!("link", psql_session_id = &psql_session_id);
     let greeting = hello_message(link_uri, &psql_session_id);
 
-    let db_info = console::mgmt::with_waiter(psql_session_id, |waiter| async {
-        // Give user a URL to spawn a new database.
-        info!(parent: &span, "sending the auth URL to the user");
-        client
-            .write_message_noflush(&Be::AuthenticationOk)?
-            .write_message_noflush(&Be::CLIENT_ENCODING)?
-            .write_message(&Be::NoticeResponse(&greeting))
-            .await?;
+    // Give user a URL to spawn a new database.
+    info!(parent: &span, "sending the auth URL to the user");
+    client
+        .write_message_noflush(&Be::AuthenticationOk)?
+        .write_message_noflush(&Be::CLIENT_ENCODING)?
+        .write_message(&Be::NoticeResponse(&greeting))
+        .await?;
 
-        // Wait for web console response (see `mgmt`).
-        info!(parent: &span, "waiting for console's reply...");
-        waiter.await?.map_err(LinkAuthError::AuthFailed)
-    })
-    .await?;
+    // Wait for web console response (see `mgmt`).
+    info!(parent: &span, "waiting for console's reply...");
+    let db_info = waiter.await.map_err(LinkAuthError::from)?;
 
     client.write_message_noflush(&Be::NoticeResponse("Connecting to database."))?;
 

proxy/src/auth/credentials.rs

@@ -7,7 +7,7 @@ use crate::{
 use itertools::Itertools;
 use pq_proto::StartupMessageParams;
 use smol_str::SmolStr;
-use std::{collections::HashSet, net::IpAddr};
+use std::{collections::HashSet, net::IpAddr, str::FromStr};
 use thiserror::Error;
 use tracing::{info, warn};
 
@@ -151,30 +151,51 @@ impl ComputeUserInfoMaybeEndpoint {
     }
 }
 
-pub fn check_peer_addr_is_in_list(peer_addr: &IpAddr, ip_list: &Vec<SmolStr>) -> bool {
-    if ip_list.is_empty() {
-        return true;
-    }
-    for ip in ip_list {
-        // We expect that all ip addresses from control plane are correct.
-        // However, if some of them are broken, we still can check the others.
-        match parse_ip_pattern(ip) {
-            Ok(pattern) => {
-                if check_ip(peer_addr, &pattern) {
-                    return true;
-                }
-            }
-            Err(err) => warn!("Cannot parse ip: {}; err: {}", ip, err),
-        }
-    }
-    false
+pub fn check_peer_addr_is_in_list(peer_addr: &IpAddr, ip_list: &[IpPattern]) -> bool {
+    ip_list.is_empty() || ip_list.iter().any(|pattern| check_ip(peer_addr, pattern))
 }
 
 #[derive(Debug, Clone, Eq, PartialEq)]
-enum IpPattern {
+pub enum IpPattern {
     Subnet(ipnet::IpNet),
     Range(IpAddr, IpAddr),
     Single(IpAddr),
+    None,
+}
+
+impl<'de> serde::de::Deserialize<'de> for IpPattern {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        struct StrVisitor;
+        impl<'de> serde::de::Visitor<'de> for StrVisitor {
+            type Value = IpPattern;
+
+            fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
+                write!(formatter, "comma separated list with ip address, ip address range, or ip address subnet mask")
+            }
+
+            fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
+            where
+                E: serde::de::Error,
+            {
+                Ok(parse_ip_pattern(v).unwrap_or_else(|e| {
+                    warn!("Cannot parse ip pattern {v}: {e}");
+                    IpPattern::None
+                }))
+            }
+        }
+        deserializer.deserialize_str(StrVisitor)
+    }
+}
+
+impl FromStr for IpPattern {
+    type Err = anyhow::Error;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        parse_ip_pattern(s)
+    }
 }
 
 fn parse_ip_pattern(pattern: &str) -> anyhow::Result<IpPattern> {
@@ -196,6 +217,7 @@ fn check_ip(ip: &IpAddr, pattern: &IpPattern) -> bool {
         IpPattern::Subnet(subnet) => subnet.contains(ip),
         IpPattern::Range(start, end) => start <= ip && ip <= end,
         IpPattern::Single(addr) => addr == ip,
+        IpPattern::None => false,
     }
 }
 
@@ -206,6 +228,7 @@ fn project_name_valid(name: &str) -> bool {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use serde_json::json;
     use ComputeUserInfoParseError::*;
 
     #[test]
@@ -415,21 +438,17 @@ mod tests {
 
     #[test]
     fn test_check_peer_addr_is_in_list() {
-        let peer_addr = IpAddr::from([127, 0, 0, 1]);
-        assert!(check_peer_addr_is_in_list(&peer_addr, &vec![]));
-        assert!(check_peer_addr_is_in_list(
-            &peer_addr,
-            &vec!["127.0.0.1".into()]
-        ));
-        assert!(!check_peer_addr_is_in_list(
-            &peer_addr,
-            &vec!["8.8.8.8".into()]
-        ));
+        fn check(v: serde_json::Value) -> bool {
+            let peer_addr = IpAddr::from([127, 0, 0, 1]);
+            let ip_list: Vec<IpPattern> = serde_json::from_value(v).unwrap();
+            check_peer_addr_is_in_list(&peer_addr, &ip_list)
+        }
+
+        assert!(check(json!([])));
+        assert!(check(json!(["127.0.0.1"])));
+        assert!(!check(json!(["8.8.8.8"])));
         // If there is an incorrect address, it will be skipped.
-        assert!(check_peer_addr_is_in_list(
-            &peer_addr,
-            &vec!["88.8.8".into(), "127.0.0.1".into()]
-        ));
+        assert!(check(json!(["88.8.8", "127.0.0.1"])));
     }
     #[test]
     fn test_parse_ip_v4() -> anyhow::Result<()> {

proxy/src/bin/proxy.rs

@@ -249,12 +249,18 @@ async fn main() -> anyhow::Result<()> {
     }
 
     if let auth::BackendType::Console(api, _) = &config.auth_backend {
-        let cache = api.caches.project_info.clone();
-        if let Some(url) = args.redis_notifications {
-            info!("Starting redis notifications listener ({url})");
-            maintenance_tasks.spawn(notifications::task_main(url.to_owned(), cache.clone()));
+        match &**api {
+            proxy::console::provider::ConsoleBackend::Console(api) => {
+                let cache = api.caches.project_info.clone();
+                if let Some(url) = args.redis_notifications {
+                    info!("Starting redis notifications listener ({url})");
+                    maintenance_tasks
+                        .spawn(notifications::task_main(url.to_owned(), cache.clone()));
+                }
+            }
+            #[cfg(feature = "testing")]
+            proxy::console::provider::ConsoleBackend::Postgres(_) => {}
         }
-        maintenance_tasks.spawn(async move { cache.clone().gc_worker().await });
     }
 
     let maintenance = loop {
@@ -351,13 +357,15 @@ fn build_config(args: &ProxyCliArgs) -> anyhow::Result<&'static ProxyConfig> {
             let endpoint = http::Endpoint::new(url, http::new_client(rate_limiter_config));
 
             let api = console::provider::neon::Api::new(endpoint, caches, locks);
+            let api = console::provider::ConsoleBackend::Console(api);
             auth::BackendType::Console(Cow::Owned(api), ())
         }
         #[cfg(feature = "testing")]
         AuthBackend::Postgres => {
             let url = args.auth_endpoint.parse()?;
             let api = console::provider::mock::Api::new(url);
-            auth::BackendType::Postgres(Cow::Owned(api), ())
+            let api = console::provider::ConsoleBackend::Postgres(api);
+            auth::BackendType::Console(Cow::Owned(api), ())
         }
         AuthBackend::Link => {
             let url = args.uri.parse()?;

proxy/src/cache/project_info.rs

@@ -1,17 +1,17 @@
 use std::{
-    collections::HashSet,
-    convert::Infallible,
     sync::{atomic::AtomicU64, Arc},
     time::Duration,
 };
 
 use dashmap::DashMap;
-use rand::{thread_rng, Rng};
+use hashlink::LruCache;
+use parking_lot::Mutex;
+use smallvec::SmallVec;
 use smol_str::SmolStr;
 use tokio::time::Instant;
-use tracing::{debug, info};
+use tracing::info;
 
-use crate::{config::ProjectInfoCacheOptions, console::AuthSecret};
+use crate::{auth::IpPattern, config::ProjectInfoCacheOptions, console::AuthSecret};
 
 use super::{Cache, Cached};
 
@@ -42,56 +42,10 @@ impl<T> From<T> for Entry<T> {
     }
 }
 
-#[derive(Default)]
-struct EndpointInfo {
-    secret: std::collections::HashMap<SmolStr, Entry<AuthSecret>>,
-    allowed_ips: Option<Entry<Arc<Vec<SmolStr>>>>,
-}
-
-impl EndpointInfo {
-    fn check_ignore_cache(ignore_cache_since: Option<Instant>, created_at: Instant) -> bool {
-        match ignore_cache_since {
-            None => false,
-            Some(t) => t < created_at,
-        }
-    }
-    pub fn get_role_secret(
-        &self,
-        role_name: &SmolStr,
-        valid_since: Instant,
-        ignore_cache_since: Option<Instant>,
-    ) -> Option<(AuthSecret, bool)> {
-        if let Some(secret) = self.secret.get(role_name) {
-            if valid_since < secret.created_at {
-                return Some((
-                    secret.value.clone(),
-                    Self::check_ignore_cache(ignore_cache_since, secret.created_at),
-                ));
-            }
-        }
-        None
-    }
-
-    pub fn get_allowed_ips(
-        &self,
-        valid_since: Instant,
-        ignore_cache_since: Option<Instant>,
-    ) -> Option<(Arc<Vec<SmolStr>>, bool)> {
-        if let Some(allowed_ips) = &self.allowed_ips {
-            if valid_since < allowed_ips.created_at {
-                return Some((
-                    allowed_ips.value.clone(),
-                    Self::check_ignore_cache(ignore_cache_since, allowed_ips.created_at),
-                ));
-            }
-        }
-        None
-    }
-    pub fn invalidate_allowed_ips(&mut self) {
-        self.allowed_ips = None;
-    }
-    pub fn invalidate_role_secret(&mut self, role_name: &SmolStr) {
-        self.secret.remove(role_name);
+fn check_ignore_cache(ignore_cache_since: Option<Instant>, created_at: Instant) -> bool {
+    match ignore_cache_since {
+        None => false,
+        Some(t) => t < created_at,
     }
 }
 
@@ -103,12 +57,33 @@ impl EndpointInfo {
 /// One may ask, why the data is stored per project, when on the user request there is only data about the endpoint available?
 /// On the cplane side updates are done per project (or per branch), so it's easier to invalidate the whole project cache.
 pub struct ProjectInfoCacheImpl {
-    cache: DashMap<SmolStr, EndpointInfo>,
+    ip_cache: Mutex<LruCache<SmolStr, Entry<Arc<Vec<IpPattern>>>>>,
+    role_cache: Mutex<LruCache<(SmolStr, SmolStr), Entry<AuthSecret>>>,
 
-    project2ep: DashMap<SmolStr, HashSet<SmolStr>>,
-    config: ProjectInfoCacheOptions,
+    // endpoints per project:
+    // P90:  1
+    // P99:  2
+    // P995: 3
+    // P999: 10
+    // P9999: 186
+    //
+    // Assuming 1 million projects with this distribution:
+    // (0.9 * 1 + 0.09 * 2 + 0.005 * 3 + 0.004 * 10 + 0.0009 * 186) * 1,000,000
+    // =~ 1,500,000 endpoints
+    //
+    // 1,000,000 * size_of(SmolStr) = 24MB
+    // 1,500,000 * size_of(SmolStr) = 36MB
+    // SmallVec inline overhead: 8B * 0.9 * 1,000,000 = 7.2MB
+    // SmallVec outline overhead: 32B * 0.1 * 1,000,000 = 3.2MB
+    //
+    // Total size: 70.4MB.
+    //
+    // We do not need to prune this hashmap and can safely
+    // keep it in memory up until 100s of millions of projects
+    project2ep: DashMap<SmolStr, SmallVec<[SmolStr; 1]>>,
 
     start_time: Instant,
+    ttl: Duration,
     ttl_disabled_since_us: AtomicU64,
 }
 
@@ -121,9 +96,7 @@ impl ProjectInfoCache for ProjectInfoCacheImpl {
             .map(|kv| kv.value().clone())
             .unwrap_or_default();
         for endpoint_id in endpoints {
-            if let Some(mut endpoint_info) = self.cache.get_mut(&endpoint_id) {
-                endpoint_info.invalidate_allowed_ips();
-            }
+            self.ip_cache.lock().remove(&endpoint_id);
         }
     }
     fn invalidate_role_secret_for_project(&self, project_id: &SmolStr, role_name: &SmolStr) {
@@ -137,9 +110,9 @@ impl ProjectInfoCache for ProjectInfoCacheImpl {
             .map(|kv| kv.value().clone())
             .unwrap_or_default();
         for endpoint_id in endpoints {
-            if let Some(mut endpoint_info) = self.cache.get_mut(&endpoint_id) {
-                endpoint_info.invalidate_role_secret(role_name);
-            }
+            self.role_cache
+                .lock()
+                .remove(&(endpoint_id, role_name.clone()));
         }
     }
     fn enable_ttl(&self) {
@@ -148,7 +121,7 @@ impl ProjectInfoCache for ProjectInfoCacheImpl {
     }
 
     fn disable_ttl(&self) {
-        let new_ttl = (self.start_time.elapsed() + self.config.ttl).as_micros() as u64;
+        let new_ttl = (self.start_time.elapsed() + self.ttl).as_micros() as u64;
         self.ttl_disabled_since_us
             .store(new_ttl, std::sync::atomic::Ordering::Relaxed);
     }
@@ -157,9 +130,10 @@ impl ProjectInfoCache for ProjectInfoCacheImpl {
 impl ProjectInfoCacheImpl {
     pub fn new(config: ProjectInfoCacheOptions) -> Self {
         Self {
-            cache: DashMap::new(),
+            ip_cache: Mutex::new(LruCache::new(config.size)),
+            role_cache: Mutex::new(LruCache::new(config.size * config.max_roles)),
             project2ep: DashMap::new(),
-            config,
+            ttl: config.ttl,
             ttl_disabled_since_us: AtomicU64::new(u64::MAX),
             start_time: Instant::now(),
         }
@@ -171,9 +145,17 @@ impl ProjectInfoCacheImpl {
         role_name: &SmolStr,
     ) -> Option<Cached<&Self, AuthSecret>> {
         let (valid_since, ignore_cache_since) = self.get_cache_times();
-        let endpoint_info = self.cache.get(endpoint_id)?;
-        let (value, ignore_cache) =
-            endpoint_info.get_role_secret(role_name, valid_since, ignore_cache_since)?;
+        let (value, ignore_cache) = {
+            let mut cache = self.role_cache.lock();
+            let secret = cache.get(&(endpoint_id.clone(), role_name.clone()))?;
+            if secret.created_at <= valid_since {
+                return None;
+            }
+            (
+                secret.value.clone(),
+                check_ignore_cache(ignore_cache_since, secret.created_at),
+            )
+        };
         if !ignore_cache {
             let cached = Cached {
                 token: Some((
@@ -186,14 +168,23 @@ impl ProjectInfoCacheImpl {
         }
         Some(Cached::new_uncached(value))
     }
+
     pub fn get_allowed_ips(
         &self,
         endpoint_id: &SmolStr,
-    ) -> Option<Cached<&Self, Arc<Vec<SmolStr>>>> {
+    ) -> Option<Cached<&Self, Arc<Vec<IpPattern>>>> {
         let (valid_since, ignore_cache_since) = self.get_cache_times();
-        let endpoint_info = self.cache.get(endpoint_id)?;
-        let value = endpoint_info.get_allowed_ips(valid_since, ignore_cache_since);
-        let (value, ignore_cache) = value?;
+        let (value, ignore_cache) = {
+            let mut cache = self.ip_cache.lock();
+            let allowed_ips = cache.get(endpoint_id)?;
+            if allowed_ips.created_at <= valid_since {
+                return None;
+            }
+            (
+                allowed_ips.value.clone(),
+                check_ignore_cache(ignore_cache_since, allowed_ips.created_at),
+            )
+        };
         if !ignore_cache {
             let cached = Cached {
                 token: Some((self, CachedLookupInfo::new_allowed_ips(endpoint_id.clone()))),
@@ -203,6 +194,7 @@ impl ProjectInfoCacheImpl {
         }
         Some(Cached::new_uncached(value))
     }
+
     pub fn insert_role_secret(
         &self,
         project_id: &SmolStr,
@@ -210,42 +202,33 @@ impl ProjectInfoCacheImpl {
         role_name: &SmolStr,
         secret: AuthSecret,
     ) {
-        if self.cache.len() >= self.config.size {
-            // If there are too many entries, wait until the next gc cycle.
-            return;
-        }
-        self.inser_project2endpoint(project_id, endpoint_id);
-        let mut entry = self.cache.entry(endpoint_id.clone()).or_default();
-        if entry.secret.len() < self.config.max_roles {
-            entry.secret.insert(role_name.clone(), secret.into());
-        }
+        self.insert_project2endpoint(project_id, endpoint_id);
+        self.role_cache
+            .lock()
+            .insert((endpoint_id.clone(), role_name.clone()), secret.into());
     }
+
     pub fn insert_allowed_ips(
         &self,
         project_id: &SmolStr,
         endpoint_id: &SmolStr,
-        allowed_ips: Arc<Vec<SmolStr>>,
+        allowed_ips: Arc<Vec<IpPattern>>,
     ) {
-        if self.cache.len() >= self.config.size {
-            // If there are too many entries, wait until the next gc cycle.
-            return;
-        }
-        self.inser_project2endpoint(project_id, endpoint_id);
-        self.cache
-            .entry(endpoint_id.clone())
+        self.insert_project2endpoint(project_id, endpoint_id);
+        self.ip_cache
+            .lock()
+            .insert(endpoint_id.clone(), allowed_ips.into());
+    }
+
+    fn insert_project2endpoint(&self, project_id: &SmolStr, endpoint_id: &SmolStr) {
+        self.project2ep
+            .entry(project_id.clone())
             .or_default()
-            .allowed_ips = Some(allowed_ips.into());
-    }
-    fn inser_project2endpoint(&self, project_id: &SmolStr, endpoint_id: &SmolStr) {
-        if let Some(mut endpoints) = self.project2ep.get_mut(project_id) {
-            endpoints.insert(endpoint_id.clone());
-        } else {
-            self.project2ep
-                .insert(project_id.clone(), HashSet::from([endpoint_id.clone()]));
-        }
+            .push(endpoint_id.clone());
     }
+
     fn get_cache_times(&self) -> (Instant, Option<Instant>) {
-        let mut valid_since = Instant::now() - self.config.ttl;
+        let mut valid_since = Instant::now() - self.ttl;
         // Only ignore cache if ttl is disabled.
         let ttl_disabled_since_us = self
             .ttl_disabled_since_us
@@ -260,37 +243,6 @@ impl ProjectInfoCacheImpl {
         };
         (valid_since, ignore_cache_since)
     }
-
-    pub async fn gc_worker(&self) -> anyhow::Result<Infallible> {
-        let mut interval =
-            tokio::time::interval(self.config.gc_interval / (self.cache.shards().len()) as u32);
-        loop {
-            interval.tick().await;
-            if self.cache.len() <= self.config.size {
-                // If there are not too many entries, wait until the next gc cycle.
-                continue;
-            }
-            self.gc();
-        }
-    }
-
-    fn gc(&self) {
-        let shard = thread_rng().gen_range(0..self.project2ep.shards().len());
-        debug!(shard, "project_info_cache: performing epoch reclamation");
-
-        // acquire a random shard lock
-        let mut removed = 0;
-        let shard = self.project2ep.shards()[shard].write();
-        for (_, endpoints) in shard.iter() {
-            for endpoint in endpoints.get().iter() {
-                self.cache.remove(endpoint);
-                removed += 1;
-            }
-        }
-        // We can drop this shard only after making sure that all endpoints are removed.
-        drop(shard);
-        info!("project_info_cache: removed {removed} endpoints");
-    }
 }
 
 /// Lookup info for project info cache.
@@ -331,14 +283,12 @@ impl Cache for ProjectInfoCacheImpl {
     fn invalidate(&self, key: &Self::LookupInfo<SmolStr>) {
         match &key.lookup_type {
             LookupType::RoleSecret(role_name) => {
-                if let Some(mut endpoint_info) = self.cache.get_mut(&key.endpoint_id) {
-                    endpoint_info.invalidate_role_secret(role_name);
-                }
+                self.role_cache
+                    .lock()
+                    .remove(&(key.endpoint_id.clone(), role_name.clone()));
             }
             LookupType::AllowedIps => {
-                if let Some(mut endpoint_info) = self.cache.get_mut(&key.endpoint_id) {
-                    endpoint_info.invalidate_allowed_ips();
-                }
+                self.ip_cache.lock().remove(&key.endpoint_id);
             }
         }
     }
@@ -356,9 +306,8 @@ mod tests {
         tokio::time::pause();
         let cache = ProjectInfoCacheImpl::new(ProjectInfoCacheOptions {
             size: 2,
-            max_roles: 2,
+            max_roles: 1,
             ttl: Duration::from_secs(1),
-            gc_interval: Duration::from_secs(600),
         });
         let project_id = "project".into();
         let endpoint_id = "endpoint".into();
@@ -366,7 +315,10 @@ mod tests {
         let user2: SmolStr = "user2".into();
         let secret1 = AuthSecret::Scram(ServerSecret::mock(user1.as_str(), [1; 32]));
         let secret2 = AuthSecret::Scram(ServerSecret::mock(user2.as_str(), [2; 32]));
-        let allowed_ips = Arc::new(vec!["allowed_ip1".into(), "allowed_ip2".into()]);
+        let allowed_ips = Arc::new(vec![
+            "127.0.0.1".parse().unwrap(),
+            "127.0.0.2".parse().unwrap(),
+        ]);
         cache.insert_role_secret(&project_id, &endpoint_id, &user1, secret1.clone());
         cache.insert_role_secret(&project_id, &endpoint_id, &user2, secret2.clone());
         cache.insert_allowed_ips(&project_id, &endpoint_id, allowed_ips.clone());
@@ -382,7 +334,7 @@ mod tests {
         let user3: SmolStr = "user3".into();
         let secret3 = AuthSecret::Scram(ServerSecret::mock(user3.as_str(), [3; 32]));
         cache.insert_role_secret(&project_id, &endpoint_id, &user3, secret3.clone());
-        assert!(cache.get_role_secret(&endpoint_id, &user3).is_none());
+        assert!(cache.get_role_secret(&endpoint_id, &user1).is_none(),);
 
         let cached = cache.get_allowed_ips(&endpoint_id).unwrap();
         assert!(cached.cached());
@@ -404,7 +356,6 @@ mod tests {
             size: 2,
             max_roles: 2,
             ttl: Duration::from_secs(1),
-            gc_interval: Duration::from_secs(600),
         }));
         cache.clone().disable_ttl();
         tokio::time::advance(Duration::from_secs(2)).await;
@@ -415,7 +366,10 @@ mod tests {
         let user2: SmolStr = "user2".into();
         let secret1 = AuthSecret::Scram(ServerSecret::mock(user1.as_str(), [1; 32]));
         let secret2 = AuthSecret::Scram(ServerSecret::mock(user2.as_str(), [2; 32]));
-        let allowed_ips = Arc::new(vec!["allowed_ip1".into(), "allowed_ip2".into()]);
+        let allowed_ips = Arc::new(vec![
+            "127.0.0.1".parse().unwrap(),
+            "127.0.0.2".parse().unwrap(),
+        ]);
         cache.insert_role_secret(&project_id, &endpoint_id, &user1, secret1.clone());
         cache.insert_role_secret(&project_id, &endpoint_id, &user2, secret2.clone());
         cache.insert_allowed_ips(&project_id, &endpoint_id, allowed_ips.clone());
@@ -452,7 +406,6 @@ mod tests {
             size: 2,
             max_roles: 2,
             ttl: Duration::from_secs(1),
-            gc_interval: Duration::from_secs(600),
         }));
 
         let project_id = "project".into();
@@ -461,7 +414,10 @@ mod tests {
         let user2: SmolStr = "user2".into();
         let secret1 = AuthSecret::Scram(ServerSecret::mock(user1.as_str(), [1; 32]));
         let secret2 = AuthSecret::Scram(ServerSecret::mock(user2.as_str(), [2; 32]));
-        let allowed_ips = Arc::new(vec!["allowed_ip1".into(), "allowed_ip2".into()]);
+        let allowed_ips = Arc::new(vec![
+            "127.0.0.1".parse().unwrap(),
+            "127.0.0.2".parse().unwrap(),
+        ]);
         cache.insert_role_secret(&project_id, &endpoint_id, &user1, secret1.clone());
         cache.clone().disable_ttl();
         tokio::time::advance(Duration::from_millis(100)).await;

proxy/src/config.rs

@@ -361,14 +361,11 @@ pub struct ProjectInfoCacheOptions {
     pub ttl: Duration,
     /// Max number of roles per endpoint.
     pub max_roles: usize,
-    /// Gc interval.
-    pub gc_interval: Duration,
 }
 
 impl ProjectInfoCacheOptions {
     /// Default options for [`crate::console::provider::NodeInfoCache`].
-    pub const CACHE_DEFAULT_OPTIONS: &'static str =
-        "size=10000,ttl=4m,max_roles=10,gc_interval=60m";
+    pub const CACHE_DEFAULT_OPTIONS: &'static str = "size=10000,ttl=4m,max_roles=5,gc_interval=60m";
 
     /// Parse cache options passed via cmdline.
     /// Example: [`Self::CACHE_DEFAULT_OPTIONS`].
@@ -376,7 +373,7 @@ impl ProjectInfoCacheOptions {
         let mut size = None;
         let mut ttl = None;
         let mut max_roles = None;
-        let mut gc_interval = None;
+        let mut _gc_interval = None;
 
         for option in options.split(',') {
             let (key, value) = option
@@ -387,7 +384,7 @@ impl ProjectInfoCacheOptions {
                 "size" => size = Some(value.parse()?),
                 "ttl" => ttl = Some(humantime::parse_duration(value)?),
                 "max_roles" => max_roles = Some(value.parse()?),
-                "gc_interval" => gc_interval = Some(humantime::parse_duration(value)?),
+                "gc_interval" => _gc_interval = Some(humantime::parse_duration(value)?),
                 unknown => bail!("unknown key: {unknown}"),
             }
         }
@@ -401,7 +398,6 @@ impl ProjectInfoCacheOptions {
             size: size.context("missing `size`")?,
             ttl: ttl.context("missing `ttl`")?,
             max_roles: max_roles.context("missing `max_roles`")?,
-            gc_interval: gc_interval.context("missing `gc_interval`")?,
         })
     }
 }

proxy/src/console/messages.rs

@@ -2,6 +2,8 @@ use serde::Deserialize;
 use smol_str::SmolStr;
 use std::fmt;
 
+use crate::auth::IpPattern;
+
 /// Generic error response with human-readable description.
 /// Note that we can't always present it to user as is.
 #[derive(Debug, Deserialize)]
@@ -14,7 +16,7 @@ pub struct ConsoleError {
 #[derive(Deserialize)]
 pub struct GetRoleSecret {
     pub role_secret: Box<str>,
-    pub allowed_ips: Option<Vec<Box<str>>>,
+    pub allowed_ips: Option<Vec<IpPattern>>,
     pub project_id: Option<Box<str>>,
 }
 

proxy/src/console/mgmt.rs

@@ -13,16 +13,10 @@ use tracing::{error, info, info_span, Instrument};
 static CPLANE_WAITERS: Lazy<Waiters<ComputeReady>> = Lazy::new(Default::default);
 
 /// Give caller an opportunity to wait for the cloud's reply.
-pub async fn with_waiter<R, T, E>(
+pub fn get_waiter(
     psql_session_id: impl Into<String>,
-    action: impl FnOnce(Waiter<'static, ComputeReady>) -> R,
-) -> Result<T, E>
-where
-    R: std::future::Future<Output = Result<T, E>>,
-    E: From<waiters::RegisterError>,
-{
-    let waiter = CPLANE_WAITERS.register(psql_session_id.into())?;
-    action(waiter).await
+) -> Result<Waiter<'static, ComputeReady>, waiters::RegisterError> {
+    CPLANE_WAITERS.register(psql_session_id.into())
 }
 
 pub fn notify(psql_session_id: &str, msg: ComputeReady) -> Result<(), waiters::NotifyError> {
@@ -77,7 +71,7 @@ async fn handle_connection(socket: TcpStream) -> Result<(), QueryError> {
 }
 
 /// A message received by `mgmt` when a compute node is ready.
-pub type ComputeReady = Result<DatabaseInfo, String>;
+pub type ComputeReady = DatabaseInfo;
 
 // TODO: replace with an http-based protocol.
 struct MgmtHandler;
@@ -102,7 +96,7 @@ fn try_process_query(pgb: &mut PostgresBackendTCP, query: &str) -> Result<(), Qu
     let _enter = span.enter();
     info!("got response: {:?}", resp.result);
 
-    match notify(resp.session_id, Ok(resp.result)) {
+    match notify(resp.session_id, resp.result) {
         Ok(()) => {
             pgb.write_message_noflush(&SINGLE_COL_ROWDESC)?
                 .write_message_noflush(&BeMessage::DataRow(&[Some(b"ok")]))?

proxy/src/console/provider.rs

@@ -4,7 +4,7 @@ pub mod neon;
 
 use super::messages::MetricsAuxInfo;
 use crate::{
-    auth::backend::ComputeUserInfo,
+    auth::{backend::ComputeUserInfo, IpPattern},
     cache::{project_info::ProjectInfoCacheImpl, Cached, TimedLru},
     compute,
     config::{CacheOptions, ProjectInfoCacheOptions},
@@ -212,7 +212,7 @@ pub enum AuthSecret {
 pub struct AuthInfo {
     pub secret: Option<AuthSecret>,
     /// List of IP addresses allowed for the autorization.
-    pub allowed_ips: Vec<SmolStr>,
+    pub allowed_ips: Vec<IpPattern>,
     /// Project ID. This is used for cache invalidation.
     pub project_id: Option<SmolStr>,
 }
@@ -236,7 +236,7 @@ pub struct NodeInfo {
 pub type NodeInfoCache = TimedLru<SmolStr, NodeInfo>;
 pub type CachedNodeInfo = Cached<&'static NodeInfoCache>;
 pub type CachedRoleSecret = Cached<&'static ProjectInfoCacheImpl, AuthSecret>;
-pub type CachedAllowedIps = Cached<&'static ProjectInfoCacheImpl, Arc<Vec<SmolStr>>>;
+pub type CachedAllowedIps = Cached<&'static ProjectInfoCacheImpl, Arc<Vec<IpPattern>>>;
 
 /// This will allocate per each call, but the http requests alone
 /// already require a few allocations, so it should be fine.
@@ -248,23 +248,75 @@ pub trait Api {
     async fn get_role_secret(
         &self,
         ctx: &mut RequestMonitoring,
-        creds: &ComputeUserInfo,
+        user_info: &ComputeUserInfo,
     ) -> Result<Option<CachedRoleSecret>, errors::GetAuthInfoError>;
 
     async fn get_allowed_ips(
         &self,
         ctx: &mut RequestMonitoring,
-        creds: &ComputeUserInfo,
+        user_info: &ComputeUserInfo,
     ) -> Result<CachedAllowedIps, errors::GetAuthInfoError>;
 
     /// Wake up the compute node and return the corresponding connection info.
     async fn wake_compute(
         &self,
         ctx: &mut RequestMonitoring,
-        creds: &ComputeUserInfo,
+        user_info: &ComputeUserInfo,
     ) -> Result<CachedNodeInfo, errors::WakeComputeError>;
 }
 
+#[derive(Clone)]
+pub enum ConsoleBackend {
+    /// Current Cloud API (V2).
+    Console(neon::Api),
+    /// Local mock of Cloud API (V2).
+    #[cfg(feature = "testing")]
+    Postgres(mock::Api),
+}
+
+#[async_trait]
+impl Api for ConsoleBackend {
+    async fn get_role_secret(
+        &self,
+        ctx: &mut RequestMonitoring,
+        user_info: &ComputeUserInfo,
+    ) -> Result<Option<CachedRoleSecret>, errors::GetAuthInfoError> {
+        use ConsoleBackend::*;
+        match self {
+            Console(api) => api.get_role_secret(ctx, user_info).await,
+            #[cfg(feature = "testing")]
+            Postgres(api) => api.get_role_secret(ctx, user_info).await,
+        }
+    }
+
+    async fn get_allowed_ips(
+        &self,
+        ctx: &mut RequestMonitoring,
+        user_info: &ComputeUserInfo,
+    ) -> Result<CachedAllowedIps, errors::GetAuthInfoError> {
+        use ConsoleBackend::*;
+        match self {
+            Console(api) => api.get_allowed_ips(ctx, user_info).await,
+            #[cfg(feature = "testing")]
+            Postgres(api) => api.get_allowed_ips(ctx, user_info).await,
+        }
+    }
+
+    async fn wake_compute(
+        &self,
+        ctx: &mut RequestMonitoring,
+        user_info: &ComputeUserInfo,
+    ) -> Result<CachedNodeInfo, errors::WakeComputeError> {
+        use ConsoleBackend::*;
+
+        match self {
+            Console(api) => api.wake_compute(ctx, user_info).await,
+            #[cfg(feature = "testing")]
+            Postgres(api) => api.wake_compute(ctx, user_info).await,
+        }
+    }
+}
+
 /// Various caches for [`console`](super).
 pub struct ApiCaches {
     /// Cache for the `wake_compute` API method.

proxy/src/console/provider/mock.rs

@@ -4,14 +4,13 @@ use super::{
     errors::{ApiError, GetAuthInfoError, WakeComputeError},
     AuthInfo, AuthSecret, CachedNodeInfo, NodeInfo,
 };
-use crate::cache::Cached;
 use crate::console::provider::{CachedAllowedIps, CachedRoleSecret};
 use crate::context::RequestMonitoring;
 use crate::{auth::backend::ComputeUserInfo, compute, error::io_error, scram, url::ApiUrl};
+use crate::{auth::IpPattern, cache::Cached};
 use async_trait::async_trait;
 use futures::TryFutureExt;
-use smol_str::SmolStr;
-use std::sync::Arc;
+use std::{str::FromStr, sync::Arc};
 use thiserror::Error;
 use tokio_postgres::{config::SslMode, Client};
 use tracing::{error, info, info_span, warn, Instrument};
@@ -88,7 +87,9 @@ impl Api {
             {
                 Some(s) => {
                     info!("got allowed_ips: {s}");
-                    s.split(',').map(String::from).collect()
+                    s.split(',')
+                        .map(|s| IpPattern::from_str(s).unwrap())
+                        .collect()
                 }
                 None => vec![],
             };
@@ -100,7 +101,7 @@ impl Api {
         .await?;
         Ok(AuthInfo {
             secret,
-            allowed_ips: allowed_ips.iter().map(SmolStr::from).collect(),
+            allowed_ips,
             project_id: None,
         })
     }

proxy/src/console/provider/neon.rs

@@ -14,7 +14,6 @@ use crate::{
 };
 use async_trait::async_trait;
 use futures::TryFutureExt;
-use itertools::Itertools;
 use smol_str::SmolStr;
 use std::sync::Arc;
 use tokio::time::Instant;
@@ -89,12 +88,7 @@ impl Api {
             let secret = scram::ServerSecret::parse(&body.role_secret)
                 .map(AuthSecret::Scram)
                 .ok_or(GetAuthInfoError::BadSecret)?;
-            let allowed_ips = body
-                .allowed_ips
-                .into_iter()
-                .flatten()
-                .map(SmolStr::from)
-                .collect_vec();
+            let allowed_ips = body.allowed_ips.unwrap_or_default();
             ALLOWED_IPS_NUMBER.observe(allowed_ips.len() as f64);
             Ok(AuthInfo {
                 secret: Some(secret),
@@ -179,21 +173,23 @@ impl super::Api for Api {
             return Ok(Some(role_secret));
         }
         let auth_info = self.do_get_auth_info(ctx, user_info).await?;
-        let project_id = auth_info.project_id.unwrap_or(ep.clone());
-        if let Some(secret) = &auth_info.secret {
-            self.caches
-                .project_info
-                .insert_role_secret(&project_id, ep, user, secret.clone())
+        if let Some(project_id) = auth_info.project_id {
+            if let Some(secret) = &auth_info.secret {
+                self.caches
+                    .project_info
+                    .insert_role_secret(&project_id, ep, user, secret.clone())
+            }
+            self.caches.project_info.insert_allowed_ips(
+                &project_id,
+                ep,
+                Arc::new(auth_info.allowed_ips),
+            );
         }
-        self.caches.project_info.insert_allowed_ips(
-            &project_id,
-            ep,
-            Arc::new(auth_info.allowed_ips),
-        );
         // When we just got a secret, we don't need to invalidate it.
         Ok(auth_info.secret.map(Cached::new_uncached))
     }
 
+    #[tracing::instrument(skip_all)]
     async fn get_allowed_ips(
         &self,
         ctx: &mut RequestMonitoring,
@@ -212,15 +208,16 @@ impl super::Api for Api {
         let auth_info = self.do_get_auth_info(ctx, user_info).await?;
         let allowed_ips = Arc::new(auth_info.allowed_ips);
         let user = &user_info.user;
-        let project_id = auth_info.project_id.unwrap_or(ep.clone());
-        if let Some(secret) = &auth_info.secret {
+        if let Some(project_id) = auth_info.project_id {
+            if let Some(secret) = &auth_info.secret {
+                self.caches
+                    .project_info
+                    .insert_role_secret(&project_id, ep, user, secret.clone())
+            }
             self.caches
                 .project_info
-                .insert_role_secret(&project_id, ep, user, secret.clone())
+                .insert_allowed_ips(&project_id, ep, allowed_ips.clone());
         }
-        self.caches
-            .project_info
-            .insert_allowed_ips(&project_id, ep, allowed_ips.clone());
         Ok(Cached::new_uncached(allowed_ips))
     }
 

proxy/src/proxy/connect_compute.rs

@@ -160,8 +160,6 @@ where
     let node_info = loop {
         let wake_res = match user_info {
             auth::BackendType::Console(api, user_info) => api.wake_compute(ctx, user_info).await,
-            #[cfg(feature = "testing")]
-            auth::BackendType::Postgres(api, user_info) => api.wake_compute(ctx, user_info).await,
             // nothing to do?
             auth::BackendType::Link(_) => return Err(err.into()),
             // test backend

proxy/src/proxy/tests.rs

@@ -6,13 +6,13 @@ use super::connect_compute::ConnectMechanism;
 use super::retry::ShouldRetry;
 use super::*;
 use crate::auth::backend::{ComputeUserInfo, TestBackend};
+use crate::auth::IpPattern;
 use crate::config::CertResolver;
 use crate::console::{self, CachedNodeInfo, NodeInfo};
 use crate::proxy::retry::{retry_after, NUM_RETRIES_CONNECT};
 use crate::{auth, http, sasl, scram};
 use async_trait::async_trait;
 use rstest::rstest;
-use smol_str::SmolStr;
 use tokio_postgres::config::SslMode;
 use tokio_postgres::tls::{MakeTlsConnect, NoTls};
 use tokio_postgres_rustls::{MakeRustlsConnect, RustlsStream};
@@ -471,7 +471,7 @@ impl TestBackend for TestConnectMechanism {
         }
     }
 
-    fn get_allowed_ips(&self) -> Result<Vec<SmolStr>, console::errors::GetAuthInfoError> {
+    fn get_allowed_ips(&self) -> Result<Vec<IpPattern>, console::errors::GetAuthInfoError> {
         unimplemented!("not used in tests")
     }
 }

safekeeper/src/bin/safekeeper.rs

@@ -8,8 +8,6 @@ use futures::future::BoxFuture;
 use futures::stream::FuturesUnordered;
 use futures::{FutureExt, StreamExt};
 use remote_storage::RemoteStorageConfig;
-use safekeeper::control_file::FileStorage;
-use safekeeper::state::TimelinePersistentState;
 use sd_notify::NotifyState;
 use tokio::runtime::Handle;
 use tokio::signal::unix::{signal, SignalKind};
@@ -32,12 +30,12 @@ use safekeeper::defaults::{
     DEFAULT_HEARTBEAT_TIMEOUT, DEFAULT_HTTP_LISTEN_ADDR, DEFAULT_MAX_OFFLOADER_LAG_BYTES,
     DEFAULT_PG_LISTEN_ADDR,
 };
+use safekeeper::wal_service;
 use safekeeper::GlobalTimelines;
 use safekeeper::SafeKeeperConf;
 use safekeeper::{broker, WAL_SERVICE_RUNTIME};
 use safekeeper::{control_file, BROKER_RUNTIME};
 use safekeeper::{http, WAL_REMOVER_RUNTIME};
-use safekeeper::{json_merge, wal_service};
 use safekeeper::{remove_wal, WAL_BACKUP_RUNTIME};
 use safekeeper::{wal_backup, HTTP_RUNTIME};
 use storage_broker::DEFAULT_ENDPOINT;
@@ -107,6 +105,9 @@ struct Args {
     /// Do not wait for changes to be written safely to disk. Unsafe.
     #[arg(short, long)]
     no_sync: bool,
+    /// Dump control file at path specified by this argument and exit.
+    #[arg(long)]
+    dump_control_file: Option<Utf8PathBuf>,
     /// Broker endpoint for storage nodes coordination in the form
     /// http[s]://host:port. In case of https schema TLS is connection is
     /// established; plaintext otherwise.
@@ -165,21 +166,6 @@ struct Args {
     /// useful for debugging.
     #[arg(long)]
     current_thread_runtime: bool,
-    /// Dump control file at path specified by this argument and exit.
-    #[arg(long)]
-    dump_control_file: Option<Utf8PathBuf>,
-    /// Patch control file at path specified by this argument and exit.
-    /// Patch is specified in --patch option and imposed over
-    /// control file as per rfc7386.
-    /// Without --write-patched the result is only printed.
-    #[arg(long, verbatim_doc_comment)]
-    patch_control_file: Option<Utf8PathBuf>,
-    /// The patch to apply to control file at --patch-control-file, in JSON.
-    #[arg(long, default_value = None)]
-    patch: Option<String>,
-    /// Write --patch-control-file result back in place.
-    #[arg(long, default_value = "false")]
-    write_patched: bool,
 }
 
 // Like PathBufValueParser, but allows empty string.
@@ -221,13 +207,7 @@ async fn main() -> anyhow::Result<()> {
     if let Some(addr) = args.dump_control_file {
         let state = control_file::FileStorage::load_control_file(addr)?;
         let json = serde_json::to_string(&state)?;
-        println!("{json}");
-        return Ok(());
-    }
-
-    if let Some(cfile_path) = args.patch_control_file {
-        let patch = args.patch.ok_or(anyhow::anyhow!("patch is missing"))?;
-        patch_control_file(cfile_path, patch, args.write_patched).await?;
+        print!("{json}");
         return Ok(());
     }
 
@@ -549,26 +529,6 @@ fn parse_remote_storage(storage_conf: &str) -> anyhow::Result<RemoteStorageConfi
     })
 }
 
-async fn patch_control_file(
-    cfile_path: Utf8PathBuf,
-    patch: String,
-    write: bool,
-) -> anyhow::Result<()> {
-    let state = control_file::FileStorage::load_control_file(&cfile_path)?;
-    // serialize to json, impose patch and deserialize back
-    let mut state_json =
-        serde_json::to_value(state).context("failed to serialize state to json")?;
-    let patch_json = serde_json::from_str(&patch).context("failed to parse patch")?;
-    json_merge(&mut state_json, patch_json);
-    let patched_state: TimelinePersistentState =
-        serde_json::from_value(state_json.clone()).context("failed to deserialize patched json")?;
-    println!("{state_json}");
-    if write {
-        FileStorage::do_persist(&patched_state, &cfile_path, true).await?;
-    }
-    return Ok(());
-}
-
 #[test]
 fn verify_cli() {
     use clap::CommandFactory;

safekeeper/src/control_file.rs

@@ -2,7 +2,7 @@
 
 use anyhow::{bail, ensure, Context, Result};
 use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
-use camino::{Utf8Path, Utf8PathBuf};
+use camino::Utf8PathBuf;
 use tokio::fs::{self, File};
 use tokio::io::AsyncWriteExt;
 
@@ -155,46 +155,6 @@ impl FileStorage {
             })?;
         Ok(state)
     }
-
-    /// Persist state s to dst_path, optionally fsyncing file.
-    pub async fn do_persist(
-        s: &TimelinePersistentState,
-        dst_path: &Utf8Path,
-        sync: bool,
-    ) -> Result<()> {
-        let mut f = File::create(&dst_path)
-            .await
-            .with_context(|| format!("failed to create partial control file at: {}", &dst_path))?;
-        let mut buf: Vec<u8> = Vec::new();
-        WriteBytesExt::write_u32::<LittleEndian>(&mut buf, SK_MAGIC)?;
-        WriteBytesExt::write_u32::<LittleEndian>(&mut buf, SK_FORMAT_VERSION)?;
-        s.ser_into(&mut buf)?;
-
-        // calculate checksum before resize
-        let checksum = crc32c::crc32c(&buf);
-        buf.extend_from_slice(&checksum.to_le_bytes());
-
-        f.write_all(&buf).await.with_context(|| {
-            format!(
-                "failed to write safekeeper state into control file at: {}",
-                dst_path
-            )
-        })?;
-        f.flush().await.with_context(|| {
-            format!(
-                "failed to flush safekeeper state into control file at: {}",
-                dst_path
-            )
-        })?;
-
-        // fsync the file
-        if sync {
-            f.sync_all()
-                .await
-                .with_context(|| format!("failed to sync partial control file at {}", dst_path))?;
-        }
-        Ok(())
-    }
 }
 
 impl Deref for FileStorage {
@@ -207,7 +167,7 @@ impl Deref for FileStorage {
 
 #[async_trait::async_trait]
 impl Storage for FileStorage {
-    /// Atomically persists state durably to the underlying storage.
+    /// Persists state durably to the underlying storage.
     ///
     /// For a description, see <https://lwn.net/Articles/457667/>.
     async fn persist(&mut self, s: &TimelinePersistentState) -> Result<()> {
@@ -215,9 +175,46 @@ impl Storage for FileStorage {
 
         // write data to safekeeper.control.partial
         let control_partial_path = self.timeline_dir.join(CONTROL_FILE_NAME_PARTIAL);
-        FileStorage::do_persist(s, &control_partial_path, !self.conf.no_sync).await?;
+        let mut control_partial = File::create(&control_partial_path).await.with_context(|| {
+            format!(
+                "failed to create partial control file at: {}",
+                &control_partial_path
+            )
+        })?;
+        let mut buf: Vec<u8> = Vec::new();
+        WriteBytesExt::write_u32::<LittleEndian>(&mut buf, SK_MAGIC)?;
+        WriteBytesExt::write_u32::<LittleEndian>(&mut buf, SK_FORMAT_VERSION)?;
+        s.ser_into(&mut buf)?;
+
+        // calculate checksum before resize
+        let checksum = crc32c::crc32c(&buf);
+        buf.extend_from_slice(&checksum.to_le_bytes());
+
+        control_partial.write_all(&buf).await.with_context(|| {
+            format!(
+                "failed to write safekeeper state into control file at: {}",
+                control_partial_path
+            )
+        })?;
+        control_partial.flush().await.with_context(|| {
+            format!(
+                "failed to flush safekeeper state into control file at: {}",
+                control_partial_path
+            )
+        })?;
+
+        // fsync the file
+        if !self.conf.no_sync {
+            control_partial.sync_all().await.with_context(|| {
+                format!(
+                    "failed to sync partial control file at {}",
+                    control_partial_path
+                )
+            })?;
+        }
 
         let control_path = self.timeline_dir.join(CONTROL_FILE_NAME);
+
         // rename should be atomic
         fs::rename(&control_partial_path, &control_path).await?;
         // this sync is not required by any standard but postgres does this (see durable_rename)

safekeeper/src/lib.rs

@@ -2,7 +2,6 @@
 use camino::Utf8PathBuf;
 use once_cell::sync::Lazy;
 use remote_storage::RemoteStorageConfig;
-use serde_json::Value;
 use tokio::runtime::Runtime;
 
 use std::time::Duration;
@@ -176,24 +175,3 @@ pub static METRICS_SHIFTER_RUNTIME: Lazy<Runtime> = Lazy::new(|| {
         .build()
         .expect("Failed to create broker runtime")
 });
-
-/// Merge json b into json a according to
-/// https://www.rfc-editor.org/rfc/rfc7396
-/// https://stackoverflow.com/a/54118457/4014587
-pub fn json_merge(a: &mut Value, b: Value) {
-    if let Value::Object(a) = a {
-        if let Value::Object(b) = b {
-            for (k, v) in b {
-                if v.is_null() {
-                    a.remove(&k);
-                } else {
-                    json_merge(a.entry(k).or_insert(Value::Null), v);
-                }
-            }
-
-            return;
-        }
-    }
-
-    *a = b;
-}

test_runner/regress/test_next_xid.py

@@ -1,10 +1,18 @@
+import json
+import os
 import time
+from pathlib import Path
 
-from fixtures.neon_fixtures import NeonEnvBuilder
+from fixtures.log_helper import log
+from fixtures.neon_fixtures import NeonEnvBuilder, wait_for_wal_insert_lsn
+from fixtures.pageserver.utils import (
+    wait_for_last_record_lsn,
+)
+from fixtures.remote_storage import RemoteStorageKind
+from fixtures.types import Lsn, TenantId, TimelineId
+from fixtures.utils import query_scalar
 
 
-# Test restarting page server, while safekeeper and compute node keep
-# running.
 def test_next_xid(neon_env_builder: NeonEnvBuilder):
     env = neon_env_builder.init_start()
 
@@ -52,3 +60,161 @@ def test_next_xid(neon_env_builder: NeonEnvBuilder):
     cur = conn.cursor()
     cur.execute("SELECT count(*) FROM t")
     assert cur.fetchone() == (iterations,)
+
+
+# Test for a bug we had, where nextXid was incorrectly updated when the
+# XID counter reached 2 billion. The nextXid tracking logic incorrectly
+# treated 0 (InvalidTransactionId) as a regular XID, and after reaching
+# 2 billion, it started to look like a very new XID, which caused nextXid
+# to be immediately advanced to the next epoch.
+#
+def test_import_at_2bil(
+    neon_env_builder: NeonEnvBuilder,
+    test_output_dir: Path,
+    pg_distrib_dir: Path,
+    pg_bin,
+    vanilla_pg,
+):
+    neon_env_builder.enable_pageserver_remote_storage(RemoteStorageKind.LOCAL_FS)
+    env = neon_env_builder.init_start()
+    ps_http = env.pageserver.http_client()
+
+    # Set LD_LIBRARY_PATH in the env properly, otherwise we may use the wrong libpq.
+    # PgBin sets it automatically, but here we need to pipe psql output to the tar command.
+    psql_env = {"LD_LIBRARY_PATH": str(pg_distrib_dir / "lib")}
+
+    # Reset the vanilla Postgres instance to somewhat before 2 billion transactions.
+    pg_resetwal_path = os.path.join(pg_bin.pg_bin_path, "pg_resetwal")
+    cmd = [pg_resetwal_path, "--next-transaction-id=2129920000", "-D", str(vanilla_pg.pgdatadir)]
+    pg_bin.run_capture(cmd, env=psql_env)
+
+    vanilla_pg.start()
+    vanilla_pg.safe_psql("create user cloud_admin with password 'postgres' superuser")
+    vanilla_pg.safe_psql(
+        """create table tt as select 'long string to consume some space' || g
+     from generate_series(1,300000) g"""
+    )
+    assert vanilla_pg.safe_psql("select count(*) from tt") == [(300000,)]
+    vanilla_pg.safe_psql("CREATE TABLE t (t text);")
+    vanilla_pg.safe_psql("INSERT INTO t VALUES ('inserted in vanilla')")
+
+    endpoint_id = "ep-import_from_vanilla"
+    tenant = TenantId.generate()
+    timeline = TimelineId.generate()
+
+    env.pageserver.tenant_create(tenant)
+
+    # Take basebackup
+    basebackup_dir = os.path.join(test_output_dir, "basebackup")
+    base_tar = os.path.join(basebackup_dir, "base.tar")
+    wal_tar = os.path.join(basebackup_dir, "pg_wal.tar")
+    os.mkdir(basebackup_dir)
+    vanilla_pg.safe_psql("CHECKPOINT")
+    pg_bin.run(
+        [
+            "pg_basebackup",
+            "-F",
+            "tar",
+            "-d",
+            vanilla_pg.connstr(),
+            "-D",
+            basebackup_dir,
+        ]
+    )
+
+    # Get start_lsn and end_lsn
+    with open(os.path.join(basebackup_dir, "backup_manifest")) as f:
+        manifest = json.load(f)
+        start_lsn = manifest["WAL-Ranges"][0]["Start-LSN"]
+        end_lsn = manifest["WAL-Ranges"][0]["End-LSN"]
+
+    def import_tar(base, wal):
+        env.neon_cli.raw_cli(
+            [
+                "timeline",
+                "import",
+                "--tenant-id",
+                str(tenant),
+                "--timeline-id",
+                str(timeline),
+                "--node-name",
+                endpoint_id,
+                "--base-lsn",
+                start_lsn,
+                "--base-tarfile",
+                base,
+                "--end-lsn",
+                end_lsn,
+                "--wal-tarfile",
+                wal,
+                "--pg-version",
+                env.pg_version,
+            ]
+        )
+
+    # Importing correct backup works
+    import_tar(base_tar, wal_tar)
+    wait_for_last_record_lsn(ps_http, tenant, timeline, Lsn(end_lsn))
+
+    endpoint = env.endpoints.create_start(
+        endpoint_id,
+        tenant_id=tenant,
+        config_lines=[
+            "log_autovacuum_min_duration = 0",
+            "autovacuum_naptime='5 s'",
+        ],
+    )
+    assert endpoint.safe_psql("select count(*) from t") == [(1,)]
+
+    # Ok, consume
+    conn = endpoint.connect()
+    cur = conn.cursor()
+
+    # Install extension containing function needed for test
+    cur.execute("CREATE EXTENSION neon_test_utils")
+
+    # Advance nextXid close to 2 billion XIDs
+    while True:
+        xid = int(query_scalar(cur, "SELECT txid_current()"))
+        log.info(f"xid now {xid}")
+        # Consume 10k transactons at a time until we get to 2^31 - 200k
+        if xid < 2 * 1024 * 1024 * 1024 - 100000:
+            cur.execute("select test_consume_xids(50000);")
+        elif xid < 2 * 1024 * 1024 * 1024 - 10000:
+            cur.execute("select test_consume_xids(5000);")
+        else:
+            break
+
+    # Run a bunch of real INSERTs to cross over the 2 billion mark
+    # Use a begin-exception block to have a separate sub-XID for each insert.
+    cur.execute(
+        """
+        do $$
+        begin
+          for i in 1..10000 loop
+            -- Use a begin-exception block to generate a new subtransaction on each iteration
+            begin
+              insert into t values (i);
+            exception when others then
+              raise 'not expected %', sqlerrm;
+            end;
+          end loop;
+        end;
+        $$;
+        """
+    )
+    # A checkpoint writes a WAL record with xl_xid=0. Many other WAL
+    # records would have the same effect.
+    cur.execute("checkpoint")
+
+    # wait until pageserver receives that data
+    wait_for_wal_insert_lsn(env, endpoint, tenant, timeline)
+
+    # Restart endpoint
+    endpoint.stop()
+    endpoint.start()
+
+    conn = endpoint.connect()
+    cur = conn.cursor()
+    cur.execute("SELECT count(*) from t")
+    assert cur.fetchone() == (10000 + 1,)

test_runner/regress/test_vm_bits.py

@@ -1,3 +1,4 @@
+import pytest
 from fixtures.log_helper import log
 from fixtures.neon_fixtures import NeonEnv, fork_at_current_lsn
 
@@ -117,6 +118,8 @@ def test_vm_bit_clear(neon_simple_env: NeonEnv):
 # Test that the ALL_FROZEN VM bit is cleared correctly at a HEAP_LOCK
 # record.
 #
+# FIXME: This test is broken
+@pytest.mark.skip("See https://github.com/neondatabase/neon/pull/6412#issuecomment-1902072541")
 def test_vm_bit_clear_on_heap_lock(neon_simple_env: NeonEnv):
     env = neon_simple_env
 

workspace_hack/Cargo.toml

@@ -66,7 +66,7 @@ rustls = { version = "0.21", features = ["dangerous_configuration"] }
 scopeguard = { version = "1" }
 serde = { version = "1", features = ["alloc", "derive"] }
 serde_json = { version = "1", features = ["raw_value"] }
-smallvec = { version = "1", default-features = false, features = ["write"] }
+smallvec = { version = "1", default-features = false, features = ["serde", "write"] }
 subtle = { version = "2" }
 time = { version = "0.3", features = ["local-offset", "macros", "serde-well-known"] }
 tokio = { version = "1", features = ["fs", "io-std", "io-util", "macros", "net", "process", "rt-multi-thread", "signal", "test-util"] }