test(storcon): add test cases for 404 passthrough

Signed-off-by: Alex Chi Z <chi@neon.tech>

.config/hakari.toml

@@ -21,14 +21,13 @@ platforms = [
     # "x86_64-apple-darwin",
     # "x86_64-pc-windows-msvc",
 ]
+
 [final-excludes]
 workspace-members = [
     # vm_monitor benefits from the same Cargo.lock as the rest of our artifacts, but
     # it is built primarly in separate repo neondatabase/autoscaling and thus is excluded
     # from depending on workspace-hack because most of the dependencies are not used.
     "vm_monitor",
-    # subzero-core is a stub crate that should be excluded from workspace-hack
-    "subzero-core",
     # All of these exist in libs and are not usually built independently.
     # Putting workspace hack there adds a bottleneck for cargo builds.
     "compute_api",

.github/actions/prepare-for-subzero/action.yml

@@ -1,28 +0,0 @@
-name: 'Prepare current job for subzero'
-description: >
-  Set git token to access `neondatabase/subzero` from cargo build,
-  and set `CARGO_NET_GIT_FETCH_WITH_CLI=true` env variable to use git CLI
-
-inputs:
-  token:
-    description: 'GitHub token with access to neondatabase/subzero'
-    required: true
-
-runs:
-  using: "composite"
-
-  steps:
-    - name: Set git token for neondatabase/subzero
-      uses: pyTooling/Actions/with-post-step@2307b526df64d55e95884e072e49aac2a00a9afa # v5.1.0
-      env:
-        SUBZERO_ACCESS_TOKEN: ${{ inputs.token }}
-      with:
-        main: |
-          git config --global url."https://x-access-token:${SUBZERO_ACCESS_TOKEN}@github.com/neondatabase/subzero".insteadOf "https://github.com/neondatabase/subzero"
-          cargo add -p proxy subzero-core --git https://github.com/neondatabase/subzero --rev 396264617e78e8be428682f87469bb25429af88a
-        post: |
-          git config --global --unset url."https://x-access-token:${SUBZERO_ACCESS_TOKEN}@github.com/neondatabase/subzero".insteadOf "https://github.com/neondatabase/subzero"
-
-    - name: Set `CARGO_NET_GIT_FETCH_WITH_CLI=true` env variable
-      shell: bash -euxo pipefail {0}
-      run: echo "CARGO_NET_GIT_FETCH_WITH_CLI=true" >> ${GITHUB_ENV}

.github/workflows/_build-and-test-locally.yml

@@ -86,10 +86,6 @@ jobs:
         with:
           submodules: true
 
-      - uses: ./.github/actions/prepare-for-subzero
-        with:
-          token: ${{ secrets.CI_ACCESS_TOKEN }}
-
       - name: Set pg 14 revision for caching
         id: pg_v14_rev
         run: echo pg_rev=$(git rev-parse HEAD:vendor/postgres-v14) >> $GITHUB_OUTPUT
@@ -120,7 +116,7 @@ jobs:
           ARCH: ${{ inputs.arch }}
           SANITIZERS: ${{ inputs.sanitizers }}
         run: |
-          CARGO_FLAGS="--locked --features testing,rest_broker"
+          CARGO_FLAGS="--locked --features testing"
           if [[ $BUILD_TYPE == "debug" && $ARCH == 'x64' ]]; then
             cov_prefix="scripts/coverage --profraw-prefix=$GITHUB_JOB --dir=/tmp/coverage run"
             CARGO_PROFILE=""

.github/workflows/_check-codestyle-rust.yml

@@ -46,10 +46,6 @@ jobs:
         uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
         with:
           submodules: true
-      
-      - uses: ./.github/actions/prepare-for-subzero
-        with:
-          token: ${{ secrets.CI_ACCESS_TOKEN }}
 
       - name: Cache cargo deps
         uses: tespkg/actions-cache@b7bf5fcc2f98a52ac6080eb0fd282c2f752074b1  # v1.8.0

.github/workflows/build-macos.yml

@@ -54,10 +54,6 @@ jobs:
         uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
         with:
           submodules: true
-      
-      - uses: ./.github/actions/prepare-for-subzero
-        with:
-          token: ${{ secrets.CI_ACCESS_TOKEN }}
 
       - name: Install build dependencies
         run: |

.github/workflows/build_and_test.yml

@@ -632,8 +632,6 @@ jobs:
             BUILD_TAG=${{ needs.meta.outputs.release-tag || needs.meta.outputs.build-tag }}
             TAG=${{ needs.build-build-tools-image.outputs.image-tag }}-bookworm
             DEBIAN_VERSION=bookworm
-          secrets: |
-            SUBZERO_ACCESS_TOKEN=${{ secrets.CI_ACCESS_TOKEN }}
           provenance: false
           push: true
           pull: true

.github/workflows/neon_extra_builds.yml

@@ -72,7 +72,6 @@ jobs:
   check-macos-build:
     needs: [ check-permissions, files-changed ]
     uses: ./.github/workflows/build-macos.yml
-    secrets: inherit
     with:
       pg_versions: ${{ needs.files-changed.outputs.postgres_changes }}
       rebuild_rust_code: ${{ fromJSON(needs.files-changed.outputs.rebuild_rust_code) }}

.gitignore

@@ -26,14 +26,9 @@ docker-compose/docker-compose-parallel.yml
 *.o
 *.so
 *.Po
-*.pid
 
 # pgindent typedef lists
 *.list
 
 # Node
 **/node_modules/
-
-# various files for local testing
-/proxy/.subzero
-local_proxy.json

Cargo.toml

@@ -49,7 +49,6 @@ members = [
     "libs/proxy/tokio-postgres2",
     "endpoint_storage",
     "pgxn/neon/communicator",
-    "proxy/subzero_core",
 ]
 
 [workspace.package]
@@ -131,7 +130,6 @@ jemalloc_pprof = { version = "0.7", features = ["symbolize", "flamegraph"] }
 jsonwebtoken = "9"
 lasso = "0.7"
 libc = "0.2"
-lock_api = "0.4.13"
 md5 = "0.7.0"
 measured = { version = "0.0.22", features=["lasso"] }
 measured-process = { version = "0.0.22" }
@@ -143,10 +141,10 @@ notify = "6.0.0"
 num_cpus = "1.15"
 num-traits = "0.2.19"
 once_cell = "1.13"
-opentelemetry = "0.30"
-opentelemetry_sdk = "0.30"
-opentelemetry-otlp = { version = "0.30", default-features = false, features = ["http-proto", "trace", "http", "reqwest-blocking-client"] }
-opentelemetry-semantic-conventions = "0.30"
+opentelemetry = "0.27"
+opentelemetry_sdk = "0.27"
+opentelemetry-otlp = { version = "0.27", default-features = false, features = ["http-proto", "trace", "http", "reqwest-client"] }
+opentelemetry-semantic-conventions = "0.27"
 parking_lot = "0.12"
 parquet = { version = "53", default-features = false, features = ["zstd"] }
 parquet_derive = "53"
@@ -158,18 +156,16 @@ procfs = "0.16"
 prometheus = {version = "0.13", default-features=false, features = ["process"]} # removes protobuf dependency
 prost = "0.13.5"
 prost-types = "0.13.5"
-rand = "0.9"
-# Remove after p256 is updated to 0.14.
-rand_core = "=0.6"
+rand = "0.8"
 redis = { version = "0.29.2", features = ["tokio-rustls-comp", "keep-alive"] }
 regex = "1.10.2"
 reqwest = { version = "0.12", default-features = false, features = ["rustls-tls"] }
-reqwest-tracing = { version = "0.5", features = ["opentelemetry_0_30"] }
+reqwest-tracing = { version = "0.5", features = ["opentelemetry_0_27"] }
 reqwest-middleware = "0.4"
 reqwest-retry = "0.7"
 routerify = "3"
 rpds = "0.13"
-rustc-hash = "2.1.1"
+rustc-hash = "1.1.0"
 rustls = { version = "0.23.16", default-features = false }
 rustls-pemfile = "2"
 rustls-pki-types = "1.11"
@@ -214,15 +210,17 @@ tonic = { version = "0.13.1", default-features = false, features = ["channel", "
 tonic-reflection = { version = "0.13.1", features = ["server"] }
 tower = { version = "0.5.2", default-features = false }
 tower-http = { version = "0.6.2", features = ["auth", "request-id", "trace"] }
-tower-otel = { version = "0.6", features = ["axum"] }
+
+# This revision uses opentelemetry 0.27. There's no tag for it.
+tower-otel = { git = "https://github.com/mattiapenati/tower-otel", rev = "56a7321053bcb72443888257b622ba0d43a11fcd" }
+
 tower-service = "0.3.3"
 tracing = "0.1"
 tracing-error = "0.2"
 tracing-log = "0.2"
-tracing-opentelemetry = "0.31"
+tracing-opentelemetry = "0.28"
 tracing-serde = "0.2.0"
 tracing-subscriber = { version = "0.3", default-features = false, features = ["smallvec", "fmt", "tracing-log", "std", "env-filter", "json"] }
-tracing-appender = "0.2.3"
 try-lock = "0.2.5"
 test-log = { version = "0.2.17", default-features = false, features = ["log"] }
 twox-hash = { version = "1.6.3", default-features = false }

Dockerfile

@@ -63,14 +63,7 @@ WORKDIR /home/nonroot
 
 COPY --chown=nonroot . .
 
-RUN --mount=type=secret,uid=1000,id=SUBZERO_ACCESS_TOKEN \
-    set -e \
-    && if [ -s /run/secrets/SUBZERO_ACCESS_TOKEN ]; then \
-        export CARGO_NET_GIT_FETCH_WITH_CLI=true && \
-        git config --global url."https://$(cat /run/secrets/SUBZERO_ACCESS_TOKEN)@github.com/neondatabase/subzero".insteadOf "https://github.com/neondatabase/subzero" && \
-        cargo add -p proxy subzero-core --git https://github.com/neondatabase/subzero --rev 396264617e78e8be428682f87469bb25429af88a; \
-    fi \
-    && cargo chef prepare --recipe-path recipe.json
+RUN cargo chef prepare --recipe-path recipe.json
 
 # Main build image
 FROM $REPOSITORY/$IMAGE:$TAG AS build
@@ -78,33 +71,20 @@ WORKDIR /home/nonroot
 ARG GIT_VERSION=local
 ARG BUILD_TAG
 ARG ADDITIONAL_RUSTFLAGS=""
-ENV CARGO_FEATURES="default"
 
 # 3. Build cargo dependencies. Note that this step doesn't depend on anything else than
 # `recipe.json`, so the layer can be reused as long as none of the dependencies change.
 COPY --from=plan     /home/nonroot/recipe.json                              recipe.json
-RUN --mount=type=secret,uid=1000,id=SUBZERO_ACCESS_TOKEN \
-    set -e \
-    && if [ -s /run/secrets/SUBZERO_ACCESS_TOKEN ]; then \
-        export CARGO_NET_GIT_FETCH_WITH_CLI=true && \
-        git config --global url."https://$(cat /run/secrets/SUBZERO_ACCESS_TOKEN)@github.com/neondatabase/subzero".insteadOf "https://github.com/neondatabase/subzero"; \
-    fi \
+RUN set -e \
     && RUSTFLAGS="-Clinker=clang -Clink-arg=-fuse-ld=mold -Clink-arg=-Wl,--no-rosegment -Cforce-frame-pointers=yes ${ADDITIONAL_RUSTFLAGS}" cargo chef cook --locked --release --recipe-path recipe.json
 
 # Perform the main build. We reuse the Postgres build artifacts from the intermediate 'pg-build'
 # layer, and the cargo dependencies built in the previous step.
 COPY --chown=nonroot --from=pg-build /home/nonroot/pg_install/ pg_install
 COPY --chown=nonroot . .
-COPY --chown=nonroot --from=plan     /home/nonroot/proxy/Cargo.toml         proxy/Cargo.toml
-COPY --chown=nonroot --from=plan     /home/nonroot/Cargo.lock               Cargo.lock
 
-RUN  --mount=type=secret,uid=1000,id=SUBZERO_ACCESS_TOKEN \
-    set -e \
-    && if [ -s /run/secrets/SUBZERO_ACCESS_TOKEN ]; then \
-        export CARGO_FEATURES="rest_broker"; \
-    fi \
+RUN set -e \
     && RUSTFLAGS="-Clinker=clang -Clink-arg=-fuse-ld=mold -Clink-arg=-Wl,--no-rosegment -Cforce-frame-pointers=yes ${ADDITIONAL_RUSTFLAGS}" cargo build \
-      --features $CARGO_FEATURES \
       --bin pg_sni_router  \
       --bin pageserver  \
       --bin pagectl  \

Makefile

@@ -2,7 +2,7 @@ ROOT_PROJECT_DIR := $(dir $(abspath $(lastword $(MAKEFILE_LIST))))
 
 # Where to install Postgres, default is ./pg_install, maybe useful for package
 # managers.
-POSTGRES_INSTALL_DIR ?= $(ROOT_PROJECT_DIR)/pg_install
+POSTGRES_INSTALL_DIR ?= $(ROOT_PROJECT_DIR)/pg_install/
 
 # Supported PostgreSQL versions
 POSTGRES_VERSIONS = v17 v16 v15 v14
@@ -14,7 +14,7 @@ POSTGRES_VERSIONS = v17 v16 v15 v14
 # it is derived from BUILD_TYPE.
 
 # All intermediate build artifacts are stored here.
-BUILD_DIR := $(ROOT_PROJECT_DIR)/build
+BUILD_DIR := build
 
 ICU_PREFIX_DIR := /usr/local/icu
 
@@ -212,7 +212,7 @@ neon-pgindent: postgres-v17-pg-bsd-indent neon-pg-ext-v17
 		FIND_TYPEDEF=$(ROOT_PROJECT_DIR)/vendor/postgres-v17/src/tools/find_typedef \
 		INDENT=$(BUILD_DIR)/v17/src/tools/pg_bsd_indent/pg_bsd_indent \
 		PGINDENT_SCRIPT=$(ROOT_PROJECT_DIR)/vendor/postgres-v17/src/tools/pgindent/pgindent \
-		-C $(BUILD_DIR)/pgxn-v17/neon \
+		-C $(BUILD_DIR)/neon-v17 \
 		-f $(ROOT_PROJECT_DIR)/pgxn/neon/Makefile pgindent
 
 

compute/compute-node.Dockerfile

@@ -133,7 +133,7 @@ RUN case $DEBIAN_VERSION in \
       # Install newer version (3.25) from backports.
       # libstdc++-10-dev is required for plv8
       bullseye) \
-        echo "deb http://archive.debian.org/debian bullseye-backports main" > /etc/apt/sources.list.d/bullseye-backports.list; \
+        echo "deb http://deb.debian.org/debian bullseye-backports main" > /etc/apt/sources.list.d/bullseye-backports.list; \
         VERSION_INSTALLS="cmake/bullseye-backports cmake-data/bullseye-backports libstdc++-10-dev"; \
       ;; \
       # Version-specific installs for Bookworm (PG17):

compute_tools/Cargo.toml

@@ -27,10 +27,7 @@ fail.workspace = true
 flate2.workspace = true
 futures.workspace = true
 http.workspace = true
-http-body-util.workspace = true
 hostname-validator = "1.1"
-hyper.workspace = true
-hyper-util.workspace = true
 indexmap.workspace = true
 itertools.workspace = true
 jsonwebtoken.workspace = true
@@ -47,7 +44,6 @@ postgres.workspace = true
 regex.workspace = true
 reqwest = { workspace = true, features = ["json"] }
 ring = "0.17"
-scopeguard.workspace = true
 serde.workspace = true
 serde_with.workspace = true
 serde_json.workspace = true
@@ -62,7 +58,6 @@ tokio-stream.workspace = true
 tonic.workspace = true
 tower-otel.workspace = true
 tracing.workspace = true
-tracing-appender.workspace = true
 tracing-opentelemetry.workspace = true
 tracing-subscriber.workspace = true
 tracing-utils.workspace = true

compute_tools/src/bin/compute_ctl.rs

@@ -51,7 +51,6 @@ use compute_tools::compute::{
 use compute_tools::extension_server::get_pg_version_string;
 use compute_tools::logger::*;
 use compute_tools::params::*;
-use compute_tools::pg_isready::get_pg_isready_bin;
 use compute_tools::spec::*;
 use rlimit::{Resource, setrlimit};
 use signal_hook::consts::{SIGINT, SIGQUIT, SIGTERM};
@@ -139,12 +138,6 @@ struct Cli {
     /// Run in development mode, skipping VM-specific operations like process termination
     #[arg(long, action = clap::ArgAction::SetTrue)]
     pub dev: bool,
-
-    #[arg(long)]
-    pub pg_init_timeout: Option<u64>,
-
-    #[arg(long, default_value_t = false, action = clap::ArgAction::Set)]
-    pub lakebase_mode: bool,
 }
 
 impl Cli {
@@ -195,12 +188,7 @@ fn main() -> Result<()> {
         .build()?;
     let _rt_guard = runtime.enter();
 
-    let mut log_dir = None;
-    if cli.lakebase_mode {
-        log_dir = std::env::var("COMPUTE_CTL_LOG_DIRECTORY").ok();
-    }
-
-    let (tracing_provider, _file_logs_guard) = init(cli.dev, log_dir)?;
+    runtime.block_on(init(cli.dev))?;
 
     // enable core dumping for all child processes
     setrlimit(Resource::CORE, rlimit::INFINITY, rlimit::INFINITY)?;
@@ -231,10 +219,6 @@ fn main() -> Result<()> {
             installed_extensions_collection_interval: Arc::new(AtomicU64::new(
                 cli.installed_extensions_collection_interval,
             )),
-            pg_init_timeout: cli.pg_init_timeout.map(Duration::from_secs),
-            pg_isready_bin: get_pg_isready_bin(&cli.pgbin),
-            instance_id: std::env::var("INSTANCE_ID").ok(),
-            lakebase_mode: cli.lakebase_mode,
         },
         config,
     )?;
@@ -243,17 +227,11 @@ fn main() -> Result<()> {
 
     scenario.teardown();
 
-    deinit_and_exit(tracing_provider, exit_code);
+    deinit_and_exit(exit_code);
 }
 
-fn init(
-    dev_mode: bool,
-    log_dir: Option<String>,
-) -> Result<(
-    Option<tracing_utils::Provider>,
-    Option<tracing_appender::non_blocking::WorkerGuard>,
-)> {
-    let (provider, file_logs_guard) = init_tracing_and_logging(DEFAULT_LOG_LEVEL, &log_dir)?;
+async fn init(dev_mode: bool) -> Result<()> {
+    init_tracing_and_logging(DEFAULT_LOG_LEVEL).await?;
 
     let mut signals = Signals::new([SIGINT, SIGTERM, SIGQUIT])?;
     thread::spawn(move || {
@@ -264,7 +242,7 @@ fn init(
 
     info!("compute build_tag: {}", &BUILD_TAG.to_string());
 
-    Ok((provider, file_logs_guard))
+    Ok(())
 }
 
 fn get_config(cli: &Cli) -> Result<ComputeConfig> {
@@ -289,27 +267,25 @@ fn get_config(cli: &Cli) -> Result<ComputeConfig> {
     }
 }
 
-fn deinit_and_exit(tracing_provider: Option<tracing_utils::Provider>, exit_code: Option<i32>) -> ! {
-    if let Some(p) = tracing_provider {
-        // Shutdown trace pipeline gracefully, so that it has a chance to send any
-        // pending traces before we exit. Shutting down OTEL tracing provider may
-        // hang for quite some time, see, for example:
-        // - https://github.com/open-telemetry/opentelemetry-rust/issues/868
-        // - and our problems with staging https://github.com/neondatabase/cloud/issues/3707#issuecomment-1493983636
-        //
-        // Yet, we want computes to shut down fast enough, as we may need a new one
-        // for the same timeline ASAP. So wait no longer than 2s for the shutdown to
-        // complete, then just error out and exit the main thread.
-        info!("shutting down tracing");
-        let (sender, receiver) = mpsc::channel();
-        let _ = thread::spawn(move || {
-            _ = p.shutdown();
-            sender.send(()).ok()
-        });
-        let shutdown_res = receiver.recv_timeout(Duration::from_millis(2000));
-        if shutdown_res.is_err() {
-            error!("timed out while shutting down tracing, exiting anyway");
-        }
+fn deinit_and_exit(exit_code: Option<i32>) -> ! {
+    // Shutdown trace pipeline gracefully, so that it has a chance to send any
+    // pending traces before we exit. Shutting down OTEL tracing provider may
+    // hang for quite some time, see, for example:
+    // - https://github.com/open-telemetry/opentelemetry-rust/issues/868
+    // - and our problems with staging https://github.com/neondatabase/cloud/issues/3707#issuecomment-1493983636
+    //
+    // Yet, we want computes to shut down fast enough, as we may need a new one
+    // for the same timeline ASAP. So wait no longer than 2s for the shutdown to
+    // complete, then just error out and exit the main thread.
+    info!("shutting down tracing");
+    let (sender, receiver) = mpsc::channel();
+    let _ = thread::spawn(move || {
+        tracing_utils::shutdown_tracing();
+        sender.send(()).ok()
+    });
+    let shutdown_res = receiver.recv_timeout(Duration::from_millis(2000));
+    if shutdown_res.is_err() {
+        error!("timed out while shutting down tracing, exiting anyway");
     }
 
     info!("shutting down");

compute_tools/src/communicator_socket_client.rs

@@ -1,98 +0,0 @@
-//! Client for making request to a running Postgres server's communicator control socket.
-//!
-//! The storage communicator process that runs inside Postgres exposes an HTTP endpoint in
-//! a Unix Domain Socket in the Postgres data directory. This provides access to it.
-
-use std::path::Path;
-
-use anyhow::Context;
-use hyper::client::conn::http1::SendRequest;
-use hyper_util::rt::TokioIo;
-
-/// Name of the socket within the Postgres data directory. This better match that in
-/// `pgxn/neon/communicator/src/lib.rs`.
-const NEON_COMMUNICATOR_SOCKET_NAME: &str = "neon-communicator.socket";
-
-/// Open a connection to the communicator's control socket, prepare to send requests to it
-/// with hyper.
-pub async fn connect_communicator_socket<B>(pgdata: &Path) -> anyhow::Result<SendRequest<B>>
-where
-    B: hyper::body::Body + 'static + Send,
-    B::Data: Send,
-    B::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
-{
-    let socket_path = pgdata.join(NEON_COMMUNICATOR_SOCKET_NAME);
-    let socket_path_len = socket_path.display().to_string().len();
-
-    // There is a limit of around 100 bytes (108 on Linux?) on the length of the path to a
-    // Unix Domain socket. The limit is on the connect(2) function used to open the
-    // socket, not on the absolute path itself. Postgres changes the current directory to
-    // the data directory and uses a relative path to bind to the socket, and the relative
-    // path "./neon-communicator.socket" is always short, but when compute_ctl needs to
-    // open the socket, we need to use a full path, which can be arbitrarily long.
-    //
-    // There are a few ways we could work around this:
-    //
-    // 1. Change the current directory to the Postgres data directory and use a relative
-    //    path in the connect(2) call. That's problematic because the current directory
-    //    applies to the whole process. We could change the current directory early in
-    //    compute_ctl startup, and that might be a good idea anyway for other reasons too:
-    //    it would be more robust if the data directory is moved around or unlinked for
-    //    some reason, and you would be less likely to accidentally litter other parts of
-    //    the filesystem with e.g. temporary files. However, that's a pretty invasive
-    //    change.
-    //
-    // 2. On Linux, you could open() the data directory, and refer to the the socket
-    //    inside it as "/proc/self/fd/<fd>/neon-communicator.socket". But that's
-    //    Linux-only.
-    //
-    // 3. Create a symbolic link to the socket with a shorter path, and use that.
-    //
-    // We use the symbolic link approach here. Hopefully the paths we use in production
-    // are shorter, so that we can open the socket directly, so that this hack is needed
-    // only in development.
-    let connect_result = if socket_path_len < 100 {
-        // We can open the path directly with no hacks.
-        tokio::net::UnixStream::connect(socket_path).await
-    } else {
-        // The path to the socket is too long. Create a symlink to it with a shorter path.
-        let short_path = std::env::temp_dir().join(format!(
-            "compute_ctl.short-socket.{}.{}",
-            std::process::id(),
-            tokio::task::id()
-        ));
-        std::os::unix::fs::symlink(&socket_path, &short_path)?;
-
-        // Delete the symlink as soon as we have connected to it. There's a small chance
-        // of leaking if the process dies before we remove it, so try to keep that window
-        // as small as possible.
-        scopeguard::defer! {
-            if let Err(err) = std::fs::remove_file(&short_path) {
-                tracing::warn!("could not remove symlink \"{}\" created for socket: {}",
-                               short_path.display(), err);
-            }
-        }
-
-        tracing::info!(
-            "created symlink \"{}\" for socket \"{}\", opening it now",
-            short_path.display(),
-            socket_path.display()
-        );
-
-        tokio::net::UnixStream::connect(&short_path).await
-    };
-
-    let stream = connect_result.context("connecting to communicator control socket")?;
-
-    let io = TokioIo::new(stream);
-    let (request_sender, connection) = hyper::client::conn::http1::handshake(io).await?;
-
-    // spawn a task to poll the connection and drive the HTTP state
-    tokio::spawn(async move {
-        if let Err(err) = connection.await {
-            eprintln!("Error in connection: {err}");
-        }
-    });
-
-    Ok(request_sender)
-}

compute_tools/src/compute.rs

@@ -113,13 +113,6 @@ pub struct ComputeNodeParams {
 
     /// Interval for installed extensions collection
     pub installed_extensions_collection_interval: Arc<AtomicU64>,
-    /// Hadron instance ID of the compute node.
-    pub instance_id: Option<String>,
-    /// Timeout of PG compute startup in the Init state.
-    pub pg_init_timeout: Option<Duration>,
-    // Path to the `pg_isready` binary.
-    pub pg_isready_bin: String,
-    pub lakebase_mode: bool,
 }
 
 type TaskHandle = Mutex<Option<JoinHandle<()>>>;
@@ -161,7 +154,6 @@ pub struct RemoteExtensionMetrics {
 #[derive(Clone, Debug)]
 pub struct ComputeState {
     pub start_time: DateTime<Utc>,
-    pub pg_start_time: Option<DateTime<Utc>>,
     pub status: ComputeStatus,
     /// Timestamp of the last Postgres activity. It could be `None` if
     /// compute wasn't used since start.
@@ -199,7 +191,6 @@ impl ComputeState {
     pub fn new() -> Self {
         Self {
             start_time: Utc::now(),
-            pg_start_time: None,
             status: ComputeStatus::Empty,
             last_active: None,
             error: None,
@@ -488,7 +479,6 @@ impl ComputeNode {
             port: this.params.external_http_port,
             config: this.compute_ctl_config.clone(),
             compute_id: this.params.compute_id.clone(),
-            instance_id: this.params.instance_id.clone(),
         }
         .launch(&this);
 
@@ -658,9 +648,6 @@ impl ComputeNode {
             };
             _this_entered = start_compute_span.enter();
 
-            // Hadron: Record postgres start time (used to enforce pg_init_timeout).
-            state_guard.pg_start_time.replace(Utc::now());
-
             state_guard.set_status(ComputeStatus::Init, &self.state_changed);
             compute_state = state_guard.clone()
         }
@@ -1454,7 +1441,7 @@ impl ComputeNode {
         })?;
 
         // Update pg_hba.conf received with basebackup.
-        update_pg_hba(pgdata_path, None)?;
+        update_pg_hba(pgdata_path)?;
 
         // Place pg_dynshmem under /dev/shm. This allows us to use
         // 'dynamic_shared_memory_type = mmap' so that the files are placed in
@@ -1759,7 +1746,6 @@ impl ComputeNode {
         }
 
         // Run migrations separately to not hold up cold starts
-        let lakebase_mode = self.params.lakebase_mode;
         let params = self.params.clone();
         tokio::spawn(async move {
             let mut conf = conf.as_ref().clone();
@@ -1772,7 +1758,7 @@ impl ComputeNode {
                             eprintln!("connection error: {e}");
                         }
                     });
-                    if let Err(e) = handle_migrations(params, &mut client, lakebase_mode).await {
+                    if let Err(e) = handle_migrations(params, &mut client).await {
                         error!("Failed to run migrations: {}", e);
                     }
                 }
@@ -1788,34 +1774,6 @@ impl ComputeNode {
         Ok::<(), anyhow::Error>(())
     }
 
-    // Signal to the configurator to refresh the configuration by pulling a new spec from the HCC.
-    // Note that this merely triggers a notification on a condition variable the configurator thread
-    // waits on. The configurator thread (in configurator.rs) pulls the new spec from the HCC and
-    // applies it.
-    pub async fn signal_refresh_configuration(&self) -> Result<()> {
-        let states_allowing_configuration_refresh = [
-            ComputeStatus::Running,
-            ComputeStatus::Failed,
-            // ComputeStatus::RefreshConfigurationPending,
-        ];
-
-        let state = self.state.lock().expect("state lock poisoned");
-        if states_allowing_configuration_refresh.contains(&state.status) {
-            // state.status = ComputeStatus::RefreshConfigurationPending;
-            self.state_changed.notify_all();
-            Ok(())
-        } else if state.status == ComputeStatus::Init {
-            // If the compute is in Init state, we can't refresh the configuration immediately,
-            // but we should be able to do that soon.
-            Ok(())
-        } else {
-            Err(anyhow::anyhow!(
-                "Cannot refresh compute configuration in state {:?}",
-                state.status
-            ))
-        }
-    }
-
     // Wrapped this around `pg_ctl reload`, but right now we don't use
     // `pg_ctl` for start / stop.
     #[instrument(skip_all)]
@@ -2492,31 +2450,14 @@ LIMIT 100",
     pub fn spawn_lfc_offload_task(self: &Arc<Self>, interval: Duration) {
         self.terminate_lfc_offload_task();
         let secs = interval.as_secs();
+        info!("spawning lfc offload worker with {secs}s interval");
         let this = self.clone();
-
-        info!("spawning LFC offload worker with {secs}s interval");
         let handle = spawn(async move {
             let mut interval = time::interval(interval);
             interval.tick().await; // returns immediately
             loop {
                 interval.tick().await;
-
-                let prewarm_state = this.state.lock().unwrap().lfc_prewarm_state.clone();
-                // Do not offload LFC state if we are currently prewarming or any issue occurred.
-                // If we'd do that, we might override the LFC state in endpoint storage with some
-                // incomplete state. Imagine a situation:
-                // 1. Endpoint started with `autoprewarm: true`
-                // 2. While prewarming is not completed, we upload the new incomplete state
-                // 3. Compute gets interrupted and restarts
-                // 4. We start again and try to prewarm with the state from 2. instead of the previous complete state
-                if matches!(
-                    prewarm_state,
-                    LfcPrewarmState::Completed
-                        | LfcPrewarmState::NotPrewarmed
-                        | LfcPrewarmState::Skipped
-                ) {
-                    this.offload_lfc_async().await;
-                }
+                this.offload_lfc_async().await;
             }
         });
         *self.lfc_offload_task.lock().unwrap() = Some(handle);

compute_tools/src/compute_prewarm.rs

@@ -89,8 +89,7 @@ impl ComputeNode {
         self.state.lock().unwrap().lfc_offload_state.clone()
     }
 
-    /// If there is a prewarm request ongoing, return `false`, `true` otherwise.
-    /// Has a failpoint "compute-prewarm"
+    /// If there is a prewarm request ongoing, return false, true otherwise
     pub fn prewarm_lfc(self: &Arc<Self>, from_endpoint: Option<String>) -> bool {
         {
             let state = &mut self.state.lock().unwrap().lfc_prewarm_state;
@@ -102,24 +101,15 @@ impl ComputeNode {
 
         let cloned = self.clone();
         spawn(async move {
-            let state = match cloned.prewarm_impl(from_endpoint).await {
-                Ok(true) => LfcPrewarmState::Completed,
-                Ok(false) => {
-                    info!(
-                        "skipping LFC prewarm because LFC state is not found in endpoint storage"
-                    );
-                    LfcPrewarmState::Skipped
-                }
-                Err(err) => {
-                    crate::metrics::LFC_PREWARM_ERRORS.inc();
-                    error!(%err, "could not prewarm LFC");
-                    LfcPrewarmState::Failed {
-                        error: format!("{err:#}"),
-                    }
-                }
+            let Err(err) = cloned.prewarm_impl(from_endpoint).await else {
+                cloned.state.lock().unwrap().lfc_prewarm_state = LfcPrewarmState::Completed;
+                return;
+            };
+            crate::metrics::LFC_PREWARM_ERRORS.inc();
+            error!(%err, "prewarming lfc");
+            cloned.state.lock().unwrap().lfc_prewarm_state = LfcPrewarmState::Failed {
+                error: err.to_string(),
             };
-
-            cloned.state.lock().unwrap().lfc_prewarm_state = state;
         });
         true
     }
@@ -130,25 +120,15 @@ impl ComputeNode {
         EndpointStoragePair::from_spec_and_endpoint(state.pspec.as_ref().unwrap(), from_endpoint)
     }
 
-    /// Request LFC state from endpoint storage and load corresponding pages into Postgres.
-    /// Returns a result with `false` if the LFC state is not found in endpoint storage.
-    async fn prewarm_impl(&self, from_endpoint: Option<String>) -> Result<bool> {
+    async fn prewarm_impl(&self, from_endpoint: Option<String>) -> Result<()> {
         let EndpointStoragePair { url, token } = self.endpoint_storage_pair(from_endpoint)?;
-
-        #[cfg(feature = "testing")]
-        fail::fail_point!("compute-prewarm", |_| {
-            bail!("prewarm configured to fail because of a failpoint")
-        });
-
         info!(%url, "requesting LFC state from endpoint storage");
+
         let request = Client::new().get(&url).bearer_auth(token);
         let res = request.send().await.context("querying endpoint storage")?;
-        match res.status() {
-            StatusCode::OK => (),
-            StatusCode::NOT_FOUND => {
-                return Ok(false);
-            }
-            status => bail!("{status} querying endpoint storage"),
+        let status = res.status();
+        if status != StatusCode::OK {
+            bail!("{status} querying endpoint storage")
         }
 
         let mut uncompressed = Vec::new();
@@ -161,8 +141,7 @@ impl ComputeNode {
             .await
             .context("decoding LFC state")?;
         let uncompressed_len = uncompressed.len();
-
-        info!(%url, "downloaded LFC state, uncompressed size {uncompressed_len}, loading into Postgres");
+        info!(%url, "downloaded LFC state, uncompressed size {uncompressed_len}, loading into postgres");
 
         ComputeNode::get_maintenance_client(&self.tokio_conn_conf)
             .await
@@ -170,9 +149,7 @@ impl ComputeNode {
             .query_one("select neon.prewarm_local_cache($1)", &[&uncompressed])
             .await
             .context("loading LFC state into postgres")
-            .map(|_| ())?;
-
-        Ok(true)
+            .map(|_| ())
     }
 
     /// If offload request is ongoing, return false, true otherwise
@@ -200,53 +177,41 @@ impl ComputeNode {
 
     async fn offload_lfc_with_state_update(&self) {
         crate::metrics::LFC_OFFLOADS.inc();
-
         let Err(err) = self.offload_lfc_impl().await else {
             self.state.lock().unwrap().lfc_offload_state = LfcOffloadState::Completed;
             return;
         };
-
         crate::metrics::LFC_OFFLOAD_ERRORS.inc();
-        error!(%err, "could not offload LFC state to endpoint storage");
+        error!(%err, "offloading lfc");
         self.state.lock().unwrap().lfc_offload_state = LfcOffloadState::Failed {
-            error: format!("{err:#}"),
+            error: err.to_string(),
         };
     }
 
     async fn offload_lfc_impl(&self) -> Result<()> {
         let EndpointStoragePair { url, token } = self.endpoint_storage_pair(None)?;
-        info!(%url, "requesting LFC state from Postgres");
+        info!(%url, "requesting LFC state from postgres");
 
-        let row = ComputeNode::get_maintenance_client(&self.tokio_conn_conf)
+        let mut compressed = Vec::new();
+        ComputeNode::get_maintenance_client(&self.tokio_conn_conf)
             .await
             .context("connecting to postgres")?
             .query_one("select neon.get_local_cache_state()", &[])
             .await
-            .context("querying LFC state")?;
-        let state = row
-            .try_get::<usize, Option<&[u8]>>(0)
-            .context("deserializing LFC state")?;
-        let Some(state) = state else {
-            info!(%url, "empty LFC state, not exporting");
-            return Ok(());
-        };
-
-        let mut compressed = Vec::new();
-        ZstdEncoder::new(state)
+            .context("querying LFC state")?
+            .try_get::<usize, &[u8]>(0)
+            .context("deserializing LFC state")
+            .map(ZstdEncoder::new)?
             .read_to_end(&mut compressed)
             .await
             .context("compressing LFC state")?;
-
         let compressed_len = compressed.len();
         info!(%url, "downloaded LFC state, compressed size {compressed_len}, writing to endpoint storage");
 
         let request = Client::new().put(url).bearer_auth(token).body(compressed);
         match request.send().await {
             Ok(res) if res.status() == StatusCode::OK => Ok(()),
-            Ok(res) => bail!(
-                "Request to endpoint storage failed with status: {}",
-                res.status()
-            ),
+            Ok(res) => bail!("Error writing to endpoint storage: {}", res.status()),
             Err(err) => Err(err).context("writing to endpoint storage"),
         }
     }

compute_tools/src/compute_promote.rs

@@ -1,12 +1,11 @@
 use crate::compute::ComputeNode;
 use anyhow::{Context, Result, bail};
-use compute_api::responses::{LfcPrewarmState, PromoteConfig, PromoteState};
-use compute_api::spec::ComputeMode;
-use itertools::Itertools;
-use std::collections::HashMap;
+use compute_api::{
+    responses::{LfcPrewarmState, PromoteState, SafekeepersLsn},
+    spec::ComputeMode,
+};
 use std::{sync::Arc, time::Duration};
 use tokio::time::sleep;
-use tracing::info;
 use utils::lsn::Lsn;
 
 impl ComputeNode {
@@ -14,22 +13,21 @@ impl ComputeNode {
     /// and http client disconnects, this does not stop promotion, and subsequent
     /// calls block until promote finishes.
     /// Called by control plane on secondary after primary endpoint is terminated
-    /// Has a failpoint "compute-promotion"
-    pub async fn promote(self: &Arc<Self>, cfg: PromoteConfig) -> PromoteState {
+    pub async fn promote(self: &Arc<Self>, safekeepers_lsn: SafekeepersLsn) -> PromoteState {
         let cloned = self.clone();
-        let promote_fn = async move || {
-            let Err(err) = cloned.promote_impl(cfg).await else {
-                return PromoteState::Completed;
-            };
-            tracing::error!(%err, "promoting");
-            PromoteState::Failed {
-                error: format!("{err:#}"),
-            }
-        };
-
         let start_promotion = || {
             let (tx, rx) = tokio::sync::watch::channel(PromoteState::NotPromoted);
-            tokio::spawn(async move { tx.send(promote_fn().await) });
+            tokio::spawn(async move {
+                tx.send(match cloned.promote_impl(safekeepers_lsn).await {
+                    Ok(_) => PromoteState::Completed,
+                    Err(err) => {
+                        tracing::error!(%err, "promoting");
+                        PromoteState::Failed {
+                            error: err.to_string(),
+                        }
+                    }
+                })
+            });
             rx
         };
 
@@ -49,7 +47,9 @@ impl ComputeNode {
         task.borrow().clone()
     }
 
-    async fn promote_impl(&self, mut cfg: PromoteConfig) -> Result<()> {
+    // Why do we have to supply safekeepers?
+    // For secondary we use primary_connection_conninfo so safekeepers field is empty
+    async fn promote_impl(&self, safekeepers_lsn: SafekeepersLsn) -> Result<()> {
         {
             let state = self.state.lock().unwrap();
             let mode = &state.pspec.as_ref().unwrap().spec.mode;
@@ -73,7 +73,7 @@ impl ComputeNode {
             .await
             .context("connecting to postgres")?;
 
-        let primary_lsn = cfg.wal_flush_lsn;
+        let primary_lsn = safekeepers_lsn.wal_flush_lsn;
         let mut last_wal_replay_lsn: Lsn = Lsn::INVALID;
         const RETRIES: i32 = 20;
         for i in 0..=RETRIES {
@@ -86,7 +86,7 @@ impl ComputeNode {
             if last_wal_replay_lsn >= primary_lsn {
                 break;
             }
-            info!("Try {i}, replica lsn {last_wal_replay_lsn}, primary lsn {primary_lsn}");
+            tracing::info!("Try {i}, replica lsn {last_wal_replay_lsn}, primary lsn {primary_lsn}");
             sleep(Duration::from_secs(1)).await;
         }
         if last_wal_replay_lsn < primary_lsn {
@@ -96,7 +96,7 @@ impl ComputeNode {
         // using $1 doesn't work with ALTER SYSTEM SET
         let safekeepers_sql = format!(
             "ALTER SYSTEM SET neon.safekeepers='{}'",
-            cfg.spec.safekeeper_connstrings.join(",")
+            safekeepers_lsn.safekeepers
         );
         client
             .query(&safekeepers_sql, &[])
@@ -106,12 +106,6 @@ impl ComputeNode {
             .query("SELECT pg_reload_conf()", &[])
             .await
             .context("reloading postgres config")?;
-
-        #[cfg(feature = "testing")]
-        fail::fail_point!("compute-promotion", |_| {
-            bail!("promotion configured to fail because of a failpoint")
-        });
-
         let row = client
             .query_one("SELECT * FROM pg_promote()", &[])
             .await
@@ -131,36 +125,8 @@ impl ComputeNode {
             bail!("replica in read only mode after promotion");
         }
 
-        {
-            let mut state = self.state.lock().unwrap();
-            let spec = &mut state.pspec.as_mut().unwrap().spec;
-            spec.mode = ComputeMode::Primary;
-            let new_conf = cfg.spec.cluster.postgresql_conf.as_mut().unwrap();
-            let existing_conf = spec.cluster.postgresql_conf.as_ref().unwrap();
-            Self::merge_spec(new_conf, existing_conf);
-        }
-        info!("applied new spec, reconfiguring as primary");
-        self.reconfigure()
-    }
-
-    /// Merge old and new Postgres conf specs to apply on secondary.
-    /// Change new spec's port and safekeepers since they are supplied
-    /// differenly
-    fn merge_spec(new_conf: &mut String, existing_conf: &str) {
-        let mut new_conf_set: HashMap<&str, &str> = new_conf
-            .split_terminator('\n')
-            .map(|e| e.split_once("=").expect("invalid item"))
-            .collect();
-        new_conf_set.remove("neon.safekeepers");
-
-        let existing_conf_set: HashMap<&str, &str> = existing_conf
-            .split_terminator('\n')
-            .map(|e| e.split_once("=").expect("invalid item"))
-            .collect();
-        new_conf_set.insert("port", existing_conf_set["port"]);
-        *new_conf = new_conf_set
-            .iter()
-            .map(|(k, v)| format!("{k}={v}"))
-            .join("\n");
+        let mut state = self.state.lock().unwrap();
+        state.pspec.as_mut().unwrap().spec.mode = ComputeMode::Primary;
+        Ok(())
     }
 }

compute_tools/src/config.rs

@@ -56,15 +56,14 @@ pub fn write_postgres_conf(
         writeln!(file, "{conf}")?;
     }
 
-    // Stripe size GUC should be defined prior to connection string
-    if let Some(stripe_size) = spec.shard_stripe_size {
-        writeln!(file, "neon.stripe_size={stripe_size}")?;
-    }
     // Add options for connecting to storage
     writeln!(file, "# Neon storage settings")?;
     if let Some(s) = &spec.pageserver_connstring {
         writeln!(file, "neon.pageserver_connstring={}", escape_conf_value(s))?;
     }
+    if let Some(stripe_size) = spec.shard_stripe_size {
+        writeln!(file, "neon.stripe_size={stripe_size}")?;
+    }
     if !spec.safekeeper_connstrings.is_empty() {
         let mut neon_safekeepers_value = String::new();
         tracing::info!(

compute_tools/src/hadron_metrics.rs

@@ -1,60 +0,0 @@
-use metrics::{
-    IntCounter, IntGaugeVec, core::Collector, proto::MetricFamily, register_int_counter,
-    register_int_gauge_vec,
-};
-use once_cell::sync::Lazy;
-
-// Counter keeping track of the number of PageStream request errors reported by Postgres.
-// An error is registered every time Postgres calls compute_ctl's /refresh_configuration API.
-// Postgres will invoke this API if it detected trouble with PageStream requests (get_page@lsn,
-// get_base_backup, etc.) it sends to any pageserver. An increase in this counter value typically
-// indicates Postgres downtime, as PageStream requests are critical for Postgres to function.
-pub static POSTGRES_PAGESTREAM_REQUEST_ERRORS: Lazy<IntCounter> = Lazy::new(|| {
-    register_int_counter!(
-        "pg_cctl_pagestream_request_errors_total",
-        "Number of PageStream request errors reported by the postgres process"
-    )
-    .expect("failed to define a metric")
-});
-
-// Counter keeping track of the number of compute configuration errors due to Postgres statement
-// timeouts. An error is registered every time `ComputeNode::reconfigure()` fails due to Postgres
-// error code 57014 (query cancelled). This statement timeout typically occurs when postgres is
-// stuck in a problematic retry loop when the PS is reject its connection requests (usually due
-// to PG pointing at the wrong PS). We should investigate the root cause when this counter value
-// increases by checking PG and PS logs.
-pub static COMPUTE_CONFIGURE_STATEMENT_TIMEOUT_ERRORS: Lazy<IntCounter> = Lazy::new(|| {
-    register_int_counter!(
-        "pg_cctl_configure_statement_timeout_errors_total",
-        "Number of compute configuration errors due to Postgres statement timeouts."
-    )
-    .expect("failed to define a metric")
-});
-
-pub static COMPUTE_ATTACHED: Lazy<IntGaugeVec> = Lazy::new(|| {
-    register_int_gauge_vec!(
-        "pg_cctl_attached",
-        "Compute node attached status (1 if attached)",
-        &[
-            "pg_compute_id",
-            "pg_instance_id",
-            "tenant_id",
-            "timeline_id"
-        ]
-    )
-    .expect("failed to define a metric")
-});
-
-pub fn collect() -> Vec<MetricFamily> {
-    let mut metrics = Vec::new();
-    metrics.extend(POSTGRES_PAGESTREAM_REQUEST_ERRORS.collect());
-    metrics.extend(COMPUTE_CONFIGURE_STATEMENT_TIMEOUT_ERRORS.collect());
-    metrics.extend(COMPUTE_ATTACHED.collect());
-    metrics
-}
-
-pub fn initialize_metrics() {
-    Lazy::force(&POSTGRES_PAGESTREAM_REQUEST_ERRORS);
-    Lazy::force(&COMPUTE_CONFIGURE_STATEMENT_TIMEOUT_ERRORS);
-    Lazy::force(&COMPUTE_ATTACHED);
-}

compute_tools/src/http/middleware/authorize.rs

@@ -16,29 +16,13 @@ use crate::http::JsonResponse;
 #[derive(Clone, Debug)]
 pub(in crate::http) struct Authorize {
     compute_id: String,
-    // BEGIN HADRON
-    // Hadron instance ID. Only set if it's a Lakebase V1 a.k.a. Hadron instance.
-    instance_id: Option<String>,
-    // END HADRON
     jwks: JwkSet,
     validation: Validation,
 }
 
 impl Authorize {
-    pub fn new(compute_id: String, instance_id: Option<String>, jwks: JwkSet) -> Self {
+    pub fn new(compute_id: String, jwks: JwkSet) -> Self {
         let mut validation = Validation::new(Algorithm::EdDSA);
-
-        // BEGIN HADRON
-        let use_rsa = jwks.keys.iter().any(|jwk| {
-            jwk.common
-                .key_algorithm
-                .is_some_and(|alg| alg == jsonwebtoken::jwk::KeyAlgorithm::RS256)
-        });
-        if use_rsa {
-            validation = Validation::new(Algorithm::RS256);
-        }
-        // END HADRON
-
         validation.validate_exp = true;
         // Unused by the control plane
         validation.validate_nbf = false;
@@ -50,7 +34,6 @@ impl Authorize {
 
         Self {
             compute_id,
-            instance_id,
             jwks,
             validation,
         }
@@ -64,20 +47,10 @@ impl AsyncAuthorizeRequest<Body> for Authorize {
 
     fn authorize(&mut self, mut request: Request<Body>) -> Self::Future {
         let compute_id = self.compute_id.clone();
-        let is_hadron_instance = self.instance_id.is_some();
         let jwks = self.jwks.clone();
         let validation = self.validation.clone();
 
         Box::pin(async move {
-            // BEGIN HADRON
-            // In Hadron deployments the "external" HTTP endpoint on compute_ctl can only be
-            // accessed by trusted components (enforced by dblet network policy), so we can bypass
-            // all auth here.
-            if is_hadron_instance {
-                return Ok(request);
-            }
-            // END HADRON
-
             let TypedHeader(Authorization(bearer)) = request
                 .extract_parts::<TypedHeader<Authorization<Bearer>>>()
                 .await

compute_tools/src/http/openapi_spec.yaml

@@ -96,7 +96,7 @@ paths:
         content:
           application/json:
             schema:
-              $ref: "#/components/schemas/ComputeSchemaWithLsn"
+                $ref: "#/components/schemas/SafekeepersLsn"
       responses:
         200:
           description: Promote succeeded or wasn't started
@@ -297,7 +297,14 @@ paths:
         content:
           application/json:
             schema:
-              $ref: "#/components/schemas/ComputeSchema"
+              type: object
+              required:
+                - spec
+              properties:
+                spec:
+                  # XXX: I don't want to explain current spec in the OpenAPI format,
+                  # as it could be changed really soon. Consider doing it later.
+                  type: object
       responses:
         200:
           description: Compute configuration finished.
@@ -584,25 +591,18 @@ components:
           type: string
           example: "1.0.0"
 
-    ComputeSchema:
+    SafekeepersLsn:
       type: object
       required:
-        - spec
-      properties:
-        spec:
-          type: object
-    ComputeSchemaWithLsn:
-      type: object
-      required:
-        - spec
+        - safekeepers
         - wal_flush_lsn
       properties:
-        spec:
-          $ref: "#/components/schemas/ComputeState"
-        wal_flush_lsn:
+        safekeepers:
+          description: Primary replica safekeepers
+          type: string
+        wal_flush_lsn:
+          description: Primary last WAL flush LSN
           type: string
-          description: "last WAL flush LSN"
-          example: "0/028F10D8"
 
     LfcPrewarmState:
       type: object
@@ -613,11 +613,11 @@ components:
         - skipped
       properties:
         status:
-          description: LFC prewarm status
-          enum: [not_prewarmed, prewarming, completed, failed, skipped]
+          description: Lfc prewarm status
+          enum: [not_prewarmed, prewarming, completed, failed]
           type: string
         error:
-          description: LFC prewarm error, if any
+          description: Lfc prewarm error, if any
           type: string
         total:
           description: Total pages processed
@@ -635,11 +635,11 @@ components:
         - status
       properties:
         status:
-          description: LFC offload status
+          description: Lfc offload status
           enum: [not_offloaded, offloading, completed, failed]
           type: string
         error:
-          description: LFC offload error, if any
+          description: Lfc offload error, if any
           type: string
 
     PromoteState:

compute_tools/src/http/routes/hadron_liveness_probe.rs

@@ -1,34 +0,0 @@
-use crate::pg_isready::pg_isready;
-use crate::{compute::ComputeNode, http::JsonResponse};
-use axum::{extract::State, http::StatusCode, response::Response};
-use std::sync::Arc;
-
-/// NOTE: NOT ENABLED YET
-/// Detect if the compute is alive.
-/// Called by the liveness probe of the compute container.
-pub(in crate::http) async fn hadron_liveness_probe(
-    State(compute): State<Arc<ComputeNode>>,
-) -> Response {
-    let port = match compute.params.connstr.port() {
-        Some(port) => port,
-        None => {
-            return JsonResponse::error(
-                StatusCode::INTERNAL_SERVER_ERROR,
-                "Failed to get the port from the connection string",
-            );
-        }
-    };
-    match pg_isready(&compute.params.pg_isready_bin, port) {
-        Ok(_) => {
-            // The connection is successful, so the compute is alive.
-            // Return a 200 OK response.
-            JsonResponse::success(StatusCode::OK, "ok")
-        }
-        Err(e) => {
-            tracing::error!("Hadron liveness probe failed: {}", e);
-            // The connection failed, so the compute is not alive.
-            // Return a 500 Internal Server Error response.
-            JsonResponse::error(StatusCode::INTERNAL_SERVER_ERROR, e)
-        }
-    }
-}

compute_tools/src/http/routes/metrics.rs

@@ -1,18 +1,10 @@
-use std::path::Path;
-use std::sync::Arc;
-
-use anyhow::Context;
 use axum::body::Body;
-use axum::extract::State;
 use axum::response::Response;
+use http::StatusCode;
 use http::header::CONTENT_TYPE;
-use http_body_util::BodyExt;
-use hyper::{Request, StatusCode};
 use metrics::proto::MetricFamily;
 use metrics::{Encoder, TextEncoder};
 
-use crate::communicator_socket_client::connect_communicator_socket;
-use crate::compute::ComputeNode;
 use crate::http::JsonResponse;
 use crate::metrics::collect;
 
@@ -39,42 +31,3 @@ pub(in crate::http) async fn get_metrics() -> Response {
         .body(Body::from(buffer))
         .unwrap()
 }
-
-/// Fetch and forward metrics from the Postgres neon extension's metrics
-/// exporter that are used by autoscaling-agent.
-///
-/// The neon extension exposes these metrics over a Unix domain socket
-/// in the data directory. That's not accessible directly from the outside
-/// world, so we have this endpoint in compute_ctl to expose it
-pub(in crate::http) async fn get_autoscaling_metrics(
-    State(compute): State<Arc<ComputeNode>>,
-) -> Result<Response, Response> {
-    let pgdata = Path::new(&compute.params.pgdata);
-
-    // Connect to the communicator process's metrics socket
-    let mut metrics_client = connect_communicator_socket(pgdata)
-        .await
-        .map_err(|e| JsonResponse::error(StatusCode::INTERNAL_SERVER_ERROR, format!("{e:#}")))?;
-
-    // Make a request for /autoscaling_metrics
-    let request = Request::builder()
-        .method("GET")
-        .uri("/autoscaling_metrics")
-        .header("Host", "localhost") // hyper requires Host, even though the server won't care
-        .body(Body::from(""))
-        .unwrap();
-    let resp = metrics_client
-        .send_request(request)
-        .await
-        .context("fetching metrics from Postgres metrics service")
-        .map_err(|e| JsonResponse::error(StatusCode::INTERNAL_SERVER_ERROR, format!("{e:#}")))?;
-
-    // Build a response that just forwards the response we got.
-    let mut response = Response::builder();
-    response = response.status(resp.status());
-    if let Some(content_type) = resp.headers().get(CONTENT_TYPE) {
-        response = response.header(CONTENT_TYPE, content_type);
-    }
-    let body = tonic::service::AxumBody::from_stream(resp.into_body().into_data_stream());
-    Ok(response.body(body).unwrap())
-}

compute_tools/src/http/routes/mod.rs

@@ -10,13 +10,11 @@ pub(in crate::http) mod extension_server;
 pub(in crate::http) mod extensions;
 pub(in crate::http) mod failpoints;
 pub(in crate::http) mod grants;
-pub(in crate::http) mod hadron_liveness_probe;
 pub(in crate::http) mod insights;
 pub(in crate::http) mod lfc;
 pub(in crate::http) mod metrics;
 pub(in crate::http) mod metrics_json;
 pub(in crate::http) mod promote;
-pub(in crate::http) mod refresh_configuration;
 pub(in crate::http) mod status;
 pub(in crate::http) mod terminate;
 

compute_tools/src/http/routes/promote.rs

@@ -1,14 +1,14 @@
 use crate::http::JsonResponse;
-use axum::extract::Json;
+use axum::Form;
 use http::StatusCode;
 
 pub(in crate::http) async fn promote(
     compute: axum::extract::State<std::sync::Arc<crate::compute::ComputeNode>>,
-    Json(cfg): Json<compute_api::responses::PromoteConfig>,
+    Form(safekeepers_lsn): Form<compute_api::responses::SafekeepersLsn>,
 ) -> axum::response::Response {
-    let state = compute.promote(cfg).await;
-    if let compute_api::responses::PromoteState::Failed { error: _ } = state {
-        return JsonResponse::create_response(StatusCode::INTERNAL_SERVER_ERROR, state);
+    let state = compute.promote(safekeepers_lsn).await;
+    if let compute_api::responses::PromoteState::Failed { error } = state {
+        return JsonResponse::error(StatusCode::INTERNAL_SERVER_ERROR, error);
     }
     JsonResponse::success(StatusCode::OK, state)
 }

compute_tools/src/http/routes/refresh_configuration.rs

@@ -1,34 +0,0 @@
-// This file is added by Hadron
-
-use std::sync::Arc;
-
-use axum::{
-    extract::State,
-    response::{IntoResponse, Response},
-};
-use http::StatusCode;
-use tracing::debug;
-
-use crate::compute::ComputeNode;
-// use crate::hadron_metrics::POSTGRES_PAGESTREAM_REQUEST_ERRORS;
-use crate::http::JsonResponse;
-
-// The /refresh_configuration POST method is used to nudge compute_ctl to pull a new spec
-// from the HCC and attempt to reconfigure Postgres with the new spec. The method does not wait
-// for the reconfiguration to complete. Rather, it simply delivers a signal that will cause
-// configuration to be reloaded in a best effort manner. Invocation of this method does not
-// guarantee that a reconfiguration will occur. The caller should consider keep sending this
-// request while it believes that the compute configuration is out of date.
-pub(in crate::http) async fn refresh_configuration(
-    State(compute): State<Arc<ComputeNode>>,
-) -> Response {
-    debug!("serving /refresh_configuration POST request");
-    // POSTGRES_PAGESTREAM_REQUEST_ERRORS.inc();
-    match compute.signal_refresh_configuration().await {
-        Ok(_) => StatusCode::OK.into_response(),
-        Err(e) => {
-            tracing::error!("error handling /refresh_configuration request: {}", e);
-            JsonResponse::error(StatusCode::INTERNAL_SERVER_ERROR, e)
-        }
-    }
-}

compute_tools/src/http/server.rs

@@ -27,7 +27,6 @@ use super::{
     },
 };
 use crate::compute::ComputeNode;
-use crate::http::routes::{hadron_liveness_probe, refresh_configuration};
 
 /// `compute_ctl` has two servers: internal and external. The internal server
 /// binds to the loopback interface and handles communication from clients on
@@ -44,7 +43,6 @@ pub enum Server {
         port: u16,
         config: ComputeCtlConfig,
         compute_id: String,
-        instance_id: Option<String>,
     },
 }
 
@@ -69,12 +67,7 @@ impl From<&Server> for Router<Arc<ComputeNode>> {
                         post(extension_server::download_extension),
                     )
                     .route("/extensions", post(extensions::install_extension))
-                    .route("/grants", post(grants::add_grant))
-                    // Hadron: Compute-initiated configuration refresh
-                    .route(
-                        "/refresh_configuration",
-                        post(refresh_configuration::refresh_configuration),
-                    );
+                    .route("/grants", post(grants::add_grant));
 
                 // Add in any testing support
                 if cfg!(feature = "testing") {
@@ -86,17 +79,10 @@ impl From<&Server> for Router<Arc<ComputeNode>> {
                 router
             }
             Server::External {
-                config,
-                compute_id,
-                instance_id,
-                ..
+                config, compute_id, ..
             } => {
-                let unauthenticated_router = Router::<Arc<ComputeNode>>::new()
-                    .route("/metrics", get(metrics::get_metrics))
-                    .route(
-                        "/autoscaling_metrics",
-                        get(metrics::get_autoscaling_metrics),
-                    );
+                let unauthenticated_router =
+                    Router::<Arc<ComputeNode>>::new().route("/metrics", get(metrics::get_metrics));
 
                 let authenticated_router = Router::<Arc<ComputeNode>>::new()
                     .route("/lfc/prewarm", get(lfc::prewarm_state).post(lfc::prewarm))
@@ -110,13 +96,8 @@ impl From<&Server> for Router<Arc<ComputeNode>> {
                     .route("/metrics.json", get(metrics_json::get_metrics))
                     .route("/status", get(status::get_status))
                     .route("/terminate", post(terminate::terminate))
-                    .route(
-                        "/hadron_liveness_probe",
-                        get(hadron_liveness_probe::hadron_liveness_probe),
-                    )
                     .layer(AsyncRequireAuthorizationLayer::new(Authorize::new(
                         compute_id.clone(),
-                        instance_id.clone(),
                         config.jwks.clone(),
                     )));
 

compute_tools/src/installed_extensions.rs

@@ -2,7 +2,6 @@ use std::collections::HashMap;
 
 use anyhow::Result;
 use compute_api::responses::{InstalledExtension, InstalledExtensions};
-use once_cell::sync::Lazy;
 use tokio_postgres::error::Error as PostgresError;
 use tokio_postgres::{Client, Config, NoTls};
 
@@ -120,7 +119,3 @@ pub async fn get_installed_extensions(
         extensions: extensions_map.into_values().collect(),
     })
 }
-
-pub fn initialize_metrics() {
-    Lazy::force(&INSTALLED_EXTENSIONS);
-}

compute_tools/src/lib.rs

@@ -4,7 +4,6 @@
 #![deny(clippy::undocumented_unsafe_blocks)]
 
 pub mod checker;
-pub mod communicator_socket_client;
 pub mod config;
 pub mod configurator;
 pub mod http;
@@ -16,7 +15,6 @@ pub mod compute_prewarm;
 pub mod compute_promote;
 pub mod disk_quota;
 pub mod extension_server;
-pub mod hadron_metrics;
 pub mod installed_extensions;
 pub mod local_proxy;
 pub mod lsn_lease;
@@ -25,7 +23,6 @@ mod migration;
 pub mod monitor;
 pub mod params;
 pub mod pg_helpers;
-pub mod pg_isready;
 pub mod pgbouncer;
 pub mod rsyslog;
 pub mod spec;

compute_tools/src/logger.rs

@@ -1,10 +1,7 @@
 use std::collections::HashMap;
-use std::sync::{LazyLock, RwLock};
-use tracing::Subscriber;
 use tracing::info;
-use tracing_appender;
+use tracing_subscriber::layer::SubscriberExt;
 use tracing_subscriber::prelude::*;
-use tracing_subscriber::{fmt, layer::SubscriberExt, registry::LookupSpan};
 
 /// Initialize logging to stderr, and OpenTelemetry tracing and exporter.
 ///
@@ -16,63 +13,31 @@ use tracing_subscriber::{fmt, layer::SubscriberExt, registry::LookupSpan};
 /// set `OTEL_EXPORTER_OTLP_ENDPOINT=http://jaeger:4318`. See
 /// `tracing-utils` package description.
 ///
-pub fn init_tracing_and_logging(
-    default_log_level: &str,
-    log_dir_opt: &Option<String>,
-) -> anyhow::Result<(
-    Option<tracing_utils::Provider>,
-    Option<tracing_appender::non_blocking::WorkerGuard>,
-)> {
+pub async fn init_tracing_and_logging(default_log_level: &str) -> anyhow::Result<()> {
     // Initialize Logging
     let env_filter = tracing_subscriber::EnvFilter::try_from_default_env()
         .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new(default_log_level));
 
-    // Standard output streams
     let fmt_layer = tracing_subscriber::fmt::layer()
         .with_ansi(false)
         .with_target(false)
         .with_writer(std::io::stderr);
 
-    // Logs with file rotation. Files in `$log_dir/pgcctl.yyyy-MM-dd`
-    let (json_to_file_layer, _file_logs_guard) = if let Some(log_dir) = log_dir_opt {
-        std::fs::create_dir_all(log_dir)?;
-        let file_logs_appender = tracing_appender::rolling::RollingFileAppender::builder()
-            .rotation(tracing_appender::rolling::Rotation::DAILY)
-            .filename_prefix("pgcctl")
-            // Lib appends to existing files, so we will keep files for up to 2 days even on restart loops.
-            // At minimum, log-daemon will have 1 day to detect and upload a file (if created right before midnight).
-            .max_log_files(2)
-            .build(log_dir)
-            .expect("Initializing rolling file appender should succeed");
-        let (file_logs_writer, _file_logs_guard) =
-            tracing_appender::non_blocking(file_logs_appender);
-        let json_to_file_layer = tracing_subscriber::fmt::layer()
-            .with_ansi(false)
-            .with_target(false)
-            .event_format(PgJsonLogShapeFormatter)
-            .with_writer(file_logs_writer);
-        (Some(json_to_file_layer), Some(_file_logs_guard))
-    } else {
-        (None, None)
-    };
-
     // Initialize OpenTelemetry
-    let provider =
-        tracing_utils::init_tracing("compute_ctl", tracing_utils::ExportConfig::default());
-    let otlp_layer = provider.as_ref().map(tracing_utils::layer);
+    let otlp_layer =
+        tracing_utils::init_tracing("compute_ctl", tracing_utils::ExportConfig::default()).await;
 
     // Put it all together
     tracing_subscriber::registry()
         .with(env_filter)
         .with(otlp_layer)
         .with(fmt_layer)
-        .with(json_to_file_layer)
         .init();
     tracing::info!("logging and tracing started");
 
     utils::logging::replace_panic_hook_with_tracing_panic_hook().forget();
 
-    Ok((provider, _file_logs_guard))
+    Ok(())
 }
 
 /// Replace all newline characters with a special character to make it
@@ -127,157 +92,3 @@ pub fn startup_context_from_env() -> Option<opentelemetry::Context> {
         None
     }
 }
-
-/// Track relevant id's
-const UNKNOWN_IDS: &str = r#""pg_instance_id": "", "pg_compute_id": """#;
-static IDS: LazyLock<RwLock<String>> = LazyLock::new(|| RwLock::new(UNKNOWN_IDS.to_string()));
-
-pub fn update_ids(instance_id: &Option<String>, compute_id: &Option<String>) -> anyhow::Result<()> {
-    let ids = format!(
-        r#""pg_instance_id": "{}", "pg_compute_id": "{}""#,
-        instance_id.as_ref().map(|s| s.as_str()).unwrap_or_default(),
-        compute_id.as_ref().map(|s| s.as_str()).unwrap_or_default()
-    );
-    let mut guard = IDS
-        .write()
-        .map_err(|e| anyhow::anyhow!("Log set id's rwlock poisoned: {}", e))?;
-    *guard = ids;
-    Ok(())
-}
-
-/// Massage compute_ctl logs into PG json log shape so we can use the same Lumberjack setup.
-struct PgJsonLogShapeFormatter;
-impl<S, N> fmt::format::FormatEvent<S, N> for PgJsonLogShapeFormatter
-where
-    S: Subscriber + for<'a> LookupSpan<'a>,
-    N: for<'a> fmt::format::FormatFields<'a> + 'static,
-{
-    fn format_event(
-        &self,
-        ctx: &fmt::FmtContext<'_, S, N>,
-        mut writer: fmt::format::Writer<'_>,
-        event: &tracing::Event<'_>,
-    ) -> std::fmt::Result {
-        // Format values from the event's metadata, and open message string
-        let metadata = event.metadata();
-        {
-            let ids_guard = IDS.read();
-            let ids = ids_guard
-                .as_ref()
-                .map(|guard| guard.as_str())
-                // Surpress so that we don't lose all uploaded/ file logs if something goes super wrong. We would notice the missing id's.
-                .unwrap_or(UNKNOWN_IDS);
-            write!(
-                &mut writer,
-                r#"{{"timestamp": "{}", "error_severity": "{}", "file_name": "{}", "backend_type": "compute_ctl_self", {}, "message": "#,
-                chrono::Utc::now().format("%Y-%m-%d %H:%M:%S%.3f GMT"),
-                metadata.level(),
-                metadata.target(),
-                ids
-            )?;
-        }
-
-        let mut message = String::new();
-        let message_writer = fmt::format::Writer::new(&mut message);
-
-        // Gather the message
-        ctx.field_format().format_fields(message_writer, event)?;
-
-        // TODO: any better options than to copy-paste this OSS span formatter?
-        // impl<S, N, T> FormatEvent<S, N> for Format<Full, T>
-        // https://docs.rs/tracing-subscriber/latest/tracing_subscriber/fmt/trait.FormatEvent.html#impl-FormatEvent%3CS,+N%3E-for-Format%3CFull,+T%3E
-
-        // write message, close bracket, and new line
-        writeln!(writer, "{}}}", serde_json::to_string(&message).unwrap())
-    }
-}
-
-#[cfg(feature = "testing")]
-#[cfg(test)]
-mod test {
-    use super::*;
-    use std::{cell::RefCell, io};
-
-    // Use thread_local! instead of Mutex for test isolation
-    thread_local! {
-        static WRITER_OUTPUT: RefCell<String> = const { RefCell::new(String::new()) };
-    }
-
-    #[derive(Clone, Default)]
-    struct StaticStringWriter;
-
-    impl io::Write for StaticStringWriter {
-        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-            let output = String::from_utf8(buf.to_vec()).expect("Invalid UTF-8 in test output");
-            WRITER_OUTPUT.with(|s| s.borrow_mut().push_str(&output));
-            Ok(buf.len())
-        }
-
-        fn flush(&mut self) -> io::Result<()> {
-            Ok(())
-        }
-    }
-
-    impl fmt::MakeWriter<'_> for StaticStringWriter {
-        type Writer = Self;
-
-        fn make_writer(&self) -> Self::Writer {
-            Self
-        }
-    }
-
-    #[test]
-    fn test_log_pg_json_shape_formatter() {
-        // Use a scoped subscriber to prevent global state pollution
-        let subscriber = tracing_subscriber::registry().with(
-            tracing_subscriber::fmt::layer()
-                .with_ansi(false)
-                .with_target(false)
-                .event_format(PgJsonLogShapeFormatter)
-                .with_writer(StaticStringWriter),
-        );
-
-        let _ = update_ids(&Some("000".to_string()), &Some("111".to_string()));
-
-        // Clear any previous test state
-        WRITER_OUTPUT.with(|s| s.borrow_mut().clear());
-
-        let messages = [
-            "test message",
-            r#"json escape check:  name="BatchSpanProcessor.Flush.ExportError" reason="Other(reqwest::Error { kind: Request, url: \"http://localhost:4318/v1/traces\", source: hyper_
-            util::client::legacy::Error(Connect, ConnectError(\"tcp connect error\", Os { code: 111, kind: ConnectionRefused, message: \"Connection refused\" })) })" Failed during the export process"#,
-        ];
-
-        tracing::subscriber::with_default(subscriber, || {
-            for message in messages {
-                tracing::info!(message);
-            }
-        });
-        tracing::info!("not test message");
-
-        // Get captured output
-        let output = WRITER_OUTPUT.with(|s| s.borrow().clone());
-
-        let json_strings: Vec<&str> = output.lines().collect();
-        assert_eq!(
-            json_strings.len(),
-            messages.len(),
-            "Log didn't have the expected number of json strings."
-        );
-
-        let json_string_shape_regex = regex::Regex::new(
-            r#"\{"timestamp": "\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\.\d{3} GMT", "error_severity": "INFO", "file_name": ".+", "backend_type": "compute_ctl_self", "pg_instance_id": "000", "pg_compute_id": "111", "message": ".+"\}"#
-        ).unwrap();
-
-        for (i, expected_message) in messages.iter().enumerate() {
-            let json_string = json_strings[i];
-            assert!(
-                json_string_shape_regex.is_match(json_string),
-                "Json log didn't match expected pattern:\n{json_string}",
-            );
-            let parsed_json: serde_json::Value = serde_json::from_str(json_string).unwrap();
-            let actual_message = parsed_json["message"].as_str().unwrap();
-            assert_eq!(*expected_message, actual_message);
-        }
-    }
-}

compute_tools/src/migration.rs

@@ -9,20 +9,15 @@ use crate::metrics::DB_MIGRATION_FAILED;
 pub(crate) struct MigrationRunner<'m> {
     client: &'m mut Client,
     migrations: &'m [&'m str],
-    lakebase_mode: bool,
 }
 
 impl<'m> MigrationRunner<'m> {
     /// Create a new migration runner
-    pub fn new(client: &'m mut Client, migrations: &'m [&'m str], lakebase_mode: bool) -> Self {
+    pub fn new(client: &'m mut Client, migrations: &'m [&'m str]) -> Self {
         // The neon_migration.migration_id::id column is a bigint, which is equivalent to an i64
         assert!(migrations.len() + 1 < i64::MAX as usize);
 
-        Self {
-            client,
-            migrations,
-            lakebase_mode,
-        }
+        Self { client, migrations }
     }
 
     /// Get the current value neon_migration.migration_id
@@ -135,13 +130,8 @@ impl<'m> MigrationRunner<'m> {
             // ID is also the next index
             let migration_id = (current_migration + 1) as i64;
             let migration = self.migrations[current_migration];
-            let migration = if self.lakebase_mode {
-                migration.replace("neon_superuser", "databricks_superuser")
-            } else {
-                migration.to_string()
-            };
 
-            match Self::run_migration(self.client, migration_id, &migration).await {
+            match Self::run_migration(self.client, migration_id, migration).await {
                 Ok(_) => {
                     info!("Finished migration id={}", migration_id);
                 }

compute_tools/src/monitor.rs

@@ -11,7 +11,6 @@ use tracing::{Level, error, info, instrument, span};
 use crate::compute::ComputeNode;
 use crate::metrics::{PG_CURR_DOWNTIME_MS, PG_TOTAL_DOWNTIME_MS};
 
-const PG_DEFAULT_INIT_TIMEOUIT: Duration = Duration::from_secs(60);
 const MONITOR_CHECK_INTERVAL: Duration = Duration::from_millis(500);
 
 /// Struct to store runtime state of the compute monitor thread.
@@ -353,47 +352,13 @@ impl ComputeMonitor {
 // Hang on condition variable waiting until the compute status is `Running`.
 fn wait_for_postgres_start(compute: &ComputeNode) {
     let mut state = compute.state.lock().unwrap();
-    let pg_init_timeout = compute
-        .params
-        .pg_init_timeout
-        .unwrap_or(PG_DEFAULT_INIT_TIMEOUIT);
-
     while state.status != ComputeStatus::Running {
         info!("compute is not running, waiting before monitoring activity");
-        if !compute.params.lakebase_mode {
-            state = compute.state_changed.wait(state).unwrap();
+        state = compute.state_changed.wait(state).unwrap();
 
-            if state.status == ComputeStatus::Running {
-                break;
-            }
-            continue;
+        if state.status == ComputeStatus::Running {
+            break;
         }
-
-        if state.pg_start_time.is_some()
-            && Utc::now()
-                .signed_duration_since(state.pg_start_time.unwrap())
-                .to_std()
-                .unwrap_or_default()
-                > pg_init_timeout
-        {
-            // If Postgres isn't up and running with working PS/SK connections within POSTGRES_STARTUP_TIMEOUT, it is
-            // possible that we started Postgres with a wrong spec (so it is talking to the wrong PS/SK nodes). To prevent
-            // deadends we simply exit (panic) the compute node so it can restart with the latest spec.
-            //
-            // NB: We skip this check if we have not attempted to start PG yet (indicated by state.pg_start_up == None).
-            // This is to make sure the more appropriate errors are surfaced if we encounter issues before we even attempt
-            // to start PG (e.g., if we can't pull the spec, can't sync safekeepers, or can't get the basebackup).
-            error!(
-                "compute did not enter Running state in {} seconds, exiting",
-                pg_init_timeout.as_secs()
-            );
-            std::process::exit(1);
-        }
-        state = compute
-            .state_changed
-            .wait_timeout(state, Duration::from_secs(5))
-            .unwrap()
-            .0;
     }
 }
 

compute_tools/src/pg_helpers.rs

@@ -11,9 +11,7 @@ use std::time::{Duration, Instant};
 
 use anyhow::{Result, bail};
 use compute_api::responses::TlsConfig;
-use compute_api::spec::{
-    Database, DatabricksSettings, GenericOption, GenericOptions, PgIdent, Role,
-};
+use compute_api::spec::{Database, GenericOption, GenericOptions, PgIdent, Role};
 use futures::StreamExt;
 use indexmap::IndexMap;
 use ini::Ini;
@@ -186,42 +184,6 @@ impl DatabaseExt for Database {
     }
 }
 
-pub trait DatabricksSettingsExt {
-    fn as_pg_settings(&self) -> String;
-}
-
-impl DatabricksSettingsExt for DatabricksSettings {
-    fn as_pg_settings(&self) -> String {
-        // Postgres GUCs rendered from DatabricksSettings
-        vec![
-            // ssl_ca_file
-            Some(format!(
-                "ssl_ca_file = '{}'",
-                self.pg_compute_tls_settings.ca_file
-            )),
-            // [Optional] databricks.workspace_url
-            Some(format!(
-                "databricks.workspace_url = '{}'",
-                &self.databricks_workspace_host
-            )),
-            // todo(vikas.jain): these are not required anymore as they are moved to static
-            // conf but keeping these to avoid image mismatch between hcc and pg.
-            // Once hcc and pg are in sync, we can remove these.
-            //
-            // databricks.enable_databricks_identity_login
-            Some("databricks.enable_databricks_identity_login = true".to_string()),
-            // databricks.enable_sql_restrictions
-            Some("databricks.enable_sql_restrictions = true".to_string()),
-        ]
-        .into_iter()
-        // Removes `None`s
-        .flatten()
-        .collect::<Vec<String>>()
-        .join("\n")
-            + "\n"
-    }
-}
-
 /// Generic trait used to provide quoting / encoding for strings used in the
 /// Postgres SQL queries and DATABASE_URL.
 pub trait Escaping {

compute_tools/src/pg_isready.rs

@@ -1,30 +0,0 @@
-use anyhow::{Context, anyhow};
-
-// Run `/usr/local/bin/pg_isready -p {port}`
-// Check the connectivity of PG
-// Success means PG is listening on the port and accepting connections
-// Note that PG does not need to authenticate the connection, nor reserve a connection quota for it.
-// See https://www.postgresql.org/docs/current/app-pg-isready.html
-pub fn pg_isready(bin: &str, port: u16) -> anyhow::Result<()> {
-    let child_result = std::process::Command::new(bin)
-        .arg("-p")
-        .arg(port.to_string())
-        .spawn();
-
-    child_result
-        .context("spawn() failed")
-        .and_then(|mut child| child.wait().context("wait() failed"))
-        .and_then(|status| match status.success() {
-            true => Ok(()),
-            false => Err(anyhow!("process exited with {status}")),
-        })
-        // wrap any prior error with the overall context that we couldn't run the command
-        .with_context(|| format!("could not run `{bin} --port {port}`"))
-}
-
-// It's safe to assume pg_isready is under the same directory with postgres,
-// because it is a PG util bin installed along with postgres
-pub fn get_pg_isready_bin(pgbin: &str) -> String {
-    let split = pgbin.split("/").collect::<Vec<&str>>();
-    split[0..split.len() - 1].join("/") + "/pg_isready"
-}

compute_tools/src/spec.rs

@@ -1,6 +1,4 @@
 use std::fs::File;
-use std::fs::{self, Permissions};
-use std::os::unix::fs::PermissionsExt;
 use std::path::Path;
 
 use anyhow::{Result, anyhow, bail};
@@ -135,25 +133,10 @@ pub fn get_config_from_control_plane(base_uri: &str, compute_id: &str) -> Result
 }
 
 /// Check `pg_hba.conf` and update if needed to allow external connections.
-pub fn update_pg_hba(pgdata_path: &Path, databricks_pg_hba: Option<&String>) -> Result<()> {
+pub fn update_pg_hba(pgdata_path: &Path) -> Result<()> {
     // XXX: consider making it a part of config.json
     let pghba_path = pgdata_path.join("pg_hba.conf");
 
-    // Update pg_hba to contains databricks specfic settings before adding neon settings
-    // PG uses the first record that matches to perform authentication, so we need to have
-    // our rules before the default ones from neon.
-    // See https://www.postgresql.org/docs/16/auth-pg-hba-conf.html
-    if let Some(databricks_pg_hba) = databricks_pg_hba {
-        if config::line_in_file(
-            &pghba_path,
-            &format!("include_if_exists {}\n", *databricks_pg_hba),
-        )? {
-            info!("updated pg_hba.conf to include databricks_pg_hba.conf");
-        } else {
-            info!("pg_hba.conf already included databricks_pg_hba.conf");
-        }
-    }
-
     if config::line_in_file(&pghba_path, PG_HBA_ALL_MD5)? {
         info!("updated pg_hba.conf to allow external connections");
     } else {
@@ -163,59 +146,6 @@ pub fn update_pg_hba(pgdata_path: &Path, databricks_pg_hba: Option<&String>) ->
     Ok(())
 }
 
-/// Check `pg_ident.conf` and update if needed to allow databricks config.
-pub fn update_pg_ident(pgdata_path: &Path, databricks_pg_ident: Option<&String>) -> Result<()> {
-    info!("checking pg_ident.conf");
-    let pghba_path = pgdata_path.join("pg_ident.conf");
-
-    // Update pg_ident to contains databricks specfic settings
-    if let Some(databricks_pg_ident) = databricks_pg_ident {
-        if config::line_in_file(
-            &pghba_path,
-            &format!("include_if_exists {}\n", *databricks_pg_ident),
-        )? {
-            info!("updated pg_ident.conf to include databricks_pg_ident.conf");
-        } else {
-            info!("pg_ident.conf already included databricks_pg_ident.conf");
-        }
-    }
-
-    Ok(())
-}
-
-/// Copy tls key_file and cert_file from k8s secret mount directory
-/// to pgdata and set private key file permissions as expected by Postgres.
-/// See this doc for expected permission <https://www.postgresql.org/docs/current/ssl-tcp.html>
-/// K8s secrets mount on dblet does not honor permission and ownership
-/// specified in the Volume or VolumeMount. So we need to explicitly copy the file and set the permissions.
-pub fn copy_tls_certificates(
-    key_file: &String,
-    cert_file: &String,
-    pgdata_path: &Path,
-) -> Result<()> {
-    let files = [cert_file, key_file];
-    for file in files.iter() {
-        let source = Path::new(file);
-        let dest = pgdata_path.join(source.file_name().unwrap());
-        if !dest.exists() {
-            std::fs::copy(source, &dest)?;
-            info!(
-                "Copying tls file: {} to {}",
-                &source.display(),
-                &dest.display()
-            );
-        }
-        if *file == key_file {
-            // Postgres requires private key to be readable only by the owner by having
-            // chmod 600 permissions.
-            let permissions = Permissions::from_mode(0o600);
-            fs::set_permissions(&dest, permissions)?;
-            info!("Setting permission on {}.", &dest.display());
-        }
-    }
-    Ok(())
-}
-
 /// Create a standby.signal file
 pub fn add_standby_signal(pgdata_path: &Path) -> Result<()> {
     // XXX: consider making it a part of config.json
@@ -240,11 +170,7 @@ pub async fn handle_neon_extension_upgrade(client: &mut Client) -> Result<()> {
 }
 
 #[instrument(skip_all)]
-pub async fn handle_migrations(
-    params: ComputeNodeParams,
-    client: &mut Client,
-    lakebase_mode: bool,
-) -> Result<()> {
+pub async fn handle_migrations(params: ComputeNodeParams, client: &mut Client) -> Result<()> {
     info!("handle migrations");
 
     // !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
@@ -308,7 +234,7 @@ pub async fn handle_migrations(
         ),
     ];
 
-    MigrationRunner::new(client, &migrations, lakebase_mode)
+    MigrationRunner::new(client, &migrations)
         .run_migrations()
         .await?;
 

compute_tools/src/spec_apply.rs

@@ -411,8 +411,7 @@ impl ComputeNode {
             .map(|limit| match limit {
                 0..10 => limit,
                 10..30 => 10,
-                30..300 => limit / 3,
-                300.. => 100,
+                30.. => limit / 3,
             })
             // If we didn't find max_connections, default to 10 concurrent connections.
             .unwrap_or(10)

control_plane/README.md

@@ -8,10 +8,10 @@ code changes locally, but not suitable for running production systems.
 
 ## Example: Start with Postgres 16
 
-To create and start a local development environment with Postgres 16, you will need to provide `--pg-version` flag to 2 of the start-up commands.
+To create and start a local development environment with Postgres 16, you will need to provide `--pg-version` flag to 3 of the start-up commands.
 
 ```shell
-cargo neon init
+cargo neon init --pg-version 16
 cargo neon start
 cargo neon tenant create --set-default --pg-version 16
 cargo neon endpoint create main --pg-version 16

control_plane/src/bin/neon_local.rs

@@ -407,12 +407,6 @@ struct StorageControllerStartCmdArgs {
         help = "Base port for the storage controller instance idenfified by instance-id (defaults to pageserver cplane api)"
     )]
     base_port: Option<u16>,
-
-    #[clap(
-        long,
-        help = "Whether the storage controller should handle pageserver-reported local disk loss events."
-    )]
-    handle_ps_local_disk_loss: Option<bool>,
 }
 
 #[derive(clap::Args)]
@@ -1517,7 +1511,7 @@ async fn handle_endpoint(subcmd: &EndpointCmd, env: &local_env::LocalEnv) -> Res
             let endpoint = cplane
                 .endpoints
                 .get(endpoint_id.as_str())
-                .ok_or_else(|| anyhow!("endpoint {endpoint_id} not found"))?;
+                .ok_or_else(|| anyhow::anyhow!("endpoint {endpoint_id} not found"))?;
 
             if !args.allow_multiple {
                 cplane.check_conflicting_endpoints(
@@ -1815,7 +1809,6 @@ async fn handle_storage_controller(
                 instance_id: args.instance_id,
                 base_port: args.base_port,
                 start_timeout: args.start_timeout,
-                handle_ps_local_disk_loss: args.handle_ps_local_disk_loss,
             };
 
             if let Err(e) = svc.start(start_args).await {

control_plane/src/endpoint.rs

@@ -65,6 +65,7 @@ use jsonwebtoken::jwk::{
     OctetKeyPairParameters, OctetKeyPairType, PublicKeyUse,
 };
 use nix::sys::signal::{Signal, kill};
+use pageserver_api::shard::ShardStripeSize;
 use pem::Pem;
 use reqwest::header::CONTENT_TYPE;
 use safekeeper_api::PgMajorVersion;
@@ -76,7 +77,6 @@ use spki::{SubjectPublicKeyInfo, SubjectPublicKeyInfoRef};
 use tracing::debug;
 use url::Host;
 use utils::id::{NodeId, TenantId, TimelineId};
-use utils::shard::ShardStripeSize;
 
 use crate::local_env::LocalEnv;
 use crate::postgresql_conf::PostgresConf;

control_plane/src/storage_controller.rs

@@ -56,7 +56,6 @@ pub struct NeonStorageControllerStartArgs {
     pub instance_id: u8,
     pub base_port: Option<u16>,
     pub start_timeout: humantime::Duration,
-    pub handle_ps_local_disk_loss: Option<bool>,
 }
 
 impl NeonStorageControllerStartArgs {
@@ -65,7 +64,6 @@ impl NeonStorageControllerStartArgs {
             instance_id: 1,
             base_port: None,
             start_timeout,
-            handle_ps_local_disk_loss: None,
         }
     }
 }
@@ -671,10 +669,6 @@ impl StorageController {
 
         println!("Starting storage controller at {scheme}://{host}:{listen_port}");
 
-        if start_args.handle_ps_local_disk_loss.unwrap_or_default() {
-            args.push("--handle-ps-local-disk-loss".to_string());
-        }
-
         background_process::start_process(
             COMMAND,
             &instance_dir,

control_plane/storcon_cli/src/main.rs

@@ -76,12 +76,6 @@ enum Command {
     NodeStartDelete {
         #[arg(long)]
         node_id: NodeId,
-        /// When `force` is true, skip waiting for shards to prewarm during migration.
-        /// This can significantly speed up node deletion since prewarming all shards
-        /// can take considerable time, but may result in slower initial access to
-        /// migrated shards until they warm up naturally.
-        #[arg(long)]
-        force: bool,
     },
     /// Cancel deletion of the specified pageserver and wait for `timeout`
     /// for the operation to be canceled. May be retried.
@@ -958,14 +952,13 @@ async fn main() -> anyhow::Result<()> {
                 .dispatch::<(), ()>(Method::DELETE, format!("control/v1/node/{node_id}"), None)
                 .await?;
         }
-        Command::NodeStartDelete { node_id, force } => {
-            let query = if force {
-                format!("control/v1/node/{node_id}/delete?force=true")
-            } else {
-                format!("control/v1/node/{node_id}/delete")
-            };
+        Command::NodeStartDelete { node_id } => {
             storcon_client
-                .dispatch::<(), ()>(Method::PUT, query, None)
+                .dispatch::<(), ()>(
+                    Method::PUT,
+                    format!("control/v1/node/{node_id}/delete"),
+                    None,
+                )
                 .await?;
             println!("Delete started for {node_id}");
         }

deny.toml

@@ -35,7 +35,6 @@ reason = "The paste crate is a build-only dependency with no runtime components.
 # More documentation for the licenses section can be found here:
 # https://embarkstudios.github.io/cargo-deny/checks/licenses/cfg.html
 [licenses]
-version = 2
 allow = [
     "0BSD",
     "Apache-2.0",

endpoint_storage/src/app.rs

@@ -233,7 +233,7 @@ mod tests {
                 .unwrap()
                 .as_millis();
             use rand::Rng;
-            let random = rand::rng().random::<u32>();
+            let random = rand::thread_rng().r#gen::<u32>();
 
             let s3_config = remote_storage::S3Config {
                 bucket_name: var(REAL_S3_BUCKET).unwrap(),

libs/compute_api/src/responses.rs

@@ -46,33 +46,16 @@ pub struct ExtensionInstallResponse {
     pub version: ExtVersion,
 }
 
-/// Status of the LFC prewarm process. The same state machine is reused for
-/// both autoprewarm (prewarm after compute/Postgres start using the previously
-/// stored LFC state) and explicit prewarming via API.
 #[derive(Serialize, Default, Debug, Clone, PartialEq)]
 #[serde(tag = "status", rename_all = "snake_case")]
 pub enum LfcPrewarmState {
-    /// Default value when compute boots up.
     #[default]
     NotPrewarmed,
-    /// Prewarming thread is active and loading pages into LFC.
     Prewarming,
-    /// We found requested LFC state in the endpoint storage and
-    /// completed prewarming successfully.
     Completed,
-    /// Unexpected error happened during prewarming. Note, `Not Found 404`
-    /// response from the endpoint storage is explicitly excluded here
-    /// because it can normally happen on the first compute start,
-    /// since LFC state is not available yet.
-    Failed { error: String },
-    /// We tried to fetch the corresponding LFC state from the endpoint storage,
-    /// but received `Not Found 404`. This should normally happen only during the
-    /// first endpoint start after creation with `autoprewarm: true`.
-    ///
-    /// During the orchestrated prewarm via API, when a caller explicitly
-    /// provides the LFC state key to prewarm from, it's the caller responsibility
-    /// to handle this status as an error state in this case.
-    Skipped,
+    Failed {
+        error: String,
+    },
 }
 
 impl Display for LfcPrewarmState {
@@ -81,7 +64,6 @@ impl Display for LfcPrewarmState {
             LfcPrewarmState::NotPrewarmed => f.write_str("NotPrewarmed"),
             LfcPrewarmState::Prewarming => f.write_str("Prewarming"),
             LfcPrewarmState::Completed => f.write_str("Completed"),
-            LfcPrewarmState::Skipped => f.write_str("Skipped"),
             LfcPrewarmState::Failed { error } => write!(f, "Error({error})"),
         }
     }
@@ -108,10 +90,11 @@ pub enum PromoteState {
     Failed { error: String },
 }
 
-#[derive(Deserialize, Default, Debug)]
+#[derive(Deserialize, Serialize, Default, Debug, Clone)]
 #[serde(rename_all = "snake_case")]
-pub struct PromoteConfig {
-    pub spec: ComputeSpec,
+/// Result of /safekeepers_lsn
+pub struct SafekeepersLsn {
+    pub safekeepers: String,
     pub wal_flush_lsn: utils::lsn::Lsn,
 }
 

libs/compute_api/src/spec.rs

@@ -416,32 +416,6 @@ pub struct GenericOption {
     pub vartype: String,
 }
 
-/// Postgres compute TLS settings.
-#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
-pub struct PgComputeTlsSettings {
-    // Absolute path to the certificate file for server-side TLS.
-    pub cert_file: String,
-    // Absolute path to the private key file for server-side TLS.
-    pub key_file: String,
-    // Absolute path to the certificate authority file for verifying client certificates.
-    pub ca_file: String,
-}
-
-/// Databricks specific options for compute instance.
-/// This is used to store any other settings that needs to be propagate to Compute
-/// but should not be persisted to ComputeSpec in the database.
-#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
-pub struct DatabricksSettings {
-    pub pg_compute_tls_settings: PgComputeTlsSettings,
-    // Absolute file path to databricks_pg_hba.conf file.
-    pub databricks_pg_hba: String,
-    // Absolute file path to databricks_pg_ident.conf file.
-    pub databricks_pg_ident: String,
-    // Hostname portion of the Databricks workspace URL of the endpoint, or empty string if not known.
-    // A valid hostname is required for the compute instance to support PAT logins.
-    pub databricks_workspace_host: String,
-}
-
 /// Optional collection of `GenericOption`'s. Type alias allows us to
 /// declare a `trait` on it.
 pub type GenericOptions = Option<Vec<GenericOption>>;

libs/consumption_metrics/src/lib.rs

@@ -90,7 +90,7 @@ impl<'a> IdempotencyKey<'a> {
         IdempotencyKey {
             now: Utc::now(),
             node_id,
-            nonce: rand::rng().random_range(0..=9999),
+            nonce: rand::thread_rng().gen_range(0..=9999),
         }
     }
 

libs/desim/src/node_os.rs

@@ -41,7 +41,7 @@ impl NodeOs {
 
     /// Generate a random number in range [0, max).
     pub fn random(&self, max: u64) -> u64 {
-        self.internal.rng.lock().random_range(0..max)
+        self.internal.rng.lock().gen_range(0..max)
     }
 
     /// Append a new event to the world event log.

libs/desim/src/options.rs

@@ -32,10 +32,10 @@ impl Delay {
     /// Generate a random delay in range [min, max]. Return None if the
     /// message should be dropped.
     pub fn delay(&self, rng: &mut StdRng) -> Option<u64> {
-        if rng.random_bool(self.fail_prob) {
+        if rng.gen_bool(self.fail_prob) {
             return None;
         }
-        Some(rng.random_range(self.min..=self.max))
+        Some(rng.gen_range(self.min..=self.max))
     }
 }
 

libs/desim/src/world.rs

@@ -69,7 +69,7 @@ impl World {
     /// Create a new random number generator.
     pub fn new_rng(&self) -> StdRng {
         let mut rng = self.rng.lock();
-        StdRng::from_rng(rng.deref_mut())
+        StdRng::from_rng(rng.deref_mut()).unwrap()
     }
 
     /// Create a new node.

libs/metrics/Cargo.toml

@@ -17,5 +17,5 @@ procfs.workspace = true
 measured-process.workspace = true
 
 [dev-dependencies]
-rand.workspace = true
-rand_distr = "0.5"
+rand = "0.8"
+rand_distr = "0.4.3"

libs/metrics/src/hll.rs

@@ -260,7 +260,7 @@ mod tests {
 
     #[test]
     fn test_cardinality_small() {
-        let (actual, estimate) = test_cardinality(100, Zipf::new(100.0, 1.2f64).unwrap());
+        let (actual, estimate) = test_cardinality(100, Zipf::new(100, 1.2f64).unwrap());
 
         assert_eq!(actual, [46, 30, 32]);
         assert!(51.3 < estimate[0] && estimate[0] < 51.4);
@@ -270,7 +270,7 @@ mod tests {
 
     #[test]
     fn test_cardinality_medium() {
-        let (actual, estimate) = test_cardinality(10000, Zipf::new(10000.0, 1.2f64).unwrap());
+        let (actual, estimate) = test_cardinality(10000, Zipf::new(10000, 1.2f64).unwrap());
 
         assert_eq!(actual, [2529, 1618, 1629]);
         assert!(2309.1 < estimate[0] && estimate[0] < 2309.2);
@@ -280,8 +280,7 @@ mod tests {
 
     #[test]
     fn test_cardinality_large() {
-        let (actual, estimate) =
-            test_cardinality(1_000_000, Zipf::new(1_000_000.0, 1.2f64).unwrap());
+        let (actual, estimate) = test_cardinality(1_000_000, Zipf::new(1_000_000, 1.2f64).unwrap());
 
         assert_eq!(actual, [129077, 79579, 79630]);
         assert!(126067.2 < estimate[0] && estimate[0] < 126067.3);
@@ -291,7 +290,7 @@ mod tests {
 
     #[test]
     fn test_cardinality_small2() {
-        let (actual, estimate) = test_cardinality(100, Zipf::new(200.0, 0.8f64).unwrap());
+        let (actual, estimate) = test_cardinality(100, Zipf::new(200, 0.8f64).unwrap());
 
         assert_eq!(actual, [92, 58, 60]);
         assert!(116.1 < estimate[0] && estimate[0] < 116.2);
@@ -301,7 +300,7 @@ mod tests {
 
     #[test]
     fn test_cardinality_medium2() {
-        let (actual, estimate) = test_cardinality(10000, Zipf::new(20000.0, 0.8f64).unwrap());
+        let (actual, estimate) = test_cardinality(10000, Zipf::new(20000, 0.8f64).unwrap());
 
         assert_eq!(actual, [8201, 5131, 5051]);
         assert!(6846.4 < estimate[0] && estimate[0] < 6846.5);
@@ -311,8 +310,7 @@ mod tests {
 
     #[test]
     fn test_cardinality_large2() {
-        let (actual, estimate) =
-            test_cardinality(1_000_000, Zipf::new(2_000_000.0, 0.8f64).unwrap());
+        let (actual, estimate) = test_cardinality(1_000_000, Zipf::new(2_000_000, 0.8f64).unwrap());
 
         assert_eq!(actual, [777847, 482069, 482246]);
         assert!(699437.4 < estimate[0] && estimate[0] < 699437.5);

libs/metrics/src/lib.rs

@@ -4,14 +4,12 @@
 //! a default registry.
 #![deny(clippy::undocumented_unsafe_blocks)]
 
-use std::sync::RwLock;
-
 use measured::label::{LabelGroupSet, LabelGroupVisitor, LabelName, NoLabels};
 use measured::metric::counter::CounterState;
 use measured::metric::gauge::GaugeState;
 use measured::metric::group::Encoding;
 use measured::metric::name::{MetricName, MetricNameEncoder};
-use measured::metric::{MetricEncoding, MetricFamilyEncoding, MetricType};
+use measured::metric::{MetricEncoding, MetricFamilyEncoding};
 use measured::{FixedCardinalityLabel, LabelGroup, MetricGroup};
 use once_cell::sync::Lazy;
 use prometheus::Registry;
@@ -118,52 +116,12 @@ pub fn pow2_buckets(start: usize, end: usize) -> Vec<f64> {
         .collect()
 }
 
-pub struct InfoMetric<L: LabelGroup, M: MetricType = GaugeState> {
-    label: RwLock<L>,
-    metric: M,
-}
-
-impl<L: LabelGroup> InfoMetric<L> {
-    pub fn new(label: L) -> Self {
-        Self::with_metric(label, GaugeState::new(1))
-    }
-}
-
-impl<L: LabelGroup, M: MetricType<Metadata = ()>> InfoMetric<L, M> {
-    pub fn with_metric(label: L, metric: M) -> Self {
-        Self {
-            label: RwLock::new(label),
-            metric,
-        }
-    }
-
-    pub fn set_label(&self, label: L) {
-        *self.label.write().unwrap() = label;
-    }
-}
-
-impl<L, M, E> MetricFamilyEncoding<E> for InfoMetric<L, M>
-where
-    L: LabelGroup,
-    M: MetricEncoding<E, Metadata = ()>,
-    E: Encoding,
-{
-    fn collect_family_into(
-        &self,
-        name: impl measured::metric::name::MetricNameEncoder,
-        enc: &mut E,
-    ) -> Result<(), E::Err> {
-        M::write_type(&name, enc)?;
-        self.metric
-            .collect_into(&(), &*self.label.read().unwrap(), name, enc)
-    }
-}
-
 pub struct BuildInfo {
     pub revision: &'static str,
     pub build_tag: &'static str,
 }
 
+// todo: allow label group without the set
 impl LabelGroup for BuildInfo {
     fn visit_values(&self, v: &mut impl LabelGroupVisitor) {
         const REVISION: &LabelName = LabelName::from_str("revision");
@@ -173,6 +131,24 @@ impl LabelGroup for BuildInfo {
     }
 }
 
+impl<T: Encoding> MetricFamilyEncoding<T> for BuildInfo
+where
+    GaugeState: MetricEncoding<T>,
+{
+    fn collect_family_into(
+        &self,
+        name: impl measured::metric::name::MetricNameEncoder,
+        enc: &mut T,
+    ) -> Result<(), T::Err> {
+        enc.write_help(&name, "Build/version information")?;
+        GaugeState::write_type(&name, enc)?;
+        GaugeState {
+            count: std::sync::atomic::AtomicI64::new(1),
+        }
+        .collect_into(&(), self, name, enc)
+    }
+}
+
 #[derive(MetricGroup)]
 #[metric(new(build_info: BuildInfo))]
 pub struct NeonMetrics {
@@ -189,8 +165,8 @@ pub struct NeonMetrics {
 #[derive(MetricGroup)]
 #[metric(new(build_info: BuildInfo))]
 pub struct LibMetrics {
-    #[metric(init = InfoMetric::new(build_info))]
-    build_info: InfoMetric<BuildInfo>,
+    #[metric(init = build_info)]
+    build_info: BuildInfo,
 
     #[metric(flatten)]
     rusage: Rusage,

libs/neon-shmem/Cargo.toml

@@ -8,13 +8,6 @@ license.workspace = true
 thiserror.workspace = true
 nix.workspace=true
 workspace_hack = { version = "0.1", path = "../../workspace_hack" }
-libc.workspace = true
-lock_api.workspace = true
-rustc-hash.workspace = true
 
 [target.'cfg(target_os = "macos")'.dependencies]
 tempfile = "3.14.0"
-
-[dev-dependencies]
-rand.workspace = true
-rand_distr = "0.5.1"

libs/neon-shmem/src/hash.rs

@@ -1,583 +0,0 @@
-//! Resizable hash table implementation on top of byte-level storage (either a [`ShmemHandle`] or a fixed byte array).
-//!
-//! This hash table has two major components: the bucket array and the dictionary. Each bucket within the
-//! bucket array contains a `Option<(K, V)>` and an index of another bucket. In this way there is both an
-//! implicit freelist within the bucket array (`None` buckets point to other `None` entries) and various hash
-//! chains within the bucket array (a Some bucket will point to other Some buckets that had the same hash).
-//!
-//! Buckets are never moved unless they are within a region that is being shrunk, and so the actual hash-
-//! dependent component is done with the dictionary. When a new key is inserted into the map, a position
-//! within the dictionary is decided based on its hash, the data is inserted into an empty bucket based
-//! off of the freelist, and then the index of said bucket is placed in the dictionary.
-//!
-//! This map is resizable (if initialized on top of a [`ShmemHandle`]). Both growing and shrinking happen
-//! in-place and are at a high level achieved by expanding/reducing the bucket array and rebuilding the
-//! dictionary by rehashing all keys.
-//!
-//! Concurrency is managed very simply: the entire map is guarded by one shared-memory RwLock.
-
-use std::hash::{BuildHasher, Hash};
-use std::mem::MaybeUninit;
-
-use crate::shmem::ShmemHandle;
-use crate::{shmem, sync::*};
-
-mod core;
-pub mod entry;
-
-#[cfg(test)]
-mod tests;
-
-use core::{Bucket, CoreHashMap, INVALID_POS};
-use entry::{Entry, OccupiedEntry, PrevPos, VacantEntry};
-
-use thiserror::Error;
-
-/// Error type for a hashmap shrink operation.
-#[derive(Error, Debug)]
-pub enum HashMapShrinkError {
-    /// There was an error encountered while resizing the memory area.
-    #[error("shmem resize failed: {0}")]
-    ResizeError(shmem::Error),
-    /// Occupied entries in to-be-shrunk space were encountered beginning at the given index.
-    #[error("occupied entry in deallocated space found at {0}")]
-    RemainingEntries(usize),
-}
-
-/// This represents a hash table that (possibly) lives in shared memory.
-/// If a new process is launched with fork(), the child process inherits
-/// this struct.
-#[must_use]
-pub struct HashMapInit<'a, K, V, S = rustc_hash::FxBuildHasher> {
-    shmem_handle: Option<ShmemHandle>,
-    shared_ptr: *mut HashMapShared<'a, K, V>,
-    shared_size: usize,
-    hasher: S,
-    num_buckets: u32,
-}
-
-/// This is a per-process handle to a hash table that (possibly) lives in shared memory.
-/// If a child process is launched with fork(), the child process should
-/// get its own HashMapAccess by calling HashMapInit::attach_writer/reader().
-///
-/// XXX: We're not making use of it at the moment, but this struct could
-/// hold process-local information in the future.
-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<K: Sync, V: Sync, S> Sync for HashMapAccess<'_, K, V, S> {}
-unsafe impl<K: Send, V: Send, S> Send for HashMapAccess<'_, K, V, S> {}
-
-impl<'a, K: Clone + Hash + Eq, V, S> HashMapInit<'a, K, V, S> {
-    /// Change the 'hasher' used by the hash table.
-    ///
-    /// NOTE: This must be called right after creating the hash table,
-    /// before inserting any entries and before calling attach_writer/reader.
-    /// Otherwise different accessors could be using different hash function,
-    /// with confusing results.
-    pub fn with_hasher<T: BuildHasher>(self, hasher: T) -> HashMapInit<'a, K, V, T> {
-        HashMapInit {
-            hasher,
-            shmem_handle: self.shmem_handle,
-            shared_ptr: self.shared_ptr,
-            shared_size: self.shared_size,
-            num_buckets: self.num_buckets,
-        }
-    }
-
-    /// Loosely (over)estimate the size needed to store a hash table with `num_buckets` buckets.
-    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
-    }
-
-    fn new(
-        num_buckets: u32,
-        shmem_handle: Option<ShmemHandle>,
-        area_ptr: *mut u8,
-        area_size: usize,
-        hasher: S,
-    ) -> Self {
-        let mut ptr: *mut u8 = area_ptr;
-        let end_ptr: *mut u8 = unsafe { ptr.add(area_size) };
-
-        // carve out area for the One Big Lock (TM) and the HashMapShared.
-        ptr = unsafe { ptr.add(ptr.align_offset(align_of::<libc::pthread_rwlock_t>())) };
-        let raw_lock_ptr = ptr;
-        ptr = unsafe { ptr.add(size_of::<libc::pthread_rwlock_t>()) };
-        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::Bucket<K, V>>())) };
-        let buckets_ptr = ptr;
-        ptr = unsafe { ptr.add(size_of::<core::Bucket<K, V>>() * 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 buckets =
-            unsafe { std::slice::from_raw_parts_mut(buckets_ptr.cast(), num_buckets as usize) };
-        let dictionary = unsafe {
-            std::slice::from_raw_parts_mut(dictionary_ptr.cast(), dictionary_size as usize)
-        };
-
-        let hashmap = CoreHashMap::new(buckets, dictionary);
-        unsafe {
-            let lock = RwLock::from_raw(PthreadRwLock::new(raw_lock_ptr.cast()), hashmap);
-            std::ptr::write(shared_ptr, lock);
-        }
-
-        Self {
-            num_buckets,
-            shmem_handle,
-            shared_ptr,
-            shared_size: area_size,
-            hasher,
-        }
-    }
-
-    /// Attach to a hash table for writing.
-    pub fn attach_writer(self) -> HashMapAccess<'a, K, V, S> {
-        HashMapAccess {
-            shmem_handle: self.shmem_handle,
-            shared_ptr: self.shared_ptr,
-            hasher: self.hasher,
-        }
-    }
-
-    /// Initialize a table for reading. Currently identical to [`HashMapInit::attach_writer`].
-    ///
-    /// This is a holdover from a previous implementation and is being kept around for
-    /// backwards compatibility reasons.
-    pub fn attach_reader(self) -> HashMapAccess<'a, K, V, S> {
-        self.attach_writer()
-    }
-}
-
-/// Hash table data that is actually 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:
-///
-/// [`libc::pthread_rwlock_t`]
-/// [`HashMapShared`]
-/// buckets
-/// dictionary
-///
-/// In between the above parts, there can be padding bytes to align the parts correctly.
-type HashMapShared<'a, K, V> = RwLock<CoreHashMap<'a, K, V>>;
-
-impl<'a, K, V> HashMapInit<'a, K, V, rustc_hash::FxBuildHasher>
-where
-    K: Clone + Hash + Eq,
-{
-    /// Place the hash table within a user-supplied fixed memory area.
-    pub fn with_fixed(num_buckets: u32, area: &'a mut [MaybeUninit<u8>]) -> Self {
-        Self::new(
-            num_buckets,
-            None,
-            area.as_mut_ptr().cast(),
-            area.len(),
-            rustc_hash::FxBuildHasher,
-        )
-    }
-
-    /// Place a new hash map in the given shared memory area
-    ///
-    /// # Panics
-    /// Will panic on failure to resize area to expected map size.
-    pub fn with_shmem(num_buckets: u32, shmem: ShmemHandle) -> Self {
-        let size = Self::estimate_size(num_buckets);
-        shmem
-            .set_size(size)
-            .expect("could not resize shared memory area");
-        let ptr = shmem.data_ptr.as_ptr().cast();
-        Self::new(
-            num_buckets,
-            Some(shmem),
-            ptr,
-            size,
-            rustc_hash::FxBuildHasher,
-        )
-    }
-
-    /// Make a resizable hash map within a new shared memory area with the given name.
-    pub fn new_resizeable_named(num_buckets: u32, max_buckets: u32, name: &str) -> Self {
-        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");
-        let ptr = shmem.data_ptr.as_ptr().cast();
-
-        Self::new(
-            num_buckets,
-            Some(shmem),
-            ptr,
-            size,
-            rustc_hash::FxBuildHasher,
-        )
-    }
-
-    /// Make a resizable hash map within a new anonymous shared memory area.
-    pub fn new_resizeable(num_buckets: u32, max_buckets: u32) -> Self {
-        use std::sync::atomic::{AtomicUsize, Ordering};
-        static 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,
-{
-    /// Hash a key using the map's hasher.
-    #[inline]
-    fn get_hash_value(&self, key: &K) -> u64 {
-        self.hasher.hash_one(key)
-    }
-
-    fn entry_with_hash(&self, key: K, hash: u64) -> Entry<'a, '_, K, V> {
-        let mut map = unsafe { self.shared_ptr.as_ref() }.unwrap().write();
-        let dict_pos = hash as usize % map.dictionary.len();
-        let first = map.dictionary[dict_pos];
-        if first == INVALID_POS {
-            // no existing entry
-            return Entry::Vacant(VacantEntry {
-                map,
-                key,
-                dict_pos: dict_pos as u32,
-            });
-        }
-
-        let mut prev_pos = PrevPos::First(dict_pos as u32);
-        let mut next = first;
-        loop {
-            let bucket = &mut map.buckets[next as usize];
-            let (bucket_key, _bucket_value) = bucket.inner.as_mut().expect("entry is in use");
-            if *bucket_key == key {
-                // found existing entry
-                return Entry::Occupied(OccupiedEntry {
-                    map,
-                    _key: key,
-                    prev_pos,
-                    bucket_pos: next,
-                });
-            }
-
-            if bucket.next == INVALID_POS {
-                // No existing entry
-                return Entry::Vacant(VacantEntry {
-                    map,
-                    key,
-                    dict_pos: dict_pos as u32,
-                });
-            }
-            prev_pos = PrevPos::Chained(next);
-            next = bucket.next;
-        }
-    }
-
-    /// Get a reference to the corresponding value for a key.
-    pub fn get<'e>(&'e self, key: &K) -> Option<ValueReadGuard<'e, V>> {
-        let hash = self.get_hash_value(key);
-        let map = unsafe { self.shared_ptr.as_ref() }.unwrap().read();
-        RwLockReadGuard::try_map(map, |m| m.get_with_hash(key, hash)).ok()
-    }
-
-    /// Get a reference to the entry containing a key.
-    ///
-    /// NB: THis takes a write lock as there's no way to distinguish whether the intention
-    /// is to use the entry for reading or for writing in advance.
-    pub fn entry(&self, key: K) -> Entry<'a, '_, K, V> {
-        let hash = self.get_hash_value(&key);
-        self.entry_with_hash(key, hash)
-    }
-
-    /// Remove a key given its hash. Returns the associated value if it existed.
-    pub fn remove(&self, key: &K) -> Option<V> {
-        let hash = self.get_hash_value(key);
-        match self.entry_with_hash(key.clone(), hash) {
-            Entry::Occupied(e) => Some(e.remove()),
-            Entry::Vacant(_) => None,
-        }
-    }
-
-    /// Insert/update a key. Returns the previous associated value if it existed.
-    ///
-    /// # Errors
-    /// Will return [`core::FullError`] if there is no more space left in the map.
-    pub fn insert(&self, key: K, value: V) -> Result<Option<V>, core::FullError> {
-        let hash = self.get_hash_value(&key);
-        match self.entry_with_hash(key.clone(), hash) {
-            Entry::Occupied(mut e) => Ok(Some(e.insert(value))),
-            Entry::Vacant(e) => {
-                _ = e.insert(value)?;
-                Ok(None)
-            }
-        }
-    }
-
-    /// Optionally return the entry for a bucket at a given index if it exists.
-    ///
-    /// Has more overhead than one would intuitively expect: performs both a clone of the key
-    /// due to the [`OccupiedEntry`] type owning the key and also a hash of the key in order
-    /// to enable repairing the hash chain if the entry is removed.
-    pub fn entry_at_bucket(&self, pos: usize) -> Option<OccupiedEntry<'a, '_, K, V>> {
-        let map = unsafe { self.shared_ptr.as_mut() }.unwrap().write();
-        if pos >= map.buckets.len() {
-            return None;
-        }
-
-        let entry = map.buckets[pos].inner.as_ref();
-        match entry {
-            Some((key, _)) => Some(OccupiedEntry {
-                _key: key.clone(),
-                bucket_pos: pos as u32,
-                prev_pos: entry::PrevPos::Unknown(self.get_hash_value(key)),
-                map,
-            }),
-            _ => None,
-        }
-    }
-
-    /// Returns the number of buckets in the table.
-    pub fn get_num_buckets(&self) -> usize {
-        let map = unsafe { self.shared_ptr.as_ref() }.unwrap().read();
-        map.get_num_buckets()
-    }
-
-    /// Return the key and value stored in bucket with given index. This can be used to
-    /// iterate through the hash map.
-    // TODO: 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<ValueReadGuard<(K, V)>> {
-        let map = unsafe { self.shared_ptr.as_ref() }.unwrap().read();
-        if pos >= map.buckets.len() {
-            return None;
-        }
-        RwLockReadGuard::try_map(map, |m| m.buckets[pos].inner.as_ref()).ok()
-    }
-
-    /// Returns the index of the bucket a given value corresponds to.
-    pub fn get_bucket_for_value(&self, val_ptr: *const V) -> usize {
-        let map = unsafe { self.shared_ptr.as_ref() }.unwrap().read();
-
-        let origin = map.buckets.as_ptr();
-        let idx = (val_ptr as usize - origin as usize) / size_of::<Bucket<K, V>>();
-        assert!(idx < map.buckets.len());
-
-        idx
-    }
-
-    /// Returns the number of occupied buckets in the table.
-    pub fn get_num_buckets_in_use(&self) -> usize {
-        let map = unsafe { self.shared_ptr.as_ref() }.unwrap().read();
-        map.buckets_in_use as usize
-    }
-
-    /// Clears all entries in a table. Does not reset any shrinking operations.
-    pub fn clear(&self) {
-        let mut map = unsafe { self.shared_ptr.as_mut() }.unwrap().write();
-        map.clear();
-    }
-
-    /// Perform an in-place rehash of some region (0..`rehash_buckets`) of the table and reset
-    /// the `buckets` and `dictionary` slices to be as long as `num_buckets`. Resets the freelist
-    /// in the process.
-    fn rehash_dict(
-        &self,
-        inner: &mut CoreHashMap<'a, K, V>,
-        buckets_ptr: *mut core::Bucket<K, V>,
-        end_ptr: *mut u8,
-        num_buckets: u32,
-        rehash_buckets: u32,
-    ) {
-        inner.free_head = INVALID_POS;
-
-        let buckets;
-        let dictionary;
-        unsafe {
-            let buckets_end_ptr = buckets_ptr.add(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>();
-
-            buckets = std::slice::from_raw_parts_mut(buckets_ptr, num_buckets as usize);
-            dictionary = std::slice::from_raw_parts_mut(dictionary_ptr, dictionary_size);
-        }
-        for e in dictionary.iter_mut() {
-            *e = INVALID_POS;
-        }
-
-        for (i, bucket) in buckets.iter_mut().enumerate().take(rehash_buckets as usize) {
-            if bucket.inner.is_none() {
-                bucket.next = inner.free_head;
-                inner.free_head = i as u32;
-                continue;
-            }
-
-            let hash = self.hasher.hash_one(&bucket.inner.as_ref().unwrap().0);
-            let pos: usize = (hash % dictionary.len() as u64) as usize;
-            bucket.next = dictionary[pos];
-            dictionary[pos] = i as u32;
-        }
-
-        inner.dictionary = dictionary;
-        inner.buckets = buckets;
-    }
-
-    /// Rehash the map without growing or shrinking.
-    pub fn shuffle(&self) {
-        let mut map = unsafe { self.shared_ptr.as_mut() }.unwrap().write();
-        let num_buckets = map.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.byte_add(size_bytes).cast() };
-        let buckets_ptr = map.buckets.as_mut_ptr();
-        self.rehash_dict(&mut map, buckets_ptr, end_ptr, num_buckets, num_buckets);
-    }
-
-    /// Grow the number of buckets within the table.
-    ///
-    /// 1. Grows the underlying shared memory area
-    /// 2. Initializes new buckets and overwrites the current dictionary
-    /// 3. Rehashes the dictionary
-    ///
-    /// # Panics
-    /// Panics if called on a map initialized with [`HashMapInit::with_fixed`].
-    ///
-    /// # Errors
-    /// Returns an [`shmem::Error`] if any errors occur resizing the memory region.
-    pub fn grow(&self, num_buckets: u32) -> Result<(), shmem::Error> {
-        let mut map = unsafe { self.shared_ptr.as_mut() }.unwrap().write();
-        let old_num_buckets = map.buckets.len() as u32;
-
-        assert!(
-            num_buckets >= old_num_buckets,
-            "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 buckets_ptr = map.buckets.as_mut_ptr();
-        unsafe {
-            for i in old_num_buckets..num_buckets {
-                let bucket = buckets_ptr.add(i as usize);
-                bucket.write(core::Bucket {
-                    next: if i < num_buckets - 1 {
-                        i + 1
-                    } else {
-                        map.free_head
-                    },
-                    inner: None,
-                });
-            }
-        }
-
-        self.rehash_dict(&mut map, buckets_ptr, end_ptr, num_buckets, old_num_buckets);
-        map.free_head = old_num_buckets;
-
-        Ok(())
-    }
-
-    /// Begin a shrink, limiting all new allocations to be in buckets with index below `num_buckets`.
-    ///
-    /// # Panics
-    /// Panics if called on a map initialized with [`HashMapInit::with_fixed`] or if `num_buckets` is
-    /// greater than the number of buckets in the map.
-    pub fn begin_shrink(&mut self, num_buckets: u32) {
-        let mut map = unsafe { self.shared_ptr.as_mut() }.unwrap().write();
-        assert!(
-            num_buckets <= map.get_num_buckets() as u32,
-            "shrink called with a larger number of buckets"
-        );
-        _ = self
-            .shmem_handle
-            .as_ref()
-            .expect("shrink called on a fixed-size hash table");
-        map.alloc_limit = num_buckets;
-    }
-
-    /// If a shrink operation is underway, returns the target size of the map. Otherwise, returns None.
-    pub fn shrink_goal(&self) -> Option<usize> {
-        let map = unsafe { self.shared_ptr.as_mut() }.unwrap().read();
-        let goal = map.alloc_limit;
-        if goal == INVALID_POS {
-            None
-        } else {
-            Some(goal as usize)
-        }
-    }
-
-    /// Complete a shrink after caller has evicted entries, removing the unused buckets and rehashing.
-    ///
-    /// # Panics
-    /// The following cases result in a panic:
-    /// - Calling this function on a map initialized with [`HashMapInit::with_fixed`].
-    /// - Calling this function on a map when no shrink operation is in progress.
-    pub fn finish_shrink(&self) -> Result<(), HashMapShrinkError> {
-        let mut map = unsafe { self.shared_ptr.as_mut() }.unwrap().write();
-        assert!(
-            map.alloc_limit != INVALID_POS,
-            "called finish_shrink when no shrink is in progress"
-        );
-
-        let num_buckets = map.alloc_limit;
-
-        if map.get_num_buckets() == num_buckets as usize {
-            return Ok(());
-        }
-
-        assert!(
-            map.buckets_in_use <= num_buckets,
-            "called finish_shrink before enough entries were removed"
-        );
-
-        for i in (num_buckets as usize)..map.buckets.len() {
-            if map.buckets[i].inner.is_some() {
-                return Err(HashMapShrinkError::RemainingEntries(i));
-            }
-        }
-
-        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);
-        if let Err(e) = shmem_handle.set_size(size_bytes) {
-            return Err(HashMapShrinkError::ResizeError(e));
-        }
-        let end_ptr: *mut u8 = unsafe { shmem_handle.data_ptr.as_ptr().add(size_bytes) };
-        let buckets_ptr = map.buckets.as_mut_ptr();
-        self.rehash_dict(&mut map, buckets_ptr, end_ptr, num_buckets, num_buckets);
-        map.alloc_limit = INVALID_POS;
-
-        Ok(())
-    }
-}

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

@@ -1,174 +0,0 @@
-//! Simple hash table with chaining.
-
-use std::hash::Hash;
-use std::mem::MaybeUninit;
-
-use crate::hash::entry::*;
-
-/// Invalid position within the map (either within the dictionary or bucket array).
-pub(crate) const INVALID_POS: u32 = u32::MAX;
-
-/// Fundamental storage unit within the hash table. Either empty or contains a key-value pair.
-/// Always part of a chain of some kind (either a freelist if empty or a hash chain if full).
-pub(crate) struct Bucket<K, V> {
-    /// Index of next bucket in the chain.
-    pub(crate) next: u32,
-    /// Key-value pair contained within bucket.
-    pub(crate) inner: Option<(K, V)>,
-}
-
-/// Core hash table implementation.
-pub(crate) struct CoreHashMap<'a, K, V> {
-    /// Dictionary used to map hashes to bucket indices.
-    pub(crate) dictionary: &'a mut [u32],
-    /// Buckets containing key-value pairs.
-    pub(crate) buckets: &'a mut [Bucket<K, V>],
-    /// Head of the freelist.
-    pub(crate) free_head: u32,
-    /// Maximum index of a bucket allowed to be allocated. [`INVALID_POS`] if no limit.
-    pub(crate) alloc_limit: u32,
-    /// The number of currently occupied buckets.
-    pub(crate) buckets_in_use: u32,
-}
-
-/// Error for when there are no empty buckets left but one is needed.
-#[derive(Debug, PartialEq)]
-pub struct FullError;
-
-impl<'a, K: Clone + Hash + Eq, V> CoreHashMap<'a, K, V> {
-    const FILL_FACTOR: f32 = 0.60;
-
-    /// Estimate the size of data contained within the the hash map.
-    pub fn estimate_size(num_buckets: u32) -> usize {
-        let mut size = 0;
-
-        // buckets
-        size += size_of::<Bucket<K, 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(
-        buckets: &'a mut [MaybeUninit<Bucket<K, V>>],
-        dictionary: &'a mut [MaybeUninit<u32>],
-    ) -> Self {
-        // Initialize the buckets
-        for i in 0..buckets.len() {
-            buckets[i].write(Bucket {
-                next: if i < buckets.len() - 1 {
-                    i as u32 + 1
-                } else {
-                    INVALID_POS
-                },
-                inner: None,
-            });
-        }
-
-        // Initialize the dictionary
-        for e in dictionary.iter_mut() {
-            e.write(INVALID_POS);
-        }
-
-        // TODO: use std::slice::assume_init_mut() once it stabilizes
-        let buckets =
-            unsafe { std::slice::from_raw_parts_mut(buckets.as_mut_ptr().cast(), buckets.len()) };
-        let dictionary = unsafe {
-            std::slice::from_raw_parts_mut(dictionary.as_mut_ptr().cast(), dictionary.len())
-        };
-
-        Self {
-            dictionary,
-            buckets,
-            free_head: 0,
-            buckets_in_use: 0,
-            alloc_limit: INVALID_POS,
-        }
-    }
-
-    /// Get the value associated with a key (if it exists) given its hash.
-    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 bucket = &self.buckets[next as usize];
-            let (bucket_key, bucket_value) = bucket.inner.as_ref().expect("entry is in use");
-            if bucket_key == key {
-                return Some(bucket_value);
-            }
-            next = bucket.next;
-        }
-    }
-
-    /// Get number of buckets in map.
-    pub fn get_num_buckets(&self) -> usize {
-        self.buckets.len()
-    }
-
-    /// 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.buckets.len() {
-            self.buckets[i] = Bucket {
-                next: if i < self.buckets.len() - 1 {
-                    i as u32 + 1
-                } else {
-                    INVALID_POS
-                },
-                inner: None,
-            }
-        }
-        for i in 0..self.dictionary.len() {
-            self.dictionary[i] = INVALID_POS;
-        }
-
-        self.free_head = 0;
-        self.buckets_in_use = 0;
-    }
-
-    /// 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 bucket = &mut self.buckets[pos as usize];
-            prev = PrevPos::Chained(pos);
-            pos = bucket.next;
-        }
-        if pos == INVALID_POS {
-            return Err(FullError);
-        }
-
-        // Repair the freelist.
-        match prev {
-            PrevPos::First(_) => {
-                let next_pos = self.buckets[pos as usize].next;
-                self.free_head = next_pos;
-            }
-            PrevPos::Chained(p) => {
-                if p != INVALID_POS {
-                    let next_pos = self.buckets[pos as usize].next;
-                    self.buckets[p as usize].next = next_pos;
-                }
-            }
-            _ => unreachable!(),
-        }
-
-        // Initialize the bucket.
-        let bucket = &mut self.buckets[pos as usize];
-        self.buckets_in_use += 1;
-        bucket.next = INVALID_POS;
-        bucket.inner = Some((key, value));
-
-        Ok(pos)
-    }
-}

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

@@ -1,130 +0,0 @@
-//! Equivalent of [`std::collections::hash_map::Entry`] for this hashmap.
-
-use crate::hash::core::{CoreHashMap, FullError, INVALID_POS};
-use crate::sync::{RwLockWriteGuard, ValueWriteGuard};
-
-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>),
-}
-
-/// Enum representing the previous position within a 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(u64),
-}
-
-pub struct OccupiedEntry<'a, 'b, K, V> {
-    /// Mutable reference to the map containing this entry.
-    pub(crate) map: RwLockWriteGuard<'b, 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`] bucket array.
-    pub(crate) bucket_pos: u32,
-}
-
-impl<K, V> OccupiedEntry<'_, '_, K, V> {
-    pub fn get(&self) -> &V {
-        &self.map.buckets[self.bucket_pos as usize]
-            .inner
-            .as_ref()
-            .unwrap()
-            .1
-    }
-
-    pub fn get_mut(&mut self) -> &mut V {
-        &mut self.map.buckets[self.bucket_pos as usize]
-            .inner
-            .as_mut()
-            .unwrap()
-            .1
-    }
-
-    /// Inserts a value into the entry, replacing (and returning) the existing value.
-    pub fn insert(&mut self, value: V) -> V {
-        let bucket = &mut self.map.buckets[self.bucket_pos as usize];
-        // This assumes inner is Some, which it must be for an OccupiedEntry
-        mem::replace(&mut bucket.inner.as_mut().unwrap().1, value)
-    }
-
-    /// Removes the entry from the hash map, returning the value originally stored within it.
-    ///
-    /// This may result in multiple bucket accesses if the entry was obtained by index as the
-    /// previous chain entry needs to be discovered in this case.
-    pub fn remove(mut self) -> V {
-        // If this bucket was queried by index, go ahead and follow its chain from the start.
-        let prev = if let PrevPos::Unknown(hash) = self.prev_pos {
-            let dict_idx = hash as usize % self.map.dictionary.len();
-            let mut prev = PrevPos::First(dict_idx as u32);
-            let mut curr = self.map.dictionary[dict_idx];
-            while curr != self.bucket_pos {
-                assert!(curr != INVALID_POS);
-                prev = PrevPos::Chained(curr);
-                curr = self.map.buckets[curr as usize].next;
-            }
-            prev
-        } else {
-            self.prev_pos
-        };
-
-        // CoreHashMap::remove returns Option<(K, V)>. We know it's Some for an OccupiedEntry.
-        let bucket = &mut self.map.buckets[self.bucket_pos as usize];
-
-        // unlink it from the chain
-        match prev {
-            PrevPos::First(dict_pos) => {
-                self.map.dictionary[dict_pos as usize] = bucket.next;
-            }
-            PrevPos::Chained(bucket_pos) => {
-                self.map.buckets[bucket_pos as usize].next = bucket.next;
-            }
-            _ => unreachable!(),
-        }
-
-        // and add it to the freelist
-        let free = self.map.free_head;
-        let bucket = &mut self.map.buckets[self.bucket_pos as usize];
-        let old_value = bucket.inner.take();
-        bucket.next = free;
-        self.map.free_head = self.bucket_pos;
-        self.map.buckets_in_use -= 1;
-
-        old_value.unwrap().1
-    }
-}
-
-/// An abstract view into a vacant entry within the map.
-pub struct VacantEntry<'a, 'b, K, V> {
-    /// Mutable reference to the map containing this entry.
-    pub(crate) map: RwLockWriteGuard<'b, CoreHashMap<'a, K, V>>,
-    /// The key to be inserted into this entry.
-    pub(crate) key: K,
-    /// The position within the dictionary corresponding to the key's hash.
-    pub(crate) dict_pos: u32,
-}
-
-impl<'b, K: Clone + Hash + Eq, V> VacantEntry<'_, 'b, K, V> {
-    /// Insert a value into the vacant entry, finding and populating an empty bucket in the process.
-    ///
-    /// # Errors
-    /// Will return [`FullError`] if there are no unoccupied buckets in the map.
-    pub fn insert(mut self, value: V) -> Result<ValueWriteGuard<'b, V>, FullError> {
-        let pos = self.map.alloc_bucket(self.key, value)?;
-        self.map.buckets[pos as usize].next = self.map.dictionary[self.dict_pos as usize];
-        self.map.dictionary[self.dict_pos as usize] = pos;
-
-        Ok(RwLockWriteGuard::map(self.map, |m| {
-            &mut m.buckets[pos as usize].inner.as_mut().unwrap().1
-        }))
-    }
-}

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

@@ -1,428 +0,0 @@
-use std::collections::BTreeMap;
-use std::collections::HashSet;
-use std::fmt::Debug;
-use std::mem::MaybeUninit;
-
-use crate::hash::Entry;
-use crate::hash::HashMapAccess;
-use crate::hash::HashMapInit;
-use crate::hash::core::FullError;
-
-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 w = HashMapInit::<TestKey, usize>::new_resizeable_named(100000, 120000, "test_inserts")
-        .attach_writer();
-
-    for (idx, k) in keys.iter().enumerate() {
-        let res = w.entry((*k).into());
-        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 x = w.get(&(*k).into());
-        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.contains(&key) {
-                continue;
-            }
-            used_keys.insert(key);
-            keys.push(key.into());
-            break;
-        }
-    }
-    test_inserts(&keys);
-}
-
-#[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 entry = map.entry(op.0);
-    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 _ in 0..num_ops {
-        let (k, _) = shadow.pop_first().unwrap();
-        writer.remove(&k);
-    }
-}
-
-fn do_shrink(
-    writer: &mut HashMapAccess<TestKey, usize>,
-    shadow: &mut BTreeMap<TestKey, usize>,
-    from: u32,
-    to: u32,
-) {
-    assert!(writer.shrink_goal().is_none());
-    writer.begin_shrink(to);
-    assert_eq!(writer.shrink_goal(), Some(to as usize));
-    for i in to..from {
-        if let Some(entry) = writer.entry_at_bucket(i as usize) {
-            shadow.remove(&entry._key);
-            entry.remove();
-        }
-    }
-    let old_usage = writer.get_num_buckets_in_use();
-    writer.finish_shrink().unwrap();
-    assert!(writer.shrink_goal().is_none());
-    assert_eq!(writer.get_num_buckets_in_use(), old_usage);
-}
-
-#[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);
-    }
-}
-
-#[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);
-    let old_usage = writer.get_num_buckets_in_use();
-    writer.grow(1500).unwrap();
-    assert_eq!(writer.get_num_buckets_in_use(), old_usage);
-    assert_eq!(writer.get_num_buckets(), 1500);
-    do_random_ops(10000, 1500, 0.75, &mut writer, &mut shadow, &mut rng);
-}
-
-#[test]
-fn test_clear() {
-    let mut writer = HashMapInit::<TestKey, usize>::new_resizeable_named(1500, 2000, "test_clear")
-        .attach_writer();
-    let mut shadow: std::collections::BTreeMap<TestKey, usize> = BTreeMap::new();
-    let mut rng = rand::rng();
-    do_random_ops(2000, 1500, 0.75, &mut writer, &mut shadow, &mut rng);
-    writer.clear();
-    assert_eq!(writer.get_num_buckets_in_use(), 0);
-    assert_eq!(writer.get_num_buckets(), 1500);
-    while let Some((key, _)) = shadow.pop_first() {
-        assert!(writer.get(&key).is_none());
-    }
-    do_random_ops(2000, 1500, 0.75, &mut writer, &mut shadow, &mut rng);
-    for i in 0..(1500 - writer.get_num_buckets_in_use()) {
-        writer.insert((1500 + i as u128).into(), 0).unwrap();
-    }
-    assert_eq!(writer.insert(5000.into(), 0), Err(FullError {}));
-    writer.clear();
-    assert!(writer.insert(5000.into(), 0).is_ok());
-}
-
-#[test]
-fn test_idx_remove() {
-    let mut writer = HashMapInit::<TestKey, usize>::new_resizeable_named(1500, 2000, "test_clear")
-        .attach_writer();
-    let mut shadow: std::collections::BTreeMap<TestKey, usize> = BTreeMap::new();
-    let mut rng = rand::rng();
-    do_random_ops(2000, 1500, 0.25, &mut writer, &mut shadow, &mut rng);
-    for _ in 0..100 {
-        let idx = (rng.next_u32() % 1500) as usize;
-        if let Some(e) = writer.entry_at_bucket(idx) {
-            shadow.remove(&e._key);
-            e.remove();
-        }
-    }
-    while let Some((key, val)) = shadow.pop_first() {
-        assert_eq!(*writer.get(&key).unwrap(), val);
-    }
-}
-
-#[test]
-fn test_idx_get() {
-    let mut writer = HashMapInit::<TestKey, usize>::new_resizeable_named(1500, 2000, "test_clear")
-        .attach_writer();
-    let mut shadow: std::collections::BTreeMap<TestKey, usize> = BTreeMap::new();
-    let mut rng = rand::rng();
-    do_random_ops(2000, 1500, 0.25, &mut writer, &mut shadow, &mut rng);
-    for _ in 0..100 {
-        let idx = (rng.next_u32() % 1500) as usize;
-        if let Some(pair) = writer.get_at_bucket(idx) {
-            {
-                let v: *const usize = &pair.1;
-                assert_eq!(writer.get_bucket_for_value(v), idx);
-            }
-            {
-                let v: *const usize = &pair.1;
-                assert_eq!(writer.get_bucket_for_value(v), idx);
-            }
-        }
-    }
-}
-
-#[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);
-    assert_eq!(writer.get_num_buckets(), 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");
-    while shadow.len() > 100 {
-        do_deletes(1, &mut writer, &mut shadow);
-    }
-    do_shrink(&mut writer, &mut shadow, 1500, 200);
-    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 writer = HashMapInit::<TestKey, usize>::new_resizeable_named(1000, 1200, "test_bucket_ops")
-        .attach_writer();
-    match writer.entry(1.into()) {
-        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(&1.into()).unwrap(), 2);
-    let pos = match writer.entry(1.into()) {
-        Entry::Occupied(e) => {
-            assert_eq!(e._key, 1.into());
-            e.bucket_pos as usize
-        }
-        Entry::Vacant(_) => {
-            panic!("Insert didn't affect entry");
-        }
-    };
-    assert_eq!(writer.entry_at_bucket(pos).unwrap()._key, 1.into());
-    assert_eq!(*writer.get_at_bucket(pos).unwrap(), (1.into(), 2));
-    {
-        let ptr: *const usize = &*writer.get(&1.into()).unwrap();
-        assert_eq!(writer.get_bucket_for_value(ptr), pos);
-    }
-    writer.remove(&1.into());
-    assert!(writer.get(&1.into()).is_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 entry = writer.entry(1.into());
-    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(1.into());
-    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 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 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,3 +1 @@
-pub mod hash;
 pub mod shmem;
-pub mod sync;

libs/neon-shmem/src/sync.rs

@@ -1,111 +0,0 @@
-//! Simple utilities akin to what's in [`std::sync`] but designed to work with shared memory.
-
-use std::mem::MaybeUninit;
-use std::ptr::NonNull;
-
-use nix::errno::Errno;
-
-pub type RwLock<T> = lock_api::RwLock<PthreadRwLock, T>;
-pub type RwLockReadGuard<'a, T> = lock_api::RwLockReadGuard<'a, PthreadRwLock, T>;
-pub type RwLockWriteGuard<'a, T> = lock_api::RwLockWriteGuard<'a, PthreadRwLock, T>;
-pub type ValueReadGuard<'a, T> = lock_api::MappedRwLockReadGuard<'a, PthreadRwLock, T>;
-pub type ValueWriteGuard<'a, T> = lock_api::MappedRwLockWriteGuard<'a, PthreadRwLock, T>;
-
-/// Shared memory read-write lock.
-pub struct PthreadRwLock(Option<NonNull<libc::pthread_rwlock_t>>);
-
-/// Simple macro that calls a function in the libc namespace and panics if return value is nonzero.
-macro_rules! libc_checked {
-    ($fn_name:ident ( $($arg:expr),* )) => {{
-        let res = libc::$fn_name($($arg),*);
-        if res != 0 {
-            panic!("{} failed with {}", stringify!($fn_name), Errno::from_raw(res));
-        }
-    }};
-}
-
-impl PthreadRwLock {
-    /// Creates a new `PthreadRwLock` on top of a pointer to a pthread rwlock.
-    ///
-    /// # Safety
-    /// `lock` must be non-null. Every unsafe operation will panic in the event of an error.
-    pub unsafe fn new(lock: *mut libc::pthread_rwlock_t) -> Self {
-        unsafe {
-            let mut attrs = MaybeUninit::uninit();
-            libc_checked!(pthread_rwlockattr_init(attrs.as_mut_ptr()));
-            libc_checked!(pthread_rwlockattr_setpshared(
-                attrs.as_mut_ptr(),
-                libc::PTHREAD_PROCESS_SHARED
-            ));
-            libc_checked!(pthread_rwlock_init(lock, attrs.as_mut_ptr()));
-            // Safety: POSIX specifies that "any function affecting the attributes
-            // object (including destruction) shall not affect any previously
-            // initialized read-write locks".
-            libc_checked!(pthread_rwlockattr_destroy(attrs.as_mut_ptr()));
-            Self(Some(NonNull::new_unchecked(lock)))
-        }
-    }
-
-    fn inner(&self) -> NonNull<libc::pthread_rwlock_t> {
-        match self.0 {
-            None => {
-                panic!("PthreadRwLock constructed badly - something likely used RawRwLock::INIT")
-            }
-            Some(x) => x,
-        }
-    }
-}
-
-unsafe impl lock_api::RawRwLock for PthreadRwLock {
-    type GuardMarker = lock_api::GuardSend;
-    const INIT: Self = Self(None);
-
-    fn try_lock_shared(&self) -> bool {
-        unsafe {
-            let res = libc::pthread_rwlock_tryrdlock(self.inner().as_ptr());
-            match res {
-                0 => true,
-                libc::EAGAIN => false,
-                _ => panic!(
-                    "pthread_rwlock_tryrdlock failed with {}",
-                    Errno::from_raw(res)
-                ),
-            }
-        }
-    }
-
-    fn try_lock_exclusive(&self) -> bool {
-        unsafe {
-            let res = libc::pthread_rwlock_trywrlock(self.inner().as_ptr());
-            match res {
-                0 => true,
-                libc::EAGAIN => false,
-                _ => panic!("try_wrlock failed with {}", Errno::from_raw(res)),
-            }
-        }
-    }
-
-    fn lock_shared(&self) {
-        unsafe {
-            libc_checked!(pthread_rwlock_rdlock(self.inner().as_ptr()));
-        }
-    }
-
-    fn lock_exclusive(&self) {
-        unsafe {
-            libc_checked!(pthread_rwlock_wrlock(self.inner().as_ptr()));
-        }
-    }
-
-    unsafe fn unlock_exclusive(&self) {
-        unsafe {
-            libc_checked!(pthread_rwlock_unlock(self.inner().as_ptr()));
-        }
-    }
-
-    unsafe fn unlock_shared(&self) {
-        unsafe {
-            libc_checked!(pthread_rwlock_unlock(self.inner().as_ptr()));
-        }
-    }
-}

libs/pageserver_api/src/config.rs

@@ -394,7 +394,7 @@ impl From<&OtelExporterConfig> for tracing_utils::ExportConfig {
         tracing_utils::ExportConfig {
             endpoint: Some(val.endpoint.clone()),
             protocol: val.protocol.into(),
-            timeout: Some(val.timeout),
+            timeout: val.timeout,
         }
     }
 }

libs/pageserver_api/src/controller_api.rs

@@ -596,7 +596,6 @@ pub struct TimelineImportRequest {
     pub timeline_id: TimelineId,
     pub start_lsn: Lsn,
     pub sk_set: Vec<NodeId>,
-    pub force_upsert: bool,
 }
 
 #[derive(serde::Serialize, serde::Deserialize, Clone)]

libs/pageserver_api/src/key.rs

@@ -981,12 +981,12 @@ mod tests {
         let mut rng = rand::rngs::StdRng::seed_from_u64(42);
 
         let key = Key {
-            field1: rng.random(),
-            field2: rng.random(),
-            field3: rng.random(),
-            field4: rng.random(),
-            field5: rng.random(),
-            field6: rng.random(),
+            field1: rng.r#gen(),
+            field2: rng.r#gen(),
+            field3: rng.r#gen(),
+            field4: rng.r#gen(),
+            field5: rng.r#gen(),
+            field6: rng.r#gen(),
         };
 
         assert_eq!(key, Key::from_str(&format!("{key}")).unwrap());

libs/pageserver_api/src/models.rs

@@ -443,9 +443,9 @@ pub struct ImportPgdataIdempotencyKey(pub String);
 impl ImportPgdataIdempotencyKey {
     pub fn random() -> Self {
         use rand::Rng;
-        use rand::distr::Alphanumeric;
+        use rand::distributions::Alphanumeric;
         Self(
-            rand::rng()
+            rand::thread_rng()
                 .sample_iter(&Alphanumeric)
                 .take(20)
                 .map(char::from)
@@ -1500,7 +1500,6 @@ pub struct TimelineArchivalConfigRequest {
 #[derive(Serialize, Deserialize, PartialEq, Eq, Clone)]
 pub struct TimelinePatchIndexPartRequest {
     pub rel_size_migration: Option<RelSizeMigration>,
-    pub rel_size_migrated_at: Option<Lsn>,
     pub gc_compaction_last_completed_lsn: Option<Lsn>,
     pub applied_gc_cutoff_lsn: Option<Lsn>,
     #[serde(default)]
@@ -1534,10 +1533,10 @@ pub enum RelSizeMigration {
     /// `None` is the same as `Some(RelSizeMigration::Legacy)`.
     Legacy,
     /// The tenant is migrating to the new rel_size format. Both old and new rel_size format are
-    /// persisted in the storage. The read path will read both formats and validate them.
+    /// persisted in the index part. The read path will read both formats and merge them.
     Migrating,
     /// The tenant has migrated to the new rel_size format. Only the new rel_size format is persisted
-    /// in the storage, and the read path will not read the old format.
+    /// in the index part, and the read path will not read the old format.
     Migrated,
 }
 
@@ -1620,7 +1619,6 @@ pub struct TimelineInfo {
 
     /// The status of the rel_size migration.
     pub rel_size_migration: Option<RelSizeMigration>,
-    pub rel_size_migrated_at: Option<Lsn>,
 
     /// Whether the timeline is invisible in synthetic size calculations.
     pub is_invisible: Option<bool>,

libs/pageserver_api/src/shard.rs

@@ -69,6 +69,22 @@ impl Hash for ShardIdentity {
     }
 }
 
+/// Stripe size in number of pages
+#[derive(Clone, Copy, Serialize, Deserialize, Eq, PartialEq, Debug)]
+pub struct ShardStripeSize(pub u32);
+
+impl Default for ShardStripeSize {
+    fn default() -> Self {
+        DEFAULT_STRIPE_SIZE
+    }
+}
+
+impl std::fmt::Display for ShardStripeSize {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
 /// Layout version: for future upgrades where we might change how the key->shard mapping works
 #[derive(Clone, Copy, Serialize, Deserialize, Eq, PartialEq, Hash, Debug)]
 pub struct ShardLayout(u8);

libs/pageserver_api/src/upcall_api.rs

@@ -21,14 +21,6 @@ pub struct ReAttachRequest {
     /// if the node already has a node_id set.
     #[serde(skip_serializing_if = "Option::is_none", default)]
     pub register: Option<NodeRegisterRequest>,
-
-    /// Hadron: Optional flag to indicate whether the node is starting with an empty local disk.
-    /// Will be set to true if the node couldn't find any local tenant data on startup, could be
-    /// due to the node starting for the first time or due to a local SSD failure/disk wipe event.
-    /// The flag may be used by the storage controller to update its observed state of the world
-    /// to make sure that it sends explicit location_config calls to the node following the
-    /// re-attach request.
-    pub empty_local_disk: Option<bool>,
 }
 
 #[derive(Serialize, Deserialize, Debug)]

libs/postgres_backend/src/lib.rs

@@ -749,18 +749,7 @@ impl<IO: AsyncRead + AsyncWrite + Unpin> PostgresBackend<IO> {
                 trace!("got query {query_string:?}");
                 if let Err(e) = handler.process_query(self, query_string).await {
                     match e {
-                        err @ QueryError::Shutdown => {
-                            // Notify postgres of the connection shutdown at the libpq
-                            // protocol level. This avoids postgres having to tell apart
-                            // from an idle connection and a stale one, which is bug prone.
-                            let shutdown_error = short_error(&err);
-                            self.write_message_noflush(&BeMessage::ErrorResponse(
-                                &shutdown_error,
-                                Some(err.pg_error_code()),
-                            ))?;
-
-                            return Ok(ProcessMsgResult::Break);
-                        }
+                        QueryError::Shutdown => return Ok(ProcessMsgResult::Break),
                         QueryError::SimulatedConnectionError => {
                             return Err(QueryError::SimulatedConnectionError);
                         }

libs/pq_proto/src/lib.rs

@@ -203,12 +203,12 @@ impl fmt::Display for CancelKeyData {
     }
 }
 
-use rand::distr::{Distribution, StandardUniform};
-impl Distribution<CancelKeyData> for StandardUniform {
+use rand::distributions::{Distribution, Standard};
+impl Distribution<CancelKeyData> for Standard {
     fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> CancelKeyData {
         CancelKeyData {
-            backend_pid: rng.random(),
-            cancel_key: rng.random(),
+            backend_pid: rng.r#gen(),
+            cancel_key: rng.r#gen(),
         }
     }
 }

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

@@ -155,10 +155,10 @@ pub struct ScramSha256 {
 
 fn nonce() -> String {
     // rand 0.5's ThreadRng is cryptographically secure
-    let mut rng = rand::rng();
+    let mut rng = rand::thread_rng();
     (0..NONCE_LENGTH)
         .map(|_| {
-            let mut v = rng.random_range(0x21u8..0x7e);
+            let mut v = rng.gen_range(0x21u8..0x7e);
             if v == 0x2c {
                 v = 0x7e
             }

libs/proxy/postgres-protocol2/src/message/backend.rs

@@ -74,6 +74,7 @@ impl Header {
 }
 
 /// An enum representing Postgres backend messages.
+#[non_exhaustive]
 pub enum Message {
     AuthenticationCleartextPassword,
     AuthenticationGss,
@@ -144,7 +145,16 @@ impl Message {
             PARSE_COMPLETE_TAG => Message::ParseComplete,
             BIND_COMPLETE_TAG => Message::BindComplete,
             CLOSE_COMPLETE_TAG => Message::CloseComplete,
-            NOTIFICATION_RESPONSE_TAG => Message::NotificationResponse(NotificationResponseBody {}),
+            NOTIFICATION_RESPONSE_TAG => {
+                let process_id = buf.read_i32::<BigEndian>()?;
+                let channel = buf.read_cstr()?;
+                let message = buf.read_cstr()?;
+                Message::NotificationResponse(NotificationResponseBody {
+                    process_id,
+                    channel,
+                    message,
+                })
+            }
             COPY_DONE_TAG => Message::CopyDone,
             COMMAND_COMPLETE_TAG => {
                 let tag = buf.read_cstr()?;
@@ -533,7 +543,28 @@ impl NoticeResponseBody {
     }
 }
 
-pub struct NotificationResponseBody {}
+pub struct NotificationResponseBody {
+    process_id: i32,
+    channel: Bytes,
+    message: Bytes,
+}
+
+impl NotificationResponseBody {
+    #[inline]
+    pub fn process_id(&self) -> i32 {
+        self.process_id
+    }
+
+    #[inline]
+    pub fn channel(&self) -> io::Result<&str> {
+        get_str(&self.channel)
+    }
+
+    #[inline]
+    pub fn message(&self) -> io::Result<&str> {
+        get_str(&self.message)
+    }
+}
 
 pub struct ParameterDescriptionBody {
     storage: Bytes,

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

@@ -28,7 +28,7 @@ const SCRAM_DEFAULT_SALT_LEN: usize = 16;
 /// special characters that would require escaping in an SQL command.
 pub async fn scram_sha_256(password: &[u8]) -> String {
     let mut salt: [u8; SCRAM_DEFAULT_SALT_LEN] = [0; SCRAM_DEFAULT_SALT_LEN];
-    let mut rng = rand::rng();
+    let mut rng = rand::thread_rng();
     rng.fill_bytes(&mut salt);
     scram_sha_256_salt(password, salt).await
 }

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

@@ -13,7 +13,7 @@ use serde::{Deserialize, Serialize};
 use tokio::sync::mpsc;
 
 use crate::cancel_token::RawCancelToken;
-use crate::codec::{BackendMessages, FrontendMessage, RecordNotices};
+use crate::codec::{BackendMessages, FrontendMessage};
 use crate::config::{Host, SslMode};
 use crate::query::RowStream;
 use crate::simple_query::SimpleQueryStream;
@@ -221,18 +221,6 @@ impl Client {
         &mut self.inner
     }
 
-    pub fn record_notices(&mut self, limit: usize) -> mpsc::UnboundedReceiver<Box<str>> {
-        let (tx, rx) = mpsc::unbounded_channel();
-
-        let notices = RecordNotices { sender: tx, limit };
-        self.inner
-            .sender
-            .send(FrontendMessage::RecordNotices(notices))
-            .ok();
-
-        rx
-    }
-
     /// Pass text directly to the Postgres backend to allow it to sort out typing itself and
     /// to save a roundtrip
     pub async fn query_raw_txt<S, I>(
@@ -292,32 +280,8 @@ impl Client {
         simple_query::batch_execute(self.inner_mut(), query).await
     }
 
-    /// Similar to `discard_all`, but it does not clear any query plans
-    ///
-    /// This runs in the background, so it can be executed without `await`ing.
-    pub fn reset_session_background(&mut self) -> Result<(), Error> {
-        // "CLOSE ALL": closes any cursors
-        // "SET SESSION AUTHORIZATION DEFAULT": resets the current_user back to the session_user
-        // "RESET ALL": resets any GUCs back to their session defaults.
-        // "DEALLOCATE ALL": deallocates any prepared statements
-        // "UNLISTEN *": stops listening on all channels
-        // "SELECT pg_advisory_unlock_all();": unlocks all advisory locks
-        // "DISCARD TEMP;": drops all temporary tables
-        // "DISCARD SEQUENCES;": deallocates all cached sequence state
-
-        let _responses = self.inner_mut().send_simple_query(
-            "ROLLBACK;
-            CLOSE ALL;
-            SET SESSION AUTHORIZATION DEFAULT;
-            RESET ALL;
-            DEALLOCATE ALL;
-            UNLISTEN *;
-            SELECT pg_advisory_unlock_all();
-            DISCARD TEMP;
-            DISCARD SEQUENCES;",
-        )?;
-
-        Ok(())
+    pub async fn discard_all(&mut self) -> Result<ReadyForQueryStatus, Error> {
+        self.batch_execute("discard all").await
     }
 
     /// Begins a new database transaction.

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

@@ -3,17 +3,10 @@ use std::io;
 use bytes::{Bytes, BytesMut};
 use fallible_iterator::FallibleIterator;
 use postgres_protocol2::message::backend;
-use tokio::sync::mpsc::UnboundedSender;
 use tokio_util::codec::{Decoder, Encoder};
 
 pub enum FrontendMessage {
     Raw(Bytes),
-    RecordNotices(RecordNotices),
-}
-
-pub struct RecordNotices {
-    pub sender: UnboundedSender<Box<str>>,
-    pub limit: usize,
 }
 
 pub enum BackendMessage {
@@ -40,11 +33,14 @@ impl FallibleIterator for BackendMessages {
 
 pub struct PostgresCodec;
 
-impl Encoder<Bytes> for PostgresCodec {
+impl Encoder<FrontendMessage> for PostgresCodec {
     type Error = io::Error;
 
-    fn encode(&mut self, item: Bytes, dst: &mut BytesMut) -> io::Result<()> {
-        dst.extend_from_slice(&item);
+    fn encode(&mut self, item: FrontendMessage, dst: &mut BytesMut) -> io::Result<()> {
+        match item {
+            FrontendMessage::Raw(buf) => dst.extend_from_slice(&buf),
+        }
+
         Ok(())
     }
 }

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

@@ -11,8 +11,9 @@ use tokio::io::{AsyncRead, AsyncWrite};
 use tokio::net::TcpStream;
 
 use crate::connect::connect;
-use crate::connect_raw::{self, StartupStream};
+use crate::connect_raw::{RawConnection, connect_raw};
 use crate::connect_tls::connect_tls;
+use crate::maybe_tls_stream::MaybeTlsStream;
 use crate::tls::{MakeTlsConnect, TlsConnect, TlsStream};
 use crate::{Client, Connection, Error};
 
@@ -243,26 +244,24 @@ impl Config {
         &self,
         stream: S,
         tls: T,
-    ) -> Result<StartupStream<S, T::Stream>, Error>
+    ) -> Result<RawConnection<S, T::Stream>, Error>
     where
         S: AsyncRead + AsyncWrite + Unpin,
         T: TlsConnect<S>,
     {
         let stream = connect_tls(stream, self.ssl_mode, tls).await?;
-        let mut stream = StartupStream::new(stream);
-        connect_raw::startup(&mut stream, self).await?;
-        connect_raw::authenticate(&mut stream, self).await?;
-
-        Ok(stream)
+        connect_raw(stream, self).await
     }
 
-    pub async fn authenticate<S, T>(&self, stream: &mut StartupStream<S, T>) -> Result<(), Error>
+    pub async fn authenticate<S, T>(
+        &self,
+        stream: MaybeTlsStream<S, T>,
+    ) -> Result<RawConnection<S, T>, Error>
     where
         S: AsyncRead + AsyncWrite + Unpin,
         T: TlsStream + Unpin,
     {
-        connect_raw::startup(stream, self).await?;
-        connect_raw::authenticate(stream, self).await
+        connect_raw(stream, self).await
     }
 }
 

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

@@ -1,17 +1,17 @@
 use std::net::IpAddr;
 
-use futures_util::TryStreamExt;
 use postgres_protocol2::message::backend::Message;
-use tokio::io::{AsyncRead, AsyncWrite};
 use tokio::net::TcpStream;
 use tokio::sync::mpsc;
 
 use crate::client::SocketConfig;
+use crate::codec::BackendMessage;
 use crate::config::Host;
-use crate::connect_raw::StartupStream;
+use crate::connect_raw::connect_raw;
 use crate::connect_socket::connect_socket;
+use crate::connect_tls::connect_tls;
 use crate::tls::{MakeTlsConnect, TlsConnect};
-use crate::{Client, Config, Connection, Error};
+use crate::{Client, Config, Connection, Error, RawConnection};
 
 pub async fn connect<T>(
     tls: &T,
@@ -45,8 +45,14 @@ where
     T: TlsConnect<TcpStream>,
 {
     let socket = connect_socket(host_addr, host, port, config.connect_timeout).await?;
-    let mut stream = config.tls_and_authenticate(socket, tls).await?;
-    let (process_id, secret_key) = wait_until_ready(&mut stream).await?;
+    let stream = connect_tls(socket, config.ssl_mode, tls).await?;
+    let RawConnection {
+        stream,
+        parameters,
+        delayed_notice,
+        process_id,
+        secret_key,
+    } = connect_raw(stream, config).await?;
 
     let socket_config = SocketConfig {
         host_addr,
@@ -66,32 +72,13 @@ where
         secret_key,
     );
 
-    let stream = stream.into_framed();
-    let connection = Connection::new(stream, conn_tx, conn_rx);
+    // delayed notices are always sent as "Async" messages.
+    let delayed = delayed_notice
+        .into_iter()
+        .map(|m| BackendMessage::Async(Message::NoticeResponse(m)))
+        .collect();
+
+    let connection = Connection::new(stream, delayed, parameters, conn_tx, conn_rx);
 
     Ok((client, connection))
 }
-
-async fn wait_until_ready<S, T>(stream: &mut StartupStream<S, T>) -> Result<(i32, i32), Error>
-where
-    S: AsyncRead + AsyncWrite + Unpin,
-    T: AsyncRead + AsyncWrite + Unpin,
-{
-    let mut process_id = 0;
-    let mut secret_key = 0;
-
-    loop {
-        match stream.try_next().await.map_err(Error::io)? {
-            Some(Message::BackendKeyData(body)) => {
-                process_id = body.process_id();
-                secret_key = body.secret_key();
-            }
-            // These values are currently not used by `Client`/`Connection`. Ignore them.
-            Some(Message::ParameterStatus(_)) | Some(Message::NoticeResponse(_)) => {}
-            Some(Message::ReadyForQuery(_)) => return Ok((process_id, secret_key)),
-            Some(Message::ErrorResponse(body)) => return Err(Error::db(body)),
-            Some(_) => return Err(Error::unexpected_message()),
-            None => return Err(Error::closed()),
-        }
-    }
-}

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

@@ -1,29 +1,31 @@
+use std::collections::HashMap;
 use std::io;
 use std::pin::Pin;
-use std::task::{Context, Poll, ready};
+use std::task::{Context, Poll};
 
-use bytes::{Bytes, BytesMut};
+use bytes::BytesMut;
 use fallible_iterator::FallibleIterator;
-use futures_util::{Sink, SinkExt, Stream, TryStreamExt};
+use futures_util::{Sink, SinkExt, Stream, TryStreamExt, ready};
 use postgres_protocol2::authentication::sasl;
 use postgres_protocol2::authentication::sasl::ScramSha256;
-use postgres_protocol2::message::backend::{AuthenticationSaslBody, Message};
+use postgres_protocol2::message::backend::{AuthenticationSaslBody, Message, NoticeResponseBody};
 use postgres_protocol2::message::frontend;
-use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
-use tokio_util::codec::{Framed, FramedParts, FramedWrite};
+use tokio::io::{AsyncRead, AsyncWrite};
+use tokio_util::codec::Framed;
 
 use crate::Error;
-use crate::codec::PostgresCodec;
+use crate::codec::{BackendMessage, BackendMessages, FrontendMessage, PostgresCodec};
 use crate::config::{self, AuthKeys, Config};
 use crate::maybe_tls_stream::MaybeTlsStream;
 use crate::tls::TlsStream;
 
 pub struct StartupStream<S, T> {
-    inner: FramedWrite<MaybeTlsStream<S, T>, PostgresCodec>,
-    read_buf: BytesMut,
+    inner: Framed<MaybeTlsStream<S, T>, PostgresCodec>,
+    buf: BackendMessages,
+    delayed_notice: Vec<NoticeResponseBody>,
 }
 
-impl<S, T> Sink<Bytes> for StartupStream<S, T>
+impl<S, T> Sink<FrontendMessage> for StartupStream<S, T>
 where
     S: AsyncRead + AsyncWrite + Unpin,
     T: AsyncRead + AsyncWrite + Unpin,
@@ -34,7 +36,7 @@ where
         Pin::new(&mut self.inner).poll_ready(cx)
     }
 
-    fn start_send(mut self: Pin<&mut Self>, item: Bytes) -> io::Result<()> {
+    fn start_send(mut self: Pin<&mut Self>, item: FrontendMessage) -> io::Result<()> {
         Pin::new(&mut self.inner).start_send(item)
     }
 
@@ -54,93 +56,63 @@ where
 {
     type Item = io::Result<Message>;
 
-    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        // read 1 byte tag, 4 bytes length.
-        let header = ready!(self.as_mut().poll_fill_buf_exact(cx, 5)?);
+    fn poll_next(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<io::Result<Message>>> {
+        loop {
+            match self.buf.next() {
+                Ok(Some(message)) => return Poll::Ready(Some(Ok(message))),
+                Ok(None) => {}
+                Err(e) => return Poll::Ready(Some(Err(e))),
+            }
 
-        let len = u32::from_be_bytes(header[1..5].try_into().unwrap());
-        if len < 4 {
-            return Poll::Ready(Some(Err(std::io::Error::other(
-                "postgres message too small",
-            ))));
+            match ready!(Pin::new(&mut self.inner).poll_next(cx)) {
+                Some(Ok(BackendMessage::Normal { messages, .. })) => self.buf = messages,
+                Some(Ok(BackendMessage::Async(message))) => return Poll::Ready(Some(Ok(message))),
+                Some(Err(e)) => return Poll::Ready(Some(Err(e))),
+                None => return Poll::Ready(None),
+            }
         }
-        if len >= 65536 {
-            return Poll::Ready(Some(Err(std::io::Error::other(
-                "postgres message too large",
-            ))));
-        }
-
-        // the tag is an additional byte.
-        let _message = ready!(self.as_mut().poll_fill_buf_exact(cx, len as usize + 1)?);
-
-        // Message::parse will remove the all the bytes from the buffer.
-        Poll::Ready(Message::parse(&mut self.read_buf).transpose())
     }
 }
 
-impl<S, T> StartupStream<S, T>
+pub struct RawConnection<S, T> {
+    pub stream: Framed<MaybeTlsStream<S, T>, PostgresCodec>,
+    pub parameters: HashMap<String, String>,
+    pub delayed_notice: Vec<NoticeResponseBody>,
+    pub process_id: i32,
+    pub secret_key: i32,
+}
+
+pub async fn connect_raw<S, T>(
+    stream: MaybeTlsStream<S, T>,
+    config: &Config,
+) -> Result<RawConnection<S, T>, Error>
 where
     S: AsyncRead + AsyncWrite + Unpin,
-    T: AsyncRead + AsyncWrite + Unpin,
+    T: TlsStream + Unpin,
 {
-    /// Fill the buffer until it's the exact length provided. No additional data will be read from the socket.
-    ///
-    /// If the current buffer length is greater, nothing happens.
-    fn poll_fill_buf_exact(
-        self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-        len: usize,
-    ) -> Poll<Result<&[u8], std::io::Error>> {
-        let this = self.get_mut();
-        let mut stream = Pin::new(this.inner.get_mut());
+    let mut stream = StartupStream {
+        inner: Framed::new(stream, PostgresCodec),
+        buf: BackendMessages::empty(),
+        delayed_notice: Vec::new(),
+    };
 
-        let mut n = this.read_buf.len();
-        while n < len {
-            this.read_buf.resize(len, 0);
+    startup(&mut stream, config).await?;
+    authenticate(&mut stream, config).await?;
+    let (process_id, secret_key, parameters) = read_info(&mut stream).await?;
 
-            let mut buf = ReadBuf::new(&mut this.read_buf[..]);
-            buf.set_filled(n);
-
-            if stream.as_mut().poll_read(cx, &mut buf)?.is_pending() {
-                this.read_buf.truncate(n);
-                return Poll::Pending;
-            }
-
-            if buf.filled().len() == n {
-                return Poll::Ready(Err(std::io::Error::new(
-                    std::io::ErrorKind::UnexpectedEof,
-                    "early eof",
-                )));
-            }
-            n = buf.filled().len();
-
-            this.read_buf.truncate(n);
-        }
-
-        Poll::Ready(Ok(&this.read_buf[..len]))
-    }
-
-    pub fn into_framed(mut self) -> Framed<MaybeTlsStream<S, T>, PostgresCodec> {
-        let write_buf = std::mem::take(self.inner.write_buffer_mut());
-        let io = self.inner.into_inner();
-        let mut parts = FramedParts::new(io, PostgresCodec);
-        parts.read_buf = self.read_buf;
-        parts.write_buf = write_buf;
-        Framed::from_parts(parts)
-    }
-
-    pub fn new(io: MaybeTlsStream<S, T>) -> Self {
-        Self {
-            inner: FramedWrite::new(io, PostgresCodec),
-            read_buf: BytesMut::new(),
-        }
-    }
+    Ok(RawConnection {
+        stream: stream.inner,
+        parameters,
+        delayed_notice: stream.delayed_notice,
+        process_id,
+        secret_key,
+    })
 }
 
-pub(crate) async fn startup<S, T>(
-    stream: &mut StartupStream<S, T>,
-    config: &Config,
-) -> Result<(), Error>
+async fn startup<S, T>(stream: &mut StartupStream<S, T>, config: &Config) -> Result<(), Error>
 where
     S: AsyncRead + AsyncWrite + Unpin,
     T: AsyncRead + AsyncWrite + Unpin,
@@ -148,13 +120,13 @@ where
     let mut buf = BytesMut::new();
     frontend::startup_message(&config.server_params, &mut buf).map_err(Error::encode)?;
 
-    stream.send(buf.freeze()).await.map_err(Error::io)
+    stream
+        .send(FrontendMessage::Raw(buf.freeze()))
+        .await
+        .map_err(Error::io)
 }
 
-pub(crate) async fn authenticate<S, T>(
-    stream: &mut StartupStream<S, T>,
-    config: &Config,
-) -> Result<(), Error>
+async fn authenticate<S, T>(stream: &mut StartupStream<S, T>, config: &Config) -> Result<(), Error>
 where
     S: AsyncRead + AsyncWrite + Unpin,
     T: TlsStream + Unpin,
@@ -219,7 +191,10 @@ where
     let mut buf = BytesMut::new();
     frontend::password_message(password, &mut buf).map_err(Error::encode)?;
 
-    stream.send(buf.freeze()).await.map_err(Error::io)
+    stream
+        .send(FrontendMessage::Raw(buf.freeze()))
+        .await
+        .map_err(Error::io)
 }
 
 async fn authenticate_sasl<S, T>(
@@ -278,7 +253,10 @@ where
 
     let mut buf = BytesMut::new();
     frontend::sasl_initial_response(mechanism, scram.message(), &mut buf).map_err(Error::encode)?;
-    stream.send(buf.freeze()).await.map_err(Error::io)?;
+    stream
+        .send(FrontendMessage::Raw(buf.freeze()))
+        .await
+        .map_err(Error::io)?;
 
     let body = match stream.try_next().await.map_err(Error::io)? {
         Some(Message::AuthenticationSaslContinue(body)) => body,
@@ -294,7 +272,10 @@ where
 
     let mut buf = BytesMut::new();
     frontend::sasl_response(scram.message(), &mut buf).map_err(Error::encode)?;
-    stream.send(buf.freeze()).await.map_err(Error::io)?;
+    stream
+        .send(FrontendMessage::Raw(buf.freeze()))
+        .await
+        .map_err(Error::io)?;
 
     let body = match stream.try_next().await.map_err(Error::io)? {
         Some(Message::AuthenticationSaslFinal(body)) => body,
@@ -309,3 +290,35 @@ where
 
     Ok(())
 }
+
+async fn read_info<S, T>(
+    stream: &mut StartupStream<S, T>,
+) -> Result<(i32, i32, HashMap<String, String>), Error>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+    T: AsyncRead + AsyncWrite + Unpin,
+{
+    let mut process_id = 0;
+    let mut secret_key = 0;
+    let mut parameters = HashMap::new();
+
+    loop {
+        match stream.try_next().await.map_err(Error::io)? {
+            Some(Message::BackendKeyData(body)) => {
+                process_id = body.process_id();
+                secret_key = body.secret_key();
+            }
+            Some(Message::ParameterStatus(body)) => {
+                parameters.insert(
+                    body.name().map_err(Error::parse)?.to_string(),
+                    body.value().map_err(Error::parse)?.to_string(),
+                );
+            }
+            Some(Message::NoticeResponse(body)) => stream.delayed_notice.push(body),
+            Some(Message::ReadyForQuery(_)) => return Ok((process_id, secret_key, parameters)),
+            Some(Message::ErrorResponse(body)) => return Err(Error::db(body)),
+            Some(_) => return Err(Error::unexpected_message()),
+            None => return Err(Error::closed()),
+        }
+    }
+}

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

@@ -1,23 +1,22 @@
+use std::collections::{HashMap, VecDeque};
 use std::future::Future;
 use std::pin::Pin;
 use std::task::{Context, Poll};
 
 use bytes::BytesMut;
-use fallible_iterator::FallibleIterator;
-use futures_util::{Sink, StreamExt, ready};
-use postgres_protocol2::message::backend::{Message, NoticeResponseBody};
+use futures_util::{Sink, Stream, ready};
+use postgres_protocol2::message::backend::Message;
 use postgres_protocol2::message::frontend;
 use tokio::io::{AsyncRead, AsyncWrite};
 use tokio::sync::mpsc;
 use tokio_util::codec::Framed;
 use tokio_util::sync::PollSender;
-use tracing::trace;
+use tracing::{info, trace};
 
-use crate::Error;
-use crate::codec::{
-    BackendMessage, BackendMessages, FrontendMessage, PostgresCodec, RecordNotices,
-};
+use crate::codec::{BackendMessage, BackendMessages, FrontendMessage, PostgresCodec};
+use crate::error::DbError;
 use crate::maybe_tls_stream::MaybeTlsStream;
+use crate::{AsyncMessage, Error, Notification};
 
 #[derive(PartialEq, Debug)]
 enum State {
@@ -34,18 +33,18 @@ enum State {
 /// occurred, or because its associated `Client` has dropped and all outstanding work has completed.
 #[must_use = "futures do nothing unless polled"]
 pub struct Connection<S, T> {
-    stream: Framed<MaybeTlsStream<S, T>, PostgresCodec>,
+    /// HACK: we need this in the Neon Proxy.
+    pub stream: Framed<MaybeTlsStream<S, T>, PostgresCodec>,
+    /// HACK: we need this in the Neon Proxy to forward params.
+    pub parameters: HashMap<String, String>,
 
     sender: PollSender<BackendMessages>,
     receiver: mpsc::UnboundedReceiver<FrontendMessage>,
-    notices: Option<RecordNotices>,
 
-    pending_response: Option<BackendMessages>,
+    pending_responses: VecDeque<BackendMessage>,
     state: State,
 }
 
-pub enum Never {}
-
 impl<S, T> Connection<S, T>
 where
     S: AsyncRead + AsyncWrite + Unpin,
@@ -53,44 +52,72 @@ where
 {
     pub(crate) fn new(
         stream: Framed<MaybeTlsStream<S, T>, PostgresCodec>,
+        pending_responses: VecDeque<BackendMessage>,
+        parameters: HashMap<String, String>,
         sender: mpsc::Sender<BackendMessages>,
         receiver: mpsc::UnboundedReceiver<FrontendMessage>,
     ) -> Connection<S, T> {
         Connection {
             stream,
+            parameters,
             sender: PollSender::new(sender),
             receiver,
-            notices: None,
-            pending_response: None,
+            pending_responses,
             state: State::Active,
         }
     }
 
+    fn poll_response(
+        &mut self,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<BackendMessage, Error>>> {
+        if let Some(message) = self.pending_responses.pop_front() {
+            trace!("retrying pending response");
+            return Poll::Ready(Some(Ok(message)));
+        }
+
+        Pin::new(&mut self.stream)
+            .poll_next(cx)
+            .map(|o| o.map(|r| r.map_err(Error::io)))
+    }
+
     /// Read and process messages from the connection to postgres.
     /// client <- postgres
-    fn poll_read(&mut self, cx: &mut Context<'_>) -> Poll<Result<Never, Error>> {
+    fn poll_read(&mut self, cx: &mut Context<'_>) -> Poll<Result<AsyncMessage, Error>> {
         loop {
-            let messages = match self.pending_response.take() {
-                Some(messages) => messages,
-                None => {
-                    let message = match self.stream.poll_next_unpin(cx) {
-                        Poll::Pending => return Poll::Pending,
-                        Poll::Ready(None) => return Poll::Ready(Err(Error::closed())),
-                        Poll::Ready(Some(Err(e))) => return Poll::Ready(Err(Error::io(e))),
-                        Poll::Ready(Some(Ok(message))) => message,
-                    };
-
-                    match message {
-                        BackendMessage::Async(Message::NoticeResponse(body)) => {
-                            self.handle_notice(body)?;
-                            continue;
-                        }
-                        BackendMessage::Async(_) => continue,
-                        BackendMessage::Normal { messages } => messages,
-                    }
+            let message = match self.poll_response(cx)? {
+                Poll::Ready(Some(message)) => message,
+                Poll::Ready(None) => return Poll::Ready(Err(Error::closed())),
+                Poll::Pending => {
+                    trace!("poll_read: waiting on response");
+                    return Poll::Pending;
                 }
             };
 
+            let messages = match message {
+                BackendMessage::Async(Message::NoticeResponse(body)) => {
+                    let error = DbError::parse(&mut body.fields()).map_err(Error::parse)?;
+                    return Poll::Ready(Ok(AsyncMessage::Notice(error)));
+                }
+                BackendMessage::Async(Message::NotificationResponse(body)) => {
+                    let notification = Notification {
+                        process_id: body.process_id(),
+                        channel: body.channel().map_err(Error::parse)?.to_string(),
+                        payload: body.message().map_err(Error::parse)?.to_string(),
+                    };
+                    return Poll::Ready(Ok(AsyncMessage::Notification(notification)));
+                }
+                BackendMessage::Async(Message::ParameterStatus(body)) => {
+                    self.parameters.insert(
+                        body.name().map_err(Error::parse)?.to_string(),
+                        body.value().map_err(Error::parse)?.to_string(),
+                    );
+                    continue;
+                }
+                BackendMessage::Async(_) => unreachable!(),
+                BackendMessage::Normal { messages } => messages,
+            };
+
             match self.sender.poll_reserve(cx) {
                 Poll::Ready(Ok(())) => {
                     let _ = self.sender.send_item(messages);
@@ -99,7 +126,8 @@ where
                     return Poll::Ready(Err(Error::closed()));
                 }
                 Poll::Pending => {
-                    self.pending_response = Some(messages);
+                    self.pending_responses
+                        .push_back(BackendMessage::Normal { messages });
                     trace!("poll_read: waiting on sender");
                     return Poll::Pending;
                 }
@@ -107,31 +135,6 @@ where
         }
     }
 
-    fn handle_notice(&mut self, body: NoticeResponseBody) -> Result<(), Error> {
-        let Some(notices) = &mut self.notices else {
-            return Ok(());
-        };
-
-        let mut fields = body.fields();
-        while let Some(field) = fields.next().map_err(Error::parse)? {
-            // loop until we find the message field
-            if field.type_() == b'M' {
-                // if the message field is within the limit, send it.
-                if let Some(new_limit) = notices.limit.checked_sub(field.value().len()) {
-                    match notices.sender.send(field.value().into()) {
-                        // set the new limit.
-                        Ok(()) => notices.limit = new_limit,
-                        // closed.
-                        Err(_) => self.notices = None,
-                    }
-                }
-                break;
-            }
-        }
-
-        Ok(())
-    }
-
     /// Fetch the next client request and enqueue the response sender.
     fn poll_request(&mut self, cx: &mut Context<'_>) -> Poll<Option<FrontendMessage>> {
         if self.receiver.is_closed() {
@@ -165,23 +168,21 @@ where
 
             match self.poll_request(cx) {
                 // send the message to postgres
-                Poll::Ready(Some(FrontendMessage::Raw(request))) => {
+                Poll::Ready(Some(request)) => {
                     Pin::new(&mut self.stream)
                         .start_send(request)
                         .map_err(Error::io)?;
                 }
-                Poll::Ready(Some(FrontendMessage::RecordNotices(notices))) => {
-                    self.notices = Some(notices)
-                }
                 // No more messages from the client, and no more responses to wait for.
                 // Send a terminate message to postgres
                 Poll::Ready(None) => {
                     trace!("poll_write: at eof, terminating");
                     let mut request = BytesMut::new();
                     frontend::terminate(&mut request);
+                    let request = FrontendMessage::Raw(request.freeze());
 
                     Pin::new(&mut self.stream)
-                        .start_send(request.freeze())
+                        .start_send(request)
                         .map_err(Error::io)?;
 
                     trace!("poll_write: sent eof, closing");
@@ -230,17 +231,34 @@ where
         }
     }
 
-    fn poll_message(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<Never, Error>>> {
+    /// Returns the value of a runtime parameter for this connection.
+    pub fn parameter(&self, name: &str) -> Option<&str> {
+        self.parameters.get(name).map(|s| &**s)
+    }
+
+    /// Polls for asynchronous messages from the server.
+    ///
+    /// The server can send notices as well as notifications asynchronously to the client. Applications that wish to
+    /// examine those messages should use this method to drive the connection rather than its `Future` implementation.
+    pub fn poll_message(
+        &mut self,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<AsyncMessage, Error>>> {
         if self.state != State::Closing {
             // if the state is still active, try read from and write to postgres.
-            let Poll::Pending = self.poll_read(cx)?;
-            if self.poll_write(cx)?.is_ready() {
+            let message = self.poll_read(cx)?;
+            let closing = self.poll_write(cx)?;
+            if let Poll::Ready(()) = closing {
                 self.state = State::Closing;
             }
 
+            if let Poll::Ready(message) = message {
+                return Poll::Ready(Some(Ok(message)));
+            }
+
             // poll_read returned Pending.
-            // poll_write returned Pending or Ready(()).
-            // if poll_write returned Ready(()), then we are waiting to read more data from postgres.
+            // poll_write returned Pending or Ready(WriteReady::WaitingOnRead).
+            // if poll_write returned Ready(WriteReady::WaitingOnRead), then we are waiting to read more data from postgres.
             if self.state != State::Closing {
                 return Poll::Pending;
             }
@@ -262,9 +280,11 @@ where
     type Output = Result<(), Error>;
 
     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
-        match self.poll_message(cx)? {
-            Poll::Ready(None) => Poll::Ready(Ok(())),
-            Poll::Pending => Poll::Pending,
+        while let Some(message) = ready!(self.poll_message(cx)?) {
+            if let AsyncMessage::Notice(notice) = message {
+                info!("{}: {}", notice.severity(), notice.message());
+            }
         }
+        Poll::Ready(Ok(()))
     }
 }

libs/proxy/tokio-postgres2/src/error/mod.rs

@@ -452,16 +452,16 @@ impl Error {
         Error(Box::new(ErrorInner { kind, cause }))
     }
 
-    pub fn closed() -> Error {
+    pub(crate) fn closed() -> Error {
         Error::new(Kind::Closed, None)
     }
 
-    pub fn unexpected_message() -> Error {
+    pub(crate) fn unexpected_message() -> Error {
         Error::new(Kind::UnexpectedMessage, None)
     }
 
     #[allow(clippy::needless_pass_by_value)]
-    pub fn db(error: ErrorResponseBody) -> Error {
+    pub(crate) fn db(error: ErrorResponseBody) -> Error {
         match DbError::parse(&mut error.fields()) {
             Ok(e) => Error::new(Kind::Db, Some(Box::new(e))),
             Err(e) => Error::new(Kind::Parse, Some(Box::new(e))),
@@ -493,7 +493,7 @@ impl Error {
         Error::new(Kind::Tls, Some(e))
     }
 
-    pub fn io(e: io::Error) -> Error {
+    pub(crate) fn io(e: io::Error) -> Error {
         Error::new(Kind::Io, Some(Box::new(e)))
     }
 

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

@@ -6,7 +6,9 @@ use postgres_protocol2::message::backend::ReadyForQueryBody;
 pub use crate::cancel_token::{CancelToken, RawCancelToken};
 pub use crate::client::{Client, SocketConfig};
 pub use crate::config::Config;
+pub use crate::connect_raw::RawConnection;
 pub use crate::connection::Connection;
+use crate::error::DbError;
 pub use crate::error::Error;
 pub use crate::generic_client::GenericClient;
 pub use crate::query::RowStream;
@@ -49,7 +51,7 @@ mod client;
 mod codec;
 pub mod config;
 mod connect;
-pub mod connect_raw;
+mod connect_raw;
 mod connect_socket;
 mod connect_tls;
 mod connection;
@@ -91,6 +93,21 @@ impl Notification {
     }
 }
 
+/// An asynchronous message from the server.
+#[allow(clippy::large_enum_variant)]
+#[derive(Debug, Clone)]
+#[non_exhaustive]
+pub enum AsyncMessage {
+    /// A notice.
+    ///
+    /// Notices use the same format as errors, but aren't "errors" per-se.
+    Notice(DbError),
+    /// A notification.
+    ///
+    /// Connections can subscribe to notifications with the `LISTEN` command.
+    Notification(Notification),
+}
+
 /// Message returned by the `SimpleQuery` stream.
 #[derive(Debug)]
 #[non_exhaustive]

libs/remote_storage/Cargo.toml

@@ -43,7 +43,7 @@ itertools.workspace = true
 sync_wrapper = { workspace = true, features = ["futures"] }
 
 byteorder = "1.4"
-rand.workspace = true
+rand = "0.8.5"
 
 [dev-dependencies]
 camino-tempfile.workspace = true

libs/remote_storage/src/simulate_failures.rs

@@ -81,7 +81,7 @@ impl UnreliableWrapper {
     ///
     fn attempt(&self, op: RemoteOp) -> anyhow::Result<u64> {
         let mut attempts = self.attempts.lock().unwrap();
-        let mut rng = rand::rng();
+        let mut rng = rand::thread_rng();
 
         match attempts.entry(op) {
             Entry::Occupied(mut e) => {
@@ -94,7 +94,7 @@ impl UnreliableWrapper {
                 /* BEGIN_HADRON */
                 // If there are more attempts to fail, fail the request by probability.
                 if (attempts_before_this < self.attempts_to_fail)
-                    && (rng.random_range(0..=100) < self.attempt_failure_probability)
+                    && (rng.gen_range(0..=100) < self.attempt_failure_probability)
                 {
                     let error =
                         anyhow::anyhow!("simulated failure of remote operation {:?}", e.key());

libs/remote_storage/tests/test_real_azure.rs

@@ -208,7 +208,7 @@ async fn create_azure_client(
         .as_millis();
 
     // because nanos can be the same for two threads so can millis, add randomness
-    let random = rand::rng().random::<u32>();
+    let random = rand::thread_rng().r#gen::<u32>();
 
     let remote_storage_config = RemoteStorageConfig {
         storage: RemoteStorageKind::AzureContainer(AzureConfig {

libs/remote_storage/tests/test_real_s3.rs

@@ -385,7 +385,7 @@ async fn create_s3_client(
         .as_millis();
 
     // because nanos can be the same for two threads so can millis, add randomness
-    let random = rand::rng().random::<u32>();
+    let random = rand::thread_rng().r#gen::<u32>();
 
     let remote_storage_config = RemoteStorageConfig {
         storage: RemoteStorageKind::AwsS3(S3Config {

libs/safekeeper_api/src/models.rs

@@ -301,12 +301,7 @@ pub struct PullTimelineRequest {
     pub tenant_id: TenantId,
     pub timeline_id: TimelineId,
     pub http_hosts: Vec<String>,
-    /// Membership configuration to switch to after pull.
-    /// It guarantees that if pull_timeline returns successfully, the timeline will
-    /// not be deleted by request with an older generation.
-    /// Storage controller always sets this field.
-    /// None is only allowed for manual pull_timeline requests.
-    pub mconf: Option<Configuration>,
+    pub ignore_tombstone: Option<bool>,
 }
 
 #[derive(Debug, Serialize, Deserialize)]

libs/tracing-utils/Cargo.toml

@@ -8,7 +8,7 @@ license.workspace = true
 hyper0.workspace = true
 opentelemetry = { workspace = true, features = ["trace"] }
 opentelemetry_sdk = { workspace = true, features = ["rt-tokio"] }
-opentelemetry-otlp = { workspace = true, default-features = false, features = ["http-proto", "trace", "http", "reqwest-blocking-client"] }
+opentelemetry-otlp = { workspace = true, default-features = false, features = ["http-proto", "trace", "http", "reqwest-client"] }
 opentelemetry-semantic-conventions.workspace = true
 tokio = { workspace = true, features = ["rt", "rt-multi-thread"] }
 tracing.workspace = true

libs/tracing-utils/src/lib.rs

@@ -1,5 +1,11 @@
 //! Helper functions to set up OpenTelemetry tracing.
 //!
+//! This comes in two variants, depending on whether you have a Tokio runtime available.
+//! If you do, call `init_tracing()`. It sets up the trace processor and exporter to use
+//! the current tokio runtime. If you don't have a runtime available, or you don't want
+//! to share the runtime with the tracing tasks, call `init_tracing_without_runtime()`
+//! instead. It sets up a dedicated single-threaded Tokio runtime for the tracing tasks.
+//!
 //! Example:
 //!
 //! ```rust,no_run
@@ -15,8 +21,7 @@
 //!         .with_writer(std::io::stderr);
 //!
 //!     // Initialize OpenTelemetry. Exports tracing spans as OpenTelemetry traces
-//!     let provider = tracing_utils::init_tracing("my_application", tracing_utils::ExportConfig::default());
-//!     let otlp_layer = provider.as_ref().map(tracing_utils::layer);
+//!     let otlp_layer = tracing_utils::init_tracing("my_application", tracing_utils::ExportConfig::default()).await;
 //!
 //!     // Put it all together
 //!     tracing_subscriber::registry()
@@ -31,18 +36,16 @@
 pub mod http;
 pub mod perf_span;
 
+use opentelemetry::KeyValue;
 use opentelemetry::trace::TracerProvider;
 use opentelemetry_otlp::WithExportConfig;
 pub use opentelemetry_otlp::{ExportConfig, Protocol};
-use opentelemetry_sdk::trace::SdkTracerProvider;
 use tracing::level_filters::LevelFilter;
 use tracing::{Dispatch, Subscriber};
 use tracing_subscriber::Layer;
 use tracing_subscriber::layer::SubscriberExt;
 use tracing_subscriber::registry::LookupSpan;
 
-pub type Provider = SdkTracerProvider;
-
 /// Set up OpenTelemetry exporter, using configuration from environment variables.
 ///
 /// `service_name` is set as the OpenTelemetry 'service.name' resource (see
@@ -67,7 +70,16 @@ pub type Provider = SdkTracerProvider;
 /// If you need some other setting, please test if it works first. And perhaps
 /// add a comment in the list above to save the effort of testing for the next
 /// person.
-pub fn init_tracing(service_name: &str, export_config: ExportConfig) -> Option<Provider> {
+///
+/// This doesn't block, but is marked as 'async' to hint that this must be called in
+/// asynchronous execution context.
+pub async fn init_tracing<S>(
+    service_name: &str,
+    export_config: ExportConfig,
+) -> Option<impl Layer<S>>
+where
+    S: Subscriber + for<'span> LookupSpan<'span>,
+{
     if std::env::var("OTEL_SDK_DISABLED") == Ok("true".to_string()) {
         return None;
     };
@@ -77,14 +89,52 @@ pub fn init_tracing(service_name: &str, export_config: ExportConfig) -> Option<P
     ))
 }
 
-pub fn layer<S>(p: &Provider) -> impl Layer<S>
+/// Like `init_tracing`, but creates a separate tokio Runtime for the tracing
+/// tasks.
+pub fn init_tracing_without_runtime<S>(
+    service_name: &str,
+    export_config: ExportConfig,
+) -> Option<impl Layer<S>>
 where
     S: Subscriber + for<'span> LookupSpan<'span>,
 {
-    tracing_opentelemetry::layer().with_tracer(p.tracer("global"))
+    if std::env::var("OTEL_SDK_DISABLED") == Ok("true".to_string()) {
+        return None;
+    };
+
+    // The opentelemetry batch processor and the OTLP exporter needs a Tokio
+    // runtime. Create a dedicated runtime for them. One thread should be
+    // enough.
+    //
+    // (Alternatively, instead of batching, we could use the "simple
+    // processor", which doesn't need Tokio, and use "reqwest-blocking"
+    // feature for the OTLP exporter, which also doesn't need Tokio.  However,
+    // batching is considered best practice, and also I have the feeling that
+    // the non-Tokio codepaths in the opentelemetry crate are less used and
+    // might be more buggy, so better to stay on the well-beaten path.)
+    //
+    // We leak the runtime so that it keeps running after we exit the
+    // function.
+    let runtime = Box::leak(Box::new(
+        tokio::runtime::Builder::new_multi_thread()
+            .enable_all()
+            .thread_name("otlp runtime thread")
+            .worker_threads(1)
+            .build()
+            .unwrap(),
+    ));
+    let _guard = runtime.enter();
+
+    Some(init_tracing_internal(
+        service_name.to_string(),
+        export_config,
+    ))
 }
 
-fn init_tracing_internal(service_name: String, export_config: ExportConfig) -> Provider {
+fn init_tracing_internal<S>(service_name: String, export_config: ExportConfig) -> impl Layer<S>
+where
+    S: Subscriber + for<'span> LookupSpan<'span>,
+{
     // Sets up exporter from the provided [`ExportConfig`] parameter.
     // If the endpoint is not specified, it is loaded from the
     // OTEL_EXPORTER_OTLP_ENDPOINT environment variable.
@@ -103,14 +153,22 @@ fn init_tracing_internal(service_name: String, export_config: ExportConfig) -> P
         opentelemetry_sdk::propagation::TraceContextPropagator::new(),
     );
 
-    Provider::builder()
-        .with_batch_exporter(exporter)
-        .with_resource(
-            opentelemetry_sdk::Resource::builder()
-                .with_service_name(service_name)
-                .build(),
-        )
+    let tracer = opentelemetry_sdk::trace::TracerProvider::builder()
+        .with_batch_exporter(exporter, opentelemetry_sdk::runtime::Tokio)
+        .with_resource(opentelemetry_sdk::Resource::new(vec![KeyValue::new(
+            opentelemetry_semantic_conventions::resource::SERVICE_NAME,
+            service_name,
+        )]))
         .build()
+        .tracer("global");
+
+    tracing_opentelemetry::layer().with_tracer(tracer)
+}
+
+// Shutdown trace pipeline gracefully, so that it has a chance to send any
+// pending traces before we exit.
+pub fn shutdown_tracing() {
+    opentelemetry::global::shutdown_tracer_provider();
 }
 
 pub enum OtelEnablement {
@@ -118,17 +176,17 @@ pub enum OtelEnablement {
     Enabled {
         service_name: String,
         export_config: ExportConfig,
+        runtime: &'static tokio::runtime::Runtime,
     },
 }
 
 pub struct OtelGuard {
-    provider: Provider,
     pub dispatch: Dispatch,
 }
 
 impl Drop for OtelGuard {
     fn drop(&mut self) {
-        _ = self.provider.shutdown();
+        shutdown_tracing();
     }
 }
 
@@ -141,19 +199,22 @@ impl Drop for OtelGuard {
 /// The lifetime of the guard should match taht of the application. On drop, it tears down the
 /// OTEL infra.
 pub fn init_performance_tracing(otel_enablement: OtelEnablement) -> Option<OtelGuard> {
-    match otel_enablement {
+    let otel_subscriber = match otel_enablement {
         OtelEnablement::Disabled => None,
         OtelEnablement::Enabled {
             service_name,
             export_config,
+            runtime,
         } => {
-            let provider = init_tracing(&service_name, export_config)?;
-
-            let otel_layer = layer(&provider).with_filter(LevelFilter::INFO);
+            let otel_layer = runtime
+                .block_on(init_tracing(&service_name, export_config))
+                .with_filter(LevelFilter::INFO);
             let otel_subscriber = tracing_subscriber::registry().with(otel_layer);
-            let dispatch = Dispatch::new(otel_subscriber);
+            let otel_dispatch = Dispatch::new(otel_subscriber);
 
-            Some(OtelGuard { dispatch, provider })
+            Some(otel_dispatch)
         }
-    }
+    };
+
+    otel_subscriber.map(|dispatch| OtelGuard { dispatch })
 }

libs/utils/Cargo.toml

@@ -47,7 +47,6 @@ tracing-subscriber = { workspace = true, features = ["json", "registry"] }
 tracing-utils.workspace = true
 rand.workspace = true
 scopeguard.workspace = true
-uuid.workspace = true
 strum.workspace = true
 strum_macros.workspace = true
 walkdir.workspace = true

libs/utils/src/auth.rs

@@ -12,8 +12,7 @@ use jsonwebtoken::{
     Algorithm, DecodingKey, EncodingKey, Header, TokenData, Validation, decode, encode,
 };
 use pem::Pem;
-use serde::{Deserialize, Deserializer, Serialize, de::DeserializeOwned};
-use uuid::Uuid;
+use serde::{Deserialize, Serialize, de::DeserializeOwned};
 
 use crate::id::TenantId;
 
@@ -26,11 +25,6 @@ pub enum Scope {
     /// Provides access to all data for a specific tenant (specified in `struct Claims` below)
     // TODO: join these two?
     Tenant,
-    /// Provides access to all data for a specific tenant, but based on endpoint ID. This token scope
-    /// is only used by compute to fetch the spec for a specific endpoint. The spec contains a Tenant-scoped
-    /// token authorizing access to all data of a tenant, so the spec-fetch API requires a TenantEndpoint
-    /// scope token to ensure that untrusted compute nodes can't fetch spec for arbitrary endpoints.
-    TenantEndpoint,
     /// Provides blanket access to all tenants on the pageserver plus pageserver-wide APIs.
     /// Should only be used e.g. for status check/tenant creation/list.
     PageServerApi,
@@ -57,43 +51,17 @@ pub enum Scope {
     ControllerPeer,
 }
 
-fn deserialize_empty_string_as_none_uuid<'de, D>(deserializer: D) -> Result<Option<Uuid>, D::Error>
-where
-    D: Deserializer<'de>,
-{
-    let opt = Option::<String>::deserialize(deserializer)?;
-    match opt.as_deref() {
-        Some("") => Ok(None),
-        Some(s) => Uuid::parse_str(s)
-            .map(Some)
-            .map_err(serde::de::Error::custom),
-        None => Ok(None),
-    }
-}
-
 /// JWT payload. See docs/authentication.md for the format
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
 pub struct Claims {
     #[serde(default)]
     pub tenant_id: Option<TenantId>,
-    #[serde(
-        default,
-        skip_serializing_if = "Option::is_none",
-        // Neon control plane includes this field as empty in the claims.
-        // Consider it None in those cases.
-        deserialize_with = "deserialize_empty_string_as_none_uuid"
-    )]
-    pub endpoint_id: Option<Uuid>,
     pub scope: Scope,
 }
 
 impl Claims {
     pub fn new(tenant_id: Option<TenantId>, scope: Scope) -> Self {
-        Self {
-            tenant_id,
-            scope,
-            endpoint_id: None,
-        }
+        Self { tenant_id, scope }
     }
 }
 
@@ -244,7 +212,6 @@ MC4CAQAwBQYDK2VwBCIEID/Drmc1AA6U/znNRWpF3zEGegOATQxfkdWxitcOMsIH
         let expected_claims = Claims {
             tenant_id: Some(TenantId::from_str("3d1f7595b468230304e0b73cecbcb081").unwrap()),
             scope: Scope::Tenant,
-            endpoint_id: None,
         };
 
         // A test token containing the following payload, signed using TEST_PRIV_KEY_ED25519:
@@ -273,7 +240,6 @@ MC4CAQAwBQYDK2VwBCIEID/Drmc1AA6U/znNRWpF3zEGegOATQxfkdWxitcOMsIH
         let claims = Claims {
             tenant_id: Some(TenantId::from_str("3d1f7595b468230304e0b73cecbcb081").unwrap()),
             scope: Scope::Tenant,
-            endpoint_id: None,
         };
 
         let pem = pem::parse(TEST_PRIV_KEY_ED25519).unwrap();

libs/utils/src/id.rs

@@ -104,7 +104,7 @@ impl Id {
 
     pub fn generate() -> Self {
         let mut tli_buf = [0u8; 16];
-        rand::rng().fill(&mut tli_buf);
+        rand::thread_rng().fill(&mut tli_buf);
         Id::from(tli_buf)
     }
 

libs/utils/src/lsn.rs

@@ -364,37 +364,42 @@ impl MonotonicCounter<Lsn> for RecordLsn {
     }
 }
 
-/// Implements  [`rand::distr::uniform::UniformSampler`] so we can sample [`Lsn`]s.
+/// Implements  [`rand::distributions::uniform::UniformSampler`] so we can sample [`Lsn`]s.
 ///
 /// This is used by the `pagebench` pageserver benchmarking tool.
-pub struct LsnSampler(<u64 as rand::distr::uniform::SampleUniform>::Sampler);
+pub struct LsnSampler(<u64 as rand::distributions::uniform::SampleUniform>::Sampler);
 
-impl rand::distr::uniform::SampleUniform for Lsn {
+impl rand::distributions::uniform::SampleUniform for Lsn {
     type Sampler = LsnSampler;
 }
 
-impl rand::distr::uniform::UniformSampler for LsnSampler {
+impl rand::distributions::uniform::UniformSampler for LsnSampler {
     type X = Lsn;
 
-    fn new<B1, B2>(low: B1, high: B2) -> Result<Self, rand::distr::uniform::Error>
+    fn new<B1, B2>(low: B1, high: B2) -> Self
     where
-        B1: rand::distr::uniform::SampleBorrow<Self::X> + Sized,
-        B2: rand::distr::uniform::SampleBorrow<Self::X> + Sized,
+        B1: rand::distributions::uniform::SampleBorrow<Self::X> + Sized,
+        B2: rand::distributions::uniform::SampleBorrow<Self::X> + Sized,
     {
-        <u64 as rand::distr::uniform::SampleUniform>::Sampler::new(low.borrow().0, high.borrow().0)
-            .map(Self)
+        Self(
+            <u64 as rand::distributions::uniform::SampleUniform>::Sampler::new(
+                low.borrow().0,
+                high.borrow().0,
+            ),
+        )
     }
 
-    fn new_inclusive<B1, B2>(low: B1, high: B2) -> Result<Self, rand::distr::uniform::Error>
+    fn new_inclusive<B1, B2>(low: B1, high: B2) -> Self
     where
-        B1: rand::distr::uniform::SampleBorrow<Self::X> + Sized,
-        B2: rand::distr::uniform::SampleBorrow<Self::X> + Sized,
+        B1: rand::distributions::uniform::SampleBorrow<Self::X> + Sized,
+        B2: rand::distributions::uniform::SampleBorrow<Self::X> + Sized,
     {
-        <u64 as rand::distr::uniform::SampleUniform>::Sampler::new_inclusive(
-            low.borrow().0,
-            high.borrow().0,
+        Self(
+            <u64 as rand::distributions::uniform::SampleUniform>::Sampler::new_inclusive(
+                low.borrow().0,
+                high.borrow().0,
+            ),
         )
-        .map(Self)
     }
 
     fn sample<R: rand::prelude::Rng + ?Sized>(&self, rng: &mut R) -> Self::X {

libs/utils/src/shard.rs

@@ -25,12 +25,6 @@ pub struct ShardIndex {
     pub shard_count: ShardCount,
 }
 
-/// Stripe size as number of pages.
-///
-/// NB: don't implement Default, so callers don't lazily use it by mistake. See DEFAULT_STRIPE_SIZE.
-#[derive(Clone, Copy, Serialize, Deserialize, Eq, PartialEq, Debug)]
-pub struct ShardStripeSize(pub u32);
-
 /// Formatting helper, for generating the `shard_id` label in traces.
 pub struct ShardSlug<'a>(&'a TenantShardId);
 
@@ -183,12 +177,6 @@ impl std::fmt::Display for ShardCount {
     }
 }
 
-impl std::fmt::Display for ShardStripeSize {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        self.0.fmt(f)
-    }
-}
-
 impl std::fmt::Display for ShardSlug<'_> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(

libs/walproposer/src/api_bindings.rs

@@ -429,11 +429,9 @@ pub fn empty_shmem() -> crate::bindings::WalproposerShmemState {
     };
 
     let empty_wal_rate_limiter = crate::bindings::WalRateLimiter {
-        effective_max_wal_bytes_per_second: crate::bindings::pg_atomic_uint32 { value: 0 },
         should_limit: crate::bindings::pg_atomic_uint32 { value: 0 },
         sent_bytes: 0,
-        batch_start_time_us: crate::bindings::pg_atomic_uint64 { value: 0 },
-        batch_end_time_us: crate::bindings::pg_atomic_uint64 { value: 0 },
+        last_recorded_time_us: crate::bindings::pg_atomic_uint64 { value: 0 },
     };
 
     crate::bindings::WalproposerShmemState {

pageserver/benches/bench_layer_map.rs

@@ -11,8 +11,7 @@ use pageserver::tenant::layer_map::LayerMap;
 use pageserver::tenant::storage_layer::{LayerName, PersistentLayerDesc};
 use pageserver_api::key::Key;
 use pageserver_api::shard::TenantShardId;
-use rand::prelude::{SeedableRng, StdRng};
-use rand::seq::IndexedRandom;
+use rand::prelude::{SeedableRng, SliceRandom, StdRng};
 use utils::id::{TenantId, TimelineId};
 use utils::lsn::Lsn;
 

pageserver/client/src/mgmt_api.rs

@@ -873,22 +873,6 @@ impl Client {
             .map_err(Error::ReceiveBody)
     }
 
-    pub async fn reset_alert_gauges(&self) -> Result<()> {
-        let uri = format!(
-            "{}/hadron-internal/reset_alert_gauges",
-            self.mgmt_api_endpoint
-        );
-        self.start_request(Method::POST, uri)
-            .send()
-            .await
-            .map_err(Error::SendRequest)?
-            .error_from_body()
-            .await?
-            .json()
-            .await
-            .map_err(Error::ReceiveBody)
-    }
-
     pub async fn wait_lsn(
         &self,
         tenant_shard_id: TenantShardId,

pageserver/client_grpc/src/client.rs

@@ -16,9 +16,10 @@ use crate::pool::{ChannelPool, ClientGuard, ClientPool, StreamGuard, StreamPool}
 use crate::retry::Retry;
 use crate::split::GetPageSplitter;
 use compute_api::spec::PageserverProtocol;
+use pageserver_api::shard::ShardStripeSize;
 use pageserver_page_api as page_api;
 use utils::id::{TenantId, TimelineId};
-use utils::shard::{ShardCount, ShardIndex, ShardNumber, ShardStripeSize};
+use utils::shard::{ShardCount, ShardIndex, ShardNumber};
 
 /// Max number of concurrent clients per channel (i.e. TCP connection). New channels will be spun up
 /// when full.
@@ -140,8 +141,8 @@ impl PageserverClient {
         if !old.count.is_unsharded() && shard_spec.stripe_size != old.stripe_size {
             return Err(anyhow!(
                 "can't change stripe size from {} to {}",
-                old.stripe_size.expect("always Some when sharded"),
-                shard_spec.stripe_size.expect("always Some when sharded")
+                old.stripe_size,
+                shard_spec.stripe_size
             ));
         }
 
@@ -156,6 +157,23 @@ impl PageserverClient {
         Ok(())
     }
 
+    /// Returns whether a relation exists.
+    #[instrument(skip_all, fields(rel=%req.rel, lsn=%req.read_lsn))]
+    pub async fn check_rel_exists(
+        &self,
+        req: page_api::CheckRelExistsRequest,
+    ) -> tonic::Result<page_api::CheckRelExistsResponse> {
+        debug!("sending request: {req:?}");
+        let resp = Self::with_retries(CALL_TIMEOUT, async |_| {
+            // Relation metadata is only available on shard 0.
+            let mut client = self.shards.load_full().get_zero().client().await?;
+            Self::with_timeout(REQUEST_TIMEOUT, client.check_rel_exists(req)).await
+        })
+        .await?;
+        debug!("received response: {resp:?}");
+        Ok(resp)
+    }
+
     /// Returns the total size of a database, as # of bytes.
     #[instrument(skip_all, fields(db_oid=%req.db_oid, lsn=%req.read_lsn))]
     pub async fn get_db_size(
@@ -231,15 +249,13 @@ impl PageserverClient {
         // Fast path: request is for a single shard.
         if let Some(shard_id) =
             GetPageSplitter::for_single_shard(&req, shards.count, shards.stripe_size)
-                .map_err(|err| tonic::Status::internal(err.to_string()))?
         {
             return Self::get_page_with_shard(req, shards.get(shard_id)?).await;
         }
 
         // Request spans multiple shards. Split it, dispatch concurrent per-shard requests, and
         // reassemble the responses.
-        let mut splitter = GetPageSplitter::split(req, shards.count, shards.stripe_size)
-            .map_err(|err| tonic::Status::internal(err.to_string()))?;
+        let mut splitter = GetPageSplitter::split(req, shards.count, shards.stripe_size);
 
         let mut shard_requests = FuturesUnordered::new();
         for (shard_id, shard_req) in splitter.drain_requests() {
@@ -249,14 +265,10 @@ impl PageserverClient {
         }
 
         while let Some((shard_id, shard_response)) = shard_requests.next().await.transpose()? {
-            splitter
-                .add_response(shard_id, shard_response)
-                .map_err(|err| tonic::Status::internal(err.to_string()))?;
+            splitter.add_response(shard_id, shard_response)?;
         }
 
-        splitter
-            .get_response()
-            .map_err(|err| tonic::Status::internal(err.to_string()))
+        splitter.get_response()
     }
 
     /// Fetches pages on the given shard. Does not retry internally.
@@ -384,14 +396,12 @@ pub struct ShardSpec {
     /// NB: this is 0 for unsharded tenants, following `ShardIndex::unsharded()` convention.
     count: ShardCount,
     /// The stripe size for these shards.
-    ///
-    /// INVARIANT: None for unsharded tenants, Some for sharded.
-    stripe_size: Option<ShardStripeSize>,
+    stripe_size: ShardStripeSize,
 }
 
 impl ShardSpec {
     /// Creates a new shard spec with the given URLs and stripe size. All shards must be given.
-    /// The stripe size must be Some for sharded tenants, or None for unsharded tenants.
+    /// The stripe size may be omitted for unsharded tenants.
     pub fn new(
         urls: HashMap<ShardIndex, String>,
         stripe_size: Option<ShardStripeSize>,
@@ -404,13 +414,11 @@ impl ShardSpec {
             n => ShardCount::new(n as u8),
         };
 
-        // Validate the stripe size.
+        // Determine the stripe size. It doesn't matter for unsharded tenants.
         if stripe_size.is_none() && !count.is_unsharded() {
             return Err(anyhow!("stripe size must be given for sharded tenants"));
         }
-        if stripe_size.is_some() && count.is_unsharded() {
-            return Err(anyhow!("stripe size can't be given for unsharded tenants"));
-        }
+        let stripe_size = stripe_size.unwrap_or_default();
 
         // Validate the shard spec.
         for (shard_id, url) in &urls {
@@ -450,10 +458,8 @@ struct Shards {
     ///
     /// NB: this is 0 for unsharded tenants, following `ShardIndex::unsharded()` convention.
     count: ShardCount,
-    /// The stripe size.
-    ///
-    /// INVARIANT: None for unsharded tenants, Some for sharded.
-    stripe_size: Option<ShardStripeSize>,
+    /// The stripe size. Only used for sharded tenants.
+    stripe_size: ShardStripeSize,
 }
 
 impl Shards {

pageserver/client_grpc/src/split.rs

@@ -1,12 +1,11 @@
 use std::collections::HashMap;
 
-use anyhow::anyhow;
 use bytes::Bytes;
 
 use pageserver_api::key::rel_block_to_key;
-use pageserver_api::shard::key_to_shard_number;
+use pageserver_api::shard::{ShardStripeSize, key_to_shard_number};
 use pageserver_page_api as page_api;
-use utils::shard::{ShardCount, ShardIndex, ShardStripeSize};
+use utils::shard::{ShardCount, ShardIndex, ShardNumber};
 
 /// Splits GetPageRequests that straddle shard boundaries and assembles the responses.
 /// TODO: add tests for this.
@@ -26,54 +25,43 @@ impl GetPageSplitter {
     pub fn for_single_shard(
         req: &page_api::GetPageRequest,
         count: ShardCount,
-        stripe_size: Option<ShardStripeSize>,
-    ) -> anyhow::Result<Option<ShardIndex>> {
+        stripe_size: ShardStripeSize,
+    ) -> Option<ShardIndex> {
         // Fast path: unsharded tenant.
         if count.is_unsharded() {
-            return Ok(Some(ShardIndex::unsharded()));
+            return Some(ShardIndex::unsharded());
         }
 
-        let Some(stripe_size) = stripe_size else {
-            return Err(anyhow!("stripe size must be given for sharded tenants"));
-        };
-
-        // Find the first page's shard, for comparison.
+        // Find the first page's shard, for comparison. If there are no pages, just return the first
+        // shard (caller likely checked already, otherwise the server will reject it).
         let Some(&first_page) = req.block_numbers.first() else {
-            return Err(anyhow!("no block numbers in request"));
+            return Some(ShardIndex::new(ShardNumber(0), count));
         };
         let key = rel_block_to_key(req.rel, first_page);
         let shard_number = key_to_shard_number(count, stripe_size, &key);
 
-        Ok(req
-            .block_numbers
+        req.block_numbers
             .iter()
             .skip(1) // computed above
             .all(|&blkno| {
                 let key = rel_block_to_key(req.rel, blkno);
                 key_to_shard_number(count, stripe_size, &key) == shard_number
             })
-            .then_some(ShardIndex::new(shard_number, count)))
+            .then_some(ShardIndex::new(shard_number, count))
     }
 
     /// Splits the given request.
     pub fn split(
         req: page_api::GetPageRequest,
         count: ShardCount,
-        stripe_size: Option<ShardStripeSize>,
-    ) -> anyhow::Result<Self> {
+        stripe_size: ShardStripeSize,
+    ) -> Self {
         // The caller should make sure we don't split requests unnecessarily.
         debug_assert!(
-            Self::for_single_shard(&req, count, stripe_size)?.is_none(),
+            Self::for_single_shard(&req, count, stripe_size).is_none(),
             "unnecessary request split"
         );
 
-        if count.is_unsharded() {
-            return Err(anyhow!("unsharded tenant, no point in splitting request"));
-        }
-        let Some(stripe_size) = stripe_size else {
-            return Err(anyhow!("stripe size must be given for sharded tenants"));
-        };
-
         // Split the requests by shard index.
         let mut requests = HashMap::with_capacity(2); // common case
         let mut block_shards = Vec::with_capacity(req.block_numbers.len());
@@ -115,11 +103,11 @@ impl GetPageSplitter {
                 .collect(),
         };
 
-        Ok(Self {
+        Self {
             requests,
             response,
             block_shards,
-        })
+        }
     }
 
     /// Drains the per-shard requests, moving them out of the splitter to avoid extra allocations.
@@ -136,30 +124,21 @@ impl GetPageSplitter {
         &mut self,
         shard_id: ShardIndex,
         response: page_api::GetPageResponse,
-    ) -> anyhow::Result<()> {
+    ) -> tonic::Result<()> {
         // The caller should already have converted status codes into tonic::Status.
         if response.status_code != page_api::GetPageStatusCode::Ok {
-            return Err(anyhow!(
+            return Err(tonic::Status::internal(format!(
                 "unexpected non-OK response for shard {shard_id}: {} {}",
                 response.status_code,
                 response.reason.unwrap_or_default()
-            ));
+            )));
         }
 
         if response.request_id != self.response.request_id {
-            return Err(anyhow!(
+            return Err(tonic::Status::internal(format!(
                 "response ID mismatch for shard {shard_id}: expected {}, got {}",
-                self.response.request_id,
-                response.request_id
-            ));
-        }
-
-        if response.request_id != self.response.request_id {
-            return Err(anyhow!(
-                "response ID mismatch for shard {shard_id}: expected {}, got {}",
-                self.response.request_id,
-                response.request_id
-            ));
+                self.response.request_id, response.request_id
+            )));
         }
 
         // Place the shard response pages into the assembled response, in request order.
@@ -171,26 +150,27 @@ impl GetPageSplitter {
             }
 
             let Some(slot) = self.response.pages.get_mut(i) else {
-                return Err(anyhow!("no block_shards slot {i} for shard {shard_id}"));
+                return Err(tonic::Status::internal(format!(
+                    "no block_shards slot {i} for shard {shard_id}"
+                )));
             };
             let Some(page) = pages.next() else {
-                return Err(anyhow!(
+                return Err(tonic::Status::internal(format!(
                     "missing page {} in shard {shard_id} response",
                     slot.block_number
-                ));
+                )));
             };
             if page.block_number != slot.block_number {
-                return Err(anyhow!(
+                return Err(tonic::Status::internal(format!(
                     "shard {shard_id} returned wrong page at index {i}, expected {} got {}",
-                    slot.block_number,
-                    page.block_number
-                ));
+                    slot.block_number, page.block_number
+                )));
             }
             if !slot.image.is_empty() {
-                return Err(anyhow!(
+                return Err(tonic::Status::internal(format!(
                     "shard {shard_id} returned duplicate page {} at index {i}",
                     slot.block_number
-                ));
+                )));
             }
 
             *slot = page;
@@ -198,10 +178,10 @@ impl GetPageSplitter {
 
         // Make sure we've consumed all pages from the shard response.
         if let Some(extra_page) = pages.next() {
-            return Err(anyhow!(
+            return Err(tonic::Status::internal(format!(
                 "shard {shard_id} returned extra page: {}",
                 extra_page.block_number
-            ));
+            )));
         }
 
         Ok(())
@@ -209,18 +189,18 @@ impl GetPageSplitter {
 
     /// Fetches the final, assembled response.
     #[allow(clippy::result_large_err)]
-    pub fn get_response(self) -> anyhow::Result<page_api::GetPageResponse> {
+    pub fn get_response(self) -> tonic::Result<page_api::GetPageResponse> {
         // Check that the response is complete.
         for (i, page) in self.response.pages.iter().enumerate() {
             if page.image.is_empty() {
-                return Err(anyhow!(
+                return Err(tonic::Status::internal(format!(
                     "missing page {} for shard {}",
                     page.block_number,
                     self.block_shards
                         .get(i)
                         .map(|s| s.to_string())
                         .unwrap_or_else(|| "?".to_string())
-                ));
+                )));
             }
         }
 

pageserver/compaction/src/bin/compaction-simulator.rs

@@ -89,7 +89,7 @@ async fn simulate(cmd: &SimulateCmd, results_path: &Path) -> anyhow::Result<()>
             let cold_key_range = splitpoint..key_range.end;
 
             for i in 0..cmd.num_records {
-                let chosen_range = if rand::rng().random_bool(0.9) {
+                let chosen_range = if rand::thread_rng().gen_bool(0.9) {
                     &hot_key_range
                 } else {
                     &cold_key_range

pageserver/compaction/src/simulator.rs

@@ -300,9 +300,9 @@ impl MockTimeline {
         key_range: &Range<Key>,
     ) -> anyhow::Result<()> {
         crate::helpers::union_to_keyspace(&mut self.keyspace, vec![key_range.clone()]);
-        let mut rng = rand::rng();
+        let mut rng = rand::thread_rng();
         for _ in 0..num_records {
-            self.ingest_record(rng.random_range(key_range.clone()), len);
+            self.ingest_record(rng.gen_range(key_range.clone()), len);
             self.wal_ingested += len;
         }
         Ok(())

pageserver/ctl/src/key.rs

@@ -4,7 +4,7 @@ use anyhow::Context;
 use clap::Parser;
 use pageserver_api::key::Key;
 use pageserver_api::reltag::{BlockNumber, RelTag, SlruKind};
-use pageserver_api::shard::{DEFAULT_STRIPE_SIZE, ShardCount, ShardStripeSize};
+use pageserver_api::shard::{ShardCount, ShardStripeSize};
 
 #[derive(Parser)]
 pub(super) struct DescribeKeyCommand {
@@ -128,9 +128,7 @@ impl DescribeKeyCommand {
             // seeing the sharding placement might be confusing, so leave it out unless shard
             // count was given.
 
-            let stripe_size = stripe_size
-                .map(ShardStripeSize)
-                .unwrap_or(DEFAULT_STRIPE_SIZE);
+            let stripe_size = stripe_size.map(ShardStripeSize).unwrap_or_default();
             println!(
                 "# placement with shard_count: {} and stripe_size: {}:",
                 shard_count.0, stripe_size.0