Add skeleton of parallel xxHash implementation

The new implementation lives in a separately allocated shared memory
area, which could be resized. Resizing it isn't actually implemented
yet, though. It would require some co-operation from the LFC code.

Cargo.lock

@@ -753,7 +753,6 @@ dependencies = [
  "axum",
  "axum-core",
  "bytes",
- "form_urlencoded",
  "futures-util",
  "headers",
  "http 1.1.0",
@@ -762,8 +761,6 @@ dependencies = [
  "mime",
  "pin-project-lite",
  "serde",
- "serde_html_form",
- "serde_path_to_error",
  "tower 0.5.2",
  "tower-layer",
  "tower-service",
@@ -903,6 +900,12 @@ version = "0.13.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8"
 
+[[package]]
+name = "base64"
+version = "0.20.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0ea22880d78093b0cbe17c89f64a7d457941e65759157ec6cb31a31d652b05e5"
+
 [[package]]
 name = "base64"
 version = "0.21.7"
@@ -1083,6 +1086,25 @@ version = "0.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5"
 
+[[package]]
+name = "cbindgen"
+version = "0.28.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "eadd868a2ce9ca38de7eeafdcec9c7065ef89b42b32f0839278d55f35c54d1ff"
+dependencies = [
+ "clap",
+ "heck 0.4.1",
+ "indexmap 2.9.0",
+ "log",
+ "proc-macro2",
+ "quote",
+ "serde",
+ "serde_json",
+ "syn 2.0.100",
+ "tempfile",
+ "toml",
+]
+
 [[package]]
 name = "cc"
 version = "1.2.16"
@@ -1209,7 +1231,7 @@ version = "4.5.18"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "4ac6a0c7b1a9e9a5186361f67dfa1b88213572f427fb9ab038efb2bd8c582dab"
 dependencies = [
- "heck",
+ "heck 0.5.0",
  "proc-macro2",
  "quote",
  "syn 2.0.100",
@@ -1267,6 +1289,14 @@ dependencies = [
  "unicode-width",
 ]
 
+[[package]]
+name = "communicator"
+version = "0.1.0"
+dependencies = [
+ "cbindgen",
+ "neon-shmem",
+]
+
 [[package]]
 name = "compute_api"
 version = "0.1.0"
@@ -1294,7 +1324,7 @@ dependencies = [
  "aws-smithy-types",
  "axum",
  "axum-extra",
- "base64 0.22.1",
+ "base64 0.13.1",
  "bytes",
  "camino",
  "cfg-if",
@@ -1420,7 +1450,7 @@ name = "control_plane"
 version = "0.1.0"
 dependencies = [
  "anyhow",
- "base64 0.22.1",
+ "base64 0.13.1",
  "camino",
  "clap",
  "comfy-table",
@@ -1442,7 +1472,6 @@ dependencies = [
  "regex",
  "reqwest",
  "safekeeper_api",
- "safekeeper_client",
  "scopeguard",
  "serde",
  "serde_json",
@@ -1934,7 +1963,7 @@ checksum = "0892a17df262a24294c382f0d5997571006e7a4348b4327557c4ff1cd4a8bccc"
 dependencies = [
  "darling",
  "either",
- "heck",
+ "heck 0.5.0",
  "proc-macro2",
  "quote",
  "syn 2.0.100",
@@ -2052,7 +2081,6 @@ dependencies = [
  "axum-extra",
  "camino",
  "camino-tempfile",
- "clap",
  "futures",
  "http-body-util",
  "itertools 0.10.5",
@@ -2499,6 +2527,18 @@ dependencies = [
  "wasm-bindgen",
 ]
 
+[[package]]
+name = "getrandom"
+version = "0.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "26145e563e54f2cadc477553f1ec5ee650b00862f0a58bcd12cbdc5f0ea2d2f4"
+dependencies = [
+ "cfg-if",
+ "libc",
+ "r-efi",
+ "wasi 0.14.2+wasi-0.2.4",
+]
+
 [[package]]
 name = "gettid"
 version = "0.1.3"
@@ -2711,6 +2751,12 @@ dependencies = [
  "http 1.1.0",
 ]
 
+[[package]]
+name = "heck"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8"
+
 [[package]]
 name = "heck"
 version = "0.5.0"
@@ -3647,7 +3693,7 @@ version = "0.0.22"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b9e6777fc80a575f9503d908c8b498782a6c3ee88a06cb416dc3941401e43b94"
 dependencies = [
- "heck",
+ "heck 0.5.0",
  "proc-macro2",
  "quote",
  "syn 2.0.100",
@@ -3709,7 +3755,7 @@ dependencies = [
  "procfs",
  "prometheus",
  "rand 0.8.5",
- "rand_distr",
+ "rand_distr 0.4.3",
  "twox-hash",
 ]
 
@@ -3797,7 +3843,11 @@ checksum = "e5ce46fe64a9d73be07dcbe690a38ce1b293be448fd8ce1e6c1b8062c9f72c6a"
 name = "neon-shmem"
 version = "0.1.0"
 dependencies = [
+ "criterion",
  "nix 0.30.1",
+ "rand 0.9.1",
+ "rand_distr 0.5.1",
+ "rustc-hash 1.1.0",
  "tempfile",
  "thiserror 1.0.69",
  "workspace_hack",
@@ -4812,7 +4862,7 @@ dependencies = [
 name = "postgres-protocol2"
 version = "0.1.0"
 dependencies = [
- "base64 0.22.1",
+ "base64 0.20.0",
  "byteorder",
  "bytes",
  "fallible-iterator",
@@ -5091,7 +5141,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "22505a5c94da8e3b7c2996394d1c933236c4d743e81a410bcca4e6989fc066a4"
 dependencies = [
  "bytes",
- "heck",
+ "heck 0.5.0",
  "itertools 0.12.1",
  "log",
  "multimap",
@@ -5112,7 +5162,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0c1318b19085f08681016926435853bbf7858f9c082d0999b80550ff5d9abe15"
 dependencies = [
  "bytes",
- "heck",
+ "heck 0.5.0",
  "itertools 0.12.1",
  "log",
  "multimap",
@@ -5184,7 +5234,7 @@ dependencies = [
  "aws-config",
  "aws-sdk-iam",
  "aws-sigv4",
- "base64 0.22.1",
+ "base64 0.13.1",
  "bstr",
  "bytes",
  "camino",
@@ -5237,7 +5287,7 @@ dependencies = [
  "postgres_backend",
  "pq_proto",
  "rand 0.8.5",
- "rand_distr",
+ "rand_distr 0.4.3",
  "rcgen",
  "redis",
  "regex",
@@ -5341,6 +5391,12 @@ dependencies = [
  "proc-macro2",
 ]
 
+[[package]]
+name = "r-efi"
+version = "5.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "74765f6d916ee2faa39bc8e68e4f3ed8949b48cccdac59983d287a7cb71ce9c5"
+
 [[package]]
 name = "rand"
 version = "0.7.3"
@@ -5365,6 +5421,16 @@ dependencies = [
  "rand_core 0.6.4",
 ]
 
+[[package]]
+name = "rand"
+version = "0.9.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9fbfd9d094a40bf3ae768db9361049ace4c0e04a4fd6b359518bd7b73a73dd97"
+dependencies = [
+ "rand_chacha 0.9.0",
+ "rand_core 0.9.3",
+]
+
 [[package]]
 name = "rand_chacha"
 version = "0.2.2"
@@ -5385,6 +5451,16 @@ dependencies = [
  "rand_core 0.6.4",
 ]
 
+[[package]]
+name = "rand_chacha"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb"
+dependencies = [
+ "ppv-lite86",
+ "rand_core 0.9.3",
+]
+
 [[package]]
 name = "rand_core"
 version = "0.5.1"
@@ -5403,6 +5479,15 @@ dependencies = [
  "getrandom 0.2.11",
 ]
 
+[[package]]
+name = "rand_core"
+version = "0.9.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "99d9a13982dcf210057a8a78572b2217b667c3beacbf3a0d8b454f6f82837d38"
+dependencies = [
+ "getrandom 0.3.3",
+]
+
 [[package]]
 name = "rand_distr"
 version = "0.4.3"
@@ -5413,6 +5498,16 @@ dependencies = [
  "rand 0.8.5",
 ]
 
+[[package]]
+name = "rand_distr"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463"
+dependencies = [
+ "num-traits",
+ "rand 0.9.1",
+]
+
 [[package]]
 name = "rand_hc"
 version = "0.2.0"
@@ -6419,19 +6514,6 @@ dependencies = [
  "syn 2.0.100",
 ]
 
-[[package]]
-name = "serde_html_form"
-version = "0.2.7"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9d2de91cf02bbc07cde38891769ccd5d4f073d22a40683aa4bc7a95781aaa2c4"
-dependencies = [
- "form_urlencoded",
- "indexmap 2.9.0",
- "itoa",
- "ryu",
- "serde",
-]
-
 [[package]]
 name = "serde_json"
 version = "1.0.125"
@@ -6488,17 +6570,15 @@ dependencies = [
 
 [[package]]
 name = "serde_with"
-version = "3.12.0"
+version = "2.3.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d6b6f7f2fcb69f747921f79f3926bd1e203fce4fef62c268dd3abfb6d86029aa"
+checksum = "07ff71d2c147a7b57362cead5e22f772cd52f6ab31cfcd9edcd7f6aeb2a0afbe"
 dependencies = [
- "base64 0.22.1",
+ "base64 0.13.1",
  "chrono",
  "hex",
  "indexmap 1.9.3",
- "indexmap 2.9.0",
  "serde",
- "serde_derive",
  "serde_json",
  "serde_with_macros",
  "time",
@@ -6506,9 +6586,9 @@ dependencies = [
 
 [[package]]
 name = "serde_with_macros"
-version = "3.12.0"
+version = "2.3.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8d00caa5193a3c8362ac2b73be6b9e768aa5a4b2f721d8f4b339600c3cb51f8e"
+checksum = "881b6f881b17d13214e5d494c939ebab463d01264ce1811e9d4ac3a882e7695f"
 dependencies = [
  "darling",
  "proc-macro2",
@@ -6914,7 +6994,7 @@ version = "0.26.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "4c6bee85a5a24955dc440386795aa378cd9cf82acd5f764469152d2270e581be"
 dependencies = [
- "heck",
+ "heck 0.5.0",
  "proc-macro2",
  "quote",
  "rustversion",
@@ -8213,6 +8293,15 @@ version = "0.11.0+wasi-snapshot-preview1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423"
 
+[[package]]
+name = "wasi"
+version = "0.14.2+wasi-0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9683f9a5a998d873c0d21fcbe3c083009670149a8fab228644b8bd36b2c48cb3"
+dependencies = [
+ "wit-bindgen-rt",
+]
+
 [[package]]
 name = "wasite"
 version = "0.1.0"
@@ -8570,6 +8659,15 @@ dependencies = [
  "windows-sys 0.48.0",
 ]
 
+[[package]]
+name = "wit-bindgen-rt"
+version = "0.39.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6f42320e61fe2cfd34354ecb597f86f413484a798ba44a8ca1165c58d42da6c1"
+dependencies = [
+ "bitflags 2.8.0",
+]
+
 [[package]]
 name = "workspace_hack"
 version = "0.1.0"
@@ -8579,6 +8677,7 @@ dependencies = [
  "anyhow",
  "axum",
  "axum-core",
+ "base64 0.13.1",
  "base64 0.21.7",
  "base64ct",
  "bytes",

Cargo.toml

@@ -44,6 +44,7 @@ members = [
     "libs/proxy/postgres-types2",
     "libs/proxy/tokio-postgres2",
     "endpoint_storage",
+    "pgxn/neon/communicator",
 ]
 
 [workspace.package]
@@ -71,8 +72,8 @@ aws-credential-types = "1.2.0"
 aws-sigv4 = { version = "1.2", features = ["sign-http"] }
 aws-types = "1.3"
 axum = { version = "0.8.1", features = ["ws"] }
-axum-extra = { version = "0.10.0", features = ["typed-header", "query"] }
-base64 = "0.22"
+axum-extra = { version = "0.10.0", features = ["typed-header"] }
+base64 = "0.13.0"
 bincode = "1.3"
 bindgen = "0.71"
 bit_field = "0.10.2"
@@ -171,7 +172,7 @@ sentry = { version = "0.37", default-features = false, features = ["backtrace",
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1"
 serde_path_to_error = "0.1"
-serde_with = { version = "3", features = [ "base64" ] }
+serde_with = { version = "2.0", features = [ "base64" ] }
 serde_assert = "0.5.0"
 sha2 = "0.10.2"
 signal-hook = "0.3"
@@ -251,6 +252,7 @@ desim = { version = "0.1", path = "./libs/desim" }
 endpoint_storage = { version = "0.0.1", path = "./endpoint_storage/" }
 http-utils = { version = "0.1", path = "./libs/http-utils/" }
 metrics = { version = "0.1", path = "./libs/metrics/" }
+neon-shmem = { version = "0.1", path = "./libs/neon-shmem/" }
 pageserver = { path = "./pageserver" }
 pageserver_api = { version = "0.1", path = "./libs/pageserver_api/" }
 pageserver_client = { path = "./pageserver/client" }
@@ -278,6 +280,7 @@ walproposer = { version = "0.1", path = "./libs/walproposer/" }
 workspace_hack = { version = "0.1", path = "./workspace_hack/" }
 
 ## Build dependencies
+cbindgen = "0.28.0"
 criterion = "0.5.1"
 rcgen = "0.13"
 rstest = "0.18"

Dockerfile

@@ -110,19 +110,6 @@ RUN set -e \
 	# System postgres for use with client libraries (e.g. in storage controller)
         postgresql-15 \
         openssl \
-        unzip \
-        curl \
-    && ARCH=$(uname -m) \
-    && if [ "$ARCH" = "x86_64" ]; then \
-        curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip"; \
-    elif [ "$ARCH" = "aarch64" ]; then \
-        curl "https://awscli.amazonaws.com/awscli-exe-linux-aarch64.zip" -o "awscliv2.zip"; \
-    else \
-        echo "Unsupported architecture: $ARCH" && exit 1; \
-    fi \
-    && unzip awscliv2.zip \
-    && ./aws/install \
-    && rm -rf aws awscliv2.zip \
     && rm -f /etc/apt/apt.conf.d/80-retries \
     && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* \
     && useradd -d /data neon \

Makefile

@@ -18,10 +18,12 @@ ifeq ($(BUILD_TYPE),release)
 	PG_LDFLAGS = $(LDFLAGS)
 	# Unfortunately, `--profile=...` is a nightly feature
 	CARGO_BUILD_FLAGS += --release
+	NEON_CARGO_ARTIFACT_TARGET_DIR = $(ROOT_PROJECT_DIR)/target/release
 else ifeq ($(BUILD_TYPE),debug)
 	PG_CONFIGURE_OPTS = --enable-debug --with-openssl --enable-cassert --enable-depend
 	PG_CFLAGS += -O0 -g3 $(CFLAGS)
 	PG_LDFLAGS = $(LDFLAGS)
+	NEON_CARGO_ARTIFACT_TARGET_DIR = $(ROOT_PROJECT_DIR)/target/debug
 else
 	$(error Bad build type '$(BUILD_TYPE)', see Makefile for options)
 endif
@@ -180,11 +182,16 @@ postgres-check-%: postgres-%
 
 .PHONY: neon-pg-ext-%
 neon-pg-ext-%: postgres-%
+	+@echo "Compiling communicator $*"
+	$(CARGO_CMD_PREFIX) cargo build -p communicator $(CARGO_BUILD_FLAGS)
+
 	+@echo "Compiling neon $*"
 	mkdir -p $(POSTGRES_INSTALL_DIR)/build/neon-$*
 	$(MAKE) PG_CONFIG=$(POSTGRES_INSTALL_DIR)/$*/bin/pg_config COPT='$(COPT)' \
+		LIBCOMMUNICATOR_PATH=$(NEON_CARGO_ARTIFACT_TARGET_DIR) \
 		-C $(POSTGRES_INSTALL_DIR)/build/neon-$* \
 		-f $(ROOT_PROJECT_DIR)/pgxn/neon/Makefile install
+
 	+@echo "Compiling neon_walredo $*"
 	mkdir -p $(POSTGRES_INSTALL_DIR)/build/neon-walredo-$*
 	$(MAKE) PG_CONFIG=$(POSTGRES_INSTALL_DIR)/$*/bin/pg_config COPT='$(COPT)' \

compute_tools/src/compute.rs

@@ -785,7 +785,7 @@ impl ComputeNode {
         self.spawn_extension_stats_task();
 
         if pspec.spec.autoprewarm {
-            self.prewarm_lfc(None);
+            self.prewarm_lfc();
         }
         Ok(())
     }

compute_tools/src/compute_prewarm.rs

@@ -25,16 +25,11 @@ struct EndpointStoragePair {
 }
 
 const KEY: &str = "lfc_state";
-impl EndpointStoragePair {
-    /// endpoint_id is set to None while prewarming from other endpoint, see replica promotion
-    /// If not None, takes precedence over pspec.spec.endpoint_id
-    fn from_spec_and_endpoint(
-        pspec: &crate::compute::ParsedSpec,
-        endpoint_id: Option<String>,
-    ) -> Result<Self> {
-        let endpoint_id = endpoint_id.as_ref().or(pspec.spec.endpoint_id.as_ref());
-        let Some(ref endpoint_id) = endpoint_id else {
-            bail!("pspec.endpoint_id missing, other endpoint_id not provided")
+impl TryFrom<&crate::compute::ParsedSpec> for EndpointStoragePair {
+    type Error = anyhow::Error;
+    fn try_from(pspec: &crate::compute::ParsedSpec) -> Result<Self, Self::Error> {
+        let Some(ref endpoint_id) = pspec.spec.endpoint_id else {
+            bail!("pspec.endpoint_id missing")
         };
         let Some(ref base_uri) = pspec.endpoint_storage_addr else {
             bail!("pspec.endpoint_storage_addr missing")
@@ -89,7 +84,7 @@ impl ComputeNode {
     }
 
     /// Returns false if there is a prewarm request ongoing, true otherwise
-    pub fn prewarm_lfc(self: &Arc<Self>, from_endpoint: Option<String>) -> bool {
+    pub fn prewarm_lfc(self: &Arc<Self>) -> bool {
         crate::metrics::LFC_PREWARM_REQUESTS.inc();
         {
             let state = &mut self.state.lock().unwrap().lfc_prewarm_state;
@@ -102,7 +97,7 @@ impl ComputeNode {
 
         let cloned = self.clone();
         spawn(async move {
-            let Err(err) = cloned.prewarm_impl(from_endpoint).await else {
+            let Err(err) = cloned.prewarm_impl().await else {
                 cloned.state.lock().unwrap().lfc_prewarm_state = LfcPrewarmState::Completed;
                 return;
             };
@@ -114,14 +109,13 @@ impl ComputeNode {
         true
     }
 
-    /// from_endpoint: None for endpoint managed by this compute_ctl
-    fn endpoint_storage_pair(&self, from_endpoint: Option<String>) -> Result<EndpointStoragePair> {
+    fn endpoint_storage_pair(&self) -> Result<EndpointStoragePair> {
         let state = self.state.lock().unwrap();
-        EndpointStoragePair::from_spec_and_endpoint(state.pspec.as_ref().unwrap(), from_endpoint)
+        state.pspec.as_ref().unwrap().try_into()
     }
 
-    async fn prewarm_impl(&self, from_endpoint: Option<String>) -> Result<()> {
-        let EndpointStoragePair { url, token } = self.endpoint_storage_pair(from_endpoint)?;
+    async fn prewarm_impl(&self) -> Result<()> {
+        let EndpointStoragePair { url, token } = self.endpoint_storage_pair()?;
         info!(%url, "requesting LFC state from endpoint storage");
 
         let request = Client::new().get(&url).bearer_auth(token);
@@ -179,7 +173,7 @@ impl ComputeNode {
     }
 
     async fn offload_lfc_impl(&self) -> Result<()> {
-        let EndpointStoragePair { url, token } = self.endpoint_storage_pair(None)?;
+        let EndpointStoragePair { url, token } = self.endpoint_storage_pair()?;
         info!(%url, "requesting LFC state from postgres");
 
         let mut compressed = Vec::new();

compute_tools/src/http/routes/lfc.rs

@@ -2,7 +2,6 @@ use crate::compute_prewarm::LfcPrewarmStateWithProgress;
 use crate::http::JsonResponse;
 use axum::response::{IntoResponse, Response};
 use axum::{Json, http::StatusCode};
-use axum_extra::extract::OptionalQuery;
 use compute_api::responses::LfcOffloadState;
 type Compute = axum::extract::State<std::sync::Arc<crate::compute::ComputeNode>>;
 
@@ -17,16 +16,8 @@ pub(in crate::http) async fn offload_state(compute: Compute) -> Json<LfcOffloadS
     Json(compute.lfc_offload_state())
 }
 
-#[derive(serde::Deserialize)]
-pub struct PrewarmQuery {
-    pub from_endpoint: String,
-}
-
-pub(in crate::http) async fn prewarm(
-    compute: Compute,
-    OptionalQuery(query): OptionalQuery<PrewarmQuery>,
-) -> Response {
-    if compute.prewarm_lfc(query.map(|q| q.from_endpoint)) {
+pub(in crate::http) async fn prewarm(compute: Compute) -> Response {
+    if compute.prewarm_lfc() {
         StatusCode::ACCEPTED.into_response()
     } else {
         JsonResponse::error(

control_plane/Cargo.toml

@@ -36,7 +36,6 @@ pageserver_api.workspace = true
 pageserver_client.workspace = true
 postgres_backend.workspace = true
 safekeeper_api.workspace = true
-safekeeper_client.workspace = true
 postgres_connection.workspace = true
 storage_broker.workspace = true
 http-utils.workspace = true

control_plane/src/bin/neon_local.rs

@@ -45,7 +45,7 @@ use pageserver_api::models::{
 use pageserver_api::shard::{DEFAULT_STRIPE_SIZE, ShardCount, ShardStripeSize, TenantShardId};
 use postgres_backend::AuthType;
 use postgres_connection::parse_host_port;
-use safekeeper_api::membership::{SafekeeperGeneration, SafekeeperId};
+use safekeeper_api::membership::SafekeeperGeneration;
 use safekeeper_api::{
     DEFAULT_HTTP_LISTEN_PORT as DEFAULT_SAFEKEEPER_HTTP_PORT,
     DEFAULT_PG_LISTEN_PORT as DEFAULT_SAFEKEEPER_PG_PORT,
@@ -1255,45 +1255,6 @@ async fn handle_timeline(cmd: &TimelineCmd, env: &mut local_env::LocalEnv) -> Re
             pageserver
                 .timeline_import(tenant_id, timeline_id, base, pg_wal, args.pg_version)
                 .await?;
-            if env.storage_controller.timelines_onto_safekeepers {
-                println!("Creating timeline on safekeeper ...");
-                let timeline_info = pageserver
-                    .timeline_info(
-                        TenantShardId::unsharded(tenant_id),
-                        timeline_id,
-                        pageserver_client::mgmt_api::ForceAwaitLogicalSize::No,
-                    )
-                    .await?;
-                let default_sk = SafekeeperNode::from_env(env, env.safekeepers.first().unwrap());
-                let default_host = default_sk
-                    .conf
-                    .listen_addr
-                    .clone()
-                    .unwrap_or_else(|| "localhost".to_string());
-                let mconf = safekeeper_api::membership::Configuration {
-                    generation: SafekeeperGeneration::new(1),
-                    members: safekeeper_api::membership::MemberSet {
-                        m: vec![SafekeeperId {
-                            host: default_host,
-                            id: default_sk.conf.id,
-                            pg_port: default_sk.conf.pg_port,
-                        }],
-                    },
-                    new_members: None,
-                };
-                let pg_version = args.pg_version * 10000;
-                let req = safekeeper_api::models::TimelineCreateRequest {
-                    tenant_id,
-                    timeline_id,
-                    mconf,
-                    pg_version,
-                    system_id: None,
-                    wal_seg_size: None,
-                    start_lsn: timeline_info.last_record_lsn,
-                    commit_lsn: None,
-                };
-                default_sk.create_timeline(&req).await?;
-            }
             env.register_branch_mapping(branch_name.to_string(), tenant_id, timeline_id)?;
             println!("Done");
         }

control_plane/src/endpoint.rs

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

control_plane/src/pageserver.rs

@@ -635,16 +635,4 @@ impl PageServerNode {
 
         Ok(())
     }
-    pub async fn timeline_info(
-        &self,
-        tenant_shard_id: TenantShardId,
-        timeline_id: TimelineId,
-        force_await_logical_size: mgmt_api::ForceAwaitLogicalSize,
-    ) -> anyhow::Result<TimelineInfo> {
-        let timeline_info = self
-            .http_client
-            .timeline_info(tenant_shard_id, timeline_id, force_await_logical_size)
-            .await?;
-        Ok(timeline_info)
-    }
 }

control_plane/src/safekeeper.rs

@@ -6,6 +6,7 @@
 //!   .neon/safekeepers/<safekeeper id>
 //! ```
 use std::error::Error as _;
+use std::future::Future;
 use std::io::Write;
 use std::path::PathBuf;
 use std::time::Duration;
@@ -13,9 +14,9 @@ use std::{io, result};
 
 use anyhow::Context;
 use camino::Utf8PathBuf;
+use http_utils::error::HttpErrorBody;
 use postgres_connection::PgConnectionConfig;
-use safekeeper_api::models::TimelineCreateRequest;
-use safekeeper_client::mgmt_api;
+use reqwest::{IntoUrl, Method};
 use thiserror::Error;
 use utils::auth::{Claims, Scope};
 use utils::id::NodeId;
@@ -34,14 +35,25 @@ pub enum SafekeeperHttpError {
 
 type Result<T> = result::Result<T, SafekeeperHttpError>;
 
-fn err_from_client_err(err: mgmt_api::Error) -> SafekeeperHttpError {
-    use mgmt_api::Error::*;
-    match err {
-        ApiError(_, str) => SafekeeperHttpError::Response(str),
-        Cancelled => SafekeeperHttpError::Response("Cancelled".to_owned()),
-        ReceiveBody(err) => SafekeeperHttpError::Transport(err),
-        ReceiveErrorBody(err) => SafekeeperHttpError::Response(err),
-        Timeout(str) => SafekeeperHttpError::Response(format!("timeout: {str}")),
+pub(crate) trait ResponseErrorMessageExt: Sized {
+    fn error_from_body(self) -> impl Future<Output = Result<Self>> + Send;
+}
+
+impl ResponseErrorMessageExt for reqwest::Response {
+    async fn error_from_body(self) -> Result<Self> {
+        let status = self.status();
+        if !(status.is_client_error() || status.is_server_error()) {
+            return Ok(self);
+        }
+
+        // reqwest does not export its error construction utility functions, so let's craft the message ourselves
+        let url = self.url().to_owned();
+        Err(SafekeeperHttpError::Response(
+            match self.json::<HttpErrorBody>().await {
+                Ok(err_body) => format!("Error: {}", err_body.msg),
+                Err(_) => format!("Http error ({}) at {}.", status.as_u16(), url),
+            },
+        ))
     }
 }
 
@@ -58,8 +70,9 @@ pub struct SafekeeperNode {
 
     pub pg_connection_config: PgConnectionConfig,
     pub env: LocalEnv,
-    pub http_client: mgmt_api::Client,
+    pub http_client: reqwest::Client,
     pub listen_addr: String,
+    pub http_base_url: String,
 }
 
 impl SafekeeperNode {
@@ -69,14 +82,13 @@ impl SafekeeperNode {
         } else {
             "127.0.0.1".to_string()
         };
-        let jwt = None;
-        let http_base_url = format!("http://{}:{}", listen_addr, conf.http_port);
         SafekeeperNode {
             id: conf.id,
             conf: conf.clone(),
             pg_connection_config: Self::safekeeper_connection_config(&listen_addr, conf.pg_port),
             env: env.clone(),
-            http_client: mgmt_api::Client::new(env.create_http_client(), http_base_url, jwt),
+            http_client: env.create_http_client(),
+            http_base_url: format!("http://{}:{}/v1", listen_addr, conf.http_port),
             listen_addr,
         }
     }
@@ -266,19 +278,20 @@ impl SafekeeperNode {
         )
     }
 
-    pub async fn check_status(&self) -> Result<()> {
-        self.http_client
-            .status()
-            .await
-            .map_err(err_from_client_err)?;
-        Ok(())
+    fn http_request<U: IntoUrl>(&self, method: Method, url: U) -> reqwest::RequestBuilder {
+        // TODO: authentication
+        //if self.env.auth_type == AuthType::NeonJWT {
+        //    builder = builder.bearer_auth(&self.env.safekeeper_auth_token)
+        //}
+        self.http_client.request(method, url)
     }
 
-    pub async fn create_timeline(&self, req: &TimelineCreateRequest) -> Result<()> {
-        self.http_client
-            .create_timeline(req)
-            .await
-            .map_err(err_from_client_err)?;
+    pub async fn check_status(&self) -> Result<()> {
+        self.http_request(Method::GET, format!("{}/{}", self.http_base_url, "status"))
+            .send()
+            .await?
+            .error_from_body()
+            .await?;
         Ok(())
     }
 }

control_plane/storcon_cli/src/main.rs

@@ -61,16 +61,10 @@ enum Command {
         #[arg(long)]
         scheduling: Option<NodeSchedulingPolicy>,
     },
-    // Set a node status as deleted.
     NodeDelete {
         #[arg(long)]
         node_id: NodeId,
     },
-    /// Delete a tombstone of node from the storage controller.
-    NodeDeleteTombstone {
-        #[arg(long)]
-        node_id: NodeId,
-    },
     /// Modify a tenant's policies in the storage controller
     TenantPolicy {
         #[arg(long)]
@@ -88,8 +82,6 @@ enum Command {
     },
     /// List nodes known to the storage controller
     Nodes {},
-    /// List soft deleted nodes known to the storage controller
-    NodeTombstones {},
     /// List tenants known to the storage controller
     Tenants {
         /// If this field is set, it will list the tenants on a specific node
@@ -908,39 +900,6 @@ async fn main() -> anyhow::Result<()> {
                 .dispatch::<(), ()>(Method::DELETE, format!("control/v1/node/{node_id}"), None)
                 .await?;
         }
-        Command::NodeDeleteTombstone { node_id } => {
-            storcon_client
-                .dispatch::<(), ()>(
-                    Method::DELETE,
-                    format!("debug/v1/tombstone/{node_id}"),
-                    None,
-                )
-                .await?;
-        }
-        Command::NodeTombstones {} => {
-            let mut resp = storcon_client
-                .dispatch::<(), Vec<NodeDescribeResponse>>(
-                    Method::GET,
-                    "debug/v1/tombstone".to_string(),
-                    None,
-                )
-                .await?;
-
-            resp.sort_by(|a, b| a.listen_http_addr.cmp(&b.listen_http_addr));
-
-            let mut table = comfy_table::Table::new();
-            table.set_header(["Id", "Hostname", "AZ", "Scheduling", "Availability"]);
-            for node in resp {
-                table.add_row([
-                    format!("{}", node.id),
-                    node.listen_http_addr,
-                    node.availability_zone_id,
-                    format!("{:?}", node.scheduling),
-                    format!("{:?}", node.availability),
-                ]);
-            }
-            println!("{table}");
-        }
         Command::TenantSetTimeBasedEviction {
             tenant_id,
             period,

endpoint_storage/Cargo.toml

@@ -8,7 +8,6 @@ anyhow.workspace = true
 axum-extra.workspace = true
 axum.workspace = true
 camino.workspace = true
-clap.workspace = true
 futures.workspace = true
 jsonwebtoken.workspace = true
 prometheus.workspace = true

endpoint_storage/src/main.rs

@@ -4,8 +4,6 @@
 //! for large computes.
 mod app;
 use anyhow::Context;
-use clap::Parser;
-use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 use tracing::info;
 use utils::logging;
 
@@ -14,26 +12,9 @@ const fn max_upload_file_limit() -> usize {
     100 * 1024 * 1024
 }
 
-const fn listen() -> SocketAddr {
-    SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 51243)
-}
-
-#[derive(Parser)]
-struct Args {
-    #[arg(exclusive = true)]
-    config_file: Option<String>,
-    #[arg(long, default_value = "false", requires = "config")]
-    /// to allow testing k8s helm chart where we don't have s3 credentials
-    no_s3_check_on_startup: bool,
-    #[arg(long, value_name = "FILE")]
-    /// inline config mode for k8s helm chart
-    config: Option<String>,
-}
-
 #[derive(serde::Deserialize)]
 #[serde(tag = "type")]
 struct Config {
-    #[serde(default = "listen")]
     listen: std::net::SocketAddr,
     pemfile: camino::Utf8PathBuf,
     #[serde(flatten)]
@@ -50,18 +31,13 @@ async fn main() -> anyhow::Result<()> {
         logging::Output::Stdout,
     )?;
 
-    let args = Args::parse();
-    let config: Config = if let Some(config_path) = args.config_file {
-        info!("Reading config from {config_path}");
-        let config = std::fs::read_to_string(config_path)?;
-        serde_json::from_str(&config).context("parsing config")?
-    } else if let Some(config) = args.config {
-        info!("Reading inline config");
-        serde_json::from_str(&config).context("parsing config")?
-    } else {
-        anyhow::bail!("Supply either config file path or --config=inline-config");
-    };
-
+    let config: String = std::env::args().skip(1).take(1).collect();
+    if config.is_empty() {
+        anyhow::bail!("Usage: endpoint_storage config.json")
+    }
+    info!("Reading config from {config}");
+    let config = std::fs::read_to_string(config.clone())?;
+    let config: Config = serde_json::from_str(&config).context("parsing config")?;
     info!("Reading pemfile from {}", config.pemfile.clone());
     let pemfile = std::fs::read(config.pemfile.clone())?;
     info!("Loading public key from {}", config.pemfile.clone());
@@ -72,9 +48,7 @@ async fn main() -> anyhow::Result<()> {
 
     let storage = remote_storage::GenericRemoteStorage::from_config(&config.storage_config).await?;
     let cancel = tokio_util::sync::CancellationToken::new();
-    if !args.no_s3_check_on_startup {
-        app::check_storage_permissions(&storage, cancel.clone()).await?;
-    }
+    app::check_storage_permissions(&storage, cancel.clone()).await?;
 
     let proxy = std::sync::Arc::new(endpoint_storage::Storage {
         auth,

libs/neon-shmem/Cargo.toml

@@ -6,8 +6,20 @@ license.workspace = true
 
 [dependencies]
 thiserror.workspace = true
-nix.workspace=true
+nix.workspace = true
 workspace_hack = { version = "0.1", path = "../../workspace_hack" }
+rustc-hash = { version = "2.1.1" }
+
+[dev-dependencies]
+criterion = { workspace = true, features = ["html_reports"] }
+rand = "0.9.1"
+rand_distr = "0.5.1"
+xxhash-rust = { version = "0.8.15", features = ["xxh3"] }
+ahash.workspace = true
 
 [target.'cfg(target_os = "macos")'.dependencies]
 tempfile = "3.14.0"
+
+[[bench]]
+name = "hmap_resize"
+harness = false

libs/neon-shmem/src/hash.rs

@@ -0,0 +1,438 @@
+//! Hash table implementation on top of 'shmem'
+//!
+//! Features required in the long run by the communicator project:
+//!
+//! [X] Accessible from both Postgres processes and rust threads in the communicator process
+//! [X] Low latency
+//! [ ] Scalable to lots of concurrent accesses (currently relies on caller for locking)
+//! [ ] Resizable
+
+use std::fmt::Debug;
+use std::hash::{Hash, Hasher, BuildHasher};
+use std::mem::MaybeUninit;
+
+use rustc_hash::FxBuildHasher;
+
+use crate::shmem::ShmemHandle;
+
+mod core;
+pub mod entry;
+
+#[cfg(test)]
+mod tests;
+
+mod optim;
+
+use core::{CoreHashMap, INVALID_POS};
+use entry::{Entry, OccupiedEntry};
+
+
+pub struct HashMapInit<'a, K, V, S = rustc_hash::FxBuildHasher> {
+    // Hash table can be allocated in a fixed memory area, or in a resizeable ShmemHandle.
+    shmem_handle: Option<ShmemHandle>,
+    shared_ptr: *mut HashMapShared<'a, K, V>,
+	shared_size: usize,
+	hasher: S,
+	num_buckets: u32,
+}
+
+pub struct HashMapAccess<'a, K, V, S = rustc_hash::FxBuildHasher> {
+    shmem_handle: Option<ShmemHandle>,
+    shared_ptr: *mut HashMapShared<'a, K, V>,
+	hasher: S,
+}
+
+unsafe impl<'a, K: Sync, V: Sync, S> Sync for HashMapAccess<'a, K, V, S> {}
+unsafe impl<'a, K: Send, V: Send, S> Send for HashMapAccess<'a, K, V, S> {}
+
+impl<'a, K: Clone + Hash + Eq, V, S> HashMapInit<'a, K, V, S> {
+	pub fn with_hasher(self, hasher: S) -> HashMapInit<'a, K, V, S> {
+		Self { hasher, ..self }
+	}
+	
+	pub fn estimate_size(num_buckets: u32) -> usize {
+        // add some margin to cover alignment etc.
+        CoreHashMap::<K, V>::estimate_size(num_buckets) + size_of::<HashMapShared<K, V>>() + 1000
+    }
+	
+    pub fn attach_writer(self) -> HashMapAccess<'a, K, V, S> {
+		let mut ptr: *mut u8 = self.shared_ptr.cast();
+        let end_ptr: *mut u8 = unsafe { ptr.add(self.shared_size) };
+        ptr = unsafe { ptr.add(ptr.align_offset(align_of::<HashMapShared<K, V>>())) };
+        let shared_ptr: *mut HashMapShared<K, V> = ptr.cast();
+        ptr = unsafe { ptr.add(size_of::<HashMapShared<K, V>>()) };
+ 
+        // carve out the buckets
+        ptr = unsafe { ptr.byte_add(ptr.align_offset(align_of::<core::LinkedKey<K>>())) };
+        let keys_ptr = ptr;
+        ptr = unsafe { ptr.add(size_of::<core::LinkedKey<K>>() * self.num_buckets as usize) };
+		
+        ptr = unsafe { ptr.byte_add(ptr.align_offset(align_of::<Option<V>>())) };
+        let vals_ptr = ptr;
+        ptr = unsafe { ptr.add(size_of::<Option<V>>() * self.num_buckets as usize) };
+		
+        // use remaining space for the dictionary
+        ptr = unsafe { ptr.byte_add(ptr.align_offset(align_of::<u32>())) };
+        assert!(ptr.addr() < end_ptr.addr());
+        let dictionary_ptr = ptr;
+        let dictionary_size = unsafe { end_ptr.byte_offset_from(ptr) / size_of::<u32>() as isize };
+        assert!(dictionary_size > 0);
+
+        let keys =
+            unsafe { std::slice::from_raw_parts_mut(keys_ptr.cast(), self.num_buckets as usize) };
+		let vals =
+            unsafe { std::slice::from_raw_parts_mut(vals_ptr.cast(), self.num_buckets as usize) };
+        let dictionary = unsafe {
+            std::slice::from_raw_parts_mut(dictionary_ptr.cast(), dictionary_size as usize)
+        };
+        let hashmap = CoreHashMap::new(keys, vals, dictionary);
+        unsafe {
+            std::ptr::write(shared_ptr, HashMapShared { inner: hashmap });
+        }
+		
+        HashMapAccess {
+            shmem_handle: self.shmem_handle,
+            shared_ptr: self.shared_ptr,
+			hasher: self.hasher,
+        }
+    }
+
+    pub fn attach_reader(self) -> HashMapAccess<'a, K, V, S> {
+        // no difference to attach_writer currently
+         self.attach_writer()
+    }
+}
+
+/// This is stored in the shared memory area
+///
+/// NOTE: We carve out the parts from a contiguous chunk. Growing and shrinking the hash table
+/// relies on the memory layout! The data structures are laid out in the contiguous shared memory
+/// area as follows:
+///
+/// HashMapShared
+/// [buckets]
+/// [dictionary]
+///
+/// In between the above parts, there can be padding bytes to align the parts correctly.
+struct HashMapShared<'a, K, V> {
+    inner: CoreHashMap<'a, K, V>	
+}
+
+impl<'a, K, V> HashMapInit<'a, K, V, rustc_hash::FxBuildHasher>
+where
+	K: Clone + Hash + Eq
+{
+	pub fn with_fixed(
+		num_buckets: u32,
+        area: &'a mut [MaybeUninit<u8>],
+    ) -> HashMapInit<'a, K, V> {
+		Self {
+			num_buckets,
+			shmem_handle: None,
+			shared_ptr: area.as_mut_ptr().cast(),
+			shared_size: area.len(),
+			hasher: rustc_hash::FxBuildHasher::default(),
+		}		
+    }
+
+    /// Initialize a new hash map in the given shared memory area
+    pub fn with_shmem(num_buckets: u32, shmem: ShmemHandle) -> HashMapInit<'a, K, V> {
+		let size = Self::estimate_size(num_buckets);
+		shmem
+            .set_size(size)
+            .expect("could not resize shared memory area");
+		Self {
+			num_buckets,
+			shared_ptr: shmem.data_ptr.as_ptr().cast(),
+			shmem_handle: Some(shmem),
+			shared_size: size,
+			hasher: rustc_hash::FxBuildHasher::default()
+		}
+    }
+
+	pub fn new_resizeable_named(num_buckets: u32, max_buckets: u32, name: &str) -> HashMapInit<'a, K, V> {
+		let size = Self::estimate_size(num_buckets);
+		let max_size = Self::estimate_size(max_buckets);
+		let shmem = ShmemHandle::new(name, size, max_size)
+			.expect("failed to make shared memory area");
+		
+		Self {
+			num_buckets,
+			shared_ptr: shmem.data_ptr.as_ptr().cast(),
+			shmem_handle: Some(shmem),
+			shared_size: size,
+			hasher: rustc_hash::FxBuildHasher::default()
+		}
+	}
+
+	pub fn new_resizeable(num_buckets: u32, max_buckets: u32) -> HashMapInit<'a, K, V> {
+		use std::sync::atomic::{AtomicUsize, Ordering};
+		const COUNTER: AtomicUsize = AtomicUsize::new(0);
+		let val = COUNTER.fetch_add(1, Ordering::Relaxed);
+		let name = format!("neon_shmem_hmap{}", val);
+		Self::new_resizeable_named(num_buckets, max_buckets, &name)
+	}
+}
+
+impl<'a, K, V, S: BuildHasher> HashMapAccess<'a, K, V, S>
+where
+    K: Clone + Hash + Eq,
+{
+    pub fn get_hash_value(&self, key: &K) -> u64 {
+		self.hasher.hash_one(key)        
+    }
+
+    pub fn get_with_hash<'e>(&'e self, key: &K, hash: u64) -> Option<&'e V> {
+        let map = unsafe { self.shared_ptr.as_ref() }.unwrap();
+
+        map.inner.get_with_hash(key, hash)
+    }
+
+    pub fn entry_with_hash(&mut self, key: K, hash: u64) -> Entry<'a, '_, K, V> {
+        let map = unsafe { self.shared_ptr.as_mut() }.unwrap();
+
+        map.inner.entry_with_hash(key, hash)
+    }
+
+    pub fn remove_with_hash(&mut self, key: &K, hash: u64) {
+        let map = unsafe { self.shared_ptr.as_mut() }.unwrap();
+
+        match map.inner.entry_with_hash(key.clone(), hash) {
+            Entry::Occupied(e) => {
+                e.remove();
+            }
+            Entry::Vacant(_) => {}
+        };
+    }
+
+    pub fn entry_at_bucket(&mut self, pos: usize) -> Option<OccupiedEntry<'a, '_, K, V>> {
+        let map = unsafe { self.shared_ptr.as_mut() }.unwrap();
+        map.inner.entry_at_bucket(pos)
+    }
+
+    pub fn get_num_buckets(&self) -> usize {
+        let map = unsafe { self.shared_ptr.as_ref() }.unwrap();
+        map.inner.get_num_buckets()
+    }
+
+    /// Return the key and value stored in bucket with given index. This can be used to
+    /// iterate through the hash map. (An Iterator might be nicer. The communicator's
+    /// clock algorithm needs to _slowly_ iterate through all buckets with its clock hand,
+    /// without holding a lock. If we switch to an Iterator, it must not hold the lock.)
+    pub fn get_at_bucket(&self, pos: usize) -> Option<(&K, &V)> {
+        let map = unsafe { self.shared_ptr.as_ref() }.unwrap();
+
+        if pos >= map.inner.keys.len() {
+            return None;
+        }
+        let key = &map.inner.keys[pos];
+		key.inner.as_ref().map(|k| (k, map.inner.vals[pos].as_ref().unwrap()))
+    }
+
+    pub fn get_bucket_for_value(&self, val_ptr: *const V) -> usize {
+        let map = unsafe { self.shared_ptr.as_ref() }.unwrap();
+
+        let origin = map.inner.vals.as_ptr();
+        let idx = (val_ptr as usize - origin as usize) / (size_of::<V>() as usize);
+        assert!(idx < map.inner.vals.len());
+
+        idx
+    }
+
+    // for metrics
+    pub fn get_num_buckets_in_use(&self) -> usize {
+        let map = unsafe { self.shared_ptr.as_ref() }.unwrap();
+        map.inner.buckets_in_use as usize
+    }
+
+	pub fn clear(&mut self) {
+		let map = unsafe { self.shared_ptr.as_mut() }.unwrap();
+        let inner = &mut map.inner;
+        inner.clear()
+	}
+	
+	/// Helper function that abstracts the common logic between growing and shrinking.
+	/// The only significant difference in the rehashing step is how many buckets to rehash.
+	fn rehash_dict(
+		&mut self,
+		inner: &mut CoreHashMap<'a, K, V>,
+		keys_ptr: *mut core::LinkedKey<K>,
+		end_ptr: *mut u8,
+		num_buckets: u32,
+		rehash_buckets: u32,
+	) {
+		inner.free_head = INVALID_POS;
+		
+		// Recalculate the dictionary
+        let keys;
+        let dictionary;
+        unsafe {
+            let keys_end_ptr = keys_ptr.add(num_buckets as usize);
+            let buckets_end_ptr: *mut u8 = (keys_end_ptr as *mut u8)
+				.add(size_of::<Option<V>>() * num_buckets as usize);
+			let dictionary_ptr: *mut u32 = buckets_end_ptr
+				.byte_add(buckets_end_ptr.align_offset(align_of::<u32>()))
+                .cast();
+            let dictionary_size: usize =
+                end_ptr.byte_offset_from(buckets_end_ptr) as usize / size_of::<u32>();
+
+            keys = std::slice::from_raw_parts_mut(keys_ptr, num_buckets as usize);
+            dictionary = std::slice::from_raw_parts_mut(dictionary_ptr, dictionary_size);
+        }
+        for i in 0..dictionary.len() {
+            dictionary[i] = INVALID_POS;
+        }
+
+        for i in 0..rehash_buckets as usize {
+			if keys[i].inner.is_none() {
+				keys[i].next = inner.free_head;
+				inner.free_head = i as u32;
+				continue;
+			}
+
+			let hash = self.hasher.hash_one(&keys[i].inner.as_ref().unwrap());
+            let pos: usize = (hash % dictionary.len() as u64) as usize;
+            keys[i].next = dictionary[pos];
+            dictionary[pos] = i as u32;
+        }
+
+        // Finally, update the CoreHashMap struct
+        inner.dictionary = dictionary;
+        inner.keys = keys;
+	}
+
+	/// Rehash the map. Intended for benchmarking only.
+	pub fn shuffle(&mut self) {
+		let map = unsafe { self.shared_ptr.as_mut() }.unwrap();
+        let inner = &mut map.inner;
+		let num_buckets = inner.get_num_buckets() as u32;
+		let size_bytes = HashMapInit::<K, V, S>::estimate_size(num_buckets);
+		let end_ptr: *mut u8 = unsafe { (self.shared_ptr as *mut u8).add(size_bytes) };
+        let keys_ptr = inner.keys.as_mut_ptr();
+		self.rehash_dict(inner, keys_ptr, end_ptr, num_buckets, num_buckets);
+	}
+
+	
+    // /// Grow
+    // ///
+    // /// 1. grow the underlying shared memory area
+    // /// 2. Initialize new buckets. This overwrites the current dictionary
+    // /// 3. Recalculate the dictionary
+    // pub fn grow(&mut self, num_buckets: u32) -> Result<(), crate::shmem::Error> {
+    //     let map = unsafe { self.shared_ptr.as_mut() }.unwrap();
+    //     let inner = &mut map.inner;
+    //     let old_num_buckets = inner.buckets.len() as u32;
+    //     if num_buckets < old_num_buckets {
+    //         panic!("grow called with a smaller number of buckets");
+    //     }
+    //     if num_buckets == old_num_buckets {
+    //         return Ok(());
+    //     }
+    //     let shmem_handle = self
+    //         .shmem_handle
+    //         .as_ref()
+    //         .expect("grow called on a fixed-size hash table");
+
+    //     let size_bytes = HashMapInit::<K, V, S>::estimate_size(num_buckets);
+    //     shmem_handle.set_size(size_bytes)?;
+    //     let end_ptr: *mut u8 = unsafe { shmem_handle.data_ptr.as_ptr().add(size_bytes) };
+
+    //     // Initialize new buckets. The new buckets are linked to the free list. NB: This overwrites
+    //     // the dictionary!
+    //     let keys_ptr = inner.keys.as_mut_ptr();
+    //     unsafe {
+    //         for i in old_num_buckets..num_buckets {
+    //             let bucket_ptr = buckets_ptr.add(i as usize);
+    //             bucket_ptr.write(core::Bucket {
+    //                 next: if i < num_buckets-1 {
+    //                     i as u32 + 1
+    //                 } else {
+    //                     inner.free_head
+    //                 },
+	// 				prev: if i > 0 {
+	// 					PrevPos::Chained(i as u32 - 1)
+	// 				} else {
+	// 					PrevPos::First(INVALID_POS)
+	// 				},
+    //                 inner: None,
+    //             });
+    //         }
+    //     }
+	// 	self.rehash_dict(inner, keys_ptr, end_ptr, num_buckets, old_num_buckets);
+    //     inner.free_head = old_num_buckets;
+
+    //     Ok(())
+    // }
+
+	// /// Begin a shrink, limiting all new allocations to be in buckets with index less than `num_buckets`. 
+// 	pub fn begin_shrink(&mut self, num_buckets: u32) {
+// 		let map = unsafe { self.shared_ptr.as_mut() }.unwrap();
+// 		if num_buckets > map.inner.get_num_buckets() as u32 {
+//             panic!("shrink called with a larger number of buckets");
+//         }
+// 		_ = self
+//             .shmem_handle
+//             .as_ref()
+//             .expect("shrink called on a fixed-size hash table");
+// 		map.inner.alloc_limit = num_buckets;
+// 	}
+
+// 	/// Complete a shrink after caller has evicted entries, removing the unused buckets and rehashing.
+// 	pub fn finish_shrink(&mut self) -> Result<(), crate::shmem::Error> {
+// 		let map = unsafe { self.shared_ptr.as_mut() }.unwrap();
+// 		let inner = &mut map.inner;
+// 		if !inner.is_shrinking() {
+// 			panic!("called finish_shrink when no shrink is in progress");
+// 		}
+
+// 		let num_buckets = inner.alloc_limit; 
+
+// 		if inner.get_num_buckets() == num_buckets as usize {
+//             return Ok(());
+//         }
+		
+// 		for i in (num_buckets as usize)..inner.buckets.len() {
+// 			if inner.buckets[i].inner.is_some() {
+// 				// TODO(quantumish) Do we want to treat this as a violation of an invariant
+// 				// or a legitimate error the caller can run into? Originally I thought this
+// 				// could return something like a UnevictedError(index) as soon as it runs
+// 				// into something (that way a caller could clear their soon-to-be-shrinked 
+// 				// buckets by repeatedly trying to call `finish_shrink`). 
+// 				//
+// 				// Would require making a wider error type enum with this and shmem errors.
+// 				panic!("unevicted entries in shrinked space")
+// 			}
+// 			match inner.buckets[i].prev {
+// 				PrevPos::First(_) => {
+// 					let next_pos = inner.buckets[i].next;
+// 					inner.free_head = next_pos;
+// 					if next_pos != INVALID_POS {
+// 						inner.buckets[next_pos as usize].prev = PrevPos::First(INVALID_POS);
+// 					}
+// 				},
+// 				PrevPos::Chained(j) => {
+// 					let next_pos = inner.buckets[i].next;
+// 					inner.buckets[j as usize].next = next_pos;
+// 					if next_pos != INVALID_POS {
+// 						inner.buckets[next_pos as usize].prev = PrevPos::Chained(j);
+// 					}
+// 				}
+// 			}
+// 		}
+
+//         let shmem_handle = self
+//             .shmem_handle
+//             .as_ref()
+//             .expect("shrink called on a fixed-size hash table");
+
+// 		let size_bytes = HashMapInit::<K, V, S>::estimate_size(num_buckets);
+//         shmem_handle.set_size(size_bytes)?;
+//         let end_ptr: *mut u8 = unsafe { shmem_handle.data_ptr.as_ptr().add(size_bytes) };
+// 		let buckets_ptr = inner.buckets.as_mut_ptr();
+// 		self.rehash_dict(inner, buckets_ptr, end_ptr, num_buckets, num_buckets);
+// 		inner.alloc_limit = INVALID_POS;
+		
+// 		Ok(())
+// 	}
+
+}

libs/neon-shmem/src/hash/core.rs

@@ -0,0 +1,247 @@
+//! Simple hash table with chaining
+//!
+//! # Resizing
+//!
+
+use std::hash::Hash;
+use std::mem::MaybeUninit;
+
+use crate::hash::entry::{Entry, OccupiedEntry, PrevPos, VacantEntry};
+
+pub(crate) const INVALID_POS: u32 = u32::MAX;
+
+pub(crate) struct LinkedKey<K> {
+	pub(crate) inner: Option<K>,
+	pub(crate) next: u32,	
+}
+
+pub(crate) struct CoreHashMap<'a, K, V> {
+	/// Dictionary used to map hashes to bucket indices.	
+    pub(crate) dictionary: &'a mut [u32],
+    pub(crate) keys: &'a mut [LinkedKey<K>],
+	pub(crate) vals: &'a mut [Option<V>],
+	/// Head of the freelist.
+    pub(crate) free_head: u32,
+
+    pub(crate) _user_list_head: u32,
+	/// Maximum index of a bucket allowed to be allocated. INVALID_POS if no limit.
+	pub(crate) alloc_limit: u32,
+
+    // metrics
+    pub(crate) buckets_in_use: u32,
+}
+
+#[derive(Debug)]
+pub struct FullError();
+
+impl<'a, K: Hash + Eq, V> CoreHashMap<'a, K, V>
+where
+    K: Clone + Hash + Eq,
+{
+    const FILL_FACTOR: f32 = 0.60;
+
+    pub fn estimate_size(num_buckets: u32) -> usize {
+        let mut size = 0;
+
+        // buckets
+        size += (size_of::<LinkedKey<K>>() + size_of::<Option<V>>())
+			* num_buckets as usize;
+
+        // dictionary
+        size += (f32::ceil((size_of::<u32>() * num_buckets as usize) as f32 / Self::FILL_FACTOR))
+            as usize;
+
+        size
+    }	
+
+    pub fn new(
+        keys: &'a mut [MaybeUninit<LinkedKey<K>>],
+		vals: &'a mut [MaybeUninit<Option<V>>],
+        dictionary: &'a mut [MaybeUninit<u32>],
+    ) -> CoreHashMap<'a, K, V> {
+        // Initialize the buckets
+        for i in 0..keys.len() {
+            keys[i].write(LinkedKey {
+				next: if i < keys.len() - 1 {
+                    i as u32 + 1
+                } else {
+                    INVALID_POS
+                },
+				inner: None,
+			});
+		}
+		for i in 0..vals.len() {
+            vals[i].write(None);
+		}
+
+		// Initialize the dictionary
+        for i in 0..dictionary.len() {
+            dictionary[i].write(INVALID_POS);
+        }
+
+        // TODO: use std::slice::assume_init_mut() once it stabilizes
+        let keys =
+            unsafe { std::slice::from_raw_parts_mut(keys.as_mut_ptr().cast(), keys.len()) };
+		let vals =
+            unsafe { std::slice::from_raw_parts_mut(vals.as_mut_ptr().cast(), vals.len()) };
+        let dictionary = unsafe {
+            std::slice::from_raw_parts_mut(dictionary.as_mut_ptr().cast(), dictionary.len())
+        };
+
+        CoreHashMap {
+            dictionary,
+            keys,
+			vals,
+            free_head: 0,
+            buckets_in_use: 0,
+            _user_list_head: INVALID_POS,
+			alloc_limit: INVALID_POS,
+        }
+    }
+
+    pub fn get_with_hash(&self, key: &K, hash: u64) -> Option<&V> {
+        let mut next = self.dictionary[hash as usize % self.dictionary.len()];
+        loop {
+            if next == INVALID_POS {
+                return None;
+            }
+
+            let keylink = &self.keys[next as usize];
+            let bucket_key = keylink.inner.as_ref().expect("entry is in use");
+            if bucket_key == key {
+                return Some(self.vals[next as usize].as_ref().unwrap());
+            }
+            next = keylink.next;
+        }
+    }
+
+    // all updates are done through Entry
+    pub fn entry_with_hash(&mut self, key: K, hash: u64) -> Entry<'a, '_, K, V> {
+        let dict_pos = hash as usize % self.dictionary.len();
+        let first = self.dictionary[dict_pos];
+        if first == INVALID_POS {
+            // no existing entry
+            return Entry::Vacant(VacantEntry {
+                map: self,
+                key,
+                dict_pos: dict_pos as u32,
+            });
+        }
+
+        let mut prev_pos = PrevPos::First(dict_pos as u32);
+        let mut next = first;
+        loop {
+            let keylink = &mut self.keys[next as usize];
+            let bucket_key = keylink.inner.as_mut().expect("entry is in use");
+            if *bucket_key == key {
+                // found existing entry
+                return Entry::Occupied(OccupiedEntry {
+                    map: self,
+                    _key: key,
+                    prev_pos,
+                    bucket_pos: next,
+                });
+            }
+
+            if keylink.next == INVALID_POS {
+                // No existing entry
+                return Entry::Vacant(VacantEntry {
+                    map: self,
+                    key,
+                    dict_pos: dict_pos as u32,
+                });
+            }
+            prev_pos = PrevPos::Chained(next);
+            next = keylink.next;
+        }
+    }
+
+    pub fn get_num_buckets(&self) -> usize {
+        self.keys.len()
+    }
+
+	pub fn is_shrinking(&self) -> bool {
+		self.alloc_limit != INVALID_POS
+	}
+
+	/// Clears all entries from the hashmap.
+	/// Does not reset any allocation limits, but does clear any entries beyond them.
+	pub fn clear(&mut self) {
+		for i in 0..self.keys.len() {
+            self.keys[i] = LinkedKey {
+                next: if i < self.keys.len() - 1 {
+                    i as u32 + 1
+                } else {
+                    INVALID_POS
+                },				
+                inner: None,
+            }
+        }
+		for i in 0..self.vals.len() {
+			self.vals[i] = None;
+		}
+
+        for i in 0..self.dictionary.len() {
+            self.dictionary[i] = INVALID_POS;
+        }
+
+		self.buckets_in_use = 0;
+	}
+	
+    pub fn entry_at_bucket(&mut self, pos: usize) -> Option<OccupiedEntry<'a, '_, K, V>> {
+		if pos >= self.keys.len() {
+			return None;
+		}
+
+		let entry = self.keys[pos].inner.as_ref();
+		match entry {
+			Some(key) => Some(OccupiedEntry {
+				_key: key.clone(),
+				bucket_pos: pos as u32,
+				prev_pos: PrevPos::Unknown,
+				map: self,
+			}),
+			_ => None,
+		}		
+    }
+
+	/// Find the position of an unused bucket via the freelist and initialize it. 
+    pub(crate) fn alloc_bucket(&mut self, key: K, value: V) -> Result<u32, FullError> {
+        let mut pos = self.free_head;
+
+		// Find the first bucket we're *allowed* to use.
+		let mut prev = PrevPos::First(self.free_head);
+		while pos != INVALID_POS && pos >= self.alloc_limit {
+			let keylink = &mut self.keys[pos as usize];
+			prev = PrevPos::Chained(pos);
+			pos = keylink.next;
+		}
+		if pos == INVALID_POS {
+			return Err(FullError());
+		}
+
+		// Repair the freelist.
+		match prev {
+			PrevPos::First(_) => {
+				let next_pos = self.keys[pos as usize].next;
+				self.free_head = next_pos;		
+			}
+			PrevPos::Chained(p) => if p != INVALID_POS {
+				let next_pos = self.keys[pos as usize].next;
+				self.keys[p as usize].next = next_pos;
+			},
+			PrevPos::Unknown => unreachable!()
+		}
+
+		// Initialize the bucket.
+		let keylink = &mut self.keys[pos as usize];
+		self.buckets_in_use += 1;
+        keylink.next = INVALID_POS;		
+        keylink.inner = Some(key);
+		self.vals[pos as usize] = Some(value);
+
+        return Ok(pos);
+    }
+}
+
+	

libs/neon-shmem/src/hash/entry.rs

@@ -0,0 +1,107 @@
+//! Like std::collections::hash_map::Entry;
+
+use crate::hash::core::{CoreHashMap, FullError, INVALID_POS};
+
+use std::hash::Hash;
+use std::mem;
+
+pub enum Entry<'a, 'b, K, V> {
+    Occupied(OccupiedEntry<'a, 'b, K, V>),
+    Vacant(VacantEntry<'a, 'b, K, V>),
+}
+
+/// Helper enum representing the previous position within a hashmap chain.
+#[derive(Clone, Copy)]
+pub(crate) enum PrevPos {
+	/// Starting index within the dictionary.  
+    First(u32),
+	/// Regular index within the buckets.
+    Chained(u32),
+	/// Unknown - e.g. the associated entry was retrieved by index instead of chain.
+	Unknown,
+}
+
+impl PrevPos {
+	/// Unwrap an index from a `PrevPos::First`, panicking otherwise.
+	pub fn unwrap_first(&self) -> u32 {
+		match self {
+			Self::First(i) => *i,
+			_ => panic!("not first entry in chain")
+		}
+	}
+}
+
+pub struct OccupiedEntry<'a, 'b, K, V> {
+	pub(crate) map: &'b mut CoreHashMap<'a, K, V>,
+	/// The key of the occupied entry
+    pub(crate) _key: K,
+	/// The index of the previous entry in the chain.
+    pub(crate) prev_pos: PrevPos,
+	/// The position of the bucket in the CoreHashMap's buckets array.
+    pub(crate) bucket_pos: u32, 
+}
+
+impl<'a, 'b, K, V> OccupiedEntry<'a, 'b, K, V> {
+    pub fn get(&self) -> &V {
+        self.map.vals[self.bucket_pos as usize]
+            .as_ref()
+            .unwrap()
+    }
+
+    pub fn get_mut(&mut self) -> &mut V {
+        self.map.vals[self.bucket_pos as usize]
+            .as_mut()
+            .unwrap()
+    }
+
+    pub fn insert(&mut self, value: V) -> V {
+        let bucket = &mut self.map.vals[self.bucket_pos as usize];
+        // This assumes inner is Some, which it must be for an OccupiedEntry
+        let old_value = mem::replace(bucket.as_mut().unwrap(), value);
+        old_value
+    }
+
+    pub fn remove(self) -> V {
+        // CoreHashMap::remove returns Option<(K, V)>. We know it's Some for an OccupiedEntry.
+        let keylink = &mut self.map.keys[self.bucket_pos as usize];
+
+        // unlink it from the chain
+        match self.prev_pos {
+            PrevPos::First(dict_pos) => self.map.dictionary[dict_pos as usize] = keylink.next,
+            PrevPos::Chained(bucket_pos) => {
+                self.map.keys[bucket_pos as usize].next = keylink.next
+            },
+			PrevPos::Unknown => panic!("can't safely remove entry with unknown previous entry"),
+        }
+
+        // and add it to the freelist        
+        let keylink = &mut self.map.keys[self.bucket_pos as usize];
+        keylink.inner = None;
+		keylink.next = self.map.free_head;
+		let old_value = self.map.vals[self.bucket_pos as usize].take();
+        self.map.free_head = self.bucket_pos;
+        self.map.buckets_in_use -= 1;
+
+        return old_value.unwrap();
+    }
+}
+
+pub struct VacantEntry<'a, 'b, K, V> {
+    pub(crate) map: &'b mut CoreHashMap<'a, K, V>,
+    pub(crate) key: K, // The key to insert
+    pub(crate) dict_pos: u32,
+}
+
+impl<'a, 'b, K: Clone + Hash + Eq, V> VacantEntry<'a, 'b, K, V> {
+    pub fn insert(self, value: V) -> Result<&'b mut V, FullError> {
+        let pos = self.map.alloc_bucket(self.key, value)?;
+        if pos == INVALID_POS {
+            return Err(FullError());
+        }
+        self.map.keys[pos as usize].next = self.map.dictionary[self.dict_pos as usize];
+        self.map.dictionary[self.dict_pos as usize] = pos;
+
+        let result = self.map.vals[pos as usize].as_mut().unwrap();
+        return Ok(result);
+    }
+}

libs/neon-shmem/src/hash/optim.rs

@@ -0,0 +1,85 @@
+//! Adapted from https://github.com/jsnell/parallel-xxhash (TODO: license?)
+
+use core::arch::x86::*;
+
+const PRIME32_1: u32 = 2654435761;
+const PRIME32_2: u32 = 2246822519;
+const PRIME32_3: u32 = 3266489917;
+const PRIME32_4: u32 =  668265263;
+const PRIME32_5: u32 =  374761393;
+
+#[cfg(target_arch = "x86_64")]
+#[target_feature(enable = "avx2")]
+fn mm256_rol32<const r: u32>(x: __m256i) -> __m256i {
+    return _mm256_or_si256(_mm256_slli_epi32(x, r),
+                           _mm256_srli_epi32(x, 32 - r));
+} 
+
+#[cfg(target_arch = "x86_64")]
+#[target_feature(enable = "avx2")]
+fn mm256_fmix32(mut h: __m256i) -> __m256i {
+    h = _mm256_xor_si256(h, _mm256_srli_epi32(h, 15));
+    h = _mm256_mullo_epi32(h, _mm256_set1_epi32(PRIME32_2));
+    h = _mm256_xor_si256(h, _mm256_srli_epi32(h, 13));
+    h = _mm256_mullo_epi32(h, _mm256_set1_epi32(PRIME32_3));
+    h = _mm256_xor_si256(h, _mm256_srli_epi32(h, 16));
+	h
+}
+
+#[cfg(target_arch = "x86_64")]
+#[target_feature(enable = "avx2")]
+fn mm256_round(mut seed: __m256i, input: __m256i) -> __m256i {
+	seed = _mm256_add_epi32(
+		seed,
+        _mm256_mullo_epi32(input, _mm256_set1_epi32(PRIME32_2))
+	);
+    seed = mm256_rol32::<13>(seed);
+    seed = _mm256_mullo_epi32(seed, _mm256_set1_epi32(PRIME32_1));
+	seed
+}
+
+/// Computes xxHash for 8 keys of size 4*N bytes in column-major order.
+#[cfg(target_arch = "x86_64")]
+#[target_feature(enable = "avx2")]
+fn xxhash_many<const N: usize>(keys: *const u32, seed: u32) -> [u32; 8] {
+	let mut res = [0; 8];
+	let mut h = _mm256_set1_epi32(seed + PRIME32_5);
+	if (N >= 4) {
+		let mut v1 = _mm256_set1_epi32(seed + PRIME32_1 + PRIME32_2);
+		let mut v2 = _mm256_set1_epi32(seed + PRIME32_2);
+		let mut v3 = _mm256_set1_epi32(seed);
+		let mut v4 = _mm256_set1_eip32(seed - PRIME32_1);
+		let mut i = 0;
+		while i < (N & !3) {
+			let k1 = _mm256_loadu_si256(keys.add((i + 0) * 8).cast());
+			let k2 = _mm256_loadu_si256(keys.add((i + 1) * 8).cast());
+			let k3 = _mm256_loadu_si256(keys.add((i + 2) * 8).cast());
+			let k4 = _mm256_loadu_si256(keys.add((i + 3) * 8).cast());
+			v1 = mm256_round(v1, k1);
+			v2 = mm256_round(v2, k2);
+			v3 = mm256_round(v3, k3);
+			v4 = mm256_round(v4, k4);
+			i += 4;
+		}
+		h = mm256_rol32::<1>(v1) + mm256_rol32::<7>(v2) +
+			mm256_rol32::<12>(v3) + mm256_rol32::<18>(v4);
+	}
+
+	// Unneeded, keeps bitwise parity with xxhash though.
+	h = _m256_add_epi32(h, _mm256_set1_eip32(N * 4));
+
+	for i in -(N & 3)..0 {
+        let v = _mm256_loadu_si256(keys.add((N + i) * 8));
+        h = _mm256_add_epi32(
+			h,
+            _mm256_mullo_epi32(v, _mm256_set1_epi32(PRIME32_3))
+		);
+        h = _mm256_mullo_epi32(
+			mm256_rol32::<17>(h),
+            _mm256_set1_epi32(PRIME32_4)
+		);
+    }
+
+    _mm256_storeu_si256((&mut res as *mut _).cast(), mm256_fmix32(h));
+	res
+}

libs/neon-shmem/src/hash/tests.rs

@@ -0,0 +1,382 @@
+use std::collections::BTreeMap;
+use std::collections::HashSet;
+use std::fmt::{Debug, Formatter};
+use std::mem::uninitialized;
+use std::mem::MaybeUninit;
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+use crate::hash::HashMapAccess;
+use crate::hash::HashMapInit;
+use crate::hash::Entry;
+use crate::shmem::ShmemHandle;
+
+use rand::seq::SliceRandom;
+use rand::{Rng, RngCore};
+use rand_distr::Zipf;
+
+const TEST_KEY_LEN: usize = 16;
+
+#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
+struct TestKey([u8; TEST_KEY_LEN]);
+
+impl From<&TestKey> for u128 {
+    fn from(val: &TestKey) -> u128 {
+        u128::from_be_bytes(val.0)
+    }
+}
+
+impl From<u128> for TestKey {
+    fn from(val: u128) -> TestKey {
+        TestKey(val.to_be_bytes())
+    }
+}
+
+impl<'a> From<&'a [u8]> for TestKey {
+    fn from(bytes: &'a [u8]) -> TestKey {
+        TestKey(bytes.try_into().unwrap())
+    }
+}
+
+fn test_inserts<K: Into<TestKey> + Copy>(keys: &[K]) {	
+    let mut w = HashMapInit::<TestKey, usize>::new_resizeable_named(
+		100000, 120000, "test_inserts"
+	).attach_writer();
+
+    for (idx, k) in keys.iter().enumerate() {
+		let hash = w.get_hash_value(&(*k).into());
+		let res = w.entry_with_hash((*k).into(), hash);
+		match res {
+			Entry::Occupied(mut e) => { e.insert(idx); }
+			Entry::Vacant(e) => {
+				let res = e.insert(idx);
+				assert!(res.is_ok());
+			},
+		};
+    }
+
+    for (idx, k) in keys.iter().enumerate() {
+		let hash = w.get_hash_value(&(*k).into());
+        let x = w.get_with_hash(&(*k).into(), hash);
+        let value = x.as_deref().copied();
+        assert_eq!(value, Some(idx));
+    }
+}
+
+#[test]
+fn dense() {
+    // This exercises splitting a node with prefix
+    let keys: &[u128] = &[0, 1, 2, 3, 256];
+    test_inserts(keys);
+
+    // Dense keys
+    let mut keys: Vec<u128> = (0..10000).collect();
+    test_inserts(&keys);
+
+    // Do the same in random orders
+    for _ in 1..10 {
+        keys.shuffle(&mut rand::rng());
+        test_inserts(&keys);
+    }
+}
+
+#[test]
+fn sparse() {
+    // sparse keys
+    let mut keys: Vec<TestKey> = Vec::new();
+    let mut used_keys = HashSet::new();
+    for _ in 0..10000 {
+        loop {
+            let key = rand::random::<u128>();
+            if used_keys.get(&key).is_some() {
+                continue;
+            }
+            used_keys.insert(key);
+            keys.push(key.into());
+            break;
+        }
+    }
+    test_inserts(&keys);
+}
+
+struct TestValue(AtomicUsize);
+
+impl TestValue {
+    fn new(val: usize) -> TestValue {
+        TestValue(AtomicUsize::new(val))
+    }
+
+    fn load(&self) -> usize {
+        self.0.load(Ordering::Relaxed)
+    }
+}
+
+impl Clone for TestValue {
+    fn clone(&self) -> TestValue {
+        TestValue::new(self.load())
+    }
+}
+
+impl Debug for TestValue {
+    fn fmt(&self, fmt: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
+        write!(fmt, "{:?}", self.load())
+    }
+}
+
+#[derive(Clone, Debug)]
+struct TestOp(TestKey, Option<usize>);
+
+fn apply_op(
+    op: &TestOp,
+    map: &mut HashMapAccess<TestKey, usize>,
+    shadow: &mut BTreeMap<TestKey, usize>,
+) {
+    // apply the change to the shadow tree first
+    let shadow_existing = if let Some(v) = op.1 {
+        shadow.insert(op.0, v)
+    } else {
+        shadow.remove(&op.0)
+    };
+
+	let hash = map.get_hash_value(&op.0);
+	let entry = map.entry_with_hash(op.0, hash);
+    let hash_existing = match op.1 {
+		Some(new) => {
+			match entry {
+				Entry::Occupied(mut e) => Some(e.insert(new)),
+				Entry::Vacant(e) => { e.insert(new).unwrap(); None },
+			}
+		},
+		None => {
+			match entry {
+				Entry::Occupied(e) => Some(e.remove()),
+				Entry::Vacant(_) => None,
+			}
+		},
+	};
+
+	assert_eq!(shadow_existing, hash_existing);
+}
+
+fn do_random_ops(
+	num_ops: usize,
+	size: u32,
+	del_prob: f64,
+	writer: &mut HashMapAccess<TestKey, usize>,
+	shadow: &mut BTreeMap<TestKey, usize>,
+	rng: &mut rand::rngs::ThreadRng,
+) {
+	for i in 0..num_ops {
+        let key: TestKey = ((rng.next_u32() % size) as u128).into();
+        let op = TestOp(key, if rng.random_bool(del_prob) { Some(i) } else { None });
+        apply_op(&op, writer, shadow);
+    }
+}
+
+fn do_deletes(
+	num_ops: usize,
+	writer: &mut HashMapAccess<TestKey, usize>,
+	shadow: &mut BTreeMap<TestKey, usize>,
+) {
+	for i in 0..num_ops {
+		let (k, _) = shadow.pop_first().unwrap();
+		let hash = writer.get_hash_value(&k);
+		writer.remove_with_hash(&k, hash);
+	}
+}
+
+fn do_shrink(
+	writer: &mut HashMapAccess<TestKey, usize>,
+	shadow: &mut BTreeMap<TestKey, usize>,
+	from: u32,
+	to: u32
+) {
+	writer.begin_shrink(to);
+	while writer.get_num_buckets_in_use() > to as usize {
+		let (k, _) = shadow.pop_first().unwrap();
+		let hash = writer.get_hash_value(&k);
+		let entry = writer.entry_with_hash(k, hash);
+		if let Entry::Occupied(mut e) = entry {
+			e.remove();
+		}
+	}
+	writer.finish_shrink().unwrap();
+}
+
+#[test]
+fn random_ops() {
+	let mut writer = HashMapInit::<TestKey, usize>::new_resizeable_named(
+		100000, 120000, "test_random"
+	).attach_writer();
+    let mut shadow: std::collections::BTreeMap<TestKey, usize> = BTreeMap::new();
+	
+    let distribution = Zipf::new(u128::MAX as f64, 1.1).unwrap();
+    let mut rng = rand::rng();
+    for i in 0..100000 {
+        let key: TestKey = (rng.sample(distribution) as u128).into();
+
+        let op = TestOp(key, if rng.random_bool(0.75) { Some(i) } else { None });
+
+        apply_op(&op, &mut writer, &mut shadow);
+
+        if i % 1000 == 0 {
+            eprintln!("{i} ops processed");
+        }
+    }
+}
+
+
+#[test]
+fn test_shuffle() {
+    let mut writer = HashMapInit::<TestKey, usize>::new_resizeable_named(
+		1000, 1200, "test_shuf"
+	).attach_writer();
+    let mut shadow: std::collections::BTreeMap<TestKey, usize> = BTreeMap::new();
+    let mut rng = rand::rng();
+
+    do_random_ops(10000, 1000, 0.75, &mut writer, &mut shadow, &mut rng);
+    writer.shuffle();
+	do_random_ops(10000, 1000, 0.75, &mut writer, &mut shadow, &mut rng);
+}
+
+#[test]
+fn test_grow() {
+    let mut writer = HashMapInit::<TestKey, usize>::new_resizeable_named(
+		1000, 2000, "test_grow"
+	).attach_writer();
+    let mut shadow: std::collections::BTreeMap<TestKey, usize> = BTreeMap::new();
+    let mut rng = rand::rng();
+
+    do_random_ops(10000, 1000, 0.75, &mut writer, &mut shadow, &mut rng);
+    writer.grow(1500).unwrap();
+	do_random_ops(10000, 1500, 0.75, &mut writer, &mut shadow, &mut rng);
+}
+
+#[test]
+fn test_shrink() {
+    let mut writer = HashMapInit::<TestKey, usize>::new_resizeable_named(
+		1500, 2000, "test_shrink"
+	).attach_writer();
+    let mut shadow: std::collections::BTreeMap<TestKey, usize> = BTreeMap::new();
+    let mut rng = rand::rng();
+	
+    do_random_ops(10000, 1500, 0.75, &mut writer, &mut shadow, &mut rng);
+    do_shrink(&mut writer, &mut shadow, 1500, 1000);
+	do_deletes(500, &mut writer, &mut shadow);
+	do_random_ops(10000, 500, 0.75, &mut writer, &mut shadow, &mut rng);
+	assert!(writer.get_num_buckets_in_use() <= 1000);
+}
+
+#[test]
+fn test_shrink_grow_seq() {
+    let mut writer = HashMapInit::<TestKey, usize>::new_resizeable_named(
+		1000, 20000, "test_grow_seq"
+	).attach_writer();
+    let mut shadow: std::collections::BTreeMap<TestKey, usize> = BTreeMap::new();
+    let mut rng = rand::rng();
+
+    do_random_ops(500, 1000, 0.1, &mut writer, &mut shadow, &mut rng);
+	eprintln!("Shrinking to 750");
+    do_shrink(&mut writer, &mut shadow, 1000, 750);
+	do_random_ops(200, 1000, 0.5, &mut writer, &mut shadow, &mut rng);
+	eprintln!("Growing to 1500");
+	writer.grow(1500).unwrap();
+	do_random_ops(600, 1500, 0.1, &mut writer, &mut shadow, &mut rng);
+	eprintln!("Shrinking to 200");
+	do_shrink(&mut writer, &mut shadow, 1500, 200);
+	do_deletes(100, &mut writer, &mut shadow);
+	do_random_ops(50, 1500, 0.25, &mut writer, &mut shadow, &mut rng);
+	eprintln!("Growing to 10k");
+	writer.grow(10000).unwrap();
+	do_random_ops(10000, 5000, 0.25, &mut writer, &mut shadow, &mut rng);
+}
+
+#[test]
+fn test_bucket_ops() {
+	let mut writer = HashMapInit::<TestKey, usize>::new_resizeable_named(
+		1000, 1200, "test_bucket_ops"
+	).attach_writer();
+	let hash = writer.get_hash_value(&1.into());
+	match writer.entry_with_hash(1.into(), hash) {
+		Entry::Occupied(mut e) => { e.insert(2); },
+		Entry::Vacant(e) => { e.insert(2).unwrap(); },
+	}
+	assert_eq!(writer.get_num_buckets_in_use(), 1);
+	assert_eq!(writer.get_num_buckets(), 1000);
+	assert_eq!(writer.get_with_hash(&1.into(), hash), Some(&2));
+	match writer.entry_with_hash(1.into(), hash) {
+		Entry::Occupied(e) => {
+			assert_eq!(e._key, 1.into());
+			let pos = e.bucket_pos as usize;
+			assert_eq!(writer.entry_at_bucket(pos).unwrap()._key, 1.into());
+			assert_eq!(writer.get_at_bucket(pos), Some(&(1.into(), 2)));
+		},
+		Entry::Vacant(_) => { panic!("Insert didn't affect entry"); },
+	}
+	writer.remove_with_hash(&1.into(), hash);
+	assert_eq!(writer.get_with_hash(&1.into(), hash), None);
+}
+
+#[test]
+fn test_shrink_zero() {
+	let mut writer = HashMapInit::<TestKey, usize>::new_resizeable_named(
+		1500, 2000, "test_shrink_zero"
+	).attach_writer();
+	writer.begin_shrink(0);
+	for i in 0..1500 {
+		writer.entry_at_bucket(i).map(|x| x.remove());
+	}
+	writer.finish_shrink().unwrap();
+	assert_eq!(writer.get_num_buckets_in_use(), 0);
+	let hash = writer.get_hash_value(&1.into());
+	let entry = writer.entry_with_hash(1.into(), hash);
+	if let Entry::Vacant(v) = entry {
+		assert!(v.insert(2).is_err());
+	} else {
+		panic!("Somehow got non-vacant entry in empty map.")
+	}
+	writer.grow(50).unwrap();
+	let entry = writer.entry_with_hash(1.into(), hash);
+	if let Entry::Vacant(v) = entry {
+		assert!(v.insert(2).is_ok());
+	} else {
+		panic!("Somehow got non-vacant entry in empty map.")
+	}
+	assert_eq!(writer.get_num_buckets_in_use(), 1);
+}
+
+#[test]
+#[should_panic]
+fn test_grow_oom() {
+    let mut writer = HashMapInit::<TestKey, usize>::new_resizeable_named(
+		1500, 2000, "test_grow_oom"
+	).attach_writer();
+	writer.grow(20000).unwrap();
+}
+
+#[test]
+#[should_panic]
+fn test_shrink_bigger() {
+    let mut writer = HashMapInit::<TestKey, usize>::new_resizeable_named(
+		1500, 2500, "test_shrink_bigger"
+	).attach_writer();
+	writer.begin_shrink(2000);
+}
+
+#[test]
+#[should_panic]
+fn test_shrink_early_finish() {
+    let mut writer = HashMapInit::<TestKey, usize>::new_resizeable_named(
+		1500, 2500, "test_shrink_early_finish"
+	).attach_writer();
+	writer.finish_shrink().unwrap();
+}
+
+#[test]
+#[should_panic]
+fn test_shrink_fixed_size() {
+	let mut area = [MaybeUninit::uninit(); 10000];
+    let init_struct = HashMapInit::<TestKey, usize>::with_fixed(3, &mut area);
+    let mut writer = init_struct.attach_writer();
+	writer.begin_shrink(1);
+}
+

libs/neon-shmem/src/lib.rs

@@ -1,418 +1,4 @@
 //! Shared memory utilities for neon communicator
 
-use std::num::NonZeroUsize;
-use std::os::fd::{AsFd, BorrowedFd, OwnedFd};
-use std::ptr::NonNull;
-use std::sync::atomic::{AtomicUsize, Ordering};
-
-use nix::errno::Errno;
-use nix::sys::mman::MapFlags;
-use nix::sys::mman::ProtFlags;
-use nix::sys::mman::mmap as nix_mmap;
-use nix::sys::mman::munmap as nix_munmap;
-use nix::unistd::ftruncate as nix_ftruncate;
-
-/// ShmemHandle represents a shared memory area that can be shared by processes over fork().
-/// Unlike shared memory allocated by Postgres, this area is resizable, up to 'max_size' that's
-/// specified at creation.
-///
-/// The area is backed by an anonymous file created with memfd_create(). The full address space for
-/// 'max_size' is reserved up-front with mmap(), but whenever you call [`ShmemHandle::set_size`],
-/// the underlying file is resized. Do not access the area beyond the current size. Currently, that
-/// will cause the file to be expanded, but we might use mprotect() etc. to enforce that in the
-/// future.
-pub struct ShmemHandle {
-    /// memfd file descriptor
-    fd: OwnedFd,
-
-    max_size: usize,
-
-    // Pointer to the beginning of the shared memory area. The header is stored there.
-    shared_ptr: NonNull<SharedStruct>,
-
-    // Pointer to the beginning of the user data
-    pub data_ptr: NonNull<u8>,
-}
-
-/// This is stored at the beginning in the shared memory area.
-struct SharedStruct {
-    max_size: usize,
-
-    /// Current size of the backing file. The high-order bit is used for the RESIZE_IN_PROGRESS flag
-    current_size: AtomicUsize,
-}
-
-const RESIZE_IN_PROGRESS: usize = 1 << 63;
-
-const HEADER_SIZE: usize = std::mem::size_of::<SharedStruct>();
-
-/// Error type returned by the ShmemHandle functions.
-#[derive(thiserror::Error, Debug)]
-#[error("{msg}: {errno}")]
-pub struct Error {
-    pub msg: String,
-    pub errno: Errno,
-}
-
-impl Error {
-    fn new(msg: &str, errno: Errno) -> Error {
-        Error {
-            msg: msg.to_string(),
-            errno,
-        }
-    }
-}
-
-impl ShmemHandle {
-    /// Create a new shared memory area. To communicate between processes, the processes need to be
-    /// fork()'d after calling this, so that the ShmemHandle is inherited by all processes.
-    ///
-    /// If the ShmemHandle is dropped, the memory is unmapped from the current process. Other
-    /// processes can continue using it, however.
-    pub fn new(name: &str, initial_size: usize, max_size: usize) -> Result<ShmemHandle, Error> {
-        // create the backing anonymous file.
-        let fd = create_backing_file(name)?;
-
-        Self::new_with_fd(fd, initial_size, max_size)
-    }
-
-    fn new_with_fd(
-        fd: OwnedFd,
-        initial_size: usize,
-        max_size: usize,
-    ) -> Result<ShmemHandle, Error> {
-        // We reserve the high-order bit for the RESIZE_IN_PROGRESS flag, and the actual size
-        // is a little larger than this because of the SharedStruct header. Make the upper limit
-        // somewhat smaller than that, because with anything close to that, you'll run out of
-        // memory anyway.
-        if max_size >= 1 << 48 {
-            panic!("max size {} too large", max_size);
-        }
-        if initial_size > max_size {
-            panic!("initial size {initial_size} larger than max size {max_size}");
-        }
-
-        // The actual initial / max size is the one given by the caller, plus the size of
-        // 'SharedStruct'.
-        let initial_size = HEADER_SIZE + initial_size;
-        let max_size = NonZeroUsize::new(HEADER_SIZE + max_size).unwrap();
-
-        // Reserve address space for it with mmap
-        //
-        // TODO: Use MAP_HUGETLB if possible
-        let start_ptr = unsafe {
-            nix_mmap(
-                None,
-                max_size,
-                ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
-                MapFlags::MAP_SHARED,
-                &fd,
-                0,
-            )
-        }
-        .map_err(|e| Error::new("mmap failed: {e}", e))?;
-
-        // Reserve space for the initial size
-        enlarge_file(fd.as_fd(), initial_size as u64)?;
-
-        // Initialize the header
-        let shared: NonNull<SharedStruct> = start_ptr.cast();
-        unsafe {
-            shared.write(SharedStruct {
-                max_size: max_size.into(),
-                current_size: AtomicUsize::new(initial_size),
-            })
-        };
-
-        // The user data begins after the header
-        let data_ptr = unsafe { start_ptr.cast().add(HEADER_SIZE) };
-
-        Ok(ShmemHandle {
-            fd,
-            max_size: max_size.into(),
-            shared_ptr: shared,
-            data_ptr,
-        })
-    }
-
-    // return reference to the header
-    fn shared(&self) -> &SharedStruct {
-        unsafe { self.shared_ptr.as_ref() }
-    }
-
-    /// Resize the shared memory area. 'new_size' must not be larger than the 'max_size' specified
-    /// when creating the area.
-    ///
-    /// This may only be called from one process/thread concurrently. We detect that case
-    /// and return an Error.
-    pub fn set_size(&self, new_size: usize) -> Result<(), Error> {
-        let new_size = new_size + HEADER_SIZE;
-        let shared = self.shared();
-
-        if new_size > self.max_size {
-            panic!(
-                "new size ({} is greater than max size ({})",
-                new_size, self.max_size
-            );
-        }
-        assert_eq!(self.max_size, shared.max_size);
-
-        // Lock the area by setting the bit in 'current_size'
-        //
-        // Ordering::Relaxed would probably be sufficient here, as we don't access any other memory
-        // and the posix_fallocate/ftruncate call is surely a synchronization point anyway. But
-        // since this is not performance-critical, better safe than sorry .
-        let mut old_size = shared.current_size.load(Ordering::Acquire);
-        loop {
-            if (old_size & RESIZE_IN_PROGRESS) != 0 {
-                return Err(Error::new(
-                    "concurrent resize detected",
-                    Errno::UnknownErrno,
-                ));
-            }
-            match shared.current_size.compare_exchange(
-                old_size,
-                new_size,
-                Ordering::Acquire,
-                Ordering::Relaxed,
-            ) {
-                Ok(_) => break,
-                Err(x) => old_size = x,
-            }
-        }
-
-        // Ok, we got the lock.
-        //
-        // NB: If anything goes wrong, we *must* clear the bit!
-        let result = {
-            use std::cmp::Ordering::{Equal, Greater, Less};
-            match new_size.cmp(&old_size) {
-                Less => nix_ftruncate(&self.fd, new_size as i64).map_err(|e| {
-                    Error::new("could not shrink shmem segment, ftruncate failed: {e}", e)
-                }),
-                Equal => Ok(()),
-                Greater => enlarge_file(self.fd.as_fd(), new_size as u64),
-            }
-        };
-
-        // Unlock
-        shared.current_size.store(
-            if result.is_ok() { new_size } else { old_size },
-            Ordering::Release,
-        );
-
-        result
-    }
-
-    /// Returns the current user-visible size of the shared memory segment.
-    ///
-    /// NOTE: a concurrent set_size() call can change the size at any time. It is the caller's
-    /// responsibility not to access the area beyond the current size.
-    pub fn current_size(&self) -> usize {
-        let total_current_size =
-            self.shared().current_size.load(Ordering::Relaxed) & !RESIZE_IN_PROGRESS;
-        total_current_size - HEADER_SIZE
-    }
-}
-
-impl Drop for ShmemHandle {
-    fn drop(&mut self) {
-        // SAFETY: The pointer was obtained from mmap() with the given size.
-        // We unmap the entire region.
-        let _ = unsafe { nix_munmap(self.shared_ptr.cast(), self.max_size) };
-        // The fd is dropped automatically by OwnedFd.
-    }
-}
-
-/// Create a "backing file" for the shared memory area. On Linux, use memfd_create(), to create an
-/// anonymous in-memory file. One macos, fall back to a regular file. That's good enough for
-/// development and testing, but in production we want the file to stay in memory.
-///
-/// disable 'unused_variables' warnings, because in the macos path, 'name' is unused.
-#[allow(unused_variables)]
-fn create_backing_file(name: &str) -> Result<OwnedFd, Error> {
-    #[cfg(not(target_os = "macos"))]
-    {
-        nix::sys::memfd::memfd_create(name, nix::sys::memfd::MFdFlags::empty())
-            .map_err(|e| Error::new("memfd_create failed: {e}", e))
-    }
-    #[cfg(target_os = "macos")]
-    {
-        let file = tempfile::tempfile().map_err(|e| {
-            Error::new(
-                "could not create temporary file to back shmem area: {e}",
-                nix::errno::Errno::from_raw(e.raw_os_error().unwrap_or(0)),
-            )
-        })?;
-        Ok(OwnedFd::from(file))
-    }
-}
-
-fn enlarge_file(fd: BorrowedFd, size: u64) -> Result<(), Error> {
-    // Use posix_fallocate() to enlarge the file. It reserves the space correctly, so that
-    // we don't get a segfault later when trying to actually use it.
-    #[cfg(not(target_os = "macos"))]
-    {
-        nix::fcntl::posix_fallocate(fd, 0, size as i64).map_err(|e| {
-            Error::new(
-                "could not grow shmem segment, posix_fallocate failed: {e}",
-                e,
-            )
-        })
-    }
-    // As a fallback on macos, which doesn't have posix_fallocate, use plain 'fallocate'
-    #[cfg(target_os = "macos")]
-    {
-        nix::unistd::ftruncate(fd, size as i64)
-            .map_err(|e| Error::new("could not grow shmem segment, ftruncate failed: {e}", e))
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    use nix::unistd::ForkResult;
-    use std::ops::Range;
-
-    /// check that all bytes in given range have the expected value.
-    fn assert_range(ptr: *const u8, expected: u8, range: Range<usize>) {
-        for i in range {
-            let b = unsafe { *(ptr.add(i)) };
-            assert_eq!(expected, b, "unexpected byte at offset {}", i);
-        }
-    }
-
-    /// Write 'b' to all bytes in the given range
-    fn write_range(ptr: *mut u8, b: u8, range: Range<usize>) {
-        unsafe { std::ptr::write_bytes(ptr.add(range.start), b, range.end - range.start) };
-    }
-
-    // simple single-process test of growing and shrinking
-    #[test]
-    fn test_shmem_resize() -> Result<(), Error> {
-        let max_size = 1024 * 1024;
-        let init_struct = ShmemHandle::new("test_shmem_resize", 0, max_size)?;
-
-        assert_eq!(init_struct.current_size(), 0);
-
-        // Initial grow
-        let size1 = 10000;
-        init_struct.set_size(size1).unwrap();
-        assert_eq!(init_struct.current_size(), size1);
-
-        // Write some data
-        let data_ptr = init_struct.data_ptr.as_ptr();
-        write_range(data_ptr, 0xAA, 0..size1);
-        assert_range(data_ptr, 0xAA, 0..size1);
-
-        // Shrink
-        let size2 = 5000;
-        init_struct.set_size(size2).unwrap();
-        assert_eq!(init_struct.current_size(), size2);
-
-        // Grow again
-        let size3 = 20000;
-        init_struct.set_size(size3).unwrap();
-        assert_eq!(init_struct.current_size(), size3);
-
-        // Try to read it. The area that was shrunk and grown again should read as all zeros now
-        assert_range(data_ptr, 0xAA, 0..5000);
-        assert_range(data_ptr, 0, 5000..size1);
-
-        // Try to grow beyond max_size
-        //let size4 = max_size + 1;
-        //assert!(init_struct.set_size(size4).is_err());
-
-        // Dropping init_struct should unmap the memory
-        drop(init_struct);
-
-        Ok(())
-    }
-
-    /// This is used in tests to coordinate between test processes. It's like std::sync::Barrier,
-    /// but is stored in the shared memory area and works across processes. It's implemented by
-    /// polling, because e.g. standard rust mutexes are not guaranteed to work across processes.
-    struct SimpleBarrier {
-        num_procs: usize,
-        count: AtomicUsize,
-    }
-
-    impl SimpleBarrier {
-        unsafe fn init(ptr: *mut SimpleBarrier, num_procs: usize) {
-            unsafe {
-                *ptr = SimpleBarrier {
-                    num_procs,
-                    count: AtomicUsize::new(0),
-                }
-            }
-        }
-
-        pub fn wait(&self) {
-            let old = self.count.fetch_add(1, Ordering::Relaxed);
-
-            let generation = old / self.num_procs;
-
-            let mut current = old + 1;
-            while current < (generation + 1) * self.num_procs {
-                std::thread::sleep(std::time::Duration::from_millis(10));
-                current = self.count.load(Ordering::Relaxed);
-            }
-        }
-    }
-
-    #[test]
-    fn test_multi_process() {
-        // Initialize
-        let max_size = 1_000_000_000_000;
-        let init_struct = ShmemHandle::new("test_multi_process", 0, max_size).unwrap();
-        let ptr = init_struct.data_ptr.as_ptr();
-
-        // Store the SimpleBarrier in the first 1k of the area.
-        init_struct.set_size(10000).unwrap();
-        let barrier_ptr: *mut SimpleBarrier = unsafe {
-            ptr.add(ptr.align_offset(std::mem::align_of::<SimpleBarrier>()))
-                .cast()
-        };
-        unsafe { SimpleBarrier::init(barrier_ptr, 2) };
-        let barrier = unsafe { barrier_ptr.as_ref().unwrap() };
-
-        // Fork another test process. The code after this runs in both processes concurrently.
-        let fork_result = unsafe { nix::unistd::fork().unwrap() };
-
-        // In the parent, fill bytes between 1000..2000. In the child, between 2000..3000
-        if fork_result.is_parent() {
-            write_range(ptr, 0xAA, 1000..2000);
-        } else {
-            write_range(ptr, 0xBB, 2000..3000);
-        }
-        barrier.wait();
-        // Verify the contents. (in both processes)
-        assert_range(ptr, 0xAA, 1000..2000);
-        assert_range(ptr, 0xBB, 2000..3000);
-
-        // Grow, from the child this time
-        let size = 10_000_000;
-        if !fork_result.is_parent() {
-            init_struct.set_size(size).unwrap();
-        }
-        barrier.wait();
-
-        // make some writes at the end
-        if fork_result.is_parent() {
-            write_range(ptr, 0xAA, (size - 10)..size);
-        } else {
-            write_range(ptr, 0xBB, (size - 20)..(size - 10));
-        }
-        barrier.wait();
-
-        // Verify the contents. (This runs in both processes)
-        assert_range(ptr, 0, (size - 1000)..(size - 20));
-        assert_range(ptr, 0xBB, (size - 20)..(size - 10));
-        assert_range(ptr, 0xAA, (size - 10)..size);
-
-        if let ForkResult::Parent { child } = fork_result {
-            nix::sys::wait::waitpid(child, None).unwrap();
-        }
-    }
-}
+pub mod hash;
+pub mod shmem;

libs/neon-shmem/src/shmem.rs

@@ -0,0 +1,418 @@
+//! Dynamically resizable contiguous chunk of shared memory
+
+use std::num::NonZeroUsize;
+use std::os::fd::{AsFd, BorrowedFd, OwnedFd};
+use std::ptr::NonNull;
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+use nix::errno::Errno;
+use nix::sys::mman::MapFlags;
+use nix::sys::mman::ProtFlags;
+use nix::sys::mman::mmap as nix_mmap;
+use nix::sys::mman::munmap as nix_munmap;
+use nix::unistd::ftruncate as nix_ftruncate;
+
+/// ShmemHandle represents a shared memory area that can be shared by processes over fork().
+/// Unlike shared memory allocated by Postgres, this area is resizable, up to 'max_size' that's
+/// specified at creation.
+///
+/// The area is backed by an anonymous file created with memfd_create(). The full address space for
+/// 'max_size' is reserved up-front with mmap(), but whenever you call [`ShmemHandle::set_size`],
+/// the underlying file is resized. Do not access the area beyond the current size. Currently, that
+/// will cause the file to be expanded, but we might use mprotect() etc. to enforce that in the
+/// future.
+pub struct ShmemHandle {
+    /// memfd file descriptor
+    fd: OwnedFd,
+
+    max_size: usize,
+
+    // Pointer to the beginning of the shared memory area. The header is stored there.
+    shared_ptr: NonNull<SharedStruct>,
+
+    // Pointer to the beginning of the user data
+    pub data_ptr: NonNull<u8>,
+}
+
+/// This is stored at the beginning in the shared memory area.
+struct SharedStruct {
+    max_size: usize,
+
+    /// Current size of the backing file. The high-order bit is used for the RESIZE_IN_PROGRESS flag
+    current_size: AtomicUsize,
+}
+
+const RESIZE_IN_PROGRESS: usize = 1 << 63;
+
+const HEADER_SIZE: usize = std::mem::size_of::<SharedStruct>();
+
+/// Error type returned by the ShmemHandle functions.
+#[derive(thiserror::Error, Debug)]
+#[error("{msg}: {errno}")]
+pub struct Error {
+    pub msg: String,
+    pub errno: Errno,
+}
+
+impl Error {
+    fn new(msg: &str, errno: Errno) -> Error {
+        Error {
+            msg: msg.to_string(),
+            errno,
+        }
+    }
+}
+
+impl ShmemHandle {
+    /// Create a new shared memory area. To communicate between processes, the processes need to be
+    /// fork()'d after calling this, so that the ShmemHandle is inherited by all processes.
+    ///
+    /// If the ShmemHandle is dropped, the memory is unmapped from the current process. Other
+    /// processes can continue using it, however.
+    pub fn new(name: &str, initial_size: usize, max_size: usize) -> Result<ShmemHandle, Error> {
+        // create the backing anonymous file.
+        let fd = create_backing_file(name)?;
+
+        Self::new_with_fd(fd, initial_size, max_size)
+    }
+
+    fn new_with_fd(
+        fd: OwnedFd,
+        initial_size: usize,
+        max_size: usize,
+    ) -> Result<ShmemHandle, Error> {
+        // We reserve the high-order bit for the RESIZE_IN_PROGRESS flag, and the actual size
+        // is a little larger than this because of the SharedStruct header. Make the upper limit
+        // somewhat smaller than that, because with anything close to that, you'll run out of
+        // memory anyway.
+        if max_size >= 1 << 48 {
+            panic!("max size {} too large", max_size);
+        }
+        if initial_size > max_size {
+            panic!("initial size {initial_size} larger than max size {max_size}");
+        }
+
+        // The actual initial / max size is the one given by the caller, plus the size of
+        // 'SharedStruct'.
+        let initial_size = HEADER_SIZE + initial_size;
+        let max_size = NonZeroUsize::new(HEADER_SIZE + max_size).unwrap();
+
+        // Reserve address space for it with mmap
+        //
+        // TODO: Use MAP_HUGETLB if possible
+        let start_ptr = unsafe {
+            nix_mmap(
+                None,
+                max_size,
+                ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
+                MapFlags::MAP_SHARED,
+                &fd,
+                0,
+            )
+        }
+        .map_err(|e| Error::new("mmap failed: {e}", e))?;
+
+        // Reserve space for the initial size
+        enlarge_file(fd.as_fd(), initial_size as u64)?;
+
+        // Initialize the header
+        let shared: NonNull<SharedStruct> = start_ptr.cast();
+        unsafe {
+            shared.write(SharedStruct {
+                max_size: max_size.into(),
+                current_size: AtomicUsize::new(initial_size),
+            })
+        };
+
+        // The user data begins after the header
+        let data_ptr = unsafe { start_ptr.cast().add(HEADER_SIZE) };
+
+        Ok(ShmemHandle {
+            fd,
+            max_size: max_size.into(),
+            shared_ptr: shared,
+            data_ptr,
+        })
+    }
+
+    // return reference to the header
+    fn shared(&self) -> &SharedStruct {
+        unsafe { self.shared_ptr.as_ref() }
+    }
+
+    /// Resize the shared memory area. 'new_size' must not be larger than the 'max_size' specified
+    /// when creating the area.
+    ///
+    /// This may only be called from one process/thread concurrently. We detect that case
+    /// and return an Error.
+    pub fn set_size(&self, new_size: usize) -> Result<(), Error> {
+        let new_size = new_size + HEADER_SIZE;
+        let shared = self.shared();
+
+        if new_size > self.max_size {
+            panic!(
+                "new size ({} is greater than max size ({})",
+                new_size, self.max_size
+            );
+        }
+        assert_eq!(self.max_size, shared.max_size);
+
+        // Lock the area by setting the bit in 'current_size'
+        //
+        // Ordering::Relaxed would probably be sufficient here, as we don't access any other memory
+        // and the posix_fallocate/ftruncate call is surely a synchronization point anyway. But
+        // since this is not performance-critical, better safe than sorry .
+        let mut old_size = shared.current_size.load(Ordering::Acquire);
+        loop {
+            if (old_size & RESIZE_IN_PROGRESS) != 0 {
+                return Err(Error::new(
+                    "concurrent resize detected",
+                    Errno::UnknownErrno,
+                ));
+            }
+            match shared.current_size.compare_exchange(
+                old_size,
+                new_size,
+                Ordering::Acquire,
+                Ordering::Relaxed,
+            ) {
+                Ok(_) => break,
+                Err(x) => old_size = x,
+            }
+        }
+
+        // Ok, we got the lock.
+        //
+        // NB: If anything goes wrong, we *must* clear the bit!
+        let result = {
+            use std::cmp::Ordering::{Equal, Greater, Less};
+            match new_size.cmp(&old_size) {
+                Less => nix_ftruncate(&self.fd, new_size as i64).map_err(|e| {
+                    Error::new("could not shrink shmem segment, ftruncate failed: {e}", e)
+                }),
+                Equal => Ok(()),
+                Greater => enlarge_file(self.fd.as_fd(), new_size as u64),
+            }
+        };
+
+        // Unlock
+        shared.current_size.store(
+            if result.is_ok() { new_size } else { old_size },
+            Ordering::Release,
+        );
+
+        result
+    }
+
+    /// Returns the current user-visible size of the shared memory segment.
+    ///
+    /// NOTE: a concurrent set_size() call can change the size at any time. It is the caller's
+    /// responsibility not to access the area beyond the current size.
+    pub fn current_size(&self) -> usize {
+        let total_current_size =
+            self.shared().current_size.load(Ordering::Relaxed) & !RESIZE_IN_PROGRESS;
+        total_current_size - HEADER_SIZE
+    }
+}
+
+impl Drop for ShmemHandle {
+    fn drop(&mut self) {
+        // SAFETY: The pointer was obtained from mmap() with the given size.
+        // We unmap the entire region.
+        let _ = unsafe { nix_munmap(self.shared_ptr.cast(), self.max_size) };
+        // The fd is dropped automatically by OwnedFd.
+    }
+}
+
+/// Create a "backing file" for the shared memory area. On Linux, use memfd_create(), to create an
+/// anonymous in-memory file. One macos, fall back to a regular file. That's good enough for
+/// development and testing, but in production we want the file to stay in memory.
+///
+/// disable 'unused_variables' warnings, because in the macos path, 'name' is unused.
+#[allow(unused_variables)]
+fn create_backing_file(name: &str) -> Result<OwnedFd, Error> {
+    #[cfg(not(target_os = "macos"))]
+    {
+        nix::sys::memfd::memfd_create(name, nix::sys::memfd::MFdFlags::empty())
+            .map_err(|e| Error::new("memfd_create failed: {e}", e))
+    }
+    #[cfg(target_os = "macos")]
+    {
+        let file = tempfile::tempfile().map_err(|e| {
+            Error::new(
+                "could not create temporary file to back shmem area: {e}",
+                nix::errno::Errno::from_raw(e.raw_os_error().unwrap_or(0)),
+            )
+        })?;
+        Ok(OwnedFd::from(file))
+    }
+}
+
+fn enlarge_file(fd: BorrowedFd, size: u64) -> Result<(), Error> {
+    // Use posix_fallocate() to enlarge the file. It reserves the space correctly, so that
+    // we don't get a segfault later when trying to actually use it.
+    #[cfg(not(target_os = "macos"))]
+    {
+        nix::fcntl::posix_fallocate(fd, 0, size as i64).map_err(|e| {
+            Error::new(
+                "could not grow shmem segment, posix_fallocate failed: {e}",
+                e,
+            )
+        })
+    }
+    // As a fallback on macos, which doesn't have posix_fallocate, use plain 'fallocate'
+    #[cfg(target_os = "macos")]
+    {
+        nix::unistd::ftruncate(fd, size as i64)
+            .map_err(|e| Error::new("could not grow shmem segment, ftruncate failed: {e}", e))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use nix::unistd::ForkResult;
+    use std::ops::Range;
+
+    /// check that all bytes in given range have the expected value.
+    fn assert_range(ptr: *const u8, expected: u8, range: Range<usize>) {
+        for i in range {
+            let b = unsafe { *(ptr.add(i)) };
+            assert_eq!(expected, b, "unexpected byte at offset {}", i);
+        }
+    }
+
+    /// Write 'b' to all bytes in the given range
+    fn write_range(ptr: *mut u8, b: u8, range: Range<usize>) {
+        unsafe { std::ptr::write_bytes(ptr.add(range.start), b, range.end - range.start) };
+    }
+
+    // simple single-process test of growing and shrinking
+    #[test]
+    fn test_shmem_resize() -> Result<(), Error> {
+        let max_size = 1024 * 1024;
+        let init_struct = ShmemHandle::new("test_shmem_resize", 0, max_size)?;
+
+        assert_eq!(init_struct.current_size(), 0);
+
+        // Initial grow
+        let size1 = 10000;
+        init_struct.set_size(size1).unwrap();
+        assert_eq!(init_struct.current_size(), size1);
+
+        // Write some data
+        let data_ptr = init_struct.data_ptr.as_ptr();
+        write_range(data_ptr, 0xAA, 0..size1);
+        assert_range(data_ptr, 0xAA, 0..size1);
+
+        // Shrink
+        let size2 = 5000;
+        init_struct.set_size(size2).unwrap();
+        assert_eq!(init_struct.current_size(), size2);
+
+        // Grow again
+        let size3 = 20000;
+        init_struct.set_size(size3).unwrap();
+        assert_eq!(init_struct.current_size(), size3);
+
+        // Try to read it. The area that was shrunk and grown again should read as all zeros now
+        assert_range(data_ptr, 0xAA, 0..5000);
+        assert_range(data_ptr, 0, 5000..size1);
+
+        // Try to grow beyond max_size
+        //let size4 = max_size + 1;
+        //assert!(init_struct.set_size(size4).is_err());
+
+        // Dropping init_struct should unmap the memory
+        drop(init_struct);
+
+        Ok(())
+    }
+
+    /// This is used in tests to coordinate between test processes. It's like std::sync::Barrier,
+    /// but is stored in the shared memory area and works across processes. It's implemented by
+    /// polling, because e.g. standard rust mutexes are not guaranteed to work across processes.
+    struct SimpleBarrier {
+        num_procs: usize,
+        count: AtomicUsize,
+    }
+
+    impl SimpleBarrier {
+        unsafe fn init(ptr: *mut SimpleBarrier, num_procs: usize) {
+            unsafe {
+                *ptr = SimpleBarrier {
+                    num_procs,
+                    count: AtomicUsize::new(0),
+                }
+            }
+        }
+
+        pub fn wait(&self) {
+            let old = self.count.fetch_add(1, Ordering::Relaxed);
+
+            let generation = old / self.num_procs;
+
+            let mut current = old + 1;
+            while current < (generation + 1) * self.num_procs {
+                std::thread::sleep(std::time::Duration::from_millis(10));
+                current = self.count.load(Ordering::Relaxed);
+            }
+        }
+    }
+
+    #[test]
+    fn test_multi_process() {
+        // Initialize
+        let max_size = 1_000_000_000_000;
+        let init_struct = ShmemHandle::new("test_multi_process", 0, max_size).unwrap();
+        let ptr = init_struct.data_ptr.as_ptr();
+
+        // Store the SimpleBarrier in the first 1k of the area.
+        init_struct.set_size(10000).unwrap();
+        let barrier_ptr: *mut SimpleBarrier = unsafe {
+            ptr.add(ptr.align_offset(std::mem::align_of::<SimpleBarrier>()))
+                .cast()
+        };
+        unsafe { SimpleBarrier::init(barrier_ptr, 2) };
+        let barrier = unsafe { barrier_ptr.as_ref().unwrap() };
+
+        // Fork another test process. The code after this runs in both processes concurrently.
+        let fork_result = unsafe { nix::unistd::fork().unwrap() };
+
+        // In the parent, fill bytes between 1000..2000. In the child, between 2000..3000
+        if fork_result.is_parent() {
+            write_range(ptr, 0xAA, 1000..2000);
+        } else {
+            write_range(ptr, 0xBB, 2000..3000);
+        }
+        barrier.wait();
+        // Verify the contents. (in both processes)
+        assert_range(ptr, 0xAA, 1000..2000);
+        assert_range(ptr, 0xBB, 2000..3000);
+
+        // Grow, from the child this time
+        let size = 10_000_000;
+        if !fork_result.is_parent() {
+            init_struct.set_size(size).unwrap();
+        }
+        barrier.wait();
+
+        // make some writes at the end
+        if fork_result.is_parent() {
+            write_range(ptr, 0xAA, (size - 10)..size);
+        } else {
+            write_range(ptr, 0xBB, (size - 20)..(size - 10));
+        }
+        barrier.wait();
+
+        // Verify the contents. (This runs in both processes)
+        assert_range(ptr, 0, (size - 1000)..(size - 20));
+        assert_range(ptr, 0xBB, (size - 20)..(size - 10));
+        assert_range(ptr, 0xAA, (size - 10)..size);
+
+        if let ForkResult::Parent { child } = fork_result {
+            nix::sys::wait::waitpid(child, None).unwrap();
+        }
+    }
+}

libs/pageserver_api/src/controller_api.rs

@@ -344,35 +344,6 @@ impl Default for ShardSchedulingPolicy {
     }
 }
 
-#[derive(Serialize, Deserialize, Clone, Copy, Eq, PartialEq, Debug)]
-pub enum NodeLifecycle {
-    Active,
-    Deleted,
-}
-
-impl FromStr for NodeLifecycle {
-    type Err = anyhow::Error;
-
-    fn from_str(s: &str) -> Result<Self, Self::Err> {
-        match s {
-            "active" => Ok(Self::Active),
-            "deleted" => Ok(Self::Deleted),
-            _ => Err(anyhow::anyhow!("Unknown node lifecycle '{s}'")),
-        }
-    }
-}
-
-impl From<NodeLifecycle> for String {
-    fn from(value: NodeLifecycle) -> String {
-        use NodeLifecycle::*;
-        match value {
-            Active => "active",
-            Deleted => "deleted",
-        }
-        .to_string()
-    }
-}
-
 #[derive(Serialize, Deserialize, Clone, Copy, Eq, PartialEq, Debug)]
 pub enum NodeSchedulingPolicy {
     Active,

libs/pageserver_api/src/upcall_api.rs

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

libs/posthog_client_lite/src/background_loop.rs

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

libs/posthog_client_lite/src/lib.rs

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

libs/proxy/postgres-protocol2/Cargo.toml

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

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

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

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

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

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

@@ -10,7 +10,7 @@ use crate::{Error, cancel_query_raw, connect_socket};
 pub(crate) async fn cancel_query<T>(
     config: Option<SocketConfig>,
     ssl_mode: SslMode,
-    tls: T,
+    mut tls: T,
     process_id: i32,
     secret_key: i32,
 ) -> Result<(), Error>

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

@@ -17,6 +17,7 @@ use crate::{Client, Connection, Error};
 
 /// TLS configuration.
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[non_exhaustive]
 pub enum SslMode {
     /// Do not use TLS.
     Disable,
@@ -230,7 +231,7 @@ impl Config {
     /// Requires the `runtime` Cargo feature (enabled by default).
     pub async fn connect<T>(
         &self,
-        tls: &T,
+        tls: T,
     ) -> Result<(Client, Connection<TcpStream, T::Stream>), Error>
     where
         T: MakeTlsConnect<TcpStream>,

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

@@ -13,7 +13,7 @@ use crate::tls::{MakeTlsConnect, TlsConnect};
 use crate::{Client, Config, Connection, Error, RawConnection};
 
 pub async fn connect<T>(
-    tls: &T,
+    mut tls: T,
     config: &Config,
 ) -> Result<(Client, Connection<TcpStream, T::Stream>), Error>
 where

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

@@ -47,7 +47,7 @@ pub trait MakeTlsConnect<S> {
     /// Creates a new `TlsConnect`or.
     ///
     /// The domain name is provided for certificate verification and SNI.
-    fn make_tls_connect(&self, domain: &str) -> Result<Self::TlsConnect, Self::Error>;
+    fn make_tls_connect(&mut self, domain: &str) -> Result<Self::TlsConnect, Self::Error>;
 }
 
 /// An asynchronous function wrapping a stream in a TLS session.
@@ -85,7 +85,7 @@ impl<S> MakeTlsConnect<S> for NoTls {
     type TlsConnect = NoTls;
     type Error = NoTlsError;
 
-    fn make_tls_connect(&self, _: &str) -> Result<NoTls, NoTlsError> {
+    fn make_tls_connect(&mut self, _: &str) -> Result<NoTls, NoTlsError> {
         Ok(NoTls)
     }
 }

libs/remote_storage/src/azure_blob.rs

@@ -10,7 +10,7 @@ use std::sync::Arc;
 use std::time::{Duration, SystemTime};
 use std::{env, io};
 
-use anyhow::{Context, Result, anyhow};
+use anyhow::{Context, Result};
 use azure_core::request_options::{IfMatchCondition, MaxResults, Metadata, Range};
 use azure_core::{Continuable, HttpClient, RetryOptions, TransportOptions};
 use azure_storage::StorageCredentials;
@@ -37,7 +37,6 @@ use crate::metrics::{AttemptOutcome, RequestKind, start_measuring_requests};
 use crate::{
     ConcurrencyLimiter, Download, DownloadError, DownloadKind, DownloadOpts, Listing, ListingMode,
     ListingObject, RemotePath, RemoteStorage, StorageMetadata, TimeTravelError, TimeoutOrCancel,
-    Version, VersionKind,
 };
 
 pub struct AzureBlobStorage {
@@ -406,39 +405,6 @@ impl AzureBlobStorage {
     pub fn container_name(&self) -> &str {
         &self.container_name
     }
-
-    async fn list_versions_with_permit(
-        &self,
-        _permit: &tokio::sync::SemaphorePermit<'_>,
-        prefix: Option<&RemotePath>,
-        mode: ListingMode,
-        max_keys: Option<NonZeroU32>,
-        cancel: &CancellationToken,
-    ) -> Result<crate::VersionListing, DownloadError> {
-        let customize_builder = |mut builder: ListBlobsBuilder| {
-            builder = builder.include_versions(true);
-            // We do not return this info back to `VersionListing` yet.
-            builder = builder.include_deleted(true);
-            builder
-        };
-        let kind = RequestKind::ListVersions;
-
-        let mut stream = std::pin::pin!(self.list_streaming_for_fn(
-            prefix,
-            mode,
-            max_keys,
-            cancel,
-            kind,
-            customize_builder
-        ));
-        let mut combined: crate::VersionListing =
-            stream.next().await.expect("At least one item required")?;
-        while let Some(list) = stream.next().await {
-            let list = list?;
-            combined.versions.extend(list.versions.into_iter());
-        }
-        Ok(combined)
-    }
 }
 
 trait ListingCollector {
@@ -522,10 +488,27 @@ impl RemoteStorage for AzureBlobStorage {
         max_keys: Option<NonZeroU32>,
         cancel: &CancellationToken,
     ) -> std::result::Result<crate::VersionListing, DownloadError> {
+        let customize_builder = |mut builder: ListBlobsBuilder| {
+            builder = builder.include_versions(true);
+            builder
+        };
         let kind = RequestKind::ListVersions;
-        let permit = self.permit(kind, cancel).await?;
-        self.list_versions_with_permit(&permit, prefix, mode, max_keys, cancel)
-            .await
+
+        let mut stream = std::pin::pin!(self.list_streaming_for_fn(
+            prefix,
+            mode,
+            max_keys,
+            cancel,
+            kind,
+            customize_builder
+        ));
+        let mut combined: crate::VersionListing =
+            stream.next().await.expect("At least one item required")?;
+        while let Some(list) = stream.next().await {
+            let list = list?;
+            combined.versions.extend(list.versions.into_iter());
+        }
+        Ok(combined)
     }
 
     async fn head_object(
@@ -820,159 +803,14 @@ impl RemoteStorage for AzureBlobStorage {
 
     async fn time_travel_recover(
         &self,
-        prefix: Option<&RemotePath>,
-        timestamp: SystemTime,
-        done_if_after: SystemTime,
-        cancel: &CancellationToken,
-        _complexity_limit: Option<NonZeroU32>,
+        _prefix: Option<&RemotePath>,
+        _timestamp: SystemTime,
+        _done_if_after: SystemTime,
+        _cancel: &CancellationToken,
     ) -> Result<(), TimeTravelError> {
-        let msg = "PLEASE NOTE: Azure Blob storage time-travel recovery may not work as expected "
-            .to_string()
-            + "for some specific files. If a file gets deleted but then overwritten and we want to recover "
-            + "to the time during the file was not present, this functionality will recover the file. Only "
-            + "use the functionality for services that can tolerate this. For example, recovering a state of the "
-            + "pageserver tenants.";
-        tracing::error!("{}", msg);
-
-        let kind = RequestKind::TimeTravel;
-        let permit = self.permit(kind, cancel).await?;
-
-        let mode = ListingMode::NoDelimiter;
-        let version_listing = self
-            .list_versions_with_permit(&permit, prefix, mode, None, cancel)
-            .await
-            .map_err(|err| match err {
-                DownloadError::Other(e) => TimeTravelError::Other(e),
-                DownloadError::Cancelled => TimeTravelError::Cancelled,
-                other => TimeTravelError::Other(other.into()),
-            })?;
-        let versions_and_deletes = version_listing.versions;
-
-        tracing::info!(
-            "Built list for time travel with {} versions and deletions",
-            versions_and_deletes.len()
-        );
-
-        // Work on the list of references instead of the objects directly,
-        // otherwise we get lifetime errors in the sort_by_key call below.
-        let mut versions_and_deletes = versions_and_deletes.iter().collect::<Vec<_>>();
-
-        versions_and_deletes.sort_by_key(|vd| (&vd.key, &vd.last_modified));
-
-        let mut vds_for_key = HashMap::<_, Vec<_>>::new();
-
-        for vd in &versions_and_deletes {
-            let Version { key, .. } = &vd;
-            let version_id = vd.version_id().map(|v| v.0.as_str());
-            if version_id == Some("null") {
-                return Err(TimeTravelError::Other(anyhow!(
-                    "Received ListVersions response for key={key} with version_id='null', \
-                        indicating either disabled versioning, or legacy objects with null version id values"
-                )));
-            }
-            tracing::trace!("Parsing version key={key} kind={:?}", vd.kind);
-
-            vds_for_key.entry(key).or_default().push(vd);
-        }
-
-        let warn_threshold = 3;
-        let max_retries = 10;
-        let is_permanent = |e: &_| matches!(e, TimeTravelError::Cancelled);
-
-        for (key, versions) in vds_for_key {
-            let last_vd = versions.last().unwrap();
-            let key = self.relative_path_to_name(key);
-            if last_vd.last_modified > done_if_after {
-                tracing::debug!("Key {key} has version later than done_if_after, skipping");
-                continue;
-            }
-            // the version we want to restore to.
-            let version_to_restore_to =
-                match versions.binary_search_by_key(&timestamp, |tpl| tpl.last_modified) {
-                    Ok(v) => v,
-                    Err(e) => e,
-                };
-            if version_to_restore_to == versions.len() {
-                tracing::debug!("Key {key} has no changes since timestamp, skipping");
-                continue;
-            }
-            let mut do_delete = false;
-            if version_to_restore_to == 0 {
-                // All versions more recent, so the key didn't exist at the specified time point.
-                tracing::debug!(
-                    "All {} versions more recent for {key}, deleting",
-                    versions.len()
-                );
-                do_delete = true;
-            } else {
-                match &versions[version_to_restore_to - 1] {
-                    Version {
-                        kind: VersionKind::Version(version_id),
-                        ..
-                    } => {
-                        let source_url = format!(
-                            "{}/{}?versionid={}",
-                            self.client
-                                .url()
-                                .map_err(|e| TimeTravelError::Other(anyhow!("{e}")))?,
-                            key,
-                            version_id.0
-                        );
-                        tracing::debug!(
-                            "Promoting old version {} for {key} at {}...",
-                            version_id.0,
-                            source_url
-                        );
-                        backoff::retry(
-                            || async {
-                                let blob_client = self.client.blob_client(key.clone());
-                                let op = blob_client.copy(Url::from_str(&source_url).unwrap());
-                                tokio::select! {
-                                    res = op => res.map_err(|e| TimeTravelError::Other(e.into())),
-                                    _ = cancel.cancelled() => Err(TimeTravelError::Cancelled),
-                                }
-                            },
-                            is_permanent,
-                            warn_threshold,
-                            max_retries,
-                            "copying object version for time_travel_recover",
-                            cancel,
-                        )
-                        .await
-                        .ok_or_else(|| TimeTravelError::Cancelled)
-                        .and_then(|x| x)?;
-                        tracing::info!(?version_id, %key, "Copied old version in Azure blob storage");
-                    }
-                    Version {
-                        kind: VersionKind::DeletionMarker,
-                        ..
-                    } => {
-                        do_delete = true;
-                    }
-                }
-            };
-            if do_delete {
-                if matches!(last_vd.kind, VersionKind::DeletionMarker) {
-                    // Key has since been deleted (but there was some history), no need to do anything
-                    tracing::debug!("Key {key} already deleted, skipping.");
-                } else {
-                    tracing::debug!("Deleting {key}...");
-
-                    self.delete(&RemotePath::from_string(&key).unwrap(), cancel)
-                        .await
-                        .map_err(|e| {
-                            // delete_oid0 will use TimeoutOrCancel
-                            if TimeoutOrCancel::caused_by_cancel(&e) {
-                                TimeTravelError::Cancelled
-                            } else {
-                                TimeTravelError::Other(e)
-                            }
-                        })?;
-                }
-            }
-        }
-
-        Ok(())
+        // TODO use Azure point in time recovery feature for this
+        // https://learn.microsoft.com/en-us/azure/storage/blobs/point-in-time-restore-overview
+        Err(TimeTravelError::Unimplemented)
     }
 }
 

libs/remote_storage/src/lib.rs

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

libs/remote_storage/src/local_fs.rs

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

libs/remote_storage/src/s3_bucket.rs

@@ -981,16 +981,22 @@ impl RemoteStorage for S3Bucket {
         timestamp: SystemTime,
         done_if_after: SystemTime,
         cancel: &CancellationToken,
-        complexity_limit: Option<NonZeroU32>,
     ) -> Result<(), TimeTravelError> {
         let kind = RequestKind::TimeTravel;
         let permit = self.permit(kind, cancel).await?;
 
         tracing::trace!("Target time: {timestamp:?}, done_if_after {done_if_after:?}");
 
+        // Limit the number of versions deletions, mostly so that we don't
+        // keep requesting forever if the list is too long, as we'd put the
+        // list in RAM.
+        // Building a list of 100k entries that reaches the limit roughly takes
+        // 40 seconds, and roughly corresponds to tenants of 2 TiB physical size.
+        const COMPLEXITY_LIMIT: Option<NonZeroU32> = NonZeroU32::new(100_000);
+
         let mode = ListingMode::NoDelimiter;
         let version_listing = self
-            .list_versions_with_permit(&permit, prefix, mode, complexity_limit, cancel)
+            .list_versions_with_permit(&permit, prefix, mode, COMPLEXITY_LIMIT, cancel)
             .await
             .map_err(|err| match err {
                 DownloadError::Other(e) => TimeTravelError::Other(e),
@@ -1016,7 +1022,6 @@ impl RemoteStorage for S3Bucket {
             let Version { key, .. } = &vd;
             let version_id = vd.version_id().map(|v| v.0.as_str());
             if version_id == Some("null") {
-                // TODO: check the behavior of using the SDK on a non-versioned container
                 return Err(TimeTravelError::Other(anyhow!(
                     "Received ListVersions response for key={key} with version_id='null', \
                     indicating either disabled versioning, or legacy objects with null version id values"

libs/remote_storage/src/simulate_failures.rs

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

libs/remote_storage/tests/test_real_s3.rs

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

libs/safekeeper_api/src/models.rs

@@ -13,7 +13,7 @@ use utils::pageserver_feedback::PageserverFeedback;
 use crate::membership::Configuration;
 use crate::{ServerInfo, Term};
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Serialize)]
 pub struct SafekeeperStatus {
     pub id: NodeId,
 }

pageserver/ctl/src/main.rs

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

pageserver/src/basebackup_cache.rs

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

pageserver/src/bin/pageserver.rs

@@ -23,7 +23,6 @@ use pageserver::deletion_queue::DeletionQueue;
 use pageserver::disk_usage_eviction_task::{self, launch_disk_usage_global_eviction_task};
 use pageserver::feature_resolver::FeatureResolver;
 use pageserver::metrics::{STARTUP_DURATION, STARTUP_IS_LOADING};
-use pageserver::page_service::GrpcPageServiceHandler;
 use pageserver::task_mgr::{
     BACKGROUND_RUNTIME, COMPUTE_REQUEST_RUNTIME, MGMT_REQUEST_RUNTIME, WALRECEIVER_RUNTIME,
 };
@@ -573,8 +572,7 @@ fn start_pageserver(
         tokio::sync::mpsc::unbounded_channel();
     let deletion_queue_client = deletion_queue.new_client();
     let background_purges = mgr::BackgroundPurges::default();
-
-    let tenant_manager = mgr::init(
+    let tenant_manager = BACKGROUND_RUNTIME.block_on(mgr::init_tenant_mgr(
         conf,
         background_purges.clone(),
         TenantSharedResources {
@@ -585,10 +583,10 @@ fn start_pageserver(
             basebackup_prepare_sender,
             feature_resolver,
         },
+        order,
         shutdown_pageserver.clone(),
-    );
+    ))?;
     let tenant_manager = Arc::new(tenant_manager);
-    BACKGROUND_RUNTIME.block_on(mgr::init_tenant_mgr(tenant_manager.clone(), order))?;
 
     let basebackup_cache = BasebackupCache::spawn(
         BACKGROUND_RUNTIME.handle(),
@@ -816,7 +814,7 @@ fn start_pageserver(
     // necessary?
     let mut page_service_grpc = None;
     if let Some(grpc_listener) = grpc_listener {
-        page_service_grpc = Some(GrpcPageServiceHandler::spawn(
+        page_service_grpc = Some(page_service::spawn_grpc(
             tenant_manager.clone(),
             grpc_auth,
             otel_guard.as_ref().map(|g| g.dispatch.clone()),

pageserver/src/feature_resolver.rs

@@ -1,6 +1,5 @@
 use std::{collections::HashMap, sync::Arc, time::Duration};
 
-use pageserver_api::config::NodeMetadata;
 use posthog_client_lite::{
     CaptureEvent, FeatureResolverBackgroundLoop, PostHogClientConfig, PostHogEvaluationError,
     PostHogFlagFilterPropertyValue,
@@ -87,35 +86,7 @@ impl FeatureResolver {
                         }
                     }
                 }
-                // TODO: move this to a background task so that we don't block startup in case of slow disk
-                let metadata_path = conf.metadata_path();
-                match std::fs::read_to_string(&metadata_path) {
-                    Ok(metadata_str) => match serde_json::from_str::<NodeMetadata>(&metadata_str) {
-                        Ok(metadata) => {
-                            properties.insert(
-                                "hostname".to_string(),
-                                PostHogFlagFilterPropertyValue::String(metadata.http_host),
-                            );
-                            if let Some(cplane_region) = metadata.other.get("region_id") {
-                                if let Some(cplane_region) = cplane_region.as_str() {
-                                    // This region contains the cell number
-                                    properties.insert(
-                                        "neon_region".to_string(),
-                                        PostHogFlagFilterPropertyValue::String(
-                                            cplane_region.to_string(),
-                                        ),
-                                    );
-                                }
-                            }
-                        }
-                        Err(e) => {
-                            tracing::warn!("Failed to parse metadata.json: {}", e);
-                        }
-                    },
-                    Err(e) => {
-                        tracing::warn!("Failed to read metadata.json: {}", e);
-                    }
-                }
+                // TODO: add pageserver URL.
                 Arc::new(properties)
             };
             let fake_tenants = {

pageserver/src/http/routes.rs

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

pageserver/src/metrics.rs

@@ -1053,15 +1053,6 @@ pub(crate) static TENANT_STATE_METRIC: Lazy<UIntGaugeVec> = Lazy::new(|| {
     .expect("Failed to register pageserver_tenant_states_count metric")
 });
 
-pub(crate) static TIMELINE_STATE_METRIC: Lazy<UIntGaugeVec> = Lazy::new(|| {
-    register_uint_gauge_vec!(
-        "pageserver_timeline_states_count",
-        "Count of timelines per state",
-        &["state"]
-    )
-    .expect("Failed to register pageserver_timeline_states_count metric")
-});
-
 /// A set of broken tenants.
 ///
 /// These are expected to be so rare that a set is fine. Set as in a new timeseries per each broken
@@ -3334,8 +3325,6 @@ impl TimelineMetrics {
                 &timeline_id,
             );
 
-        TIMELINE_STATE_METRIC.with_label_values(&["active"]).inc();
-
         TimelineMetrics {
             tenant_id,
             shard_id,
@@ -3490,8 +3479,6 @@ impl TimelineMetrics {
             return;
         }
 
-        TIMELINE_STATE_METRIC.with_label_values(&["active"]).dec();
-
         let tenant_id = &self.tenant_id;
         let timeline_id = &self.timeline_id;
         let shard_id = &self.shard_id;

pageserver/src/page_service.rs

@@ -169,6 +169,99 @@ pub fn spawn(
     Listener { cancel, task }
 }
 
+/// Spawns a gRPC server for the page service.
+///
+/// TODO: move this onto GrpcPageServiceHandler::spawn().
+/// TODO: this doesn't support TLS. We need TLS reloading via ReloadingCertificateResolver, so we
+/// need to reimplement the TCP+TLS accept loop ourselves.
+pub fn spawn_grpc(
+    tenant_manager: Arc<TenantManager>,
+    auth: Option<Arc<SwappableJwtAuth>>,
+    perf_trace_dispatch: Option<Dispatch>,
+    get_vectored_concurrent_io: GetVectoredConcurrentIo,
+    listener: std::net::TcpListener,
+) -> anyhow::Result<CancellableTask> {
+    let cancel = CancellationToken::new();
+    let ctx = RequestContextBuilder::new(TaskKind::PageRequestHandler)
+        .download_behavior(DownloadBehavior::Download)
+        .perf_span_dispatch(perf_trace_dispatch)
+        .detached_child();
+    let gate = Gate::default();
+
+    // Set up the TCP socket. We take a preconfigured TcpListener to bind the
+    // port early during startup.
+    let incoming = {
+        let _runtime = COMPUTE_REQUEST_RUNTIME.enter(); // required by TcpListener::from_std
+        listener.set_nonblocking(true)?;
+        tonic::transport::server::TcpIncoming::from(tokio::net::TcpListener::from_std(listener)?)
+            .with_nodelay(Some(GRPC_TCP_NODELAY))
+            .with_keepalive(Some(GRPC_TCP_KEEPALIVE_TIME))
+    };
+
+    // Set up the gRPC server.
+    //
+    // TODO: consider tuning window sizes.
+    let mut server = tonic::transport::Server::builder()
+        .http2_keepalive_interval(Some(GRPC_HTTP2_KEEPALIVE_INTERVAL))
+        .http2_keepalive_timeout(Some(GRPC_HTTP2_KEEPALIVE_TIMEOUT))
+        .max_concurrent_streams(Some(GRPC_MAX_CONCURRENT_STREAMS));
+
+    // Main page service stack. Uses a mix of Tonic interceptors and Tower layers:
+    //
+    // * Interceptors: can inspect and modify the gRPC request. Sync code only, runs before service.
+    //
+    // * Layers: allow async code, can run code after the service response. However, only has access
+    //   to the raw HTTP request/response, not the gRPC types.
+    let page_service_handler = GrpcPageServiceHandler {
+        tenant_manager,
+        ctx,
+        gate_guard: gate.enter().expect("gate was just created"),
+        get_vectored_concurrent_io,
+    };
+
+    let observability_layer = ObservabilityLayer;
+    let mut tenant_interceptor = TenantMetadataInterceptor;
+    let mut auth_interceptor = TenantAuthInterceptor::new(auth);
+
+    let page_service = tower::ServiceBuilder::new()
+        // Create tracing span and record request start time.
+        .layer(observability_layer)
+        // Intercept gRPC requests.
+        .layer(tonic::service::InterceptorLayer::new(move |mut req| {
+            // Extract tenant metadata.
+            req = tenant_interceptor.call(req)?;
+            // Authenticate tenant JWT token.
+            req = auth_interceptor.call(req)?;
+            Ok(req)
+        }))
+        .service(proto::PageServiceServer::new(page_service_handler));
+    let server = server.add_service(page_service);
+
+    // Reflection service for use with e.g. grpcurl.
+    let reflection_service = tonic_reflection::server::Builder::configure()
+        .register_encoded_file_descriptor_set(proto::FILE_DESCRIPTOR_SET)
+        .build_v1()?;
+    let server = server.add_service(reflection_service);
+
+    // Spawn server task.
+    let task_cancel = cancel.clone();
+    let task = COMPUTE_REQUEST_RUNTIME.spawn(task_mgr::exit_on_panic_or_error(
+        "grpc listener",
+        async move {
+            let result = server
+                .serve_with_incoming_shutdown(incoming, task_cancel.cancelled())
+                .await;
+            if result.is_ok() {
+                // TODO: revisit shutdown logic once page service is implemented.
+                gate.close().await;
+            }
+            result
+        },
+    ));
+
+    Ok(CancellableTask { task, cancel })
+}
+
 impl Listener {
     pub async fn stop_accepting(self) -> Connections {
         self.cancel.cancel();
@@ -3273,101 +3366,6 @@ pub struct GrpcPageServiceHandler {
 }
 
 impl GrpcPageServiceHandler {
-    /// Spawns a gRPC server for the page service.
-    ///
-    /// TODO: this doesn't support TLS. We need TLS reloading via ReloadingCertificateResolver, so we
-    /// need to reimplement the TCP+TLS accept loop ourselves.
-    pub fn spawn(
-        tenant_manager: Arc<TenantManager>,
-        auth: Option<Arc<SwappableJwtAuth>>,
-        perf_trace_dispatch: Option<Dispatch>,
-        get_vectored_concurrent_io: GetVectoredConcurrentIo,
-        listener: std::net::TcpListener,
-    ) -> anyhow::Result<CancellableTask> {
-        let cancel = CancellationToken::new();
-        let ctx = RequestContextBuilder::new(TaskKind::PageRequestHandler)
-            .download_behavior(DownloadBehavior::Download)
-            .perf_span_dispatch(perf_trace_dispatch)
-            .detached_child();
-        let gate = Gate::default();
-
-        // Set up the TCP socket. We take a preconfigured TcpListener to bind the
-        // port early during startup.
-        let incoming = {
-            let _runtime = COMPUTE_REQUEST_RUNTIME.enter(); // required by TcpListener::from_std
-            listener.set_nonblocking(true)?;
-            tonic::transport::server::TcpIncoming::from(tokio::net::TcpListener::from_std(
-                listener,
-            )?)
-            .with_nodelay(Some(GRPC_TCP_NODELAY))
-            .with_keepalive(Some(GRPC_TCP_KEEPALIVE_TIME))
-        };
-
-        // Set up the gRPC server.
-        //
-        // TODO: consider tuning window sizes.
-        let mut server = tonic::transport::Server::builder()
-            .http2_keepalive_interval(Some(GRPC_HTTP2_KEEPALIVE_INTERVAL))
-            .http2_keepalive_timeout(Some(GRPC_HTTP2_KEEPALIVE_TIMEOUT))
-            .max_concurrent_streams(Some(GRPC_MAX_CONCURRENT_STREAMS));
-
-        // Main page service stack. Uses a mix of Tonic interceptors and Tower layers:
-        //
-        // * Interceptors: can inspect and modify the gRPC request. Sync code only, runs before service.
-        //
-        // * Layers: allow async code, can run code after the service response. However, only has access
-        //   to the raw HTTP request/response, not the gRPC types.
-        let page_service_handler = GrpcPageServiceHandler {
-            tenant_manager,
-            ctx,
-            gate_guard: gate.enter().expect("gate was just created"),
-            get_vectored_concurrent_io,
-        };
-
-        let observability_layer = ObservabilityLayer;
-        let mut tenant_interceptor = TenantMetadataInterceptor;
-        let mut auth_interceptor = TenantAuthInterceptor::new(auth);
-
-        let page_service = tower::ServiceBuilder::new()
-            // Create tracing span and record request start time.
-            .layer(observability_layer)
-            // Intercept gRPC requests.
-            .layer(tonic::service::InterceptorLayer::new(move |mut req| {
-                // Extract tenant metadata.
-                req = tenant_interceptor.call(req)?;
-                // Authenticate tenant JWT token.
-                req = auth_interceptor.call(req)?;
-                Ok(req)
-            }))
-            // Run the page service.
-            .service(proto::PageServiceServer::new(page_service_handler));
-        let server = server.add_service(page_service);
-
-        // Reflection service for use with e.g. grpcurl.
-        let reflection_service = tonic_reflection::server::Builder::configure()
-            .register_encoded_file_descriptor_set(proto::FILE_DESCRIPTOR_SET)
-            .build_v1()?;
-        let server = server.add_service(reflection_service);
-
-        // Spawn server task.
-        let task_cancel = cancel.clone();
-        let task = COMPUTE_REQUEST_RUNTIME.spawn(task_mgr::exit_on_panic_or_error(
-            "grpc listener",
-            async move {
-                let result = server
-                    .serve_with_incoming_shutdown(incoming, task_cancel.cancelled())
-                    .await;
-                if result.is_ok() {
-                    // TODO: revisit shutdown logic once page service is implemented.
-                    gate.close().await;
-                }
-                result
-            },
-        ));
-
-        Ok(CancellableTask { task, cancel })
-    }
-
     /// Errors if the request is executed on a non-zero shard. Only shard 0 has a complete view of
     /// relations and their sizes, as well as SLRU segments and similar data.
     #[allow(clippy::result_large_err)]

pageserver/src/tenant.rs

@@ -51,7 +51,6 @@ use secondary::heatmap::{HeatMapTenant, HeatMapTimeline};
 use storage_broker::BrokerClientChannel;
 use timeline::compaction::{CompactionOutcome, GcCompactionQueue};
 use timeline::import_pgdata::ImportingTimeline;
-use timeline::layer_manager::LayerManagerLockHolder;
 use timeline::offload::{OffloadError, offload_timeline};
 use timeline::{
     CompactFlags, CompactOptions, CompactionError, PreviousHeatmap, ShutdownMode, import_pgdata,
@@ -90,8 +89,7 @@ use crate::l0_flush::L0FlushGlobalState;
 use crate::metrics::{
     BROKEN_TENANTS_SET, CIRCUIT_BREAKERS_BROKEN, CIRCUIT_BREAKERS_UNBROKEN, CONCURRENT_INITDBS,
     INITDB_RUN_TIME, INITDB_SEMAPHORE_ACQUISITION_TIME, TENANT, TENANT_OFFLOADED_TIMELINES,
-    TENANT_STATE_METRIC, TENANT_SYNTHETIC_SIZE_METRIC, TIMELINE_STATE_METRIC,
-    remove_tenant_metrics,
+    TENANT_STATE_METRIC, TENANT_SYNTHETIC_SIZE_METRIC, remove_tenant_metrics,
 };
 use crate::task_mgr::TaskKind;
 use crate::tenant::config::LocationMode;
@@ -546,28 +544,6 @@ pub struct OffloadedTimeline {
 
     /// Part of the `OffloadedTimeline` object's lifecycle: this needs to be set before we drop it
     pub deleted_from_ancestor: AtomicBool,
-
-    _metrics_guard: OffloadedTimelineMetricsGuard,
-}
-
-/// Increases the offloaded timeline count metric when created, and decreases when dropped.
-struct OffloadedTimelineMetricsGuard;
-
-impl OffloadedTimelineMetricsGuard {
-    fn new() -> Self {
-        TIMELINE_STATE_METRIC
-            .with_label_values(&["offloaded"])
-            .inc();
-        Self
-    }
-}
-
-impl Drop for OffloadedTimelineMetricsGuard {
-    fn drop(&mut self) {
-        TIMELINE_STATE_METRIC
-            .with_label_values(&["offloaded"])
-            .dec();
-    }
 }
 
 impl OffloadedTimeline {
@@ -600,8 +576,6 @@ impl OffloadedTimeline {
 
             delete_progress: timeline.delete_progress.clone(),
             deleted_from_ancestor: AtomicBool::new(false),
-
-            _metrics_guard: OffloadedTimelineMetricsGuard::new(),
         })
     }
     fn from_manifest(tenant_shard_id: TenantShardId, manifest: &OffloadedTimelineManifest) -> Self {
@@ -621,7 +595,6 @@ impl OffloadedTimeline {
             archived_at,
             delete_progress: TimelineDeleteProgress::default(),
             deleted_from_ancestor: AtomicBool::new(false),
-            _metrics_guard: OffloadedTimelineMetricsGuard::new(),
         }
     }
     fn manifest(&self) -> OffloadedTimelineManifest {
@@ -1316,7 +1289,7 @@ impl TenantShard {
                         ancestor.is_some()
                             || timeline
                                 .layers
-                                .read(LayerManagerLockHolder::LoadLayerMap)
+                                .read()
                                 .await
                                 .layer_map()
                                 .expect(
@@ -2644,7 +2617,7 @@ impl TenantShard {
         }
         let layer_names = tline
             .layers
-            .read(LayerManagerLockHolder::Testing)
+            .read()
             .await
             .layer_map()
             .unwrap()
@@ -3159,12 +3132,7 @@ impl TenantShard {
 
         for timeline in &compact {
             // Collect L0 counts. Can't await while holding lock above.
-            if let Ok(lm) = timeline
-                .layers
-                .read(LayerManagerLockHolder::Compaction)
-                .await
-                .layer_map()
-            {
+            if let Ok(lm) = timeline.layers.read().await.layer_map() {
                 l0_counts.insert(timeline.timeline_id, lm.level0_deltas().len());
             }
         }
@@ -4906,7 +4874,7 @@ impl TenantShard {
         }
         let layer_names = tline
             .layers
-            .read(LayerManagerLockHolder::Testing)
+            .read()
             .await
             .layer_map()
             .unwrap()
@@ -6976,7 +6944,7 @@ mod tests {
             .await?;
         make_some_layers(tline.as_ref(), Lsn(0x20), &ctx).await?;
 
-        let layer_map = tline.layers.read(LayerManagerLockHolder::Testing).await;
+        let layer_map = tline.layers.read().await;
         let level0_deltas = layer_map
             .layer_map()?
             .level0_deltas()
@@ -7212,7 +7180,7 @@ mod tests {
         let lsn = Lsn(0x10);
         let inserted = bulk_insert_compact_gc(&tenant, &tline, &ctx, lsn, 50, 10000).await?;
 
-        let guard = tline.layers.read(LayerManagerLockHolder::Testing).await;
+        let guard = tline.layers.read().await;
         let lm = guard.layer_map()?;
 
         lm.dump(true, &ctx).await?;
@@ -8240,23 +8208,12 @@ mod tests {
             tline.freeze_and_flush().await?; // force create a delta layer
         }
 
-        let before_num_l0_delta_files = tline
-            .layers
-            .read(LayerManagerLockHolder::Testing)
-            .await
-            .layer_map()?
-            .level0_deltas()
-            .len();
+        let before_num_l0_delta_files =
+            tline.layers.read().await.layer_map()?.level0_deltas().len();
 
         tline.compact(&cancel, EnumSet::default(), &ctx).await?;
 
-        let after_num_l0_delta_files = tline
-            .layers
-            .read(LayerManagerLockHolder::Testing)
-            .await
-            .layer_map()?
-            .level0_deltas()
-            .len();
+        let after_num_l0_delta_files = tline.layers.read().await.layer_map()?.level0_deltas().len();
 
         assert!(
             after_num_l0_delta_files < before_num_l0_delta_files,

pageserver/src/tenant/config.rs

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

pageserver/src/tenant/remote_timeline_client/upload.rs

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

pageserver/src/tenant/storage_layer/delta_layer.rs

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

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

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

pageserver/src/tenant/timeline.rs

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

pageserver/src/tenant/timeline/analysis.rs

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

pageserver/src/tenant/timeline/compaction.rs

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

pageserver/src/tenant/timeline/detach_ancestor.rs

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

pageserver/src/tenant/timeline/eviction_task.rs

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

pageserver/src/tenant/timeline/import_pgdata.rs

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

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

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

pageserver/src/tenant/timeline/layer_manager.rs

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

pgxn/neon/Makefile

@@ -1,6 +1,5 @@
 # pgxs/neon/Makefile
 
-
 MODULE_big = neon
 OBJS = \
 	$(WIN32RES) \
@@ -22,7 +21,8 @@ OBJS = \
 	walproposer.o \
 	walproposer_pg.o \
 	control_plane_connector.o \
-	walsender_hooks.o
+	walsender_hooks.o \
+	$(LIBCOMMUNICATOR_PATH)/libcommunicator.a
 
 PG_CPPFLAGS = -I$(libpq_srcdir)
 SHLIB_LINK_INTERNAL = $(libpq)

pgxn/neon/communicator.c

@@ -1092,15 +1092,13 @@ communicator_prefetch_register_bufferv(BufferTag tag, neon_request_lsns *frlsns,
 		   MyPState->ring_last <= ring_index);
 }
 
-/* Internal version. Returns the ring index of the last block (result of this function is used only
-*  when nblocks==1)
-*/
+/* internal version. Returns the ring index */
 static uint64
 prefetch_register_bufferv(BufferTag tag, neon_request_lsns *frlsns,
 						  BlockNumber nblocks, const bits8 *mask,
 						  bool is_prefetch)
 {
-	uint64		last_ring_index;
+	uint64		min_ring_index;
 	PrefetchRequest hashkey;
 #ifdef USE_ASSERT_CHECKING
 	bool		any_hits = false;
@@ -1124,12 +1122,13 @@ Retry:
 		MyPState->ring_unused - MyPState->ring_receive;
 	MyNeonCounters->getpage_prefetches_buffered =
 		MyPState->n_responses_buffered;
-	last_ring_index = UINT64_MAX;
 
+	min_ring_index = UINT64_MAX;
 	for (int i = 0; i < nblocks; i++)
 	{
 		PrefetchRequest *slot = NULL;
 		PrfHashEntry *entry = NULL;
+		uint64		ring_index;
 		neon_request_lsns *lsns;
 
 		if (PointerIsValid(mask) && BITMAP_ISSET(mask, i))
@@ -1153,12 +1152,12 @@ Retry:
 		if (entry != NULL)
 		{
 			slot = entry->slot;
-			last_ring_index = slot->my_ring_index;
-			Assert(slot == GetPrfSlot(last_ring_index));
+			ring_index = slot->my_ring_index;
+			Assert(slot == GetPrfSlot(ring_index));
 
 			Assert(slot->status != PRFS_UNUSED);
-			Assert(MyPState->ring_last <= last_ring_index &&
-				   last_ring_index < MyPState->ring_unused);
+			Assert(MyPState->ring_last <= ring_index &&
+				   ring_index < MyPState->ring_unused);
 			Assert(BufferTagsEqual(&slot->buftag, &hashkey.buftag));
 
 			/*
@@ -1170,9 +1169,9 @@ Retry:
 				if (!neon_prefetch_response_usable(lsns, slot))
 				{
 					/* Wait for the old request to finish and discard it */
-					if (!prefetch_wait_for(last_ring_index))
+					if (!prefetch_wait_for(ring_index))
 						goto Retry;
-					prefetch_set_unused(last_ring_index);
+					prefetch_set_unused(ring_index);
 					entry = NULL;
 					slot = NULL;
 					pgBufferUsage.prefetch.expired += 1;
@@ -1189,12 +1188,13 @@ Retry:
 				 */
 				if (slot->status == PRFS_TAG_REMAINS)
 				{
-					prefetch_set_unused(last_ring_index);
+					prefetch_set_unused(ring_index);
 					entry = NULL;
 					slot = NULL;
 				}
 				else
 				{
+					min_ring_index = Min(min_ring_index, ring_index);
 					/* The buffered request is good enough, return that index */
 					if (is_prefetch)
 						pgBufferUsage.prefetch.duplicates++;
@@ -1283,12 +1283,12 @@ Retry:
 		 * The next buffer pointed to by `ring_unused` is now definitely empty, so
 		 * we can insert the new request to it.
 		 */
-		last_ring_index = MyPState->ring_unused;
+		ring_index = MyPState->ring_unused;
 
-		Assert(MyPState->ring_last <= last_ring_index &&
-			   last_ring_index <= MyPState->ring_unused);
+		Assert(MyPState->ring_last <= ring_index &&
+			   ring_index <= MyPState->ring_unused);
 
-		slot = GetPrfSlotNoCheck(last_ring_index);
+		slot = GetPrfSlotNoCheck(ring_index);
 
 		Assert(slot->status == PRFS_UNUSED);
 
@@ -1298,9 +1298,11 @@ Retry:
 		 */
 		slot->buftag = hashkey.buftag;
 		slot->shard_no = get_shard_number(&tag);
-		slot->my_ring_index = last_ring_index;
+		slot->my_ring_index = ring_index;
 		slot->flags = 0;
 
+		min_ring_index = Min(min_ring_index, ring_index);
+
 		if (is_prefetch)
 			MyNeonCounters->getpage_prefetch_requests_total++;
 		else
@@ -1313,12 +1315,11 @@ Retry:
 		MyPState->ring_unused - MyPState->ring_receive;
 
 	Assert(any_hits);
-	Assert(last_ring_index != UINT64_MAX);
 
-	Assert(GetPrfSlot(last_ring_index)->status == PRFS_REQUESTED ||
-		   GetPrfSlot(last_ring_index)->status == PRFS_RECEIVED);
-	Assert(MyPState->ring_last <= last_ring_index &&
-		   last_ring_index < MyPState->ring_unused);
+	Assert(GetPrfSlot(min_ring_index)->status == PRFS_REQUESTED ||
+		   GetPrfSlot(min_ring_index)->status == PRFS_RECEIVED);
+	Assert(MyPState->ring_last <= min_ring_index &&
+		   min_ring_index < MyPState->ring_unused);
 
 	if (flush_every_n_requests > 0 &&
 		MyPState->ring_unused - MyPState->ring_flush >= flush_every_n_requests)
@@ -1334,7 +1335,7 @@ Retry:
 		MyPState->ring_flush = MyPState->ring_unused;
 	}
 
-	return last_ring_index;
+	return min_ring_index;
 }
 
 static bool

pgxn/neon/communicator/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "communicator"
+version = "0.1.0"
+edition = "2024"
+
+[lib]
+crate-type = ["staticlib"]
+
+[dependencies]
+neon-shmem.workspace = true
+
+[build-dependencies]
+cbindgen.workspace = true

pgxn/neon/communicator/README.md

@@ -0,0 +1,8 @@
+This package will evolve into a "compute-pageserver communicator"
+process and machinery. For now, it just provides wrappers on the
+neon-shmem Rust crate, to allow using it in the C implementation of
+the LFC.
+
+At compilation time, pgxn/neon/communicator/ produces a static
+library, libcommunicator.a. It is linked to the neon.so extension
+library. 

pgxn/neon/communicator/build.rs

@@ -0,0 +1,22 @@
+use std::env;
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let crate_dir = env::var("CARGO_MANIFEST_DIR").unwrap();
+
+    cbindgen::generate(crate_dir).map_or_else(
+        |error| match error {
+            cbindgen::Error::ParseSyntaxError { .. } => {
+                // This means there was a syntax error in the Rust sources. Don't panic, because
+                // we want the build to continue and the Rust compiler to hit the error. The
+                // Rust compiler produces a better error message than cbindgen.
+                eprintln!("Generating C bindings failed because of a Rust syntax error");
+            }
+            e => panic!("Unable to generate C bindings: {:?}", e),
+        },
+        |bindings| {
+            bindings.write_to_file("communicator_bindings.h");
+        },
+    );
+
+    Ok(())
+}

pgxn/neon/communicator/cbindgen.toml

@@ -0,0 +1,4 @@
+language = "C"
+
+[enum]
+prefix_with_name = true

pgxn/neon/communicator/src/file_cache_hashmap.rs

@@ -0,0 +1,240 @@
+//! Glue code to allow using the Rust shmem hash map implementation from C code
+//!
+//! For convience of adapting existing code, the interface provided somewhat resembles the dynahash
+//! interface.
+//!
+//! NOTE: The caller is responsible for locking! The caller is expected to hold the PostgreSQL
+//! LWLock, 'lfc_lock', while accessing the hash table, in shared or exclusive mode as appropriate.
+
+use std::ffi::c_void;
+use std::marker::PhantomData;
+
+use neon_shmem::hash::entry::Entry;
+use neon_shmem::hash::{HashMapAccess, HashMapInit};
+use neon_shmem::shmem::ShmemHandle;
+
+/// NB: This must match the definition of BufferTag in Postgres C headers. We could use bindgen to
+/// generate this from the C headers, but prefer to not introduce dependency on bindgen for now.
+///
+/// Note that there are no padding bytes. If the corresponding C struct has padding bytes, the C C
+/// code must clear them.
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+#[repr(C)]
+pub struct FileCacheKey {
+    pub _spc_id: u32,
+    pub _db_id: u32,
+    pub _rel_number: u32,
+    pub _fork_num: u32,
+    pub _block_num: u32,
+}
+
+/// Like with FileCacheKey, this must match the definition of FileCacheEntry in file_cache.c.  We
+/// don't look at the contents here though, it's sufficent that the size and alignment matches.
+#[derive(Clone, Debug, Default)]
+#[repr(C)]
+pub struct FileCacheEntry {
+    pub _offset: u32,
+    pub _access_count: u32,
+    pub _prev: *mut FileCacheEntry,
+    pub _next: *mut FileCacheEntry,
+    pub _state: [u32; 8],
+}
+
+/// XXX: This could be just:
+///
+/// ```ignore
+/// type FileCacheHashMapHandle = HashMapInit<'a, FileCacheKey, FileCacheEntry>
+/// ```
+///
+/// but with that, cbindgen generates a broken typedef in the C header file which doesn't
+/// compile. It apparently gets confused by the generics.
+#[repr(transparent)]
+pub struct FileCacheHashMapHandle<'a>(
+    pub *mut c_void,
+    PhantomData<HashMapInit<'a, FileCacheKey, FileCacheEntry>>,
+);
+impl<'a> From<Box<HashMapInit<'a, FileCacheKey, FileCacheEntry>>> for FileCacheHashMapHandle<'a> {
+    fn from(x: Box<HashMapInit<'a, FileCacheKey, FileCacheEntry>>) -> Self {
+        FileCacheHashMapHandle(Box::into_raw(x) as *mut c_void, PhantomData::default())
+    }
+}
+impl<'a> From<FileCacheHashMapHandle<'a>> for Box<HashMapInit<'a, FileCacheKey, FileCacheEntry>> {
+    fn from(x: FileCacheHashMapHandle) -> Self {
+        unsafe { Box::from_raw(x.0.cast()) }
+    }
+}
+
+/// XXX: same for this
+#[repr(transparent)]
+pub struct FileCacheHashMapAccess<'a>(
+    pub *mut c_void,
+    PhantomData<HashMapAccess<'a, FileCacheKey, FileCacheEntry>>,
+);
+impl<'a> From<Box<HashMapAccess<'a, FileCacheKey, FileCacheEntry>>> for FileCacheHashMapAccess<'a> {
+    fn from(x: Box<HashMapAccess<'a, FileCacheKey, FileCacheEntry>>) -> Self {
+        // Convert the Box into a raw mutable pointer to the HashMapAccess itself.
+        // This transfers ownership of the HashMapAccess (and its contained ShmemHandle)
+        // to the raw pointer. The C caller is now responsible for managing this memory.
+        FileCacheHashMapAccess(Box::into_raw(x) as *mut c_void, PhantomData::default())
+    }
+}
+impl<'a> FileCacheHashMapAccess<'a> {
+    fn as_ref(self) -> &'a HashMapAccess<'a, FileCacheKey, FileCacheEntry> {
+        let ptr: *mut HashMapAccess<'_, FileCacheKey, FileCacheEntry> = self.0.cast();
+        unsafe { ptr.as_ref().unwrap() }
+    }
+    fn as_mut(self) -> &'a mut HashMapAccess<'a, FileCacheKey, FileCacheEntry> {
+        let ptr: *mut HashMapAccess<'_, FileCacheKey, FileCacheEntry> = self.0.cast();
+        unsafe { ptr.as_mut().unwrap() }
+    }
+}
+
+/// Initialize the shared memory area at postmaster startup. The returned handle is inherited
+/// by all the backend processes across fork()
+#[unsafe(no_mangle)]
+pub extern "C" fn bcomm_file_cache_shmem_init<'a>(
+    initial_num_buckets: u32,
+    max_num_buckets: u32,
+) -> FileCacheHashMapHandle<'a> {
+    let max_bytes = HashMapInit::<FileCacheKey, FileCacheEntry>::estimate_size(max_num_buckets);
+    let shmem_handle =
+        ShmemHandle::new("lfc mapping", 0, max_bytes).expect("shmem initialization failed");
+
+    let handle = HashMapInit::<FileCacheKey, FileCacheEntry>::init_in_shmem(
+        initial_num_buckets,
+        shmem_handle,
+    );
+
+    Box::new(handle).into()
+}
+
+/// Initialize the access to the shared memory area in a backend process.
+///
+/// XXX: I'm not sure if this actually gets called in each process, or if the returned struct
+/// is also inherited across fork(). It currently works either way but if this did more
+/// initialization that needed to be done after fork(), then it would matter.
+#[unsafe(no_mangle)]
+pub extern "C" fn bcomm_file_cache_shmem_access<'a>(
+    handle: FileCacheHashMapHandle<'a>,
+) -> FileCacheHashMapAccess<'a> {
+    let handle: Box<HashMapInit<'_, FileCacheKey, FileCacheEntry>> = handle.into();
+    Box::new(handle.attach_writer()).into()
+}
+
+/// Return the current number of buckets in the hash table
+#[unsafe(no_mangle)]
+pub extern "C" fn bcomm_file_cache_get_num_buckets<'a>(
+    map: FileCacheHashMapAccess<'static>,
+) -> u32 {
+    let map = map.as_ref();
+    map.get_num_buckets().try_into().unwrap()
+}
+
+/// Look up the entry with given key and hash.
+///
+/// This is similar to dynahash's hash_search(... , HASH_FIND)
+#[unsafe(no_mangle)]
+pub extern "C" fn bcomm_file_cache_hash_find<'a>(
+    map: FileCacheHashMapAccess<'static>,
+    key: &FileCacheKey,
+    hash: u64,
+) -> Option<&'static FileCacheEntry> {
+    let map = map.as_ref();
+    map.get_with_hash(key, hash)
+}
+
+/// Look up the entry at given bucket position
+///
+/// This has no direct equivalent in the dynahash interface, but can be used to
+/// iterate through all entries in the hash table.
+#[unsafe(no_mangle)]
+pub extern "C" fn bcomm_file_cache_hash_get_at_pos<'a>(
+    map: FileCacheHashMapAccess<'static>,
+    pos: u32,
+) -> Option<&'static FileCacheEntry> {
+    let map = map.as_ref();
+    map.get_at_bucket(pos as usize).map(|(_k, v)| v)
+}
+
+/// Remove entry, given a pointer to the value.
+///
+/// This is equivalent to dynahash hash_search(entry->key, HASH_REMOVE), where 'entry'
+/// is an entry you have previously looked up
+#[unsafe(no_mangle)]
+pub extern "C" fn bcomm_file_cache_hash_remove_entry<'a, 'b>(
+    map: FileCacheHashMapAccess,
+    entry: *mut FileCacheEntry,
+) {
+    let map = map.as_mut();
+    let pos = map.get_bucket_for_value(entry);
+    match map.entry_at_bucket(pos) {
+        Some(e) => {
+            e.remove();
+        }
+        None => {
+            // todo: shouldn't happen, panic?
+        }
+    }
+}
+
+/// Compute the hash for given key
+///
+/// This is equivalent to dynahash get_hash_value() function. We use Rust's default hasher
+/// for calculating the hash though.
+#[unsafe(no_mangle)]
+pub extern "C" fn bcomm_file_cache_get_hash_value<'a, 'b>(
+    map: FileCacheHashMapAccess<'static>,
+    key: &FileCacheKey,
+) -> u64 {
+    map.as_ref().get_hash_value(key)
+}
+
+/// Insert a new entry to the hash table
+///
+/// This is equivalent to dynahash hash_search(..., HASH_ENTER).
+#[unsafe(no_mangle)]
+pub extern "C" fn bcomm_file_cache_hash_enter<'a, 'b>(
+    map: FileCacheHashMapAccess,
+    key: &FileCacheKey,
+    hash: u64,
+    found: &mut bool,
+) -> *mut FileCacheEntry {
+    match map.as_mut().entry_with_hash(key.clone(), hash) {
+        Entry::Occupied(mut e) => {
+            *found = true;
+            e.get_mut()
+        }
+        Entry::Vacant(e) => {
+            *found = false;
+            let initial_value = FileCacheEntry::default();
+            e.insert(initial_value).expect("TODO: hash table full")
+        }
+    }
+}
+
+/// Get the key for a given entry, which must be present in the hash table.
+///
+/// Dynahash requires the key to be part of the "value" struct, so you can always
+/// access the key with something like `entry->key`. The Rust implementation however
+/// stores the key separately. This function extracts the separately stored key.
+#[unsafe(no_mangle)]
+pub extern "C" fn bcomm_file_cache_hash_get_key_for_entry<'a, 'b>(
+    map: FileCacheHashMapAccess,
+    entry: *const FileCacheEntry,
+) -> Option<&FileCacheKey> {
+    let map = map.as_ref();
+    let pos = map.get_bucket_for_value(entry);
+    map.get_at_bucket(pos as usize).map(|(k, _v)| k)
+}
+
+/// Remove all entries from the hash table
+#[unsafe(no_mangle)]
+pub extern "C" fn bcomm_file_cache_hash_reset<'a, 'b>(map: FileCacheHashMapAccess) {
+    let map = map.as_mut();
+    let num_buckets = map.get_num_buckets();
+    for i in 0..num_buckets {
+        if let Some(e) = map.entry_at_bucket(i) {
+            e.remove();
+        }
+    }
+}

pgxn/neon/communicator/src/lib.rs

@@ -0,0 +1 @@
+pub mod file_cache_hashmap;

pgxn/neon/file_cache.c

@@ -21,7 +21,7 @@
 #include "access/xlog.h"
 #include "funcapi.h"
 #include "miscadmin.h"
-#include "common/hashfn.h"
+#include "common/file_utils.h"
 #include "pgstat.h"
 #include "port/pg_iovec.h"
 #include "postmaster/bgworker.h"
@@ -36,7 +36,6 @@
 #include "storage/procsignal.h"
 #include "tcop/tcopprot.h"
 #include "utils/builtins.h"
-#include "utils/dynahash.h"
 #include "utils/guc.h"
 
 #if PG_VERSION_NUM >= 150000
@@ -46,6 +45,7 @@
 #include "hll.h"
 #include "bitmap.h"
 #include "file_cache.h"
+#include "file_cache_rust_hash.h"
 #include "neon.h"
 #include "neon_lwlsncache.h"
 #include "neon_perf_counters.h"
@@ -64,7 +64,7 @@
  *
  * Cache is always reconstructed at node startup, so we do not need to save mapping somewhere and worry about
  * its consistency.
-
+ *
  *
  * ## Holes
  *
@@ -76,13 +76,15 @@
  * fallocate(FALLOC_FL_PUNCH_HOLE) call. The nominal size of the file doesn't
  * shrink, but the disk space it uses does.
  *
- * Each hole is tracked by a dummy FileCacheEntry, which are kept in the
- * 'holes' linked list. They are entered into the chunk hash table, with a
- * special key where the blockNumber is used to store the 'offset' of the
- * hole, and all other fields are zero. Holes are never looked up in the hash
- * table, we only enter them there to have a FileCacheEntry that we can keep
- * in the linked list. If the soft limit is raised again, we reuse the holes
- * before extending the nominal size of the file.
+ * Each hole is tracked in a freelist. The freelist consists of two parts: a
+ * fixed-size array in shared memory, and a linked chain of on-disk
+ * blocks. When the in-memory array fills up, it's flushed to a new on-disk
+ * chunk. If the soft limit is raised again, we reuse the holes before
+ * extending the nominal size of the file.
+ *
+ * The in-memory freelist array is protected by 'lfc_lock', while the on-disk
+ * chain is protected by a separate 'lfc_freelist_lock'.  Locking rule to
+ * avoid deadlocks: always acquire lfc_freelist_lock first, then lfc_lock.
  */
 
 /* Local file storage allocation chunk.
@@ -92,13 +94,15 @@
  *    1Mb chunks can reduce hash map size to 320Mb.
  * 2. Improve access locality, subsequent pages will be allocated together improving seqscan speed
  */
-#define MAX_BLOCKS_PER_CHUNK_LOG  7 /* 1Mb chunk */
-#define MAX_BLOCKS_PER_CHUNK	  (1 << MAX_BLOCKS_PER_CHUNK_LOG)
+#define BLOCKS_PER_CHUNK_LOG  7 /* 1Mb chunk */
+#define BLOCKS_PER_CHUNK	  (1 << BLOCKS_PER_CHUNK_LOG)
 
 #define MB					((uint64)1024*1024)
 
-#define SIZE_MB_TO_CHUNKS(size) ((uint32)((size) * MB / BLCKSZ >> lfc_chunk_size_log))
-#define BLOCK_TO_CHUNK_OFF(blkno) ((blkno) & (lfc_blocks_per_chunk-1))
+#define SIZE_MB_TO_CHUNKS(size) ((uint32)((size) * MB / BLCKSZ >> BLOCKS_PER_CHUNK_LOG))
+#define BLOCK_TO_CHUNK_OFF(blkno) ((blkno) & (BLOCKS_PER_CHUNK-1))
+
+#define INVALID_OFFSET (0xffffffff)
 
 /*
  * Blocks are read or written to LFC file outside LFC critical section.
@@ -119,15 +123,18 @@ typedef enum FileCacheBlockState
 
 typedef struct FileCacheEntry
 {
-	BufferTag	key;
-	uint32		hash;
 	uint32		offset;
 	uint32		access_count;
-	dlist_node	list_node;		/* LRU/holes list node */
-	uint32		state[FLEXIBLE_ARRAY_MEMBER]; /* two bits per block */
+	dlist_node	list_node;		/* LRU list node */
+	uint32		state[(BLOCKS_PER_CHUNK * 2 + 31) / 32]; /* two bits per block */
 } FileCacheEntry;
 
-#define FILE_CACHE_ENRTY_SIZE MAXALIGN(offsetof(FileCacheEntry, state) + (lfc_blocks_per_chunk*2+31)/32*4)
+/* Todo: alignment must be the same too */
+StaticAssertDecl(sizeof(FileCacheEntry) == sizeof(RustFileCacheEntry),
+				 "Rust and C declarations of FileCacheEntry are incompatible");
+StaticAssertDecl(sizeof(BufferTag) == sizeof(RustFileCacheKey),
+				 "Rust and C declarations of FileCacheKey are incompatible");
+
 #define GET_STATE(entry, i) (((entry)->state[(i) / 16] >> ((i) % 16 * 2)) & 3)
 #define SET_STATE(entry, i, new_state) (entry)->state[(i) / 16] = ((entry)->state[(i) / 16] & ~(3 << ((i) % 16 * 2))) | ((new_state) << ((i) % 16 * 2))
 
@@ -136,6 +143,9 @@ typedef struct FileCacheEntry
 
 #define MAX_PREWARM_WORKERS 8
 
+
+#define FREELIST_ENTRIES_PER_CHUNK (BLOCKS_PER_CHUNK * BLCKSZ / sizeof(uint32) - 2)
+
 typedef struct PrewarmWorkerState
 {
 	uint32		prewarmed_pages;
@@ -161,7 +171,6 @@ typedef struct FileCacheControl
 	uint64		evicted_pages;	/* number of evicted pages */
 	dlist_head	lru;			/* double linked list for LRU replacement
 								 * algorithm */
-	dlist_head  holes;          /* double linked list of punched holes */
 	HyperLogLogState wss_estimation; /* estimation of working set size */
 	ConditionVariable cv[N_COND_VARS]; /* turnstile of condition variables */
 	PrewarmWorkerState prewarm_workers[MAX_PREWARM_WORKERS];
@@ -172,23 +181,39 @@ typedef struct FileCacheControl
 	bool   prewarm_active;
 	bool   prewarm_canceled;
 	dsm_handle prewarm_lfc_state_handle;
+
+	/*
+	 * Free list. This is large enough to hold one chunks worth of entries.
+	 */
+	uint32		freelist_size;
+	uint32		freelist_head;
+	uint32		num_free_pages;
+	uint32		free_pages[FREELIST_ENTRIES_PER_CHUNK];
 } FileCacheControl;
 
+typedef struct FreeListChunk
+{
+	uint32		next;
+	uint32		num_free_pages;
+	uint32		free_pages[FREELIST_ENTRIES_PER_CHUNK];
+} FreeListChunk;
+
 #define FILE_CACHE_STATE_MAGIC 0xfcfcfcfc
 
 #define FILE_CACHE_STATE_BITMAP(fcs)	((uint8*)&(fcs)->chunks[(fcs)->n_chunks])
-#define FILE_CACHE_STATE_SIZE_FOR_CHUNKS(n_chunks)	(sizeof(FileCacheState) + (n_chunks)*sizeof(BufferTag) + (((n_chunks) * lfc_blocks_per_chunk)+7)/8)
+#define FILE_CACHE_STATE_SIZE_FOR_CHUNKS(n_chunks)	(sizeof(FileCacheState) + (n_chunks)*sizeof(BufferTag) + (((n_chunks) * BLOCKS_PER_CHUNK)+7)/8)
 #define FILE_CACHE_STATE_SIZE(fcs)		(sizeof(FileCacheState) + (fcs->n_chunks)*sizeof(BufferTag) + (((fcs->n_chunks) << fcs->chunk_size_log)+7)/8)
 
-static HTAB *lfc_hash;
+static FileCacheHashMapHandle lfc_hash_handle;
+static FileCacheHashMapAccess lfc_hash;
 static int	lfc_desc = -1;
 static LWLockId lfc_lock;
+static LWLockId lfc_freelist_lock;
 static int	lfc_max_size;
 static int	lfc_size_limit;
 static int	lfc_prewarm_limit;
 static int	lfc_prewarm_batch;
-static int	lfc_chunk_size_log = MAX_BLOCKS_PER_CHUNK_LOG;
-static int	lfc_blocks_per_chunk = MAX_BLOCKS_PER_CHUNK;
+static int	lfc_blocks_per_chunk_ro = BLOCKS_PER_CHUNK;
 static char *lfc_path;
 static uint64 lfc_generation;
 static FileCacheControl *lfc_ctl;
@@ -205,6 +230,11 @@ bool AmPrewarmWorker;
 
 #define LFC_ENABLED() (lfc_ctl->limit != 0)
 
+static bool freelist_push(uint32 offset);
+static bool freelist_prepare_pop(void);
+static uint32 freelist_pop(void);
+static bool freelist_is_empty(void);
+
 /*
  * Close LFC file if opened.
  * All backends should close their LFC files once LFC is disabled.
@@ -232,15 +262,9 @@ lfc_switch_off(void)
 
 	if (LFC_ENABLED())
 	{
-		HASH_SEQ_STATUS status;
-		FileCacheEntry *entry;
-
 		/* Invalidate hash */
-		hash_seq_init(&status, lfc_hash);
-		while ((entry = hash_seq_search(&status)) != NULL)
-		{
-			hash_search_with_hash_value(lfc_hash, &entry->key, entry->hash, HASH_REMOVE, NULL);
-		}
+		file_cache_hash_reset(lfc_hash);
+
 		lfc_ctl->generation += 1;
 		lfc_ctl->size = 0;
 		lfc_ctl->pinned = 0;
@@ -248,7 +272,9 @@ lfc_switch_off(void)
 		lfc_ctl->used_pages = 0;
 		lfc_ctl->limit = 0;
 		dlist_init(&lfc_ctl->lru);
-		dlist_init(&lfc_ctl->holes);
+
+		lfc_ctl->freelist_head = INVALID_OFFSET;
+		lfc_ctl->num_free_pages = 0;
 
 		/*
 		 * We need to use unlink to to avoid races in LFC write, because it is not
@@ -317,8 +343,8 @@ lfc_ensure_opened(void)
 static void
 lfc_shmem_startup(void)
 {
+	size_t		size;
 	bool		found;
-	static HASHCTL info;
 
 	if (prev_shmem_startup_hook)
 	{
@@ -327,27 +353,29 @@ lfc_shmem_startup(void)
 
 	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
 
-	lfc_ctl = (FileCacheControl *) ShmemInitStruct("lfc", sizeof(FileCacheControl), &found);
+	size = sizeof(FileCacheControl);
+
+	lfc_ctl = (FileCacheControl *) ShmemInitStruct("lfc", size, &found);
 	if (!found)
 	{
 		int			fd;
 		uint32		n_chunks = SIZE_MB_TO_CHUNKS(lfc_max_size);
 
 		lfc_lock = (LWLockId) GetNamedLWLockTranche("lfc_lock");
-		info.keysize = sizeof(BufferTag);
-		info.entrysize = FILE_CACHE_ENRTY_SIZE;
+		lfc_freelist_lock = (LWLockId) GetNamedLWLockTranche("lfc_freelist_lock");
 
 		/*
 		 * n_chunks+1 because we add new element to hash table before eviction
 		 * of victim
 		 */
-		lfc_hash = ShmemInitHash("lfc_hash",
-								 n_chunks + 1, n_chunks + 1,
-								 &info,
-								 HASH_ELEM | HASH_BLOBS);
-		memset(lfc_ctl, 0, sizeof(FileCacheControl));
+		lfc_hash_handle = file_cache_hash_shmem_init(n_chunks + 1, n_chunks + 1);
+
+		memset(lfc_ctl, 0, offsetof(FileCacheControl, free_pages));
 		dlist_init(&lfc_ctl->lru);
-		dlist_init(&lfc_ctl->holes);
+
+		lfc_ctl->freelist_size = FREELIST_ENTRIES_PER_CHUNK;
+		lfc_ctl->freelist_head = INVALID_OFFSET;
+		lfc_ctl->num_free_pages = 0;
 
 		/* Initialize hyper-log-log structure for estimating working set size */
 		initSHLL(&lfc_ctl->wss_estimation);
@@ -371,18 +399,25 @@ lfc_shmem_startup(void)
 
 	}
 	LWLockRelease(AddinShmemInitLock);
+
+	lfc_hash = file_cache_hash_shmem_access(lfc_hash_handle);
 }
 
 static void
 lfc_shmem_request(void)
 {
+	size_t		size;
+
 #if PG_VERSION_NUM>=150000
 	if (prev_shmem_request_hook)
 		prev_shmem_request_hook();
 #endif
 
-	RequestAddinShmemSpace(sizeof(FileCacheControl) + hash_estimate_size(SIZE_MB_TO_CHUNKS(lfc_max_size) + 1, FILE_CACHE_ENRTY_SIZE));
+	size = sizeof(FileCacheControl);
+
+	RequestAddinShmemSpace(size);
 	RequestNamedLWLockTranche("lfc_lock", 1);
+	RequestNamedLWLockTranche("lfc_freelist_lock", 2);
 }
 
 static bool
@@ -398,24 +433,6 @@ is_normal_backend(void)
 	return lfc_ctl && MyProc && UsedShmemSegAddr && !IsParallelWorker();
 }
 
-static bool
-lfc_check_chunk_size(int *newval, void **extra, GucSource source)
-{
-	if (*newval & (*newval - 1))
-	{
-		elog(ERROR, "LFC chunk size should be power of two");
-		return false;
-	}
-	return true;
-}
-
-static void
-lfc_change_chunk_size(int newval, void* extra)
-{
-	lfc_chunk_size_log = pg_ceil_log2_32(newval);
-}
-
-
 static bool
 lfc_check_limit_hook(int *newval, void **extra, GucSource source)
 {
@@ -435,12 +452,14 @@ lfc_change_limit_hook(int newval, void *extra)
 	if (!lfc_ctl || !is_normal_backend())
 		return;
 
+	LWLockAcquire(lfc_freelist_lock, LW_EXCLUSIVE);
 	LWLockAcquire(lfc_lock, LW_EXCLUSIVE);
 
 	/* Open LFC file only if LFC was enabled or we are going to reenable it */
 	if (newval == 0 && !LFC_ENABLED())
 	{
 		LWLockRelease(lfc_lock);
+		LWLockRelease(lfc_freelist_lock);
 		/* File should be reopened if LFC is reenabled */
 		lfc_close_file();
 		return;
@@ -449,6 +468,7 @@ lfc_change_limit_hook(int newval, void *extra)
 	if (!lfc_ensure_opened())
 	{
 		LWLockRelease(lfc_lock);
+		LWLockRelease(lfc_freelist_lock);
 		return;
 	}
 
@@ -464,35 +484,30 @@ lfc_change_limit_hook(int newval, void *extra)
 		 * returning their space to file system
 		 */
 		FileCacheEntry *victim = dlist_container(FileCacheEntry, list_node, dlist_pop_head_node(&lfc_ctl->lru));
-		FileCacheEntry *hole;
 		uint32		offset = victim->offset;
-		uint32		hash;
-		bool		found;
-		BufferTag	holetag;
 
 		CriticalAssert(victim->access_count == 0);
 #ifdef FALLOC_FL_PUNCH_HOLE
-		if (fallocate(lfc_desc, FALLOC_FL_PUNCH_HOLE | FALLOC_FL_KEEP_SIZE, (off_t) victim->offset * lfc_blocks_per_chunk * BLCKSZ, lfc_blocks_per_chunk * BLCKSZ) < 0)
+		if (fallocate(lfc_desc, FALLOC_FL_PUNCH_HOLE | FALLOC_FL_KEEP_SIZE, (off_t) victim->offset * BLOCKS_PER_CHUNK * BLCKSZ, BLOCKS_PER_CHUNK * BLCKSZ) < 0)
 			neon_log(LOG, "Failed to punch hole in file: %m");
 #endif
-		/* We remove the old entry, and re-enter a hole to the hash table */
-		for (int i = 0; i < lfc_blocks_per_chunk; i++)
+		/* We remove the entry, and enter a hole to the freelist */
+		for (int i = 0; i < BLOCKS_PER_CHUNK; i++)
 		{
 			bool is_page_cached = GET_STATE(victim, i) == AVAILABLE;
 			lfc_ctl->used_pages -= is_page_cached;
 			lfc_ctl->evicted_pages += is_page_cached;
 		}
-		hash_search_with_hash_value(lfc_hash, &victim->key, victim->hash, HASH_REMOVE, NULL);
+		file_cache_hash_remove_entry(lfc_hash, victim);
 
-		memset(&holetag, 0, sizeof(holetag));
-		holetag.blockNum = offset;
-		hash = get_hash_value(lfc_hash, &holetag);
-		hole = hash_search_with_hash_value(lfc_hash, &holetag, hash, HASH_ENTER, &found);
-		hole->hash = hash;
-		hole->offset = offset;
-		hole->access_count = 0;
-		CriticalAssert(!found);
-		dlist_push_tail(&lfc_ctl->holes, &hole->list_node);
+		if (!freelist_push(offset))
+		{
+			/* freelist_push already logged the error */
+			lfc_switch_off();
+			LWLockRelease(lfc_lock);
+			LWLockRelease(lfc_freelist_lock);
+			return;
+		}
 
 		lfc_ctl->used -= 1;
 	}
@@ -504,6 +519,7 @@ lfc_change_limit_hook(int newval, void *extra)
 	neon_log(DEBUG1, "set local file cache limit to %d", new_size);
 
 	LWLockRelease(lfc_lock);
+	LWLockRelease(lfc_freelist_lock);
 }
 
 void
@@ -579,14 +595,14 @@ lfc_init(void)
 	DefineCustomIntVariable("neon.file_cache_chunk_size",
 							"LFC chunk size in blocks (should be power of two)",
 							NULL,
-							&lfc_blocks_per_chunk,
-							MAX_BLOCKS_PER_CHUNK,
-							1,
-							MAX_BLOCKS_PER_CHUNK,
-							PGC_POSTMASTER,
+							&lfc_blocks_per_chunk_ro,
+							BLOCKS_PER_CHUNK,
+							BLOCKS_PER_CHUNK,
+							BLOCKS_PER_CHUNK,
+							PGC_INTERNAL,
 							GUC_UNIT_BLOCKS,
-							lfc_check_chunk_size,
-							lfc_change_chunk_size,
+							NULL,
+							NULL,
 							NULL);
 
 	DefineCustomIntVariable("neon.file_cache_prewarm_limit",
@@ -649,19 +665,19 @@ lfc_get_state(size_t max_entries)
 		fcs = (FileCacheState*)palloc0(state_size);
 		SET_VARSIZE(fcs, state_size);
 		fcs->magic = FILE_CACHE_STATE_MAGIC;
-		fcs->chunk_size_log = lfc_chunk_size_log;
+		fcs->chunk_size_log = BLOCKS_PER_CHUNK_LOG;
 		fcs->n_chunks = n_entries;
 		bitmap = FILE_CACHE_STATE_BITMAP(fcs);
 
 		dlist_reverse_foreach(iter, &lfc_ctl->lru)
 		{
 			FileCacheEntry *entry = dlist_container(FileCacheEntry, list_node, iter.cur);
-			fcs->chunks[i] = entry->key;
-			for (int j = 0; j < lfc_blocks_per_chunk; j++)
+			fcs->chunks[i] = *file_cache_hash_get_key_for_entry(lfc_hash, entry);
+			for (int j = 0; j < BLOCKS_PER_CHUNK; j++)
 			{
 				if (GET_STATE(entry, j) != UNAVAILABLE)
 				{
-					BITMAP_SET(bitmap, i*lfc_blocks_per_chunk + j);
+					BITMAP_SET(bitmap, i*BLOCKS_PER_CHUNK + j);
 					n_pages += 1;
 				}
 			}
@@ -670,7 +686,7 @@ lfc_get_state(size_t max_entries)
 		}
 		Assert(i == n_entries);
 		fcs->n_pages = n_pages;
-		Assert(pg_popcount((char*)bitmap, ((n_entries << lfc_chunk_size_log) + 7)/8) == n_pages);
+		Assert(pg_popcount((char*)bitmap, ((n_entries << BLOCKS_PER_CHUNK_LOG) + 7)/8) == n_pages);
 		elog(LOG, "LFC: save state of %d chunks %d pages", (int)n_entries, (int)n_pages);
 	}
 
@@ -726,7 +742,7 @@ lfc_prewarm(FileCacheState* fcs, uint32 n_workers)
 	}
 
 	fcs_chunk_size_log = fcs->chunk_size_log;
-	if (fcs_chunk_size_log > MAX_BLOCKS_PER_CHUNK_LOG)
+	if (fcs_chunk_size_log > BLOCKS_PER_CHUNK_LOG)
 	{
 		elog(ERROR, "LFC: Invalid chunk size log: %u", fcs->chunk_size_log);
 	}
@@ -945,7 +961,7 @@ lfc_invalidate(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber nblocks)
 {
 	BufferTag	tag;
 	FileCacheEntry *entry;
-	uint32		hash;
+	uint64		hash;
 
 	if (lfc_maybe_disabled())	/* fast exit if file cache is disabled */
 		return;
@@ -958,14 +974,14 @@ lfc_invalidate(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber nblocks)
 	LWLockAcquire(lfc_lock, LW_EXCLUSIVE);
 	if (LFC_ENABLED())
 	{
-		for (BlockNumber blkno = 0; blkno < nblocks; blkno += lfc_blocks_per_chunk)
+		for (BlockNumber blkno = 0; blkno < nblocks; blkno += BLOCKS_PER_CHUNK)
 		{
 			tag.blockNum = blkno;
-			hash = get_hash_value(lfc_hash, &tag);
-			entry = hash_search_with_hash_value(lfc_hash, &tag, hash, HASH_FIND, NULL);
+			hash = file_cache_hash_get_hash_value(lfc_hash, &tag);
+			entry = file_cache_hash_find(lfc_hash, &tag, hash);
 			if (entry != NULL)
 			{
-				for (int i = 0; i < lfc_blocks_per_chunk; i++)
+				for (int i = 0; i < BLOCKS_PER_CHUNK; i++)
 				{
 					if (GET_STATE(entry, i) == AVAILABLE)
 					{
@@ -990,7 +1006,7 @@ lfc_cache_contains(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno)
 	FileCacheEntry *entry;
 	int			chunk_offs = BLOCK_TO_CHUNK_OFF(blkno);
 	bool		found = false;
-	uint32		hash;
+	uint64		hash;
 
 	if (lfc_maybe_disabled())	/* fast exit if file cache is disabled */
 		return false;
@@ -1000,12 +1016,12 @@ lfc_cache_contains(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno)
 	tag.blockNum = blkno - chunk_offs;
 
 	CriticalAssert(BufTagGetRelNumber(&tag) != InvalidRelFileNumber);
-	hash = get_hash_value(lfc_hash, &tag);
+	hash = file_cache_hash_get_hash_value(lfc_hash, &tag);
 
 	LWLockAcquire(lfc_lock, LW_SHARED);
 	if (LFC_ENABLED())
 	{
-		entry = hash_search_with_hash_value(lfc_hash, &tag, hash, HASH_FIND, NULL);
+		entry = file_cache_hash_find(lfc_hash, &tag, hash);
 		found = entry != NULL && GET_STATE(entry, chunk_offs) != UNAVAILABLE;
 	}
 	LWLockRelease(lfc_lock);
@@ -1024,7 +1040,7 @@ lfc_cache_containsv(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 	FileCacheEntry *entry;
 	uint32		chunk_offs;
 	int			found = 0;
-	uint32		hash;
+	uint64		hash;
 	int			i = 0;
 
 	if (lfc_maybe_disabled())	/* fast exit if file cache is disabled */
@@ -1037,7 +1053,7 @@ lfc_cache_containsv(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 
 	chunk_offs = BLOCK_TO_CHUNK_OFF(blkno);
 	tag.blockNum = blkno - chunk_offs;
-	hash = get_hash_value(lfc_hash, &tag);
+	hash = file_cache_hash_get_hash_value(lfc_hash, &tag);
 
 	LWLockAcquire(lfc_lock, LW_SHARED);
 
@@ -1048,12 +1064,12 @@ lfc_cache_containsv(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 	}
 	while (true)
 	{
-		int		this_chunk = Min(nblocks - i, lfc_blocks_per_chunk - chunk_offs);
-		entry = hash_search_with_hash_value(lfc_hash, &tag, hash, HASH_FIND, NULL);
+		int		this_chunk = Min(nblocks - i, BLOCKS_PER_CHUNK - chunk_offs);
+		entry = file_cache_hash_find(lfc_hash, &tag, hash);
 
 		if (entry != NULL)
 		{
-			for (; chunk_offs < lfc_blocks_per_chunk && i < nblocks; chunk_offs++, i++)
+			for (; chunk_offs < BLOCKS_PER_CHUNK && i < nblocks; chunk_offs++, i++)
 			{
 				if (GET_STATE(entry, chunk_offs) != UNAVAILABLE)
 				{
@@ -1079,7 +1095,7 @@ lfc_cache_containsv(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 		 */
 		chunk_offs = BLOCK_TO_CHUNK_OFF(blkno + i);
 		tag.blockNum = (blkno + i) - chunk_offs;
-		hash = get_hash_value(lfc_hash, &tag);
+		hash = file_cache_hash_get_hash_value(lfc_hash, &tag);
 	}
 
 	LWLockRelease(lfc_lock);
@@ -1128,7 +1144,7 @@ lfc_readv_select(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 	BufferTag	tag;
 	FileCacheEntry *entry;
 	ssize_t		rc;
-	uint32		hash;
+	uint64		hash;
 	uint64		generation;
 	uint32		entry_offset;
 	int			blocks_read = 0;
@@ -1154,9 +1170,9 @@ lfc_readv_select(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 	while (nblocks > 0)
 	{
 		struct iovec iov[PG_IOV_MAX];
-		uint8	chunk_mask[MAX_BLOCKS_PER_CHUNK / 8] = {0};
+		uint8	chunk_mask[BLOCKS_PER_CHUNK / 8] = {0};
 		int		chunk_offs = BLOCK_TO_CHUNK_OFF(blkno);
-		int		blocks_in_chunk = Min(nblocks, lfc_blocks_per_chunk - chunk_offs);
+		int		blocks_in_chunk = Min(nblocks, BLOCKS_PER_CHUNK - chunk_offs);
 		int		iteration_hits = 0;
 		int		iteration_misses = 0;
 		uint64	io_time_us = 0;
@@ -1206,7 +1222,7 @@ lfc_readv_select(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 		Assert(iov_last_used - first_block_in_chunk_read >= n_blocks_to_read);
 
 		tag.blockNum = blkno - chunk_offs;
-		hash = get_hash_value(lfc_hash, &tag);
+		hash = file_cache_hash_get_hash_value(lfc_hash, &tag);
 		cv = &lfc_ctl->cv[hash % N_COND_VARS];
 
 		LWLockAcquire(lfc_lock, LW_EXCLUSIVE);
@@ -1219,13 +1235,13 @@ lfc_readv_select(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 			return blocks_read;
 		}
 
-		entry = hash_search_with_hash_value(lfc_hash, &tag, hash, HASH_FIND, NULL);
+		entry = file_cache_hash_find(lfc_hash, &tag, hash);
 
 		/* Approximate working set for the blocks assumed in this entry */
 		for (int i = 0; i < blocks_in_chunk; i++)
 		{
 			tag.blockNum = blkno + i;
-			addSHLL(&lfc_ctl->wss_estimation, hash_bytes((uint8_t const*)&tag, sizeof(tag)));
+			addSHLL(&lfc_ctl->wss_estimation, file_cache_hash_get_hash_value(lfc_hash, &tag));
 		}
 
 		if (entry == NULL)
@@ -1296,7 +1312,7 @@ lfc_readv_select(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 		if (iteration_hits != 0)
 		{
 			/* chunk offset (# of pages) into the LFC file */
-			off_t	first_read_offset = (off_t) entry_offset * lfc_blocks_per_chunk;
+			off_t	first_read_offset = (off_t) entry_offset * BLOCKS_PER_CHUNK;
 			int		nwrite = iov_last_used - first_block_in_chunk_read;
 			/* offset of first IOV */
 			first_read_offset += chunk_offs + first_block_in_chunk_read;
@@ -1373,14 +1389,14 @@ lfc_readv_select(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
  * Returns false if there are no unpinned entries and chunk can not be added.
  */
 static bool
-lfc_init_new_entry(FileCacheEntry* entry, uint32 hash)
+lfc_init_new_entry(FileCacheEntry *entry)
 {
 	/*-----------
 	 * If the chunk wasn't already in the LFC then we have these
 	 * options, in order of preference:
 	 *
 	 * Unless there is no space available, we can:
-	 *  1. Use an entry from the `holes` list, and
+	 *  1. Use an entry from the freelist, and
 	 *  2. Create a new entry.
 	 * We can always, regardless of space in the LFC:
 	 *  3. evict an entry from LRU, and
@@ -1388,17 +1404,10 @@ lfc_init_new_entry(FileCacheEntry* entry, uint32 hash)
 	 */
 	if (lfc_ctl->used < lfc_ctl->limit)
 	{
-		if (!dlist_is_empty(&lfc_ctl->holes))
+		if (!freelist_is_empty())
 		{
 			/* We can reuse a hole that was left behind when the LFC was shrunk previously */
-			FileCacheEntry *hole = dlist_container(FileCacheEntry, list_node,
-												   dlist_pop_head_node(&lfc_ctl->holes));
-			uint32 offset = hole->offset;
-			bool hole_found;
-
-			hash_search_with_hash_value(lfc_hash, &hole->key,
-										hole->hash, HASH_REMOVE, &hole_found);
-			CriticalAssert(hole_found);
+			uint32 offset = freelist_pop();
 
 			lfc_ctl->used += 1;
 			entry->offset = offset;			/* reuse the hole */
@@ -1427,7 +1436,7 @@ lfc_init_new_entry(FileCacheEntry* entry, uint32 hash)
 		FileCacheEntry *victim = dlist_container(FileCacheEntry, list_node,
 												 dlist_pop_head_node(&lfc_ctl->lru));
 
-		for (int i = 0; i < lfc_blocks_per_chunk; i++)
+		for (int i = 0; i < BLOCKS_PER_CHUNK; i++)
 		{
 			bool is_page_cached = GET_STATE(victim, i) == AVAILABLE;
 			lfc_ctl->used_pages -= is_page_cached;
@@ -1436,24 +1445,21 @@ lfc_init_new_entry(FileCacheEntry* entry, uint32 hash)
 
 		CriticalAssert(victim->access_count == 0);
 		entry->offset = victim->offset; /* grab victim's chunk */
-		hash_search_with_hash_value(lfc_hash, &victim->key,
-									victim->hash, HASH_REMOVE, NULL);
+		file_cache_hash_remove_entry(lfc_hash, victim);
 		neon_log(DEBUG2, "Swap file cache page");
 	}
 	else
 	{
 		/* Can't add this chunk - we don't have the space for it */
-		hash_search_with_hash_value(lfc_hash, &entry->key, hash,
-									HASH_REMOVE, NULL);
+		file_cache_hash_remove_entry(lfc_hash, entry);
 		lfc_ctl->prewarm_canceled = true; /* cancel prewarm if LFC limit is reached */
 		return false;
 	}
 
 	entry->access_count = 1;
-	entry->hash = hash;
 	lfc_ctl->pinned += 1;
 
-	for (int i = 0; i < lfc_blocks_per_chunk; i++)
+	for (int i = 0; i < BLOCKS_PER_CHUNK; i++)
 		SET_STATE(entry, i, UNAVAILABLE);
 
 	return true;
@@ -1490,7 +1496,7 @@ lfc_prefetch(NRelFileInfo rinfo, ForkNumber forknum, BlockNumber blkno,
 	FileCacheEntry *entry;
 	ssize_t		rc;
 	bool		found;
-	uint32		hash;
+	uint64		hash;
 	uint64		generation;
 	uint32		entry_offset;
 	instr_time io_start, io_end;
@@ -1509,9 +1515,10 @@ lfc_prefetch(NRelFileInfo rinfo, ForkNumber forknum, BlockNumber blkno,
 	CriticalAssert(BufTagGetRelNumber(&tag) != InvalidRelFileNumber);
 
 	tag.blockNum = blkno - chunk_offs;
-	hash = get_hash_value(lfc_hash, &tag);
+	hash = file_cache_hash_get_hash_value(lfc_hash, &tag);
 	cv = &lfc_ctl->cv[hash % N_COND_VARS];
 
+ retry:
 	LWLockAcquire(lfc_lock, LW_EXCLUSIVE);
 
 	if (!LFC_ENABLED() || !lfc_ensure_opened())
@@ -1520,6 +1527,9 @@ lfc_prefetch(NRelFileInfo rinfo, ForkNumber forknum, BlockNumber blkno,
 		return false;
 	}
 
+	if (!freelist_prepare_pop())
+		goto retry;
+
 	lwlsn = neon_get_lwlsn(rinfo, forknum, blkno);
 
 	if (lwlsn > lsn)
@@ -1530,12 +1540,12 @@ lfc_prefetch(NRelFileInfo rinfo, ForkNumber forknum, BlockNumber blkno,
 		return false;
 	}
 
-	entry = hash_search_with_hash_value(lfc_hash, &tag, hash, HASH_ENTER, &found);
+	entry = file_cache_hash_enter(lfc_hash, &tag, hash, &found);
 
 	if (lfc_prewarm_update_ws_estimation)
 	{
 		tag.blockNum = blkno;
-		addSHLL(&lfc_ctl->wss_estimation, hash_bytes((uint8_t const*)&tag, sizeof(tag)));
+		addSHLL(&lfc_ctl->wss_estimation, file_cache_hash_get_hash_value(lfc_hash, &tag));
 	}
 	if (found)
 	{
@@ -1557,7 +1567,7 @@ lfc_prefetch(NRelFileInfo rinfo, ForkNumber forknum, BlockNumber blkno,
 	}
 	else
 	{
-		if (!lfc_init_new_entry(entry, hash))
+		if (!lfc_init_new_entry(entry))
 		{
 			/*
 			 * We can't process this chunk due to lack of space in LFC,
@@ -1578,7 +1588,7 @@ lfc_prefetch(NRelFileInfo rinfo, ForkNumber forknum, BlockNumber blkno,
 	pgstat_report_wait_start(WAIT_EVENT_NEON_LFC_WRITE);
 	INSTR_TIME_SET_CURRENT(io_start);
 	rc = pwrite(lfc_desc, buffer, BLCKSZ,
-				((off_t) entry_offset * lfc_blocks_per_chunk + chunk_offs) * BLCKSZ);
+				((off_t) entry_offset * BLOCKS_PER_CHUNK + chunk_offs) * BLCKSZ);
 	INSTR_TIME_SET_CURRENT(io_end);
 	pgstat_report_wait_end();
 
@@ -1640,7 +1650,7 @@ lfc_writev(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 	FileCacheEntry *entry;
 	ssize_t		rc;
 	bool		found;
-	uint32		hash;
+	uint64		hash;
 	uint64		generation;
 	uint32		entry_offset;
 	int			buf_offset = 0;
@@ -1653,6 +1663,7 @@ lfc_writev(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 
 	CriticalAssert(BufTagGetRelNumber(&tag) != InvalidRelFileNumber);
 
+ retry:
 	LWLockAcquire(lfc_lock, LW_EXCLUSIVE);
 
 	if (!LFC_ENABLED() || !lfc_ensure_opened())
@@ -1662,6 +1673,9 @@ lfc_writev(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 	}
 	generation = lfc_ctl->generation;
 
+	if (!freelist_prepare_pop())
+		goto retry;
+
 	/*
 	 * For every chunk that has blocks we're interested in, we
 	 * 1. get the chunk header
@@ -1675,7 +1689,7 @@ lfc_writev(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 	{
 		struct iovec iov[PG_IOV_MAX];
 		int		chunk_offs = BLOCK_TO_CHUNK_OFF(blkno);
-		int		blocks_in_chunk = Min(nblocks, lfc_blocks_per_chunk - chunk_offs);
+		int		blocks_in_chunk = Min(nblocks, BLOCKS_PER_CHUNK - chunk_offs);
 		instr_time io_start, io_end;
 		ConditionVariable* cv;
 
@@ -1688,16 +1702,16 @@ lfc_writev(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 		}
 
 		tag.blockNum = blkno - chunk_offs;
-		hash = get_hash_value(lfc_hash, &tag);
+		hash = file_cache_hash_get_hash_value(lfc_hash, &tag);
 		cv = &lfc_ctl->cv[hash % N_COND_VARS];
 
-		entry = hash_search_with_hash_value(lfc_hash, &tag, hash, HASH_ENTER, &found);
+		entry = file_cache_hash_enter(lfc_hash, &tag, hash, &found);
 
 		/* Approximate working set for the blocks assumed in this entry */
 		for (int i = 0; i < blocks_in_chunk; i++)
 		{
 			tag.blockNum = blkno + i;
-			addSHLL(&lfc_ctl->wss_estimation, hash_bytes((uint8_t const*)&tag, sizeof(tag)));
+			addSHLL(&lfc_ctl->wss_estimation, file_cache_hash_get_hash_value(lfc_hash, &tag));
 		}
 
 		if (found)
@@ -1714,7 +1728,7 @@ lfc_writev(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 		}
 		else
 		{
-			if (!lfc_init_new_entry(entry, hash))
+			if (!lfc_init_new_entry(entry))
 			{
 				/*
 				 * We can't process this chunk due to lack of space in LFC,
@@ -1763,7 +1777,7 @@ lfc_writev(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 		pgstat_report_wait_start(WAIT_EVENT_NEON_LFC_WRITE);
 		INSTR_TIME_SET_CURRENT(io_start);
 		rc = pwritev(lfc_desc, iov, blocks_in_chunk,
-					 ((off_t) entry_offset * lfc_blocks_per_chunk + chunk_offs) * BLCKSZ);
+					 ((off_t) entry_offset * BLOCKS_PER_CHUNK + chunk_offs) * BLCKSZ);
 		INSTR_TIME_SET_CURRENT(io_end);
 		pgstat_report_wait_end();
 
@@ -1823,6 +1837,140 @@ lfc_writev(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 	LWLockRelease(lfc_lock);
 }
 
+/**** freelist management ****/
+
+
+/*
+ * Prerequisites:
+ * - The caller is holding 'lfc_lock'. XXX
+ */
+static bool
+freelist_prepare_pop(void)
+{
+	/*
+	 * If the in-memory freelist is empty, but there are more blocks available, load them.
+	 *
+	 * TODO: if there
+	 */
+	if (lfc_ctl->num_free_pages == 0 && lfc_ctl->freelist_head != INVALID_OFFSET)
+	{
+		uint32		freelist_head;
+		FreeListChunk *freelist_chunk;
+		size_t		bytes_read;
+
+		LWLockRelease(lfc_lock);
+		LWLockAcquire(lfc_freelist_lock, LW_EXCLUSIVE);
+
+		if (!(lfc_ctl->num_free_pages == 0 && lfc_ctl->freelist_head != INVALID_OFFSET))
+		{
+			/* someone else did the work for us while we were not holding the lock */
+			LWLockRelease(lfc_freelist_lock);
+			return false;
+		}
+
+		freelist_head = lfc_ctl->freelist_head;
+		freelist_chunk = palloc(BLOCKS_PER_CHUNK * BLCKSZ);
+
+		bytes_read = 0;
+		while (bytes_read < BLOCKS_PER_CHUNK * BLCKSZ)
+		{
+			ssize_t		rc;
+
+			rc = pread(lfc_desc, freelist_chunk, BLOCKS_PER_CHUNK * BLCKSZ - bytes_read, (off_t) freelist_head * BLOCKS_PER_CHUNK * BLCKSZ + bytes_read);
+			if (rc < 0)
+			{
+				lfc_disable("read freelist page");
+				return false;
+			}
+			bytes_read += rc;
+		}
+
+		LWLockAcquire(lfc_lock, LW_EXCLUSIVE);
+		if (lfc_generation != lfc_ctl->generation)
+		{
+			LWLockRelease(lfc_lock);
+			return false;
+		}
+
+		Assert(lfc_ctl->freelist_head == freelist_head);
+		Assert(lfc_ctl->num_free_pages == 0);
+		lfc_ctl->freelist_head = freelist_chunk->next;
+		lfc_ctl->num_free_pages = freelist_chunk->num_free_pages;
+		memcpy(lfc_ctl->free_pages, freelist_chunk->free_pages, lfc_ctl->num_free_pages * sizeof(uint32));
+		pfree(freelist_chunk);
+
+		LWLockRelease(lfc_lock);
+		LWLockRelease(lfc_freelist_lock);
+		return false;
+	}
+
+	return true;
+}
+
+/*
+ * Prerequisites:
+ * - The caller is holding 'lfc_lock' and 'lfc_freelist_lock'.
+ *
+ * Returns 'false' on error.
+ */
+static bool
+freelist_push(uint32 offset)
+{
+	Assert(lfc_ctl->freelist_size == FREELIST_ENTRIES_PER_CHUNK);
+	if (lfc_ctl->num_free_pages == lfc_ctl->freelist_size)
+	{
+		FreeListChunk *freelist_chunk;
+		struct iovec iov;
+		ssize_t		rc;
+
+		freelist_chunk = palloc(BLOCKS_PER_CHUNK * BLCKSZ);
+
+		/* write the existing entries to the chunk on disk */
+		freelist_chunk->next = lfc_ctl->freelist_head;
+		freelist_chunk->num_free_pages = lfc_ctl->num_free_pages;
+		memcpy(freelist_chunk->free_pages, lfc_ctl->free_pages, lfc_ctl->num_free_pages * sizeof(uint32));
+
+		/* Use the passed-in offset to hold the freelist chunk itself */
+		iov.iov_base = freelist_chunk;
+		iov.iov_len = BLOCKS_PER_CHUNK * BLCKSZ;
+		rc = pg_pwritev_with_retry(lfc_desc, &iov, 1, (off_t) offset * BLOCKS_PER_CHUNK * BLCKSZ);
+
+		pfree(freelist_chunk);
+
+		if (rc < 0)
+			return false;
+
+		lfc_ctl->freelist_head = offset;
+		lfc_ctl->num_free_pages = 0;
+	}
+	else
+	{
+		lfc_ctl->free_pages[lfc_ctl->num_free_pages] = offset;
+		lfc_ctl->num_free_pages++;
+	}
+	return true;
+}
+
+static uint32
+freelist_pop(void)
+{
+	uint32		result;
+
+	/* The caller should've checked that the list is not empty */
+	Assert(lfc_ctl->num_free_pages > 0);
+
+	result = lfc_ctl->free_pages[lfc_ctl->num_free_pages - 1];
+	lfc_ctl->num_free_pages--;
+
+	return result;
+}
+
+static bool
+freelist_is_empty(void)
+{
+	return lfc_ctl->num_free_pages == 0;
+}
+
 typedef struct
 {
 	TupleDesc	tupdesc;
@@ -1919,7 +2067,7 @@ neon_get_lfc_stats(PG_FUNCTION_ARGS)
 			break;
 		case 8:
 			key = "file_cache_chunk_size_pages";
-			value = lfc_blocks_per_chunk;
+			value = BLOCKS_PER_CHUNK;
 			break;
 		case 9:
 			key = "file_cache_chunks_pinned";
@@ -1990,7 +2138,6 @@ local_cache_pages(PG_FUNCTION_ARGS)
 
 	if (SRF_IS_FIRSTCALL())
 	{
-		HASH_SEQ_STATUS status;
 		FileCacheEntry *entry;
 		uint32		n_pages = 0;
 
@@ -2046,15 +2193,16 @@ local_cache_pages(PG_FUNCTION_ARGS)
 
 			if (LFC_ENABLED())
 			{
-				hash_seq_init(&status, lfc_hash);
-				while ((entry = hash_seq_search(&status)) != NULL)
+				uint32		num_buckets = file_cache_hash_get_num_buckets(lfc_hash);
+
+				for (uint32 pos = 0; pos < num_buckets; pos++)
 				{
-					/* Skip hole tags */
-					if (NInfoGetRelNumber(BufTagGetNRelFileInfo(entry->key)) != 0)
-					{
-						for (int i = 0; i < lfc_blocks_per_chunk; i++)
-							n_pages += GET_STATE(entry, i) == AVAILABLE;
-					}
+					entry = file_cache_hash_get_at_pos(lfc_hash, pos);
+					if (entry == NULL)
+						continue;
+
+					for (int i = 0; i < BLOCKS_PER_CHUNK; i++)
+						n_pages += GET_STATE(entry, i) == AVAILABLE;
 				}
 			}
 		}
@@ -2076,25 +2224,28 @@ local_cache_pages(PG_FUNCTION_ARGS)
 			 * in the fctx->record structure.
 			 */
 			uint32		n = 0;
+			uint32		num_buckets = file_cache_hash_get_num_buckets(lfc_hash);
 
-			hash_seq_init(&status, lfc_hash);
-			while ((entry = hash_seq_search(&status)) != NULL)
+			for (uint32 pos = 0; pos < num_buckets; pos++)
 			{
-				for (int i = 0; i < lfc_blocks_per_chunk; i++)
+				entry = file_cache_hash_get_at_pos(lfc_hash, pos);
+				if (entry == NULL)
+					continue;
+
+				for (int i = 0; i < BLOCKS_PER_CHUNK; i++)
 				{
-					if (NInfoGetRelNumber(BufTagGetNRelFileInfo(entry->key)) != 0)
+					const BufferTag *key = file_cache_hash_get_key_for_entry(lfc_hash, entry);
+
+					if (GET_STATE(entry, i) == AVAILABLE)
 					{
-						if (GET_STATE(entry, i) == AVAILABLE)
-						{
-							fctx->record[n].pageoffs = entry->offset * lfc_blocks_per_chunk + i;
-							fctx->record[n].relfilenode = NInfoGetRelNumber(BufTagGetNRelFileInfo(entry->key));
-							fctx->record[n].reltablespace = NInfoGetSpcOid(BufTagGetNRelFileInfo(entry->key));
-							fctx->record[n].reldatabase = NInfoGetDbOid(BufTagGetNRelFileInfo(entry->key));
-							fctx->record[n].forknum = entry->key.forkNum;
-							fctx->record[n].blocknum = entry->key.blockNum + i;
-							fctx->record[n].accesscount = entry->access_count;
-							n += 1;
-						}
+						fctx->record[n].pageoffs = entry->offset * BLOCKS_PER_CHUNK + i;
+						fctx->record[n].relfilenode = NInfoGetRelNumber(BufTagGetNRelFileInfo(*key));
+						fctx->record[n].reltablespace = NInfoGetSpcOid(BufTagGetNRelFileInfo(*key));
+						fctx->record[n].reldatabase = NInfoGetDbOid(BufTagGetNRelFileInfo(*key));
+						fctx->record[n].forknum = key->forkNum;
+						fctx->record[n].blocknum = key->blockNum + i;
+						fctx->record[n].accesscount = entry->access_count;
+						n += 1;
 					}
 				}
 			}

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

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

pgxn/neon/walproposer.c

@@ -1135,7 +1135,7 @@ VotesCollectedMset(WalProposer *wp, MemberSet *mset, Safekeeper **msk, StringInf
 				wp->propTermStartLsn = sk->voteResponse.flushLsn;
 				wp->donor = sk;
 			}
-			wp->truncateLsn = Max(sk->voteResponse.truncateLsn, wp->truncateLsn);
+			wp->truncateLsn = Max(wp->safekeeper[i].voteResponse.truncateLsn, wp->truncateLsn);
 
 			if (n_votes > 0)
 				appendStringInfoString(s, ", ");

poetry.lock

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

proxy/src/auth/backend/classic.rs

@@ -18,6 +18,11 @@ pub(super) async fn authenticate(
     secret: AuthSecret,
 ) -> auth::Result<ComputeCredentials> {
     let scram_keys = match secret {
+        #[cfg(any(test, feature = "testing"))]
+        AuthSecret::Md5(_) => {
+            debug!("auth endpoint chooses MD5");
+            return Err(auth::AuthError::MalformedPassword("MD5 not supported"));
+        }
         AuthSecret::Scram(secret) => {
             debug!("auth endpoint chooses SCRAM");
 

proxy/src/auth/backend/console_redirect.rs

@@ -6,17 +6,18 @@ use thiserror::Error;
 use tokio::io::{AsyncRead, AsyncWrite};
 use tracing::{info, info_span};
 
+use super::ComputeCredentialKeys;
+use crate::auth::IpPattern;
 use crate::auth::backend::ComputeUserInfo;
 use crate::cache::Cached;
-use crate::compute::AuthInfo;
 use crate::config::AuthenticationConfig;
 use crate::context::RequestContext;
 use crate::control_plane::client::cplane_proxy_v1;
 use crate::control_plane::{self, CachedNodeInfo, NodeInfo};
 use crate::error::{ReportableError, UserFacingError};
+use crate::pglb::connect_compute::ComputeConnectBackend;
 use crate::pqproto::BeMessage;
 use crate::proxy::NeonOptions;
-use crate::proxy::wake_compute::WakeComputeBackend;
 use crate::stream::PqStream;
 use crate::types::RoleName;
 use crate::{auth, compute, waiters};
@@ -97,11 +98,15 @@ impl ConsoleRedirectBackend {
         ctx: &RequestContext,
         auth_config: &'static AuthenticationConfig,
         client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
-    ) -> auth::Result<(ConsoleRedirectNodeInfo, AuthInfo, ComputeUserInfo)> {
+    ) -> auth::Result<(
+        ConsoleRedirectNodeInfo,
+        ComputeUserInfo,
+        Option<Vec<IpPattern>>,
+    )> {
         authenticate(ctx, auth_config, &self.console_uri, client)
             .await
-            .map(|(node_info, auth_info, user_info)| {
-                (ConsoleRedirectNodeInfo(node_info), auth_info, user_info)
+            .map(|(node_info, user_info, ip_allowlist)| {
+                (ConsoleRedirectNodeInfo(node_info), user_info, ip_allowlist)
             })
     }
 }
@@ -109,13 +114,17 @@ impl ConsoleRedirectBackend {
 pub struct ConsoleRedirectNodeInfo(pub(super) NodeInfo);
 
 #[async_trait]
-impl WakeComputeBackend for ConsoleRedirectNodeInfo {
+impl ComputeConnectBackend for ConsoleRedirectNodeInfo {
     async fn wake_compute(
         &self,
         _ctx: &RequestContext,
     ) -> Result<CachedNodeInfo, control_plane::errors::WakeComputeError> {
         Ok(Cached::new_uncached(self.0.clone()))
     }
+
+    fn get_keys(&self) -> &ComputeCredentialKeys {
+        &ComputeCredentialKeys::None
+    }
 }
 
 async fn authenticate(
@@ -123,7 +132,7 @@ async fn authenticate(
     auth_config: &'static AuthenticationConfig,
     link_uri: &reqwest::Url,
     client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
-) -> auth::Result<(NodeInfo, AuthInfo, ComputeUserInfo)> {
+) -> auth::Result<(NodeInfo, ComputeUserInfo, Option<Vec<IpPattern>>)> {
     ctx.set_auth_method(crate::context::AuthMethod::ConsoleRedirect);
 
     // registering waiter can fail if we get unlucky with rng.
@@ -183,24 +192,10 @@ async fn authenticate(
 
     client.write_message(BeMessage::NoticeResponse("Connecting to database."));
 
-    // Backwards compatibility. pg_sni_proxy uses "--" in domain names
-    // while direct connections do not. Once we migrate to pg_sni_proxy
-    // everywhere, we can remove this.
-    let ssl_mode = if db_info.host.contains("--") {
-        // we need TLS connection with SNI info to properly route it
-        SslMode::Require
-    } else {
-        SslMode::Disable
-    };
-
-    let conn_info = compute::ConnectInfo {
-        host: db_info.host.into(),
-        port: db_info.port,
-        ssl_mode,
-        host_addr: None,
-    };
-    let auth_info =
-        AuthInfo::for_console_redirect(&db_info.dbname, &db_info.user, db_info.password.as_deref());
+    // This config should be self-contained, because we won't
+    // take username or dbname from client's startup message.
+    let mut config = compute::ConnCfg::new(db_info.host.to_string(), db_info.port);
+    config.dbname(&db_info.dbname).user(&db_info.user);
 
     let user: RoleName = db_info.user.into();
     let user_info = ComputeUserInfo {
@@ -214,12 +209,26 @@ async fn authenticate(
     ctx.set_project(db_info.aux.clone());
     info!("woken up a compute node");
 
+    // Backwards compatibility. pg_sni_proxy uses "--" in domain names
+    // while direct connections do not. Once we migrate to pg_sni_proxy
+    // everywhere, we can remove this.
+    if db_info.host.contains("--") {
+        // we need TLS connection with SNI info to properly route it
+        config.ssl_mode(SslMode::Require);
+    } else {
+        config.ssl_mode(SslMode::Disable);
+    }
+
+    if let Some(password) = db_info.password {
+        config.password(password.as_ref());
+    }
+
     Ok((
         NodeInfo {
-            conn_info,
+            config,
             aux: db_info.aux,
         },
-        auth_info,
         user_info,
+        db_info.allowed_ips,
     ))
 }

proxy/src/auth/backend/jwt.rs

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

proxy/src/auth/backend/local.rs

@@ -1,12 +1,11 @@
 use std::net::SocketAddr;
 
 use arc_swap::ArcSwapOption;
-use postgres_client::config::SslMode;
 use tokio::sync::Semaphore;
 
 use super::jwt::{AuthRule, FetchAuthRules};
 use crate::auth::backend::jwt::FetchAuthRulesError;
-use crate::compute::ConnectInfo;
+use crate::compute::ConnCfg;
 use crate::compute_ctl::ComputeCtlApi;
 use crate::context::RequestContext;
 use crate::control_plane::NodeInfo;
@@ -30,12 +29,7 @@ impl LocalBackend {
                 api: http::Endpoint::new(compute_ctl, http::new_client()),
             },
             node_info: NodeInfo {
-                conn_info: ConnectInfo {
-                    host_addr: Some(postgres_addr.ip()),
-                    host: postgres_addr.ip().to_string().into(),
-                    port: postgres_addr.port(),
-                    ssl_mode: SslMode::Disable,
-                },
+                config: ConnCfg::new(postgres_addr.ip().to_string(), postgres_addr.port()),
                 // TODO(conrad): make this better reflect compute info rather than endpoint info.
                 aux: MetricsAuxInfo {
                     endpoint_id: EndpointIdTag::get_interner().get_or_intern("local"),

proxy/src/auth/backend/mod.rs

@@ -14,21 +14,20 @@ use serde::{Deserialize, Serialize};
 use tokio::io::{AsyncRead, AsyncWrite};
 use tracing::{debug, info};
 
-use crate::auth::{self, ComputeUserInfoMaybeEndpoint, validate_password_and_exchange};
+use crate::auth::{self, AuthError, ComputeUserInfoMaybeEndpoint, validate_password_and_exchange};
 use crate::cache::Cached;
 use crate::config::AuthenticationConfig;
 use crate::context::RequestContext;
 use crate::control_plane::client::ControlPlaneClient;
 use crate::control_plane::errors::GetAuthInfoError;
-use crate::control_plane::messages::EndpointRateLimitConfig;
 use crate::control_plane::{
     self, AccessBlockerFlags, AuthSecret, CachedNodeInfo, ControlPlaneApi, EndpointAccessControl,
     RoleAccessControl,
 };
 use crate::intern::EndpointIdInt;
+use crate::pglb::connect_compute::ComputeConnectBackend;
 use crate::pqproto::BeMessage;
 use crate::proxy::NeonOptions;
-use crate::proxy::wake_compute::WakeComputeBackend;
 use crate::rate_limiter::EndpointRateLimiter;
 use crate::stream::Stream;
 use crate::types::{EndpointCacheKey, EndpointId, RoleName};
@@ -169,6 +168,8 @@ impl ComputeUserInfo {
 
 #[cfg_attr(test, derive(Debug))]
 pub(crate) enum ComputeCredentialKeys {
+    #[cfg(any(test, feature = "testing"))]
+    Password(Vec<u8>),
     AuthKeys(AuthKeys),
     JwtPayload(Vec<u8>),
     None,
@@ -231,8 +232,11 @@ async fn auth_quirks(
         config.is_vpc_acccess_proxy,
     )?;
 
-    access_controls.connection_attempt_rate_limit(ctx, &info.endpoint, &endpoint_rate_limiter)?;
-
+    let endpoint = EndpointIdInt::from(&info.endpoint);
+    let rate_limit_config = None;
+    if !endpoint_rate_limiter.check(endpoint, rate_limit_config, 1) {
+        return Err(AuthError::too_many_connections());
+    }
     let role_access = api
         .get_role_access_control(ctx, &info.endpoint, &info.user)
         .await?;
@@ -399,23 +403,29 @@ impl Backend<'_, ComputeUserInfo> {
                 allowed_ips: Arc::new(vec![]),
                 allowed_vpce: Arc::new(vec![]),
                 flags: AccessBlockerFlags::default(),
-                rate_limits: EndpointRateLimitConfig::default(),
             }),
         }
     }
 }
 
 #[async_trait::async_trait]
-impl WakeComputeBackend for Backend<'_, ComputeUserInfo> {
+impl ComputeConnectBackend for Backend<'_, ComputeCredentials> {
     async fn wake_compute(
         &self,
         ctx: &RequestContext,
     ) -> Result<CachedNodeInfo, control_plane::errors::WakeComputeError> {
         match self {
-            Self::ControlPlane(api, info) => api.wake_compute(ctx, info).await,
+            Self::ControlPlane(api, creds) => api.wake_compute(ctx, &creds.info).await,
             Self::Local(local) => Ok(Cached::new_uncached(local.node_info.clone())),
         }
     }
+
+    fn get_keys(&self) -> &ComputeCredentialKeys {
+        match self {
+            Self::ControlPlane(_, creds) => &creds.keys,
+            Self::Local(_) => &ComputeCredentialKeys::None,
+        }
+    }
 }
 
 #[cfg(test)]
@@ -438,7 +448,6 @@ mod tests {
     use crate::auth::{ComputeUserInfoMaybeEndpoint, IpPattern};
     use crate::config::AuthenticationConfig;
     use crate::context::RequestContext;
-    use crate::control_plane::messages::EndpointRateLimitConfig;
     use crate::control_plane::{
         self, AccessBlockerFlags, CachedNodeInfo, EndpointAccessControl, RoleAccessControl,
     };
@@ -477,7 +486,6 @@ mod tests {
                 allowed_ips: Arc::new(self.ips.clone()),
                 allowed_vpce: Arc::new(self.vpc_endpoint_ids.clone()),
                 flags: self.access_blocker_flags,
-                rate_limits: EndpointRateLimitConfig::default(),
             })
         }
 

proxy/src/auth/flow.rs

@@ -169,6 +169,13 @@ pub(crate) async fn validate_password_and_exchange(
     secret: AuthSecret,
 ) -> super::Result<sasl::Outcome<ComputeCredentialKeys>> {
     match secret {
+        #[cfg(any(test, feature = "testing"))]
+        AuthSecret::Md5(_) => {
+            // test only
+            Ok(sasl::Outcome::Success(ComputeCredentialKeys::Password(
+                password.to_owned(),
+            )))
+        }
         // perform scram authentication as both client and server to validate the keys
         AuthSecret::Scram(scram_secret) => {
             let outcome = crate::scram::exchange(pool, endpoint, &scram_secret, password).await?;

proxy/src/binary/pg_sni_router.rs

@@ -28,9 +28,10 @@ use crate::context::RequestContext;
 use crate::metrics::{Metrics, ThreadPoolMetrics};
 use crate::pqproto::FeStartupPacket;
 use crate::protocol2::ConnectionInfo;
-use crate::proxy::{ErrorSource, TlsRequired, copy_bidirectional_client_compute};
+use crate::proxy::{
+    ErrorSource, TlsRequired, copy_bidirectional_client_compute, run_until_cancelled,
+};
 use crate::stream::{PqStream, Stream};
-use crate::util::run_until_cancelled;
 
 project_git_version!(GIT_VERSION);
 

proxy/src/binary/proxy.rs

@@ -11,13 +11,11 @@ use anyhow::Context;
 use anyhow::{bail, ensure};
 use arc_swap::ArcSwapOption;
 use futures::future::Either;
-use itertools::{Itertools, Position};
-use rand::{Rng, thread_rng};
 use remote_storage::RemoteStorageConfig;
 use tokio::net::TcpListener;
 use tokio::task::JoinSet;
 use tokio_util::sync::CancellationToken;
-use tracing::{Instrument, error, info, warn};
+use tracing::{Instrument, info, warn};
 use utils::sentry_init::init_sentry;
 use utils::{project_build_tag, project_git_version};
 
@@ -316,7 +314,7 @@ pub async fn run() -> anyhow::Result<()> {
     let jemalloc = match crate::jemalloc::MetricRecorder::new() {
         Ok(t) => Some(t),
         Err(e) => {
-            error!(error = ?e, "could not start jemalloc metrics loop");
+            tracing::error!(error = ?e, "could not start jemalloc metrics loop");
             None
         }
     };
@@ -522,44 +520,23 @@ pub async fn run() -> anyhow::Result<()> {
                 }
             }
 
-            // Try to connect to Redis 3 times with 1 + (0..0.1) second interval.
-            // This prevents immediate exit and pod restart,
-            // which can cause hammering of the redis in case of connection issues.
             if let Some(mut redis_kv_client) = redis_kv_client {
-                for attempt in (0..3).with_position() {
-                    match redis_kv_client.try_connect().await {
-                        Ok(()) => {
-                            info!("Connected to Redis KV client");
-                            maintenance_tasks.spawn(async move {
-                                handle_cancel_messages(
-                                    &mut redis_kv_client,
-                                    rx_cancel,
-                                    args.cancellation_batch_size,
-                                )
-                                .await?;
+                maintenance_tasks.spawn(async move {
+                    redis_kv_client.try_connect().await?;
+                    handle_cancel_messages(
+                        &mut redis_kv_client,
+                        rx_cancel,
+                        args.cancellation_batch_size,
+                    )
+                    .await?;
 
-                                drop(redis_kv_client);
+                    drop(redis_kv_client);
 
-                                // `handle_cancel_messages` was terminated due to the tx_cancel
-                                // being dropped. this is not worthy of an error, and this task can only return `Err`,
-                                // so let's wait forever instead.
-                                std::future::pending().await
-                            });
-                            break;
-                        }
-                        Err(e) => {
-                            error!("Failed to connect to Redis KV client: {e}");
-                            if matches!(attempt, Position::Last(_)) {
-                                bail!(
-                                    "Failed to connect to Redis KV client after {} attempts",
-                                    attempt.into_inner()
-                                );
-                            }
-                            let jitter = thread_rng().gen_range(0..100);
-                            tokio::time::sleep(Duration::from_millis(1000 + jitter)).await;
-                        }
-                    }
-                }
+                    // `handle_cancel_messages` was terminated due to the tx_cancel
+                    // being dropped. this is not worthy of an error, and this task can only return `Err`,
+                    // so let's wait forever instead.
+                    std::future::pending().await
+                });
             }
 
             if let Some(regional_redis_client) = regional_redis_client {

proxy/src/cache/project_info.rs

@@ -18,7 +18,6 @@ use crate::types::{EndpointId, RoleName};
 
 #[async_trait]
 pub(crate) trait ProjectInfoCache {
-    fn invalidate_endpoint_access(&self, endpoint_id: EndpointIdInt);
     fn invalidate_endpoint_access_for_project(&self, project_id: ProjectIdInt);
     fn invalidate_endpoint_access_for_org(&self, account_id: AccountIdInt);
     fn invalidate_role_secret_for_project(&self, project_id: ProjectIdInt, role_name: RoleNameInt);
@@ -101,13 +100,6 @@ pub struct ProjectInfoCacheImpl {
 
 #[async_trait]
 impl ProjectInfoCache for ProjectInfoCacheImpl {
-    fn invalidate_endpoint_access(&self, endpoint_id: EndpointIdInt) {
-        info!("invalidating endpoint access for `{endpoint_id}`");
-        if let Some(mut endpoint_info) = self.cache.get_mut(&endpoint_id) {
-            endpoint_info.invalidate_endpoint();
-        }
-    }
-
     fn invalidate_endpoint_access_for_project(&self, project_id: ProjectIdInt) {
         info!("invalidating endpoint access for project `{project_id}`");
         let endpoints = self
@@ -364,7 +356,6 @@ mod tests {
     use std::sync::Arc;
 
     use super::*;
-    use crate::control_plane::messages::EndpointRateLimitConfig;
     use crate::control_plane::{AccessBlockerFlags, AuthSecret};
     use crate::scram::ServerSecret;
     use crate::types::ProjectId;
@@ -400,7 +391,6 @@ mod tests {
                 allowed_ips: allowed_ips.clone(),
                 allowed_vpce: Arc::new(vec![]),
                 flags: AccessBlockerFlags::default(),
-                rate_limits: EndpointRateLimitConfig::default(),
             },
             RoleAccessControl {
                 secret: secret1.clone(),
@@ -416,7 +406,6 @@ mod tests {
                 allowed_ips: allowed_ips.clone(),
                 allowed_vpce: Arc::new(vec![]),
                 flags: AccessBlockerFlags::default(),
-                rate_limits: EndpointRateLimitConfig::default(),
             },
             RoleAccessControl {
                 secret: secret2.clone(),
@@ -442,7 +431,6 @@ mod tests {
                 allowed_ips: allowed_ips.clone(),
                 allowed_vpce: Arc::new(vec![]),
                 flags: AccessBlockerFlags::default(),
-                rate_limits: EndpointRateLimitConfig::default(),
             },
             RoleAccessControl {
                 secret: secret3.clone(),

proxy/src/cancellation.rs

@@ -24,6 +24,7 @@ use crate::pqproto::CancelKeyData;
 use crate::rate_limiter::LeakyBucketRateLimiter;
 use crate::redis::keys::KeyPrefix;
 use crate::redis::kv_ops::RedisKVClient;
+use crate::tls::postgres_rustls::MakeRustlsConnect;
 
 type IpSubnetKey = IpNet;
 
@@ -496,8 +497,10 @@ impl CancelClosure {
     ) -> Result<(), CancelError> {
         let socket = TcpStream::connect(self.socket_addr).await?;
 
-        let tls = <_ as MakeTlsConnect<tokio::net::TcpStream>>::make_tls_connect(
-            compute_config,
+        let mut mk_tls =
+            crate::tls::postgres_rustls::MakeRustlsConnect::new(compute_config.tls.clone());
+        let tls = <MakeRustlsConnect as MakeTlsConnect<tokio::net::TcpStream>>::make_tls_connect(
+            &mut mk_tls,
             &self.hostname,
         )
         .map_err(|e| CancelError::IO(std::io::Error::other(e.to_string())))?;

proxy/src/compute.rs

@@ -1,24 +1,21 @@
-mod tls;
-
 use std::fmt::Debug;
 use std::io;
-use std::net::{IpAddr, SocketAddr};
+use std::net::SocketAddr;
+use std::time::Duration;
 
 use futures::{FutureExt, TryFutureExt};
 use itertools::Itertools;
-use postgres_client::config::{AuthKeys, SslMode};
-use postgres_client::maybe_tls_stream::MaybeTlsStream;
 use postgres_client::tls::MakeTlsConnect;
-use postgres_client::{CancelToken, NoTls, RawConnection};
+use postgres_client::{CancelToken, RawConnection};
 use postgres_protocol::message::backend::NoticeResponseBody;
+use rustls::pki_types::InvalidDnsNameError;
 use thiserror::Error;
 use tokio::net::{TcpStream, lookup_host};
 use tracing::{debug, error, info, warn};
 
-use crate::auth::backend::{ComputeCredentialKeys, ComputeUserInfo};
+use crate::auth::backend::ComputeUserInfo;
 use crate::auth::parse_endpoint_param;
 use crate::cancellation::CancelClosure;
-use crate::compute::tls::TlsError;
 use crate::config::ComputeConfig;
 use crate::context::RequestContext;
 use crate::control_plane::client::ApiLockError;
@@ -28,6 +25,7 @@ use crate::error::{ReportableError, UserFacingError};
 use crate::metrics::{Metrics, NumDbConnectionsGuard};
 use crate::pqproto::StartupMessageParams;
 use crate::proxy::neon_option;
+use crate::tls::postgres_rustls::MakeRustlsConnect;
 use crate::types::Host;
 
 pub const COULD_NOT_CONNECT: &str = "Couldn't connect to compute node";
@@ -40,7 +38,10 @@ pub(crate) enum ConnectionError {
     Postgres(#[from] postgres_client::Error),
 
     #[error("{COULD_NOT_CONNECT}: {0}")]
-    TlsError(#[from] TlsError),
+    CouldNotConnect(#[from] io::Error),
+
+    #[error("{COULD_NOT_CONNECT}: {0}")]
+    TlsError(#[from] InvalidDnsNameError),
 
     #[error("{COULD_NOT_CONNECT}: {0}")]
     WakeComputeError(#[from] WakeComputeError),
@@ -72,7 +73,7 @@ impl UserFacingError for ConnectionError {
             ConnectionError::TooManyConnectionAttempts(_) => {
                 "Failed to acquire permit to connect to the database. Too many database connection attempts are currently ongoing.".to_owned()
             }
-            ConnectionError::TlsError(_) => COULD_NOT_CONNECT.to_owned(),
+            _ => COULD_NOT_CONNECT.to_owned(),
         }
     }
 }
@@ -84,6 +85,7 @@ impl ReportableError for ConnectionError {
                 crate::error::ErrorKind::Postgres
             }
             ConnectionError::Postgres(_) => crate::error::ErrorKind::Compute,
+            ConnectionError::CouldNotConnect(_) => crate::error::ErrorKind::Compute,
             ConnectionError::TlsError(_) => crate::error::ErrorKind::Compute,
             ConnectionError::WakeComputeError(e) => e.get_error_kind(),
             ConnectionError::TooManyConnectionAttempts(e) => e.get_error_kind(),
@@ -94,85 +96,34 @@ impl ReportableError for ConnectionError {
 /// A pair of `ClientKey` & `ServerKey` for `SCRAM-SHA-256`.
 pub(crate) type ScramKeys = postgres_client::config::ScramKeys<32>;
 
+/// A config for establishing a connection to compute node.
+/// Eventually, `postgres_client` will be replaced with something better.
+/// Newtype allows us to implement methods on top of it.
 #[derive(Clone)]
-pub enum Auth {
-    /// Only used during console-redirect.
-    Password(Vec<u8>),
-    /// Used by sql-over-http, ws, tcp.
-    Scram(Box<ScramKeys>),
-}
-
-/// A config for authenticating to the compute node.
-pub(crate) struct AuthInfo {
-    /// None for local-proxy, as we use trust-based localhost auth.
-    /// Some for sql-over-http, ws, tcp, and in most cases for console-redirect.
-    /// Might be None for console-redirect, but that's only a consequence of testing environments ATM.
-    auth: Option<Auth>,
-    server_params: StartupMessageParams,
-
-    /// Console redirect sets user and database, we shouldn't re-use those from the params.
-    skip_db_user: bool,
-}
-
-/// Contains only the data needed to establish a secure connection to compute.
-#[derive(Clone)]
-pub struct ConnectInfo {
-    pub host_addr: Option<IpAddr>,
-    pub host: Host,
-    pub port: u16,
-    pub ssl_mode: SslMode,
-}
+pub(crate) struct ConnCfg(Box<postgres_client::Config>);
 
 /// Creation and initialization routines.
-impl AuthInfo {
-    pub(crate) fn for_console_redirect(db: &str, user: &str, pw: Option<&str>) -> Self {
-        let mut server_params = StartupMessageParams::default();
-        server_params.insert("database", db);
-        server_params.insert("user", user);
-        Self {
-            auth: pw.map(|pw| Auth::Password(pw.as_bytes().to_owned())),
-            server_params,
-            skip_db_user: true,
+impl ConnCfg {
+    pub(crate) fn new(host: String, port: u16) -> Self {
+        Self(Box::new(postgres_client::Config::new(host, port)))
+    }
+
+    /// Reuse password or auth keys from the other config.
+    pub(crate) fn reuse_password(&mut self, other: Self) {
+        if let Some(password) = other.get_password() {
+            self.password(password);
+        }
+
+        if let Some(keys) = other.get_auth_keys() {
+            self.auth_keys(keys);
         }
     }
 
-    pub(crate) fn with_auth_keys(keys: ComputeCredentialKeys) -> Self {
-        Self {
-            auth: match keys {
-                ComputeCredentialKeys::AuthKeys(AuthKeys::ScramSha256(auth_keys)) => {
-                    Some(Auth::Scram(Box::new(auth_keys)))
-                }
-                ComputeCredentialKeys::JwtPayload(_) | ComputeCredentialKeys::None => None,
-            },
-            server_params: StartupMessageParams::default(),
-            skip_db_user: false,
+    pub(crate) fn get_host(&self) -> Host {
+        match self.0.get_host() {
+            postgres_client::config::Host::Tcp(s) => s.into(),
         }
     }
-}
-
-impl ConnectInfo {
-    pub fn to_postgres_client_config(&self) -> postgres_client::Config {
-        let mut config = postgres_client::Config::new(self.host.to_string(), self.port);
-        config.ssl_mode(self.ssl_mode);
-        if let Some(host_addr) = self.host_addr {
-            config.set_host_addr(host_addr);
-        }
-        config
-    }
-}
-
-impl AuthInfo {
-    fn enrich(&self, mut config: postgres_client::Config) -> postgres_client::Config {
-        match &self.auth {
-            Some(Auth::Scram(keys)) => config.auth_keys(AuthKeys::ScramSha256(**keys)),
-            Some(Auth::Password(pw)) => config.password(pw),
-            None => &mut config,
-        };
-        for (k, v) in self.server_params.iter() {
-            config.set_param(k, v);
-        }
-        config
-    }
 
     /// Apply startup message params to the connection config.
     pub(crate) fn set_startup_params(
@@ -181,26 +132,27 @@ impl AuthInfo {
         arbitrary_params: bool,
     ) {
         if !arbitrary_params {
-            self.server_params.insert("client_encoding", "UTF8");
+            self.set_param("client_encoding", "UTF8");
         }
         for (k, v) in params.iter() {
             match k {
                 // Only set `user` if it's not present in the config.
                 // Console redirect auth flow takes username from the console's response.
-                "user" | "database" if self.skip_db_user => {}
+                "user" if self.user_is_set() => {}
+                "database" if self.db_is_set() => {}
                 "options" => {
                     if let Some(options) = filtered_options(v) {
-                        self.server_params.insert(k, &options);
+                        self.set_param(k, &options);
                     }
                 }
                 "user" | "database" | "application_name" | "replication" => {
-                    self.server_params.insert(k, v);
+                    self.set_param(k, v);
                 }
 
                 // if we allow arbitrary params, then we forward them through.
                 // this is a flag for a period of backwards compatibility
                 k if arbitrary_params => {
-                    self.server_params.insert(k, v);
+                    self.set_param(k, v);
                 }
                 _ => {}
             }
@@ -208,13 +160,25 @@ impl AuthInfo {
     }
 }
 
-impl ConnectInfo {
-    /// Establish a raw TCP+TLS connection to the compute node.
-    async fn connect_raw(
-        &self,
-        config: &ComputeConfig,
-    ) -> Result<(SocketAddr, MaybeTlsStream<TcpStream, RustlsStream>), TlsError> {
-        let timeout = config.timeout;
+impl std::ops::Deref for ConnCfg {
+    type Target = postgres_client::Config;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+/// For now, let's make it easier to setup the config.
+impl std::ops::DerefMut for ConnCfg {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+impl ConnCfg {
+    /// Establish a raw TCP connection to the compute node.
+    async fn connect_raw(&self, timeout: Duration) -> io::Result<(SocketAddr, TcpStream, &str)> {
+        use postgres_client::config::Host;
 
         // wrap TcpStream::connect with timeout
         let connect_with_timeout = |addrs| {
@@ -244,32 +208,34 @@ impl ConnectInfo {
         // We can't reuse connection establishing logic from `postgres_client` here,
         // because it has no means for extracting the underlying socket which we
         // require for our business.
-        let port = self.port;
-        let host = &*self.host;
+        let port = self.0.get_port();
+        let host = self.0.get_host();
 
-        let addrs = match self.host_addr {
+        let host = match host {
+            Host::Tcp(host) => host.as_str(),
+        };
+
+        let addrs = match self.0.get_host_addr() {
             Some(addr) => vec![SocketAddr::new(addr, port)],
             None => lookup_host((host, port)).await?.collect(),
         };
 
         match connect_once(&*addrs).await {
-            Ok((sockaddr, stream)) => Ok((
-                sockaddr,
-                tls::connect_tls(stream, self.ssl_mode, config, host).await?,
-            )),
+            Ok((sockaddr, stream)) => Ok((sockaddr, stream, host)),
             Err(err) => {
                 warn!("couldn't connect to compute node at {host}:{port}: {err}");
-                Err(TlsError::Connection(err))
+                Err(err)
             }
         }
     }
 }
 
-type RustlsStream = <ComputeConfig as MakeTlsConnect<tokio::net::TcpStream>>::Stream;
+type RustlsStream = <MakeRustlsConnect as MakeTlsConnect<tokio::net::TcpStream>>::Stream;
 
 pub(crate) struct PostgresConnection {
     /// Socket connected to a compute node.
-    pub(crate) stream: MaybeTlsStream<tokio::net::TcpStream, RustlsStream>,
+    pub(crate) stream:
+        postgres_client::maybe_tls_stream::MaybeTlsStream<tokio::net::TcpStream, RustlsStream>,
     /// PostgreSQL connection parameters.
     pub(crate) params: std::collections::HashMap<String, String>,
     /// Query cancellation token.
@@ -282,23 +248,28 @@ pub(crate) struct PostgresConnection {
     _guage: NumDbConnectionsGuard<'static>,
 }
 
-impl ConnectInfo {
+impl ConnCfg {
     /// Connect to a corresponding compute node.
     pub(crate) async fn connect(
         &self,
         ctx: &RequestContext,
         aux: MetricsAuxInfo,
-        auth: &AuthInfo,
         config: &ComputeConfig,
         user_info: ComputeUserInfo,
     ) -> Result<PostgresConnection, ConnectionError> {
-        let mut tmp_config = auth.enrich(self.to_postgres_client_config());
-        // we setup SSL early in `ConnectInfo::connect_raw`.
-        tmp_config.ssl_mode(SslMode::Disable);
-
         let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
-        let (socket_addr, stream) = self.connect_raw(config).await?;
-        let connection = tmp_config.connect_raw(stream, NoTls).await?;
+        let (socket_addr, stream, host) = self.connect_raw(config.timeout).await?;
+        drop(pause);
+
+        let mut mk_tls = crate::tls::postgres_rustls::MakeRustlsConnect::new(config.tls.clone());
+        let tls = <MakeRustlsConnect as MakeTlsConnect<tokio::net::TcpStream>>::make_tls_connect(
+            &mut mk_tls,
+            host,
+        )?;
+
+        // connect_raw() will not use TLS if sslmode is "disable"
+        let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
+        let connection = self.0.connect_raw(stream, tls).await?;
         drop(pause);
 
         let RawConnection {
@@ -311,14 +282,13 @@ impl ConnectInfo {
 
         tracing::Span::current().record("pid", tracing::field::display(process_id));
         tracing::Span::current().record("compute_id", tracing::field::display(&aux.compute_id));
-        let MaybeTlsStream::Raw(stream) = stream.into_inner();
+        let stream = stream.into_inner();
 
         // TODO: lots of useful info but maybe we can move it elsewhere (eg traces?)
         info!(
             cold_start_info = ctx.cold_start_info().as_str(),
-            "connected to compute node at {} ({socket_addr}) sslmode={:?}, latency={}, query_id={}",
-            self.host,
-            self.ssl_mode,
+            "connected to compute node at {host} ({socket_addr}) sslmode={:?}, latency={}, query_id={}",
+            self.0.get_ssl_mode(),
             ctx.get_proxy_latency(),
             ctx.get_testodrome_id().unwrap_or_default(),
         );
@@ -329,11 +299,11 @@ impl ConnectInfo {
             socket_addr,
             CancelToken {
                 socket_config: None,
-                ssl_mode: self.ssl_mode,
+                ssl_mode: self.0.get_ssl_mode(),
                 process_id,
                 secret_key,
             },
-            self.host.to_string(),
+            host.to_string(),
             user_info,
         );
 

proxy/src/compute/tls.rs

@@ -1,63 +0,0 @@
-use futures::FutureExt;
-use postgres_client::config::SslMode;
-use postgres_client::maybe_tls_stream::MaybeTlsStream;
-use postgres_client::tls::{MakeTlsConnect, TlsConnect};
-use rustls::pki_types::InvalidDnsNameError;
-use thiserror::Error;
-use tokio::io::{AsyncRead, AsyncWrite};
-
-use crate::pqproto::request_tls;
-use crate::proxy::retry::CouldRetry;
-
-#[derive(Debug, Error)]
-pub enum TlsError {
-    #[error(transparent)]
-    Dns(#[from] InvalidDnsNameError),
-    #[error(transparent)]
-    Connection(#[from] std::io::Error),
-    #[error("TLS required but not provided")]
-    Required,
-}
-
-impl CouldRetry for TlsError {
-    fn could_retry(&self) -> bool {
-        match self {
-            TlsError::Dns(_) => false,
-            TlsError::Connection(err) => err.could_retry(),
-            // perhaps compute didn't realise it supports TLS?
-            TlsError::Required => true,
-        }
-    }
-}
-
-pub async fn connect_tls<S, T>(
-    mut stream: S,
-    mode: SslMode,
-    tls: &T,
-    host: &str,
-) -> Result<MaybeTlsStream<S, T::Stream>, TlsError>
-where
-    S: AsyncRead + AsyncWrite + Unpin + Send,
-    T: MakeTlsConnect<
-            S,
-            Error = InvalidDnsNameError,
-            TlsConnect: TlsConnect<S, Error = std::io::Error, Future: Send>,
-        >,
-{
-    match mode {
-        SslMode::Disable => return Ok(MaybeTlsStream::Raw(stream)),
-        SslMode::Prefer | SslMode::Require => {}
-    }
-
-    if !request_tls(&mut stream).await? {
-        if SslMode::Require == mode {
-            return Err(TlsError::Required);
-        }
-
-        return Ok(MaybeTlsStream::Raw(stream));
-    }
-
-    Ok(MaybeTlsStream::Tls(
-        tls.make_tls_connect(host)?.connect(stream).boxed().await?,
-    ))
-}

proxy/src/console_redirect_proxy.rs

@@ -11,12 +11,13 @@ use crate::config::{ProxyConfig, ProxyProtocolV2};
 use crate::context::RequestContext;
 use crate::error::ReportableError;
 use crate::metrics::{Metrics, NumClientConnectionsGuard};
+use crate::pglb::connect_compute::{TcpMechanism, connect_to_compute};
 use crate::pglb::handshake::{HandshakeData, handshake};
 use crate::pglb::passthrough::ProxyPassthrough;
 use crate::protocol2::{ConnectHeader, ConnectionInfo, read_proxy_protocol};
-use crate::proxy::connect_compute::{TcpMechanism, connect_to_compute};
-use crate::proxy::{ClientRequestError, ErrorSource, prepare_client_connection};
-use crate::util::run_until_cancelled;
+use crate::proxy::{
+    ClientRequestError, ErrorSource, prepare_client_connection, run_until_cancelled,
+};
 
 pub async fn task_main(
     config: &'static ProxyConfig,
@@ -209,20 +210,20 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
 
     ctx.set_db_options(params.clone());
 
-    let (node_info, mut auth_info, user_info) = match backend
+    let (node_info, user_info, _ip_allowlist) = match backend
         .authenticate(ctx, &config.authentication_config, &mut stream)
         .await
     {
         Ok(auth_result) => auth_result,
         Err(e) => Err(stream.throw_error(e, Some(ctx)).await)?,
     };
-    auth_info.set_startup_params(&params, true);
 
     let node = connect_to_compute(
         ctx,
         &TcpMechanism {
             user_info,
-            auth: auth_info,
+            params_compat: true,
+            params: &params,
             locks: &config.connect_compute_locks,
         },
         &node_info,

proxy/src/control_plane/client/cplane_proxy_v1.rs

@@ -146,7 +146,6 @@ impl NeonControlPlaneClient {
                     public_access_blocked: block_public_connections,
                     vpc_access_blocked: block_vpc_connections,
                 },
-                rate_limits: body.rate_limits,
             })
         }
         .inspect_err(|e| tracing::debug!(error = ?e))
@@ -262,18 +261,24 @@ impl NeonControlPlaneClient {
                 Some(_) => SslMode::Require,
                 None => SslMode::Disable,
             };
-            let host = match body.server_name {
-                Some(host) => host.into(),
-                None => host.into(),
+            let host_name = match body.server_name {
+                Some(host) => host,
+                None => host.to_owned(),
             };
 
+            // Don't set anything but host and port! This config will be cached.
+            // We'll set username and such later using the startup message.
+            // TODO: add more type safety (in progress).
+            let mut config = compute::ConnCfg::new(host_name, port);
+
+            if let Some(addr) = host_addr {
+                config.set_host_addr(addr);
+            }
+
+            config.ssl_mode(ssl_mode);
+
             let node = NodeInfo {
-                conn_info: compute::ConnectInfo {
-                    host_addr,
-                    host,
-                    port,
-                    ssl_mode,
-                },
+                config,
                 aux: body.aux,
             };
 
@@ -313,7 +318,6 @@ impl super::ControlPlaneApi for NeonControlPlaneClient {
             allowed_ips: Arc::new(auth_info.allowed_ips),
             allowed_vpce: Arc::new(auth_info.allowed_vpc_endpoint_ids),
             flags: auth_info.access_blocker_flags,
-            rate_limits: auth_info.rate_limits,
         };
         let role_control = RoleAccessControl {
             secret: auth_info.secret,
@@ -359,7 +363,6 @@ impl super::ControlPlaneApi for NeonControlPlaneClient {
             allowed_ips: Arc::new(auth_info.allowed_ips),
             allowed_vpce: Arc::new(auth_info.allowed_vpc_endpoint_ids),
             flags: auth_info.access_blocker_flags,
-            rate_limits: auth_info.rate_limits,
         };
         let role_control = RoleAccessControl {
             secret: auth_info.secret,

proxy/src/control_plane/client/mock.rs

@@ -6,7 +6,6 @@ use std::str::FromStr;
 use std::sync::Arc;
 
 use futures::TryFutureExt;
-use postgres_client::config::SslMode;
 use thiserror::Error;
 use tokio_postgres::Client;
 use tracing::{Instrument, error, info, info_span, warn};
@@ -15,20 +14,19 @@ use crate::auth::IpPattern;
 use crate::auth::backend::ComputeUserInfo;
 use crate::auth::backend::jwt::AuthRule;
 use crate::cache::Cached;
-use crate::compute::ConnectInfo;
 use crate::context::RequestContext;
 use crate::control_plane::errors::{
     ControlPlaneError, GetAuthInfoError, GetEndpointJwksError, WakeComputeError,
 };
-use crate::control_plane::messages::{EndpointRateLimitConfig, MetricsAuxInfo};
+use crate::control_plane::messages::MetricsAuxInfo;
 use crate::control_plane::{
     AccessBlockerFlags, AuthInfo, AuthSecret, CachedNodeInfo, EndpointAccessControl, NodeInfo,
     RoleAccessControl,
 };
 use crate::intern::RoleNameInt;
-use crate::scram;
 use crate::types::{BranchId, EndpointId, ProjectId, RoleName};
 use crate::url::ApiUrl;
+use crate::{compute, scram};
 
 #[derive(Debug, Error)]
 enum MockApiError {
@@ -89,7 +87,8 @@ impl MockControlPlane {
             .await?
             {
                 info!("got a secret: {entry}"); // safe since it's not a prod scenario
-                scram::ServerSecret::parse(&entry).map(AuthSecret::Scram)
+                let secret = scram::ServerSecret::parse(&entry).map(AuthSecret::Scram);
+                secret.or_else(|| parse_md5(&entry).map(AuthSecret::Md5))
             } else {
                 warn!("user '{role}' does not exist");
                 None
@@ -130,7 +129,6 @@ impl MockControlPlane {
             project_id: None,
             account_id: None,
             access_blocker_flags: AccessBlockerFlags::default(),
-            rate_limits: EndpointRateLimitConfig::default(),
         })
     }
 
@@ -172,23 +170,25 @@ impl MockControlPlane {
 
     async fn do_wake_compute(&self) -> Result<NodeInfo, WakeComputeError> {
         let port = self.endpoint.port().unwrap_or(5432);
-        let conn_info = match self.endpoint.host_str() {
-            None => ConnectInfo {
-                host_addr: Some(IpAddr::V4(Ipv4Addr::LOCALHOST)),
-                host: "localhost".into(),
-                port,
-                ssl_mode: SslMode::Disable,
-            },
-            Some(host) => ConnectInfo {
-                host_addr: IpAddr::from_str(host).ok(),
-                host: host.into(),
-                port,
-                ssl_mode: SslMode::Disable,
-            },
+        let mut config = match self.endpoint.host_str() {
+            None => {
+                let mut config = compute::ConnCfg::new("localhost".to_string(), port);
+                config.set_host_addr(IpAddr::V4(Ipv4Addr::LOCALHOST));
+                config
+            }
+            Some(host) => {
+                let mut config = compute::ConnCfg::new(host.to_string(), port);
+                if let Ok(addr) = IpAddr::from_str(host) {
+                    config.set_host_addr(addr);
+                }
+                config
+            }
         };
 
+        config.ssl_mode(postgres_client::config::SslMode::Disable);
+
         let node = NodeInfo {
-            conn_info,
+            config,
             aux: MetricsAuxInfo {
                 endpoint_id: (&EndpointId::from("endpoint")).into(),
                 project_id: (&ProjectId::from("project")).into(),
@@ -234,7 +234,6 @@ impl super::ControlPlaneApi for MockControlPlane {
             allowed_ips: Arc::new(info.allowed_ips),
             allowed_vpce: Arc::new(info.allowed_vpc_endpoint_ids),
             flags: info.access_blocker_flags,
-            rate_limits: info.rate_limits,
         })
     }
 
@@ -267,3 +266,12 @@ impl super::ControlPlaneApi for MockControlPlane {
         self.do_wake_compute().map_ok(Cached::new_uncached).await
     }
 }
+
+fn parse_md5(input: &str) -> Option<[u8; 16]> {
+    let text = input.strip_prefix("md5")?;
+
+    let mut bytes = [0u8; 16];
+    hex::decode_to_slice(text, &mut bytes).ok()?;
+
+    Some(bytes)
+}

proxy/src/control_plane/client/mod.rs

@@ -10,7 +10,6 @@ use clashmap::ClashMap;
 use tokio::time::Instant;
 use tracing::{debug, info};
 
-use super::{EndpointAccessControl, RoleAccessControl};
 use crate::auth::backend::ComputeUserInfo;
 use crate::auth::backend::jwt::{AuthRule, FetchAuthRules, FetchAuthRulesError};
 use crate::cache::endpoints::EndpointsCache;
@@ -23,6 +22,8 @@ use crate::metrics::ApiLockMetrics;
 use crate::rate_limiter::{DynamicLimiter, Outcome, RateLimiterConfig, Token};
 use crate::types::EndpointId;
 
+use super::{EndpointAccessControl, RoleAccessControl};
+
 #[non_exhaustive]
 #[derive(Clone)]
 pub enum ControlPlaneClient {

proxy/src/control_plane/messages.rs

@@ -227,35 +227,12 @@ pub(crate) struct UserFacingMessage {
 #[derive(Deserialize)]
 pub(crate) struct GetEndpointAccessControl {
     pub(crate) role_secret: Box<str>,
-
-    pub(crate) project_id: Option<ProjectIdInt>,
-    pub(crate) account_id: Option<AccountIdInt>,
-
     pub(crate) allowed_ips: Option<Vec<IpPattern>>,
     pub(crate) allowed_vpc_endpoint_ids: Option<Vec<String>>,
+    pub(crate) project_id: Option<ProjectIdInt>,
+    pub(crate) account_id: Option<AccountIdInt>,
     pub(crate) block_public_connections: Option<bool>,
     pub(crate) block_vpc_connections: Option<bool>,
-
-    #[serde(default)]
-    pub(crate) rate_limits: EndpointRateLimitConfig,
-}
-
-#[derive(Copy, Clone, Deserialize, Default)]
-pub struct EndpointRateLimitConfig {
-    pub connection_attempts: ConnectionAttemptsLimit,
-}
-
-#[derive(Copy, Clone, Deserialize, Default)]
-pub struct ConnectionAttemptsLimit {
-    pub tcp: Option<LeakyBucketSetting>,
-    pub ws: Option<LeakyBucketSetting>,
-    pub http: Option<LeakyBucketSetting>,
-}
-
-#[derive(Copy, Clone, Deserialize)]
-pub struct LeakyBucketSetting {
-    pub rps: f64,
-    pub burst: f64,
 }
 
 /// Response which holds compute node's `host:port` pair.

proxy/src/control_plane/mod.rs

@@ -11,18 +11,15 @@ pub(crate) mod errors;
 
 use std::sync::Arc;
 
-use messages::EndpointRateLimitConfig;
-
-use crate::auth::backend::ComputeUserInfo;
 use crate::auth::backend::jwt::AuthRule;
+use crate::auth::backend::{ComputeCredentialKeys, ComputeUserInfo};
 use crate::auth::{AuthError, IpPattern, check_peer_addr_is_in_list};
 use crate::cache::{Cached, TimedLru};
 use crate::config::ComputeConfig;
 use crate::context::RequestContext;
 use crate::control_plane::messages::{ControlPlaneErrorMessage, MetricsAuxInfo};
-use crate::intern::{AccountIdInt, EndpointIdInt, ProjectIdInt};
+use crate::intern::{AccountIdInt, ProjectIdInt};
 use crate::protocol2::ConnectionInfoExtra;
-use crate::rate_limiter::{EndpointRateLimiter, LeakyBucketConfig};
 use crate::types::{EndpointCacheKey, EndpointId, RoleName};
 use crate::{compute, scram};
 
@@ -42,6 +39,10 @@ pub mod mgmt;
 /// Auth secret which is managed by the cloud.
 #[derive(Clone, Eq, PartialEq, Debug)]
 pub(crate) enum AuthSecret {
+    #[cfg(any(test, feature = "testing"))]
+    /// Md5 hash of user's password.
+    Md5([u8; 16]),
+
     /// [SCRAM](crate::scram) authentication info.
     Scram(scram::ServerSecret),
 }
@@ -59,14 +60,16 @@ pub(crate) struct AuthInfo {
     pub(crate) account_id: Option<AccountIdInt>,
     /// Are public connections or VPC connections blocked?
     pub(crate) access_blocker_flags: AccessBlockerFlags,
-    /// The rate limits for this endpoint.
-    pub(crate) rate_limits: EndpointRateLimitConfig,
 }
 
 /// Info for establishing a connection to a compute node.
+/// This is what we get after auth succeeded, but not before!
 #[derive(Clone)]
 pub(crate) struct NodeInfo {
-    pub(crate) conn_info: compute::ConnectInfo,
+    /// Compute node connection params.
+    /// It's sad that we have to clone this, but this will improve
+    /// once we migrate to a bespoke connection logic.
+    pub(crate) config: compute::ConnCfg,
 
     /// Labels for proxy's metrics.
     pub(crate) aux: MetricsAuxInfo,
@@ -76,14 +79,26 @@ impl NodeInfo {
     pub(crate) async fn connect(
         &self,
         ctx: &RequestContext,
-        auth: &compute::AuthInfo,
         config: &ComputeConfig,
         user_info: ComputeUserInfo,
     ) -> Result<compute::PostgresConnection, compute::ConnectionError> {
-        self.conn_info
-            .connect(ctx, self.aux.clone(), auth, config, user_info)
+        self.config
+            .connect(ctx, self.aux.clone(), config, user_info)
             .await
     }
+
+    pub(crate) fn reuse_settings(&mut self, other: Self) {
+        self.config.reuse_password(other.config);
+    }
+
+    pub(crate) fn set_keys(&mut self, keys: &ComputeCredentialKeys) {
+        match keys {
+            #[cfg(any(test, feature = "testing"))]
+            ComputeCredentialKeys::Password(password) => self.config.password(password),
+            ComputeCredentialKeys::AuthKeys(auth_keys) => self.config.auth_keys(*auth_keys),
+            ComputeCredentialKeys::JwtPayload(_) | ComputeCredentialKeys::None => &mut self.config,
+        };
+    }
 }
 
 #[derive(Copy, Clone, Default)]
@@ -106,8 +121,6 @@ pub struct EndpointAccessControl {
     pub allowed_ips: Arc<Vec<IpPattern>>,
     pub allowed_vpce: Arc<Vec<String>>,
     pub flags: AccessBlockerFlags,
-
-    pub rate_limits: EndpointRateLimitConfig,
 }
 
 impl EndpointAccessControl {
@@ -146,36 +159,6 @@ impl EndpointAccessControl {
 
         Ok(())
     }
-
-    pub fn connection_attempt_rate_limit(
-        &self,
-        ctx: &RequestContext,
-        endpoint: &EndpointId,
-        rate_limiter: &EndpointRateLimiter,
-    ) -> Result<(), AuthError> {
-        let endpoint = EndpointIdInt::from(endpoint);
-
-        let limits = &self.rate_limits.connection_attempts;
-        let config = match ctx.protocol() {
-            crate::metrics::Protocol::Http => limits.http,
-            crate::metrics::Protocol::Ws => limits.ws,
-            crate::metrics::Protocol::Tcp => limits.tcp,
-            crate::metrics::Protocol::SniRouter => return Ok(()),
-        };
-        let config = config.and_then(|config| {
-            if config.rps <= 0.0 || config.burst <= 0.0 {
-                return None;
-            }
-
-            Some(LeakyBucketConfig::new(config.rps, config.burst))
-        });
-
-        if !rate_limiter.check(endpoint, config, 1) {
-            return Err(AuthError::too_many_connections());
-        }
-
-        Ok(())
-    }
 }
 
 /// This will allocate per each call, but the http requests alone

proxy/src/lib.rs

@@ -106,5 +106,4 @@ mod tls;
 mod types;
 mod url;
 mod usage_metrics;
-mod util;
 mod waiters;

proxy/src/metrics.rs

@@ -610,11 +610,11 @@ pub enum RedisEventsCount {
     BranchCreated,
     ProjectCreated,
     CancelSession,
-    InvalidateRole,
-    InvalidateEndpoint,
-    InvalidateProject,
-    InvalidateProjects,
-    InvalidateOrg,
+    PasswordUpdate,
+    AllowedIpsUpdate,
+    AllowedVpcEndpointIdsUpdateForProjects,
+    AllowedVpcEndpointIdsUpdateForAllProjectsInOrg,
+    BlockPublicOrVpcAccessUpdate,
 }
 
 pub struct ThreadPoolWorkers(usize);

proxy/src/pglb/connect_compute.rs

@@ -2,25 +2,26 @@ use async_trait::async_trait;
 use tokio::time;
 use tracing::{debug, info, warn};
 
-use crate::auth::backend::ComputeUserInfo;
-use crate::compute::{self, AuthInfo, COULD_NOT_CONNECT, PostgresConnection};
+use crate::auth::backend::{ComputeCredentialKeys, ComputeUserInfo};
+use crate::compute::{self, COULD_NOT_CONNECT, PostgresConnection};
 use crate::config::{ComputeConfig, RetryConfig};
 use crate::context::RequestContext;
 use crate::control_plane::errors::WakeComputeError;
 use crate::control_plane::locks::ApiLocks;
-use crate::control_plane::{self, NodeInfo};
+use crate::control_plane::{self, CachedNodeInfo, NodeInfo};
 use crate::error::ReportableError;
 use crate::metrics::{
     ConnectOutcome, ConnectionFailureKind, Metrics, RetriesMetricGroup, RetryType,
 };
+use crate::pqproto::StartupMessageParams;
 use crate::proxy::retry::{CouldRetry, ShouldRetryWakeCompute, retry_after, should_retry};
-use crate::proxy::wake_compute::{WakeComputeBackend, wake_compute};
+use crate::proxy::wake_compute::wake_compute;
 use crate::types::Host;
 
 /// If we couldn't connect, a cached connection info might be to blame
 /// (e.g. the compute node's address might've changed at the wrong time).
 /// Invalidate the cache entry (if any) to prevent subsequent errors.
-#[tracing::instrument(skip_all)]
+#[tracing::instrument(name = "invalidate_cache", skip_all)]
 pub(crate) fn invalidate_cache(node_info: control_plane::CachedNodeInfo) -> NodeInfo {
     let is_cached = node_info.cached();
     if is_cached {
@@ -47,17 +48,34 @@ pub(crate) trait ConnectMechanism {
         node_info: &control_plane::CachedNodeInfo,
         config: &ComputeConfig,
     ) -> Result<Self::Connection, Self::ConnectError>;
+
+    fn update_connect_config(&self, conf: &mut compute::ConnCfg);
 }
 
-pub(crate) struct TcpMechanism {
-    pub(crate) auth: AuthInfo,
+#[async_trait]
+pub(crate) trait ComputeConnectBackend {
+    async fn wake_compute(
+        &self,
+        ctx: &RequestContext,
+    ) -> Result<CachedNodeInfo, control_plane::errors::WakeComputeError>;
+
+    fn get_keys(&self) -> &ComputeCredentialKeys;
+}
+
+pub(crate) struct TcpMechanism<'a> {
+    pub(crate) params_compat: bool,
+
+    /// KV-dictionary with PostgreSQL connection params.
+    pub(crate) params: &'a StartupMessageParams,
+
     /// connect_to_compute concurrency lock
     pub(crate) locks: &'static ApiLocks<Host>,
+
     pub(crate) user_info: ComputeUserInfo,
 }
 
 #[async_trait]
-impl ConnectMechanism for TcpMechanism {
+impl ConnectMechanism for TcpMechanism<'_> {
     type Connection = PostgresConnection;
     type ConnectError = compute::ConnectionError;
     type Error = compute::ConnectionError;
@@ -72,18 +90,19 @@ impl ConnectMechanism for TcpMechanism {
         node_info: &control_plane::CachedNodeInfo,
         config: &ComputeConfig,
     ) -> Result<PostgresConnection, Self::Error> {
-        let permit = self.locks.get_permit(&node_info.conn_info.host).await?;
-        permit.release_result(
-            node_info
-                .connect(ctx, &self.auth, config, self.user_info.clone())
-                .await,
-        )
+        let host = node_info.config.get_host();
+        let permit = self.locks.get_permit(&host).await?;
+        permit.release_result(node_info.connect(ctx, config, self.user_info.clone()).await)
+    }
+
+    fn update_connect_config(&self, config: &mut compute::ConnCfg) {
+        config.set_startup_params(self.params, self.params_compat);
     }
 }
 
 /// Try to connect to the compute node, retrying if necessary.
 #[tracing::instrument(skip_all)]
-pub(crate) async fn connect_to_compute<M: ConnectMechanism, B: WakeComputeBackend>(
+pub(crate) async fn connect_to_compute<M: ConnectMechanism, B: ComputeConnectBackend>(
     ctx: &RequestContext,
     mechanism: &M,
     user_info: &B,
@@ -95,9 +114,12 @@ where
     M::Error: From<WakeComputeError>,
 {
     let mut num_retries = 0;
-    let node_info =
+    let mut node_info =
         wake_compute(&mut num_retries, ctx, user_info, wake_compute_retry_config).await?;
 
+    node_info.set_keys(user_info.get_keys());
+    mechanism.update_connect_config(&mut node_info.config);
+
     // try once
     let err = match mechanism.connect_once(ctx, &node_info, compute).await {
         Ok(res) => {
@@ -133,9 +155,14 @@ where
     } else {
         // if we failed to connect, it's likely that the compute node was suspended, wake a new compute node
         debug!("compute node's state has likely changed; requesting a wake-up");
-        invalidate_cache(node_info);
+        let old_node_info = invalidate_cache(node_info);
         // TODO: increment num_retries?
-        wake_compute(&mut num_retries, ctx, user_info, wake_compute_retry_config).await?
+        let mut node_info =
+            wake_compute(&mut num_retries, ctx, user_info, wake_compute_retry_config).await?;
+        node_info.reuse_settings(old_node_info);
+
+        mechanism.update_connect_config(&mut node_info.config);
+        node_info
     };
 
     // now that we have a new node, try connect to it repeatedly.

proxy/src/pglb/mod.rs

@@ -1,3 +1,4 @@
+pub mod connect_compute;
 pub mod copy_bidirectional;
 pub mod handshake;
 pub mod inprocess;

proxy/src/pqproto.rs

@@ -8,7 +8,7 @@ use std::io::{self, Cursor};
 use bytes::{Buf, BufMut};
 use itertools::Itertools;
 use rand::distributions::{Distribution, Standard};
-use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
+use tokio::io::{AsyncRead, AsyncReadExt};
 use zerocopy::{FromBytes, Immutable, IntoBytes, big_endian};
 
 pub type ErrorCode = [u8; 5];
@@ -53,28 +53,6 @@ impl fmt::Debug for ProtocolVersion {
     }
 }
 
-/// <https://github.com/postgres/postgres/blob/ca481d3c9ab7bf69ff0c8d71ad3951d407f6a33c/src/include/libpq/pqcomm.h#L118>
-const MAX_STARTUP_PACKET_LENGTH: usize = 10000;
-const RESERVED_INVALID_MAJOR_VERSION: u16 = 1234;
-/// <https://github.com/postgres/postgres/blob/ca481d3c9ab7bf69ff0c8d71ad3951d407f6a33c/src/include/libpq/pqcomm.h#L132>
-const CANCEL_REQUEST_CODE: ProtocolVersion = ProtocolVersion::new(1234, 5678);
-/// <https://github.com/postgres/postgres/blob/ca481d3c9ab7bf69ff0c8d71ad3951d407f6a33c/src/include/libpq/pqcomm.h#L166>
-const NEGOTIATE_SSL_CODE: ProtocolVersion = ProtocolVersion::new(1234, 5679);
-/// <https://github.com/postgres/postgres/blob/ca481d3c9ab7bf69ff0c8d71ad3951d407f6a33c/src/include/libpq/pqcomm.h#L167>
-const NEGOTIATE_GSS_CODE: ProtocolVersion = ProtocolVersion::new(1234, 5680);
-
-/// This first reads the startup message header, is 8 bytes.
-/// The first 4 bytes is a big-endian message length, and the next 4 bytes is a version number.
-///
-/// The length value is inclusive of the header. For example,
-/// an empty message will always have length 8.
-#[derive(Clone, Copy, FromBytes, IntoBytes, Immutable)]
-#[repr(C)]
-struct StartupHeader {
-    len: big_endian::U32,
-    version: ProtocolVersion,
-}
-
 /// read the type from the stream using zerocopy.
 ///
 /// not cancel safe.
@@ -88,38 +66,32 @@ macro_rules! read {
     }};
 }
 
-/// Returns true if TLS is supported.
-///
-/// This is not cancel safe.
-pub async fn request_tls<S>(stream: &mut S) -> io::Result<bool>
-where
-    S: AsyncRead + AsyncWrite + Unpin,
-{
-    let payload = StartupHeader {
-        len: 8.into(),
-        version: NEGOTIATE_SSL_CODE,
-    };
-    stream.write_all(payload.as_bytes()).await?;
-    stream.flush().await?;
-
-    // we expect back either `S` or `N` as a single byte.
-    let mut res = *b"0";
-    stream.read_exact(&mut res).await?;
-
-    debug_assert!(
-        res == *b"S" || res == *b"N",
-        "unexpected SSL negotiation response: {}",
-        char::from(res[0]),
-    );
-
-    // S for SSL.
-    Ok(res == *b"S")
-}
-
 pub async fn read_startup<S>(stream: &mut S) -> io::Result<FeStartupPacket>
 where
     S: AsyncRead + Unpin,
 {
+    /// <https://github.com/postgres/postgres/blob/ca481d3c9ab7bf69ff0c8d71ad3951d407f6a33c/src/include/libpq/pqcomm.h#L118>
+    const MAX_STARTUP_PACKET_LENGTH: usize = 10000;
+    const RESERVED_INVALID_MAJOR_VERSION: u16 = 1234;
+    /// <https://github.com/postgres/postgres/blob/ca481d3c9ab7bf69ff0c8d71ad3951d407f6a33c/src/include/libpq/pqcomm.h#L132>
+    const CANCEL_REQUEST_CODE: ProtocolVersion = ProtocolVersion::new(1234, 5678);
+    /// <https://github.com/postgres/postgres/blob/ca481d3c9ab7bf69ff0c8d71ad3951d407f6a33c/src/include/libpq/pqcomm.h#L166>
+    const NEGOTIATE_SSL_CODE: ProtocolVersion = ProtocolVersion::new(1234, 5679);
+    /// <https://github.com/postgres/postgres/blob/ca481d3c9ab7bf69ff0c8d71ad3951d407f6a33c/src/include/libpq/pqcomm.h#L167>
+    const NEGOTIATE_GSS_CODE: ProtocolVersion = ProtocolVersion::new(1234, 5680);
+
+    /// This first reads the startup message header, is 8 bytes.
+    /// The first 4 bytes is a big-endian message length, and the next 4 bytes is a version number.
+    ///
+    /// The length value is inclusive of the header. For example,
+    /// an empty message will always have length 8.
+    #[derive(Clone, Copy, FromBytes, IntoBytes, Immutable)]
+    #[repr(C)]
+    struct StartupHeader {
+        len: big_endian::U32,
+        version: ProtocolVersion,
+    }
+
     let header = read!(stream => StartupHeader);
 
     // <https://github.com/postgres/postgres/blob/04bcf9e19a4261fe9c7df37c777592c2e10c32a7/src/backend/tcop/backend_startup.c#L378-L382>
@@ -592,9 +564,10 @@ mod tests {
     use tokio::io::{AsyncWriteExt, duplex};
     use zerocopy::IntoBytes;
 
-    use super::ProtocolVersion;
     use crate::pqproto::{FeStartupPacket, read_message, read_startup};
 
+    use super::ProtocolVersion;
+
     #[tokio::test]
     async fn reject_large_startup() {
         // we're going to define a v3.0 startup message with far too many parameters.

proxy/src/proxy/mod.rs

@@ -1,10 +1,8 @@
 #[cfg(test)]
 mod tests;
 
-pub(crate) mod connect_compute;
 pub(crate) mod retry;
 pub(crate) mod wake_compute;
-
 use std::sync::Arc;
 
 use futures::FutureExt;
@@ -23,16 +21,15 @@ use crate::config::{ProxyConfig, ProxyProtocolV2, TlsConfig};
 use crate::context::RequestContext;
 use crate::error::{ReportableError, UserFacingError};
 use crate::metrics::{Metrics, NumClientConnectionsGuard};
+use crate::pglb::connect_compute::{TcpMechanism, connect_to_compute};
 pub use crate::pglb::copy_bidirectional::{ErrorSource, copy_bidirectional_client_compute};
 use crate::pglb::handshake::{HandshakeData, HandshakeError, handshake};
 use crate::pglb::passthrough::ProxyPassthrough;
 use crate::pqproto::{BeMessage, CancelKeyData, StartupMessageParams};
 use crate::protocol2::{ConnectHeader, ConnectionInfo, ConnectionInfoExtra, read_proxy_protocol};
-use crate::proxy::connect_compute::{TcpMechanism, connect_to_compute};
 use crate::rate_limiter::EndpointRateLimiter;
 use crate::stream::{PqStream, Stream};
 use crate::types::EndpointCacheKey;
-use crate::util::run_until_cancelled;
 use crate::{auth, compute};
 
 const ERR_INSECURE_CONNECTION: &str = "connection is insecure (try using `sslmode=require`)";
@@ -49,6 +46,21 @@ impl ReportableError for TlsRequired {
 
 impl UserFacingError for TlsRequired {}
 
+pub async fn run_until_cancelled<F: std::future::Future>(
+    f: F,
+    cancellation_token: &CancellationToken,
+) -> Option<F::Output> {
+    match futures::future::select(
+        std::pin::pin!(f),
+        std::pin::pin!(cancellation_token.cancelled()),
+    )
+    .await
+    {
+        futures::future::Either::Left((f, _)) => Some(f),
+        futures::future::Either::Right(((), _)) => None,
+    }
+}
+
 pub async fn task_main(
     config: &'static ProxyConfig,
     auth_backend: &'static auth::Backend<'static, ()>,
@@ -346,22 +358,24 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
         }
     };
 
-    let (cplane, creds) = match user_info {
-        auth::Backend::ControlPlane(cplane, creds) => (cplane, creds),
+    let compute_user_info = match &user_info {
+        auth::Backend::ControlPlane(_, info) => &info.info,
         auth::Backend::Local(_) => unreachable!("local proxy does not run tcp proxy service"),
     };
-    let params_compat = creds.info.options.get(NeonOptions::PARAMS_COMPAT).is_some();
-    let mut auth_info = compute::AuthInfo::with_auth_keys(creds.keys);
-    auth_info.set_startup_params(&params, params_compat);
+    let params_compat = compute_user_info
+        .options
+        .get(NeonOptions::PARAMS_COMPAT)
+        .is_some();
 
     let res = connect_to_compute(
         ctx,
         &TcpMechanism {
-            user_info: creds.info.clone(),
-            auth: auth_info,
+            user_info: compute_user_info.clone(),
+            params_compat,
+            params: &params,
             locks: &config.connect_compute_locks,
         },
-        &auth::Backend::ControlPlane(cplane, creds.info),
+        &user_info,
         config.wake_compute_retry_config,
         &config.connect_to_compute,
     )

proxy/src/proxy/retry.rs

@@ -100,9 +100,9 @@ impl CouldRetry for compute::ConnectionError {
     fn could_retry(&self) -> bool {
         match self {
             compute::ConnectionError::Postgres(err) => err.could_retry(),
-            compute::ConnectionError::TlsError(err) => err.could_retry(),
+            compute::ConnectionError::CouldNotConnect(err) => err.could_retry(),
             compute::ConnectionError::WakeComputeError(err) => err.could_retry(),
-            compute::ConnectionError::TooManyConnectionAttempts(_) => false,
+            _ => false,
         }
     }
 }

proxy/src/proxy/tests/mod.rs

@@ -19,14 +19,17 @@ use tracing_test::traced_test;
 
 use super::retry::CouldRetry;
 use super::*;
-use crate::auth::backend::{ComputeUserInfo, MaybeOwned};
+use crate::auth::backend::{
+    ComputeCredentialKeys, ComputeCredentials, ComputeUserInfo, MaybeOwned,
+};
 use crate::config::{ComputeConfig, RetryConfig};
 use crate::control_plane::client::{ControlPlaneClient, TestControlPlaneClient};
 use crate::control_plane::messages::{ControlPlaneErrorMessage, Details, MetricsAuxInfo, Status};
 use crate::control_plane::{self, CachedNodeInfo, NodeInfo, NodeInfoCache};
 use crate::error::ErrorKind;
-use crate::proxy::connect_compute::ConnectMechanism;
+use crate::pglb::connect_compute::ConnectMechanism;
 use crate::tls::client_config::compute_client_config_with_certs;
+use crate::tls::postgres_rustls::MakeRustlsConnect;
 use crate::tls::server_config::CertResolver;
 use crate::types::{BranchId, EndpointId, ProjectId};
 use crate::{sasl, scram};
@@ -69,14 +72,13 @@ struct ClientConfig<'a> {
     hostname: &'a str,
 }
 
-type TlsConnect<S> = <ComputeConfig as MakeTlsConnect<S>>::TlsConnect;
+type TlsConnect<S> = <MakeRustlsConnect as MakeTlsConnect<S>>::TlsConnect;
 
 impl ClientConfig<'_> {
     fn make_tls_connect(self) -> anyhow::Result<TlsConnect<DuplexStream>> {
-        Ok(crate::tls::postgres_rustls::make_tls_connect(
-            &self.config,
-            self.hostname,
-        )?)
+        let mut mk = MakeRustlsConnect::new(self.config);
+        let tls = MakeTlsConnect::<DuplexStream>::make_tls_connect(&mut mk, self.hostname)?;
+        Ok(tls)
     }
 }
 
@@ -495,6 +497,8 @@ impl ConnectMechanism for TestConnectMechanism {
             x => panic!("expecting action {x:?}, connect is called instead"),
         }
     }
+
+    fn update_connect_config(&self, _conf: &mut compute::ConnCfg) {}
 }
 
 impl TestControlPlaneClient for TestConnectMechanism {
@@ -553,12 +557,7 @@ impl TestControlPlaneClient for TestConnectMechanism {
 
 fn helper_create_cached_node_info(cache: &'static NodeInfoCache) -> CachedNodeInfo {
     let node = NodeInfo {
-        conn_info: compute::ConnectInfo {
-            host: "test".into(),
-            port: 5432,
-            ssl_mode: SslMode::Disable,
-            host_addr: None,
-        },
+        config: compute::ConnCfg::new("test".to_owned(), 5432),
         aux: MetricsAuxInfo {
             endpoint_id: (&EndpointId::from("endpoint")).into(),
             project_id: (&ProjectId::from("project")).into(),
@@ -573,13 +572,16 @@ fn helper_create_cached_node_info(cache: &'static NodeInfoCache) -> CachedNodeIn
 
 fn helper_create_connect_info(
     mechanism: &TestConnectMechanism,
-) -> auth::Backend<'static, ComputeUserInfo> {
+) -> auth::Backend<'static, ComputeCredentials> {
     auth::Backend::ControlPlane(
         MaybeOwned::Owned(ControlPlaneClient::Test(Box::new(mechanism.clone()))),
-        ComputeUserInfo {
-            endpoint: "endpoint".into(),
-            user: "user".into(),
-            options: NeonOptions::parse_options_raw(""),
+        ComputeCredentials {
+            info: ComputeUserInfo {
+                endpoint: "endpoint".into(),
+                user: "user".into(),
+                options: NeonOptions::parse_options_raw(""),
+            },
+            keys: ComputeCredentialKeys::Password("password".into()),
         },
     )
 }