lints

On December 8th, 2023, an engineering escalation (INC-110) was opened
after it was found that BYPASSRLS was being applied to all roles.

PR that introduced the issue:
https://github.com/neondatabase/neon/pull/5657
Subsequent commit on main:
ad99fa5f03

NOBYPASSRLS and INHERIT are the defaults for a Postgres role, but
because it isn't easy to know if a Postgres cluster is affected by the
issue, we need to keep the migration around for a long time, if not
indefinitely, so any cluster can be fixed.

Branching is the gift that keeps on giving...

Signed-off-by: Tristan Partin <tristan.partin@databricks.com>

Signed-off-by: Tristan Partin <tristan.partin@databricks.com>

Cargo.lock

@@ -1330,7 +1330,6 @@ dependencies = [
  "chrono",
  "clap",
  "compute_api",
- "fail",
  "flate2",
  "futures",
  "hostname-validator",
@@ -1339,6 +1338,7 @@ dependencies = [
  "itertools 0.10.5",
  "jsonwebtoken",
  "metrics",
+ "neon_failpoint",
  "nix 0.30.1",
  "notify",
  "num_cpus",
@@ -2891,13 +2891,13 @@ dependencies = [
  "arc-swap",
  "bytes",
  "camino",
- "fail",
  "futures",
  "hyper 0.14.30",
  "itertools 0.10.5",
  "jemalloc_pprof",
  "jsonwebtoken",
  "metrics",
+ "neon_failpoint",
  "once_cell",
  "pprof",
  "regex",
@@ -3852,6 +3852,23 @@ dependencies = [
  "workspace_hack",
 ]
 
+[[package]]
+name = "neon_failpoint"
+version = "0.1.0"
+dependencies = [
+ "anyhow",
+ "either",
+ "once_cell",
+ "parking_lot 0.12.1",
+ "rand 0.8.5",
+ "regex",
+ "serde",
+ "tokio",
+ "tokio-util",
+ "tracing",
+ "tracing-subscriber",
+]
+
 [[package]]
 name = "never-say-never"
 version = "6.6.666"
@@ -4357,7 +4374,6 @@ dependencies = [
  "either",
  "enum-map",
  "enumset",
- "fail",
  "futures",
  "hashlink",
  "hex",
@@ -4372,6 +4388,7 @@ dependencies = [
  "jsonwebtoken",
  "md5",
  "metrics",
+ "neon_failpoint",
  "nix 0.30.1",
  "num-traits",
  "num_cpus",
@@ -6191,7 +6208,6 @@ dependencies = [
  "criterion",
  "desim",
  "env_logger",
- "fail",
  "futures",
  "hex",
  "http 1.1.0",
@@ -6201,6 +6217,7 @@ dependencies = [
  "itertools 0.10.5",
  "jsonwebtoken",
  "metrics",
+ "neon_failpoint",
  "once_cell",
  "pageserver_api",
  "parking_lot 0.12.1",
@@ -6887,7 +6904,6 @@ dependencies = [
  "diesel",
  "diesel-async",
  "diesel_migrations",
- "fail",
  "futures",
  "governor",
  "hex",
@@ -6900,6 +6916,7 @@ dependencies = [
  "lasso",
  "measured",
  "metrics",
+ "neon_failpoint",
  "once_cell",
  "pageserver_api",
  "pageserver_client",
@@ -8164,7 +8181,7 @@ dependencies = [
  "const_format",
  "criterion",
  "diatomic-waker",
- "fail",
+ "either",
  "futures",
  "git-version",
  "hex",
@@ -8172,6 +8189,7 @@ dependencies = [
  "humantime",
  "jsonwebtoken",
  "metrics",
+ "neon_failpoint",
  "nix 0.30.1",
  "once_cell",
  "pem",

Cargo.toml

@@ -21,6 +21,7 @@ members = [
     "workspace_hack",
     "libs/compute_api",
     "libs/http-utils",
+    "libs/neon_failpoint",
     "libs/pageserver_api",
     "libs/postgres_ffi",
     "libs/postgres_ffi_types",
@@ -97,7 +98,6 @@ diatomic-waker = { version = "0.2.3" }
 either = "1.8"
 enum-map = "2.4.2"
 enumset = "1.0.12"
-fail = "0.5.0"
 fallible-iterator = "0.2"
 framed-websockets = { version = "0.1.0", git = "https://github.com/neondatabase/framed-websockets" }
 futures = "0.3"
@@ -258,6 +258,7 @@ desim = { version = "0.1", path = "./libs/desim" }
 endpoint_storage = { version = "0.0.1", path = "./endpoint_storage/" }
 http-utils = { version = "0.1", path = "./libs/http-utils/" }
 metrics = { version = "0.1", path = "./libs/metrics/" }
+neon_failpoint = { version = "0.1", path = "./libs/neon_failpoint/" }
 neon-shmem = { version = "0.1", path = "./libs/neon-shmem/" }
 pageserver = { path = "./pageserver" }
 pageserver_api = { version = "0.1", path = "./libs/pageserver_api/" }

compute_tools/Cargo.toml

@@ -7,7 +7,7 @@ license.workspace = true
 [features]
 default = []
 # Enables test specific features.
-testing = ["fail/failpoints"]
+testing = ["neon_failpoint/testing"]
 
 [dependencies]
 async-compression.workspace = true
@@ -23,7 +23,7 @@ camino.workspace = true
 chrono.workspace = true
 cfg-if.workspace = true
 clap.workspace = true
-fail.workspace = true
+neon_failpoint.workspace = true
 flate2.workspace = true
 futures.workspace = true
 http.workspace = true

compute_tools/README.md

@@ -46,11 +46,14 @@ stateDiagram-v2
   Configuration --> Failed : Failed to configure the compute
   Configuration --> Running : Compute has been configured
   Empty --> Init : Compute spec is immediately available
-  Empty --> TerminationPending : Requested termination
+  Empty --> TerminationPendingFast : Requested termination
+  Empty --> TerminationPendingImmediate : Requested termination
   Init --> Failed : Failed to start Postgres
   Init --> Running : Started Postgres
-  Running --> TerminationPending : Requested termination
-  TerminationPending --> Terminated : Terminated compute
+  Running --> TerminationPendingFast : Requested termination
+  Running --> TerminationPendingImmediate : Requested termination
+  TerminationPendingFast --> Terminated compute with 30s delay for cplane to inspect status
+  TerminationPendingImmediate --> Terminated : Terminated compute immediately
   Failed --> [*] : Compute exited
   Terminated --> [*] : Compute exited
 ```

compute_tools/src/bin/compute_ctl.rs

@@ -154,7 +154,7 @@ impl Cli {
 fn main() -> Result<()> {
     let cli = Cli::parse();
 
-    let scenario = failpoint_support::init();
+    failpoint_support::init().unwrap();
 
     // For historical reasons, the main thread that processes the config and launches postgres
     // is synchronous, but we always have this tokio runtime available and we "enter" it so
@@ -201,8 +201,6 @@ fn main() -> Result<()> {
 
     let exit_code = compute_node.run()?;
 
-    scenario.teardown();
-
     deinit_and_exit(exit_code);
 }
 

compute_tools/src/compute.rs

@@ -956,14 +956,20 @@ impl ComputeNode {
             None
         };
 
-        let mut delay_exit = false;
         let mut state = self.state.lock().unwrap();
         state.terminate_flush_lsn = lsn;
-        if let ComputeStatus::TerminationPending { mode } = state.status {
+
+        let delay_exit = state.status == ComputeStatus::TerminationPendingFast;
+        if state.status == ComputeStatus::TerminationPendingFast
+            || state.status == ComputeStatus::TerminationPendingImmediate
+        {
+            info!(
+                "Changing compute status from {} to {}",
+                state.status,
+                ComputeStatus::Terminated
+            );
             state.status = ComputeStatus::Terminated;
             self.state_changed.notify_all();
-            // we were asked to terminate gracefully, don't exit to avoid restart
-            delay_exit = mode == compute_api::responses::TerminateMode::Fast
         }
         drop(state);
 
@@ -1901,7 +1907,8 @@ impl ComputeNode {
 
                             // exit loop
                             ComputeStatus::Failed
-                            | ComputeStatus::TerminationPending { .. }
+                            | ComputeStatus::TerminationPendingFast
+                            | ComputeStatus::TerminationPendingImmediate
                             | ComputeStatus::Terminated => break 'cert_update,
 
                             // wait

compute_tools/src/compute_prewarm.rs

@@ -70,7 +70,7 @@ impl ComputeNode {
             }
         };
         let row = match client
-            .query_one("select * from get_prewarm_info()", &[])
+            .query_one("select * from neon.get_prewarm_info()", &[])
             .await
         {
             Ok(row) => row,
@@ -146,7 +146,7 @@ impl ComputeNode {
         ComputeNode::get_maintenance_client(&self.tokio_conn_conf)
             .await
             .context("connecting to postgres")?
-            .query_one("select prewarm_local_cache($1)", &[&uncompressed])
+            .query_one("select neon.prewarm_local_cache($1)", &[&uncompressed])
             .await
             .context("loading LFC state into postgres")
             .map(|_| ())
@@ -196,7 +196,7 @@ impl ComputeNode {
         ComputeNode::get_maintenance_client(&self.tokio_conn_conf)
             .await
             .context("connecting to postgres")?
-            .query_one("select get_local_cache_state()", &[])
+            .query_one("select neon.get_local_cache_state()", &[])
             .await
             .context("querying LFC state")?
             .try_get::<usize, &[u8]>(0)

compute_tools/src/http/openapi_spec.yaml

@@ -371,9 +371,28 @@ paths:
       summary: Terminate Postgres and wait for it to exit
       description: ""
       operationId: terminate
+      parameters:
+        - name: mode
+          in: query
+          description: "Terminate mode: fast (wait 30s before returning) and immediate"
+          required: false
+          schema:
+            type: string
+            enum: ["fast", "immediate"]
+            default: fast
       responses:
         200:
           description: Result
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/TerminateResponse"
+        201:
+          description: Result if compute is already terminated
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/TerminateResponse"
         412:
           description: "wrong state"
           content:
@@ -530,11 +549,14 @@ components:
       type: string
       enum:
         - empty
-        - init
-        - failed
-        - running
         - configuration_pending
+        - init
+        - running
         - configuration
+        - failed
+        - termination_pending_fast
+        - termination_pending_immediate
+        - terminated
       example: running
 
     ExtensionInstallRequest:
@@ -660,6 +682,17 @@ components:
           description: Role name.
           example: "neon"
 
+    TerminateResponse:
+      type: object
+      required:
+        - lsn
+      properties:
+        lsn:
+          type: string
+          nullable: true
+          description: "last WAL flush LSN"
+          example: "0/028F10D8"
+
     SetRoleGrantsResponse:
       type: object
       required:

compute_tools/src/http/routes/failpoints.rs

@@ -1,8 +1,9 @@
 use axum::response::{IntoResponse, Response};
 use http::StatusCode;
+use neon_failpoint::{configure_failpoint, configure_failpoint_with_context};
 use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
 use tracing::info;
-use utils::failpoint_support::apply_failpoint;
 
 pub type ConfigureFailpointsRequest = Vec<FailpointConfig>;
 
@@ -11,10 +12,16 @@ pub type ConfigureFailpointsRequest = Vec<FailpointConfig>;
 pub struct FailpointConfig {
     /// Name of the fail point
     pub name: String,
-    /// List of actions to take, using the format described in `fail::cfg`
+    /// List of actions to take, using the format described in neon_failpoint
     ///
-    /// We also support `actions = "exit"` to cause the fail point to immediately exit.
+    /// We support actions: "pause", "sleep(N)", "return", "return(value)", "exit", "off", "panic(message)"
+    /// Plus probability-based actions: "N%return(value)", "N%M*return(value)", "N%action", "N%M*action"
     pub actions: String,
+    /// Optional context matching rules for conditional failpoints
+    /// Each key-value pair specifies a context key and a regex pattern to match against
+    /// All context matchers must match for the failpoint to trigger
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub context_matchers: Option<HashMap<String, String>>,
 }
 
 use crate::http::JsonResponse;
@@ -24,7 +31,7 @@ use crate::http::extract::Json;
 pub(in crate::http) async fn configure_failpoints(
     failpoints: Json<ConfigureFailpointsRequest>,
 ) -> Response {
-    if !fail::has_failpoints() {
+    if !neon_failpoint::has_failpoints() {
         return JsonResponse::error(
             StatusCode::PRECONDITION_FAILED,
             "Cannot manage failpoints because neon was compiled without failpoints support",
@@ -32,16 +39,21 @@ pub(in crate::http) async fn configure_failpoints(
     }
 
     for fp in &*failpoints {
-        info!("cfg failpoint: {} {}", fp.name, fp.actions);
+        info!(
+            "cfg failpoint: {} {} (context: {:?})",
+            fp.name, fp.actions, fp.context_matchers
+        );
 
-        // We recognize one extra "action" that's not natively recognized
-        // by the failpoints crate: exit, to immediately kill the process
-        let cfg_result = apply_failpoint(&fp.name, &fp.actions);
+        let cfg_result = if let Some(context_matchers) = fp.context_matchers.clone() {
+            configure_failpoint_with_context(&fp.name, &fp.actions, context_matchers)
+        } else {
+            configure_failpoint(&fp.name, &fp.actions)
+        };
 
         if let Err(e) = cfg_result {
             return JsonResponse::error(
                 StatusCode::BAD_REQUEST,
-                format!("failed to configure failpoints: {e}"),
+                format!("failed to configure failpoint '{}': {e}", fp.name),
             );
         }
     }

compute_tools/src/http/routes/terminate.rs

@@ -3,7 +3,7 @@ use crate::http::JsonResponse;
 use axum::extract::State;
 use axum::response::Response;
 use axum_extra::extract::OptionalQuery;
-use compute_api::responses::{ComputeStatus, TerminateResponse};
+use compute_api::responses::{ComputeStatus, TerminateMode, TerminateResponse};
 use http::StatusCode;
 use serde::Deserialize;
 use std::sync::Arc;
@@ -12,7 +12,7 @@ use tracing::info;
 
 #[derive(Deserialize, Default)]
 pub struct TerminateQuery {
-    mode: compute_api::responses::TerminateMode,
+    mode: TerminateMode,
 }
 
 /// Terminate the compute.
@@ -24,16 +24,16 @@ pub(in crate::http) async fn terminate(
     {
         let mut state = compute.state.lock().unwrap();
         if state.status == ComputeStatus::Terminated {
-            return JsonResponse::success(StatusCode::CREATED, state.terminate_flush_lsn);
+            let response = TerminateResponse {
+                lsn: state.terminate_flush_lsn,
+            };
+            return JsonResponse::success(StatusCode::CREATED, response);
         }
 
         if !matches!(state.status, ComputeStatus::Empty | ComputeStatus::Running) {
             return JsonResponse::invalid_status(state.status);
         }
-        state.set_status(
-            ComputeStatus::TerminationPending { mode },
-            &compute.state_changed,
-        );
+        state.set_status(mode.into(), &compute.state_changed);
     }
 
     forward_termination_signal(false);

compute_tools/src/migration.rs

@@ -1,5 +1,5 @@
 use anyhow::{Context, Result};
-use fail::fail_point;
+use neon_failpoint::fail_point;
 use tokio_postgres::{Client, Transaction};
 use tracing::{error, info};
 
@@ -40,13 +40,14 @@ impl<'m> MigrationRunner<'m> {
         // middle of applying a series of migrations fails in an expected
         // manner
         if cfg!(feature = "testing") {
-            let fail = (|| {
-                fail_point!("compute-migration", |fail_migration_id| {
+            let fail = async {
+                fail_point!("compute-migration", |fail_migration_id: Option<String>| {
                     migration_id == fail_migration_id.unwrap().parse::<i64>().unwrap()
                 });
 
                 false
-            })();
+            }
+            .await;
 
             if fail {
                 return Err(anyhow::anyhow!(format!(

compute_tools/src/migrations/0002-alter_roles.sql

@@ -1,3 +1,16 @@
+-- On December 8th, 2023, an engineering escalation (INC-110) was opened after
+-- it was found that BYPASSRLS was being applied to all roles.
+--
+-- PR that introduced the issue: https://github.com/neondatabase/neon/pull/5657
+-- Subsequent commit on main: https://github.com/neondatabase/neon/commit/ad99fa5f0393e2679e5323df653c508ffa0ac072
+--
+-- NOBYPASSRLS and INHERIT are the defaults for a Postgres role, but because it
+-- isn't easy to know if a Postgres cluster is affected by the issue, we need to
+-- keep the migration around for a long time, if not indefinitely, so any
+-- cluster can be fixed.
+--
+-- Branching is the gift that keeps on giving...
+
 DO $$
 DECLARE
     role_name text;

compute_tools/src/monitor.rs

@@ -84,7 +84,8 @@ impl ComputeMonitor {
         if matches!(
             compute_status,
             ComputeStatus::Terminated
-                | ComputeStatus::TerminationPending { .. }
+                | ComputeStatus::TerminationPendingFast
+                | ComputeStatus::TerminationPendingImmediate
                 | ComputeStatus::Failed
         ) {
             info!(

control_plane/src/endpoint.rs

@@ -922,7 +922,8 @@ impl Endpoint {
                         ComputeStatus::Empty
                         | ComputeStatus::ConfigurationPending
                         | ComputeStatus::Configuration
-                        | ComputeStatus::TerminationPending { .. }
+                        | ComputeStatus::TerminationPendingFast
+                        | ComputeStatus::TerminationPendingImmediate
                         | ComputeStatus::Terminated => {
                             bail!("unexpected compute status: {:?}", state.status)
                         }

libs/compute_api/src/responses.rs

@@ -121,6 +121,15 @@ pub enum TerminateMode {
     Immediate,
 }
 
+impl From<TerminateMode> for ComputeStatus {
+    fn from(mode: TerminateMode) -> Self {
+        match mode {
+            TerminateMode::Fast => ComputeStatus::TerminationPendingFast,
+            TerminateMode::Immediate => ComputeStatus::TerminationPendingImmediate,
+        }
+    }
+}
+
 #[derive(Serialize, Clone, Copy, Debug, Deserialize, PartialEq, Eq)]
 #[serde(rename_all = "snake_case")]
 pub enum ComputeStatus {
@@ -141,7 +150,9 @@ pub enum ComputeStatus {
     // control-plane to terminate it.
     Failed,
     // Termination requested
-    TerminationPending { mode: TerminateMode },
+    TerminationPendingFast,
+    // Termination requested, without waiting 30s before returning from /terminate
+    TerminationPendingImmediate,
     // Terminated Postgres
     Terminated,
 }
@@ -160,7 +171,10 @@ impl Display for ComputeStatus {
             ComputeStatus::Running => f.write_str("running"),
             ComputeStatus::Configuration => f.write_str("configuration"),
             ComputeStatus::Failed => f.write_str("failed"),
-            ComputeStatus::TerminationPending { .. } => f.write_str("termination-pending"),
+            ComputeStatus::TerminationPendingFast => f.write_str("termination-pending-fast"),
+            ComputeStatus::TerminationPendingImmediate => {
+                f.write_str("termination-pending-immediate")
+            }
             ComputeStatus::Terminated => f.write_str("terminated"),
         }
     }

libs/http-utils/Cargo.toml

@@ -9,7 +9,7 @@ anyhow.workspace = true
 arc-swap.workspace = true
 bytes.workspace = true
 camino.workspace = true
-fail.workspace = true
+neon_failpoint.workspace = true
 futures.workspace = true
 hyper0.workspace = true
 itertools.workspace = true

libs/http-utils/src/failpoints.rs

@@ -1,7 +1,8 @@
 use hyper::{Body, Request, Response, StatusCode};
+use neon_failpoint::{configure_failpoint, configure_failpoint_with_context, has_failpoints};
 use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
 use tokio_util::sync::CancellationToken;
-use utils::failpoint_support::apply_failpoint;
 
 use crate::error::ApiError;
 use crate::json::{json_request, json_response};
@@ -13,10 +14,16 @@ pub type ConfigureFailpointsRequest = Vec<FailpointConfig>;
 pub struct FailpointConfig {
     /// Name of the fail point
     pub name: String,
-    /// List of actions to take, using the format described in `fail::cfg`
+    /// List of actions to take, using the format described in neon_failpoint
     ///
-    /// We also support `actions = "exit"` to cause the fail point to immediately exit.
+    /// We support actions: "pause", "sleep(N)", "return", "return(value)", "exit", "off", "panic(message)"
+    /// Plus probability-based actions: "N%return(value)", "N%M*return(value)", "N%action", "N%M*action"
     pub actions: String,
+    /// Optional context matching rules for conditional failpoints
+    /// Each key-value pair specifies a context key and a regex pattern to match against
+    /// All context matchers must match for the failpoint to trigger
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub context_matchers: Option<HashMap<String, String>>,
 }
 
 /// Configure failpoints through http.
@@ -24,7 +31,7 @@ pub async fn failpoints_handler(
     mut request: Request<Body>,
     _cancel: CancellationToken,
 ) -> Result<Response<Body>, ApiError> {
-    if !fail::has_failpoints() {
+    if !has_failpoints() {
         return Err(ApiError::BadRequest(anyhow::anyhow!(
             "Cannot manage failpoints because neon was compiled without failpoints support"
         )));
@@ -32,15 +39,24 @@ pub async fn failpoints_handler(
 
     let failpoints: ConfigureFailpointsRequest = json_request(&mut request).await?;
     for fp in failpoints {
-        tracing::info!("cfg failpoint: {} {}", fp.name, fp.actions);
+        tracing::info!(
+            "cfg failpoint: {} {} (context: {:?})",
+            fp.name,
+            fp.actions,
+            fp.context_matchers
+        );
 
-        // We recognize one extra "action" that's not natively recognized
-        // by the failpoints crate: exit, to immediately kill the process
-        let cfg_result = apply_failpoint(&fp.name, &fp.actions);
+        let cfg_result = if let Some(context_matchers) = fp.context_matchers {
+            configure_failpoint_with_context(&fp.name, &fp.actions, context_matchers)
+        } else {
+            configure_failpoint(&fp.name, &fp.actions)
+        };
 
-        if let Err(err_msg) = cfg_result {
+        if let Err(err) = cfg_result {
             return Err(ApiError::BadRequest(anyhow::anyhow!(
-                "Failed to configure failpoints: {err_msg}"
+                "Failed to configure failpoint '{}': {}",
+                fp.name,
+                err
             )));
         }
     }

libs/neon_failpoint/Cargo.toml

@@ -0,0 +1,27 @@
+[package]
+name = "neon_failpoint"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+tokio = { workspace = true, features = ["time", "sync", "rt-multi-thread"] }
+tokio-util = { workspace = true }
+tracing = { workspace = true }
+serde = { workspace = true, features = ["derive"] }
+anyhow = { workspace = true }
+regex = { workspace = true }
+once_cell = { workspace = true }
+parking_lot = { workspace = true }
+rand = { workspace = true }
+either = { workspace = true }
+
+[dev-dependencies]
+tracing-subscriber = { workspace = true, features = ["fmt"] }
+
+[features]
+default = []
+testing = []
+
+[[example]]
+name = "context_demo"
+required-features = ["testing"]

libs/neon_failpoint/README.md

@@ -0,0 +1,460 @@
+# Neon Failpoint Library
+
+A modern, async-first failpoint library for Neon, replacing the `fail` crate with enhanced functionality.
+
+## Features
+
+- **Async-first**: All failpoint operations are async and don't require `spawn_blocking`
+- **Context matching**: Failpoints can be configured to trigger only when specific context conditions are met
+- **Regex support**: Context values can be matched using regular expressions
+- **Cancellation support**: All operations support cancellation tokens
+- **Dynamic reconfiguration**: Paused and sleeping tasks automatically resume when failpoint configurations change
+- **Backward compatibility**: Drop-in replacement for existing `fail` crate usage
+
+## Supported Actions
+
+- `off` - Disable the failpoint
+- `pause` - Pause indefinitely until disabled, reconfigured, or cancelled
+- `sleep(N)` - Sleep for N milliseconds (can be interrupted by reconfiguration)
+- `return` - Return early (empty value)
+- `return(value)` - Return early with a specific value
+- `exit` - Exit the process immediately
+- `panic(message)` - Panic the process with a custom message
+- `N%return(value)` - Return with a specific value N% of the time (probability-based)
+- `N%M*return(value)` - Return with a specific value N% of the time, maximum M times
+- `N%action` - Execute any action N% of the time (probability-based)
+- `N%M*action` - Execute any action N% of the time, maximum M times
+
+## Probability-Based Actions
+
+The library supports probability-based failpoints that trigger only a percentage of the time:
+
+```rust
+// 50% chance to return a value
+configure_failpoint("random_failure", "50%return(error)").unwrap();
+
+// 10% chance to sleep, maximum 3 times
+configure_failpoint("occasional_delay", "10%3*sleep(1000)").unwrap();
+
+// 25% chance to panic
+configure_failpoint("rare_panic", "25%panic(critical error)").unwrap();
+```
+
+The probability system uses a counter to track how many times a probability-based action has been triggered, allowing for precise control over test scenarios.
+
+## Dynamic Behavior
+
+When a failpoint is reconfigured while tasks are waiting on it:
+
+- **Paused tasks** will immediately resume and continue normal execution
+- **Sleeping tasks** will wake up early and continue normal execution
+- **Removed failpoints** will cause all waiting tasks to resume normally
+
+The new configuration only applies to future hits of the failpoint, not to tasks that are already waiting. This allows for flexible testing scenarios where you can pause execution, inspect state, and then resume execution dynamically.
+
+## Example: Dynamic Reconfiguration
+
+```rust
+use neon_failpoint::{configure_failpoint, failpoint, FailpointResult};
+use tokio::time::Duration;
+
+// Start a task that will hit a failpoint
+let task = tokio::spawn(async {
+    println!("About to hit failpoint");
+    match failpoint("test_pause", None).await {
+        FailpointResult::Return(value) => println!("Returned: {}", value),
+        FailpointResult::Continue => println!("Continued normally"),
+        FailpointResult::Cancelled => println!("Cancelled"),
+    }
+});
+
+// Configure the failpoint to pause
+configure_failpoint("test_pause", "pause").unwrap();
+
+// Let the task hit the failpoint and pause
+tokio::time::sleep(Duration::from_millis(10)).await;
+
+// Change the failpoint configuration - this will wake up the paused task
+// The task will resume and continue normally (not apply the new config)
+configure_failpoint("test_pause", "return(not_applied)").unwrap();
+
+// The task will complete with Continue, not Return
+let result = task.await.unwrap();
+```
+
+## Basic Usage
+
+```rust
+use neon_failpoint::{configure_failpoint, failpoint, FailpointResult};
+
+// Configure a failpoint
+configure_failpoint("my_failpoint", "return(42)").unwrap();
+
+// Use the failpoint
+match failpoint("my_failpoint", None).await {
+    FailpointResult::Return(value) => {
+        println!("Failpoint returned: {}", value);
+        return value.parse().unwrap_or_default();
+    }
+    FailpointResult::Continue => {
+        // Continue normal execution
+    }
+    FailpointResult::Cancelled => {
+        // Handle cancellation
+    }
+}
+```
+
+## Context-Based Failpoint Configuration
+
+Context allows you to create **conditional failpoints** that only trigger when specific runtime conditions are met. This is particularly useful for testing scenarios where you want to inject failures only for specific tenants, operations, or other contextual conditions.
+
+### Configuring Context-Based Failpoints
+
+Use `configure_failpoint_with_context()` to set up failpoints with context matching:
+
+```rust
+use neon_failpoint::configure_failpoint_with_context;
+use std::collections::HashMap;
+
+let mut context_matchers = HashMap::new();
+context_matchers.insert("tenant_id".to_string(), "test_.*".to_string());
+context_matchers.insert("operation".to_string(), "backup".to_string());
+
+configure_failpoint_with_context(
+    "backup_operation",           // failpoint name
+    "return(simulated_failure)",  // action to take
+    context_matchers             // context matching rules
+).unwrap();
+```
+
+### Context Matching Rules
+
+The context matching system works as follows:
+
+1. **Key-Value Matching**: Each entry in `context_matchers` specifies a key that must exist in the runtime context
+2. **Regex Support**: Values in `context_matchers` are treated as regular expressions first
+3. **Fallback to Exact Match**: If the regex compilation fails, it falls back to exact string matching
+4. **ALL Must Match**: All context matchers must match for the failpoint to trigger
+
+### Runtime Context Usage
+
+When code hits a failpoint, it provides context using a `HashMap<String, String>`:
+
+```rust
+use neon_failpoint::{failpoint, FailpointResult};
+use std::collections::HashMap;
+
+let mut context = HashMap::new();
+context.insert("tenant_id".to_string(), "test_123".to_string());
+context.insert("operation".to_string(), "backup".to_string());
+context.insert("user_id".to_string(), "user_456".to_string());
+
+match failpoint("backup_operation", Some(&context)) {
+    either::Either::Left(result) => {
+        match result {
+            FailpointResult::Return(value) => {
+                // This will only trigger if ALL context matchers match
+                println!("Backup failed: {}", value);
+            }
+            FailpointResult::Continue => {
+                // Continue with normal backup operation
+            }
+            FailpointResult::Cancelled => {}
+        }
+    }
+    either::Either::Right(future) => {
+        match future.await {
+            FailpointResult::Return(value) => {
+                // This will only trigger if ALL context matchers match
+                println!("Backup failed: {}", value);
+            }
+            FailpointResult::Continue => {
+                // Continue with normal backup operation
+            }
+            FailpointResult::Cancelled => {}
+        }
+    }
+}
+```
+
+### Context Matching Examples
+
+#### Regex Matching
+```rust
+// Configure to match test tenants only
+let mut matchers = HashMap::new();
+matchers.insert("tenant_id".to_string(), "test_.*".to_string());
+
+configure_failpoint_with_context("test_failpoint", "pause", matchers).unwrap();
+
+// This will match
+let mut context = HashMap::new();
+context.insert("tenant_id".to_string(), "test_123".to_string());
+// This will NOT match
+let mut context = HashMap::new();
+context.insert("tenant_id".to_string(), "prod_123".to_string());
+```
+
+#### Multiple Conditions
+```rust
+// Must match BOTH tenant pattern AND operation
+let mut matchers = HashMap::new();
+matchers.insert("tenant_id".to_string(), "test_.*".to_string());
+matchers.insert("operation".to_string(), "backup".to_string());
+
+configure_failpoint_with_context("backup_test", "return(failed)", matchers).unwrap();
+
+// This will match (both conditions met)
+let mut context = HashMap::new();
+context.insert("tenant_id".to_string(), "test_123".to_string());
+context.insert("operation".to_string(), "backup".to_string());
+
+// This will NOT match (missing operation)
+let mut context = HashMap::new();
+context.insert("tenant_id".to_string(), "test_123".to_string());
+context.insert("operation".to_string(), "restore".to_string());
+```
+
+#### Exact String Matching
+```rust
+// If regex compilation fails, falls back to exact match
+let mut matchers = HashMap::new();
+matchers.insert("env".to_string(), "staging".to_string());
+
+configure_failpoint_with_context("env_specific", "sleep(1000)", matchers).unwrap();
+
+// This will match
+let mut context = HashMap::new();
+context.insert("env".to_string(), "staging".to_string());
+// This will NOT match
+let mut context = HashMap::new();
+context.insert("env".to_string(), "production".to_string());
+```
+
+### Benefits of Context-Based Failpoints
+
+1. **Selective Testing**: Only inject failures for specific tenants, environments, or operations
+2. **Production Safety**: Avoid accidentally triggering failpoints in production by using context filters
+3. **Complex Scenarios**: Test interactions between different components with targeted failures
+4. **Debugging**: Isolate issues to specific contexts without affecting the entire system
+
+### Context vs. Non-Context Failpoints
+
+- **Without context**: `configure_failpoint("name", "action")` - triggers for ALL hits
+- **With context**: `configure_failpoint_with_context("name", "action", matchers)` - triggers only when context matches
+
+## Context-Specific Failpoints
+
+```rust
+use neon_failpoint::{configure_failpoint_with_context, failpoint};
+use std::collections::HashMap;
+
+// Configure a failpoint that only triggers for specific tenants
+let mut context_matchers = HashMap::new();
+context_matchers.insert("tenant_id".to_string(), "test_.*".to_string());
+context_matchers.insert("operation".to_string(), "backup".to_string());
+
+configure_failpoint_with_context(
+    "backup_operation",
+    "return(simulated_failure)",
+    context_matchers
+).unwrap();
+
+// Use with context
+let mut context = HashMap::new();
+context.insert("tenant_id".to_string(), "test_123".to_string());
+context.insert("operation".to_string(), "backup".to_string());
+
+match failpoint("backup_operation", Some(&context)) {
+    either::Either::Left(result) => {
+        match result {
+            FailpointResult::Return(value) => {
+                // This will trigger for tenant_id matching "test_.*"
+                println!("Backup failed: {}", value);
+            }
+            FailpointResult::Continue => {
+                // Continue with backup
+            }
+            FailpointResult::Cancelled => {}
+        }
+    }
+    either::Either::Right(future) => {
+        match future.await {
+            FailpointResult::Return(value) => {
+                // This will trigger for tenant_id matching "test_.*"
+                println!("Backup failed: {}", value);
+            }
+            FailpointResult::Continue => {
+                // Continue with backup
+            }
+            FailpointResult::Cancelled => {}
+        }
+    }
+}
+```
+
+## Macros
+
+The library provides convenient macros for common patterns:
+
+### `fail_point!` - Basic Failpoint Macro
+
+The `fail_point!` macro has three variants:
+
+1. **Simple failpoint** - `fail_point!(name)`
+   - Just checks the failpoint and continues or returns early (no value)
+   - Panics if the failpoint is configured with `return(value)` since no closure is provided
+
+2. **Failpoint with return handler** - `fail_point!(name, closure)`
+   - Provides a closure to handle return values from the failpoint
+   - The closure receives `Option<String>` and should return the appropriate value
+
+3. **Conditional failpoint** - `fail_point!(name, condition, closure)`
+   - Only checks the failpoint if the condition is true
+   - Provides a closure to handle return values (receives `&str`)
+
+```rust
+use neon_failpoint::fail_point;
+
+// Simple failpoint - just continue or return early
+fail_point!("my_failpoint");
+
+// Failpoint with return value handling
+fail_point!("my_failpoint", |value: Option<String>| {
+    match value {
+        Some(v) => {
+            println!("Got value: {}", v);
+            return Ok(v.parse().unwrap_or_default());
+        }
+        None => return Ok(42), // Default return value
+    }
+});
+
+// Conditional failpoint - only check if condition is met
+let should_fail = some_condition();
+fail_point!("conditional_failpoint", should_fail, |value: &str| {
+    println!("Conditional failpoint triggered with: {}", value);
+    return Err(anyhow::anyhow!("Simulated failure"));
+});
+```
+
+### `fail_point_with_context!` - Context-Aware Failpoint Macro
+
+The `fail_point_with_context!` macro has three variants that mirror `fail_point!` but include context:
+
+1. **Simple with context** - `fail_point_with_context!(name, context)`
+2. **With context and return handler** - `fail_point_with_context!(name, context, closure)`
+3. **Conditional with context** - `fail_point_with_context!(name, context, condition, closure)`
+
+```rust
+use neon_failpoint::{fail_point_with_context};
+use std::collections::HashMap;
+
+let mut context = HashMap::new();
+context.insert("tenant_id".to_string(), "test_123".to_string());
+context.insert("operation".to_string(), "backup".to_string());
+
+// Simple context failpoint
+fail_point_with_context!("backup_failpoint", &context);
+
+// Context failpoint with return handler
+fail_point_with_context!("backup_failpoint", &context, |value: Option<String>| {
+    match value {
+        Some(v) => return Err(anyhow::anyhow!("Backup failed: {}", v)),
+        None => return Err(anyhow::anyhow!("Backup failed")),
+    }
+});
+
+// Conditional context failpoint
+let is_test_tenant = tenant_id.starts_with("test_");
+fail_point_with_context!("backup_failpoint", &context, is_test_tenant, |value: Option<String>| {
+    // Only triggers for test tenants
+    return Err(anyhow::anyhow!("Test tenant backup failure"));
+});
+```
+
+### Other Utility Macros
+
+```rust
+use neon_failpoint::{pausable_failpoint, sleep_millis_async};
+
+// Pausable failpoint with cancellation
+let cancel_token = CancellationToken::new();
+if let Err(()) = pausable_failpoint!("pause_here", &cancel_token).await {
+    println!("Failpoint was cancelled");
+}
+
+// Sleep failpoint
+sleep_millis_async!("sleep_here", &cancel_token).await;
+
+// Context creation helper
+let mut context = HashMap::new();
+context.insert("key1".to_string(), "value1".to_string());
+context.insert("key2".to_string(), "value2".to_string());
+```
+
+### Argument Reference
+
+- **`name`**: String literal - the name of the failpoint
+- **`context`**: Expression that evaluates to `&HashMap<String, String>` - context for matching
+- **`condition`**: Boolean expression - only check failpoint if true
+- **`closure`**: Closure that handles return values:
+  - For `fail_point!` with closure: receives `Option<String>`
+  - For conditional variants: receives `&str`
+  - For `fail_point_with_context!` with closure: receives `Option<String>`
+- **`cancel`**: `&CancellationToken` - for cancellation support
+
+## Migration from `fail` crate
+
+The library provides a compatibility layer in `libs/utils/src/failpoint_support.rs`. Most existing code should work without changes, but you can migrate to the new async APIs for better performance:
+
+### Before (with `fail` crate):
+```rust
+use utils::failpoint_support::pausable_failpoint;
+
+// This used spawn_blocking internally
+pausable_failpoint!("my_failpoint", &cancel_token).await?;
+```
+
+### After (with `neon_failpoint`):
+```rust
+use neon_failpoint::{failpoint_with_cancellation, FailpointResult};
+
+// This is fully async
+match failpoint_with_cancellation("my_failpoint", None, &cancel_token).await {
+    FailpointResult::Continue => {},
+    FailpointResult::Cancelled => return Err(()),
+    FailpointResult::Return(_) => {},
+}
+```
+
+## Environment Variable Support
+
+Failpoints can be configured via the `FAILPOINTS` environment variable:
+
+```bash
+FAILPOINTS="failpoint1=return(42);failpoint2=sleep(1000);failpoint3=exit"
+```
+
+## Testing
+
+The library includes comprehensive tests and examples. Run them with:
+
+```bash
+cargo test --features testing
+cargo run --example context_demo --features testing
+```
+
+## HTTP Configuration
+
+The library integrates with the existing HTTP failpoint configuration API. Send POST requests to `/v1/failpoints` with:
+
+```json
+[
+  {
+    "name": "my_failpoint",
+    "actions": "return(42)"
+  }
+]
+```

libs/neon_failpoint/examples/context_demo.rs

@@ -0,0 +1,82 @@
+use neon_failpoint::{configure_failpoint_with_context, failpoint, FailpointResult};
+use std::collections::HashMap;
+
+#[tokio::main]
+async fn main() {
+    // Initialize tracing for better output
+    tracing_subscriber::fmt::init();
+
+    // Set up a context-specific failpoint
+    let mut context_matchers = HashMap::new();
+    context_matchers.insert("tenant_id".to_string(), "test_.*".to_string());
+    context_matchers.insert("operation".to_string(), "backup".to_string());
+
+    configure_failpoint_with_context(
+        "backup_operation",
+        "return(simulated_failure)",
+        context_matchers,
+    )
+    .unwrap();
+
+    // Test with matching context
+    let mut context = HashMap::new();
+    context.insert("tenant_id".to_string(), "test_123".to_string());
+    context.insert("operation".to_string(), "backup".to_string());
+
+    println!("Testing with matching context...");
+    match failpoint("backup_operation", Some(&context)) {
+        either::Either::Left(result) => match result {
+            FailpointResult::Return(value) => {
+                println!("Failpoint triggered with value: {value:?}");
+            }
+            FailpointResult::Continue => {
+                println!("Failpoint not triggered");
+            }
+            FailpointResult::Cancelled => {
+                println!("Failpoint cancelled");
+            }
+        },
+        either::Either::Right(future) => match future.await {
+            FailpointResult::Return(value) => {
+                println!("Failpoint triggered with value: {value:?}");
+            }
+            FailpointResult::Continue => {
+                println!("Failpoint not triggered");
+            }
+            FailpointResult::Cancelled => {
+                println!("Failpoint cancelled");
+            }
+        },
+    }
+
+    // Test with non-matching context
+    let mut context = HashMap::new();
+    context.insert("tenant_id".to_string(), "prod_456".to_string());
+    context.insert("operation".to_string(), "backup".to_string());
+
+    println!("Testing with non-matching context...");
+    match failpoint("backup_operation", Some(&context)) {
+        either::Either::Left(result) => match result {
+            FailpointResult::Return(value) => {
+                println!("Failpoint triggered with value: {value:?}");
+            }
+            FailpointResult::Continue => {
+                println!("Failpoint not triggered (expected)");
+            }
+            FailpointResult::Cancelled => {
+                println!("Failpoint cancelled");
+            }
+        },
+        either::Either::Right(future) => match future.await {
+            FailpointResult::Return(value) => {
+                println!("Failpoint triggered with value: {value:?}");
+            }
+            FailpointResult::Continue => {
+                println!("Failpoint not triggered (expected)");
+            }
+            FailpointResult::Cancelled => {
+                println!("Failpoint cancelled");
+            }
+        },
+    }
+}

libs/neon_failpoint/src/macros.rs

@@ -0,0 +1,356 @@
+//! Macros for convenient failpoint usage
+
+/// Simple failpoint macro - async version that doesn't require spawn_blocking
+#[macro_export]
+macro_rules! fail_point {
+    ($name:literal) => {{
+        if cfg!(feature = "testing") {
+            match $crate::failpoint($name, None) {
+                $crate::either::Either::Left(result) => {
+                    match result {
+                        $crate::FailpointResult::Continue => {},
+                        $crate::FailpointResult::Return(_) => {
+                            panic!("failpoint was configured with return(X) but Rust code does not pass a closure to map X to a return value");
+                        },
+                        $crate::FailpointResult::Cancelled => {},
+                    }
+                },
+                $crate::either::Either::Right(future) => {
+                    match future.await {
+                        $crate::FailpointResult::Continue => {},
+                        $crate::FailpointResult::Return(_) => {
+                            panic!("failpoint was configured with return(X) but Rust code does not pass a closure to map X to a return value");
+                        },
+                        $crate::FailpointResult::Cancelled => {},
+                    }
+                },
+            }
+        }
+    }};
+    ($name:literal, $closure:expr) => {{
+        if cfg!(feature = "testing") {
+            match $crate::failpoint($name, None) {
+                $crate::either::Either::Left(result) => {
+                    match result {
+                        $crate::FailpointResult::Continue => {},
+                        $crate::FailpointResult::Return(value) => {
+                            let closure = $closure;
+                            return closure(value);
+                        },
+                        $crate::FailpointResult::Cancelled => {},
+                    }
+                },
+                $crate::either::Either::Right(future) => {
+                    match future.await {
+                        $crate::FailpointResult::Continue => {},
+                        $crate::FailpointResult::Return(value) => {
+                            let closure = $closure;
+                            return closure(value);
+                        },
+                        $crate::FailpointResult::Cancelled => {},
+                    }
+                },
+            }
+        }
+    }};
+    ($name:literal, $condition:expr, $closure:expr) => {{
+        if cfg!(feature = "testing") {
+            if $condition {
+                match $crate::failpoint($name, None) {
+                    $crate::either::Either::Left(result) => {
+                        match result {
+                            $crate::FailpointResult::Continue => {},
+                            $crate::FailpointResult::Return(value) => {
+                                let closure = $closure;
+                                return closure(value);
+                            },
+                            $crate::FailpointResult::Cancelled => {},
+                        }
+                    },
+                    $crate::either::Either::Right(future) => {
+                        match future.await {
+                            $crate::FailpointResult::Continue => {},
+                            $crate::FailpointResult::Return(value) => {
+                                let closure = $closure;
+                                return closure(value);
+                            },
+                            $crate::FailpointResult::Cancelled => {},
+                        }
+                    },
+                }
+            }
+        }
+    }};
+}
+
+/// Simple failpoint macro - sync version that panics if async action is triggered
+#[macro_export]
+macro_rules! fail_point_sync {
+    ($name:literal) => {{
+        if cfg!(feature = "testing") {
+            match $crate::failpoint($name, None) {
+                $crate::either::Either::Left(result) => {
+                    match result {
+                        $crate::FailpointResult::Continue => {},
+                        $crate::FailpointResult::Return(_) => {
+                            panic!("failpoint was configured with return(X) but Rust code does not pass a closure to map X to a return value");
+                        },
+                        $crate::FailpointResult::Cancelled => {},
+                    }
+                },
+                $crate::either::Either::Right(_) => {
+                    panic!("failpoint '{}' triggered an async action (sleep/pause) but fail_point_sync! was used. Use fail_point! instead.", $name);
+                },
+            }
+        }
+    }};
+    ($name:literal, $closure:expr) => {{
+        if cfg!(feature = "testing") {
+            match $crate::failpoint($name, None) {
+                $crate::either::Either::Left(result) => {
+                    match result {
+                        $crate::FailpointResult::Continue => {},
+                        $crate::FailpointResult::Return(value) => {
+                            let closure = $closure;
+                            return closure(value);
+                        },
+                        $crate::FailpointResult::Cancelled => {},
+                    }
+                },
+                $crate::either::Either::Right(_) => {
+                    panic!("failpoint '{}' triggered an async action (sleep/pause) but fail_point_sync! was used. Use fail_point! instead.", $name);
+                },
+            }
+        }
+    }};
+    ($name:literal, $condition:expr, $closure:expr) => {{
+        if cfg!(feature = "testing") {
+            if $condition {
+                match $crate::failpoint($name, None) {
+                    $crate::either::Either::Left(result) => {
+                        match result {
+                            $crate::FailpointResult::Continue => {},
+                            $crate::FailpointResult::Return(value) => {
+                                let closure = $closure;
+                                return closure(value);
+                            },
+                            $crate::FailpointResult::Cancelled => {},
+                        }
+                    },
+                    $crate::either::Either::Right(_) => {
+                        panic!("failpoint '{}' triggered an async action (sleep/pause) but fail_point_sync! was used. Use fail_point! instead.", $name);
+                    },
+                }
+            }
+        }
+    }};
+}
+
+/// Failpoint macro with context support
+#[macro_export]
+macro_rules! fail_point_with_context {
+    ($name:literal, $context:expr) => {{
+        if cfg!(feature = "testing") {
+            match $crate::failpoint($name, Some($context)) {
+                $crate::either::Either::Left(result) => {
+                    match result {
+                        $crate::FailpointResult::Continue => {},
+                        $crate::FailpointResult::Return(_) => {
+                            panic!("failpoint was configured with return(X) but Rust code does not pass a closure to map X to a return value");
+                        },
+                        $crate::FailpointResult::Cancelled => {},
+                    }
+                },
+                $crate::either::Either::Right(future) => {
+                    match future.await {
+                        $crate::FailpointResult::Continue => {},
+                        $crate::FailpointResult::Return(_) => {
+                            panic!("failpoint was configured with return(X) but Rust code does not pass a closure to map X to a return value");
+                        },
+                        $crate::FailpointResult::Cancelled => {},
+                    }
+                },
+            }
+        }
+    }};
+    ($name:literal, $context:expr, $closure:expr) => {{
+        if cfg!(feature = "testing") {
+            match $crate::failpoint($name, Some($context)) {
+                $crate::either::Either::Left(result) => {
+                    match result {
+                        $crate::FailpointResult::Continue => {},
+                        $crate::FailpointResult::Return(value) => {
+                            let closure = $closure;
+                            return closure(value);
+                        },
+                        $crate::FailpointResult::Cancelled => {},
+                    }
+                },
+                $crate::either::Either::Right(future) => {
+                    match future.await {
+                        $crate::FailpointResult::Continue => {},
+                        $crate::FailpointResult::Return(value) => {
+                            let closure = $closure;
+                            return closure(value);
+                        },
+                        $crate::FailpointResult::Cancelled => {},
+                    }
+                },
+            }
+        }
+    }};
+    ($name:literal, $context:expr, $condition:expr, $closure:expr) => {{
+        if cfg!(feature = "testing") {
+            if $condition {
+                match $crate::failpoint($name, Some($context)) {
+                    $crate::either::Either::Left(result) => {
+                        match result {
+                            $crate::FailpointResult::Continue => {},
+                            $crate::FailpointResult::Return(value) => {
+                                let closure = $closure;
+                                return closure(value);
+                            },
+                            $crate::FailpointResult::Cancelled => {},
+                        }
+                    },
+                    $crate::either::Either::Right(future) => {
+                        match future.await {
+                            $crate::FailpointResult::Continue => {},
+                            $crate::FailpointResult::Return(value) => {
+                                let closure = $closure;
+                                return closure(value);
+                            },
+                            $crate::FailpointResult::Cancelled => {},
+                        }
+                    },
+                }
+            }
+        }
+    }};
+}
+
+/// Failpoint macro with context support - sync version
+#[macro_export]
+macro_rules! fail_point_with_context_sync {
+    ($name:literal, $context:expr) => {{
+        if cfg!(feature = "testing") {
+            match $crate::failpoint($name, Some($context)) {
+                $crate::either::Either::Left(result) => {
+                    match result {
+                        $crate::FailpointResult::Continue => {},
+                        $crate::FailpointResult::Return(_) => {
+                            panic!("failpoint was configured with return(X) but Rust code does not pass a closure to map X to a return value");
+                        },
+                        $crate::FailpointResult::Cancelled => {},
+                    }
+                },
+                $crate::either::Either::Right(_) => {
+                    panic!("failpoint '{}' triggered an async action (sleep/pause) but fail_point_with_context_sync! was used. Use fail_point_with_context! instead.", $name);
+                },
+            }
+        }
+    }};
+    ($name:literal, $context:expr, $closure:expr) => {{
+        if cfg!(feature = "testing") {
+            match $crate::failpoint($name, Some($context)) {
+                $crate::either::Either::Left(result) => {
+                    match result {
+                        $crate::FailpointResult::Continue => {},
+                        $crate::FailpointResult::Return(value) => {
+                            let closure = $closure;
+                            return closure(value);
+                        },
+                        $crate::FailpointResult::Cancelled => {},
+                    }
+                },
+                $crate::either::Either::Right(_) => {
+                    panic!("failpoint '{}' triggered an async action (sleep/pause) but fail_point_with_context_sync! was used. Use fail_point_with_context! instead.", $name);
+                },
+            }
+        }
+    }};
+    ($name:literal, $context:expr, $condition:expr, $closure:expr) => {{
+        if cfg!(feature = "testing") {
+            if $condition {
+                match $crate::failpoint($name, Some($context)) {
+                    $crate::either::Either::Left(result) => {
+                        match result {
+                            $crate::FailpointResult::Continue => {},
+                            $crate::FailpointResult::Return(value) => {
+                                let closure = $closure;
+                                return closure(value);
+                            },
+                            $crate::FailpointResult::Cancelled => {},
+                        }
+                    },
+                    $crate::either::Either::Right(_) => {
+                        panic!("failpoint '{}' triggered an async action (sleep/pause) but fail_point_with_context_sync! was used. Use fail_point_with_context! instead.", $name);
+                    },
+                }
+            }
+        }
+    }};
+}
+
+/// Pausable failpoint macro - equivalent to the old pausable_failpoint
+#[macro_export]
+macro_rules! pausable_failpoint {
+    ($name:literal) => {{
+        if cfg!(feature = "testing") {
+            let cancel = ::tokio_util::sync::CancellationToken::new();
+            let _ = $crate::pausable_failpoint!($name, &cancel);
+        }
+    }};
+    ($name:literal, $cancel:expr) => {{
+        if cfg!(feature = "testing") {
+            ::tracing::info!("at failpoint {}", $name); // tests rely on this
+            match $crate::failpoint_with_cancellation($name, None, $cancel) {
+                $crate::either::Either::Left(result) => match result {
+                    $crate::FailpointResult::Continue => Ok(()),
+                    $crate::FailpointResult::Return(_) => Ok(()),
+                    $crate::FailpointResult::Cancelled => Err(()),
+                },
+                $crate::either::Either::Right(future) => match future.await {
+                    $crate::FailpointResult::Continue => Ok(()),
+                    $crate::FailpointResult::Return(_) => Ok(()),
+                    $crate::FailpointResult::Cancelled => Err(()),
+                },
+            }
+        } else {
+            Ok(())
+        }
+    }};
+}
+
+/// Sleep failpoint macro - for async sleep operations
+#[macro_export]
+macro_rules! sleep_millis_async {
+    ($name:literal) => {{
+        if cfg!(feature = "testing") {
+            match $crate::failpoint($name, None) {
+                $crate::either::Either::Left(_) => {}
+                $crate::either::Either::Right(future) => {
+                    future.await;
+                }
+            }
+        }
+    }};
+    ($name:literal, $cancel:expr) => {{
+        if cfg!(feature = "testing") {
+            match $crate::failpoint_with_cancellation($name, None, $cancel) {
+                $crate::either::Either::Left(_) => {}
+                $crate::either::Either::Right(future) => {
+                    future.await;
+                }
+            }
+        }
+    }};
+}
+
+// Re-export for convenience
+pub use fail_point;
+pub use fail_point_sync;
+pub use fail_point_with_context;
+pub use fail_point_with_context_sync;
+pub use pausable_failpoint;
+pub use sleep_millis_async;

libs/utils/Cargo.toml

@@ -9,7 +9,7 @@ default = ["rename_noreplace"]
 rename_noreplace = []
 # Enables test-only APIs, incuding failpoints. In particular, enables the `fail_point!` macro,
 # which adds some runtime cost to run tests on outage conditions
-testing = ["fail/failpoints"]
+testing = ["neon_failpoint/testing"]
 
 [dependencies]
 arc-swap.workspace = true
@@ -21,10 +21,11 @@ bytes.workspace = true
 camino.workspace = true
 chrono.workspace = true
 diatomic-waker.workspace = true
+either.workspace = true
 git-version.workspace = true
 hex = { workspace = true, features = ["serde"] }
 humantime.workspace = true
-fail.workspace = true
+neon_failpoint.workspace = true
 futures = { workspace = true }
 jsonwebtoken.workspace = true
 nix = { workspace = true, features = ["ioctl"] }

libs/utils/src/failpoint_support.rs

@@ -1,59 +1,22 @@
 //! Failpoint support code shared between pageserver and safekeepers.
+//!
+//! This module provides a compatibility layer over the new neon_failpoint crate.
 
+pub use neon_failpoint::{configure_failpoint as apply_failpoint, has_failpoints, init};
 use tokio_util::sync::CancellationToken;
 
-/// Declare a failpoint that can use to `pause` failpoint action.
-/// We don't want to block the executor thread, hence, spawn_blocking + await.
-///
-/// Optionally pass a cancellation token, and this failpoint will drop out of
-/// its pause when the cancellation token fires. This is useful for testing
-/// cases where we would like to block something, but test its clean shutdown behavior.
-/// The macro evaluates to a Result in that case, where Ok(()) is the case
-/// where the failpoint was not paused, and Err() is the case where cancellation
-/// token fired while evaluating the failpoint.
-///
-/// Remember to unpause the failpoint in the test; until that happens, one of the
-/// limited number of spawn_blocking thread pool threads is leaked.
+/// Mere forward to neon_failpoint::pausable_failpoint
 #[macro_export]
 macro_rules! pausable_failpoint {
-    ($name:literal) => {{
-        if cfg!(feature = "testing") {
-            let cancel = ::tokio_util::sync::CancellationToken::new();
-            let _ = $crate::pausable_failpoint!($name, &cancel);
-        }
-    }};
-    ($name:literal, $cancel:expr) => {{
-        if cfg!(feature = "testing") {
-            let failpoint_fut = ::tokio::task::spawn_blocking({
-                let current = ::tracing::Span::current();
-                move || {
-                    let _entered = current.entered();
-                    ::tracing::info!("at failpoint {}", $name);
-                    ::fail::fail_point!($name);
-                }
-            });
-            let cancel_fut = async move {
-                $cancel.cancelled().await;
-            };
-            ::tokio::select! {
-                res = failpoint_fut => {
-                    res.expect("spawn_blocking");
-                    // continue with execution
-                    Ok(())
-                },
-                _ = cancel_fut => {
-                    Err(())
-                }
-            }
-        } else {
-            Ok(())
-        }
-    }};
+    ($name:literal) => {
+        ::neon_failpoint::pausable_failpoint!($name)
+    };
+    ($name:literal, $cancel:expr) => {
+        ::neon_failpoint::pausable_failpoint!($name, $cancel)
+    };
 }
 
-pub use pausable_failpoint;
-
-/// use with fail::cfg("$name", "return(2000)")
+/// DEPRECATED! - use with fail::cfg("$name", "return(2000)")
 ///
 /// The effect is similar to a "sleep(2000)" action, i.e. we sleep for the
 /// specified time (in milliseconds). The main difference is that we use async
@@ -69,7 +32,7 @@ macro_rules! __failpoint_sleep_millis_async {
         // If the failpoint is used with a "return" action, set should_sleep to the
         // returned value (as string). Otherwise it's set to None.
         let should_sleep = (|| {
-            ::fail::fail_point!($name, |x| x);
+            ::neon_failpoint::fail_point_sync!($name, |x| x);
             ::std::option::Option::None
         })();
 
@@ -82,7 +45,7 @@ macro_rules! __failpoint_sleep_millis_async {
         // If the failpoint is used with a "return" action, set should_sleep to the
         // returned value (as string). Otherwise it's set to None.
         let should_sleep = (|| {
-            ::fail::fail_point!($name, |x| x);
+            ::neon_failpoint::fail_point_sync!($name, |x| x);
             ::std::option::Option::None
         })();
 
@@ -126,60 +89,3 @@ pub async fn failpoint_sleep_cancellable_helper(
     tokio::time::timeout(d, cancel.cancelled()).await.ok();
     tracing::info!("failpoint {:?}: sleep done", name);
 }
-
-/// Initialize the configured failpoints
-///
-/// You must call this function before any concurrent threads do operations.
-pub fn init() -> fail::FailScenario<'static> {
-    // The failpoints lib provides support for parsing the `FAILPOINTS` env var.
-    // We want non-default behavior for `exit`, though, so, we handle it separately.
-    //
-    // Format for FAILPOINTS is "name=actions" separated by ";".
-    let actions = std::env::var("FAILPOINTS");
-    if actions.is_ok() {
-        // SAFETY: this function should before any threads start and access env vars concurrently
-        unsafe {
-            std::env::remove_var("FAILPOINTS");
-        }
-    } else {
-        // let the library handle non-utf8, or nothing for not present
-    }
-
-    let scenario = fail::FailScenario::setup();
-
-    if let Ok(val) = actions {
-        val.split(';')
-            .enumerate()
-            .map(|(i, s)| s.split_once('=').ok_or((i, s)))
-            .for_each(|res| {
-                let (name, actions) = match res {
-                    Ok(t) => t,
-                    Err((i, s)) => {
-                        panic!(
-                            "startup failpoints: missing action on the {}th failpoint; try `{s}=return`",
-                            i + 1,
-                        );
-                    }
-                };
-                if let Err(e) = apply_failpoint(name, actions) {
-                    panic!("startup failpoints: failed to apply failpoint {name}={actions}: {e}");
-                }
-            });
-    }
-
-    scenario
-}
-
-pub fn apply_failpoint(name: &str, actions: &str) -> Result<(), String> {
-    if actions == "exit" {
-        fail::cfg_callback(name, exit_failpoint)
-    } else {
-        fail::cfg(name, actions)
-    }
-}
-
-#[inline(never)]
-fn exit_failpoint() {
-    tracing::info!("Exit requested by failpoint");
-    std::process::exit(1);
-}

libs/walproposer/src/api_bindings.rs

@@ -428,6 +428,12 @@ pub fn empty_shmem() -> crate::bindings::WalproposerShmemState {
         shard_number: 0,
     };
 
+    let empty_wal_rate_limiter = crate::bindings::WalRateLimiter {
+        should_limit: crate::bindings::pg_atomic_uint32 { value: 0 },
+        sent_bytes: 0,
+        last_recorded_time_us: 0,
+    };
+
     crate::bindings::WalproposerShmemState {
         propEpochStartLsn: crate::bindings::pg_atomic_uint64 { value: 0 },
         donor_name: [0; 64],
@@ -441,6 +447,7 @@ pub fn empty_shmem() -> crate::bindings::WalproposerShmemState {
         num_shards: 0,
         replica_promote: false,
         min_ps_feedback: empty_feedback,
+        wal_rate_limiter: empty_wal_rate_limiter,
     }
 }
 

pageserver/Cargo.toml

@@ -8,7 +8,7 @@ license.workspace = true
 default = []
 # Enables test-only APIs, incuding failpoints. In particular, enables the `fail_point!` macro,
 # which adds some runtime cost to run tests on outage conditions
-testing = ["fail/failpoints", "pageserver_api/testing", "wal_decoder/testing", "pageserver_client/testing"]
+testing = ["neon_failpoint/testing", "pageserver_api/testing", "wal_decoder/testing", "pageserver_client/testing"]
 
 fuzz-read-path = ["testing"]
 
@@ -33,7 +33,7 @@ crc32c.workspace = true
 either.workspace = true
 enum-map.workspace = true
 enumset = { workspace = true, features = ["serde"]}
-fail.workspace = true
+neon_failpoint.workspace = true
 futures.workspace = true
 hashlink.workspace = true
 hex.workspace = true

pageserver/client_grpc/src/client.rs

@@ -143,7 +143,7 @@ impl PageserverClient {
         req: page_api::CheckRelExistsRequest,
     ) -> tonic::Result<page_api::CheckRelExistsResponse> {
         self.retry
-            .with(async || {
+            .with(async |_| {
                 // Relation metadata is only available on shard 0.
                 let mut client = self.shards.load_full().get_zero().client().await?;
                 client.check_rel_exists(req).await
@@ -158,7 +158,7 @@ impl PageserverClient {
         req: page_api::GetDbSizeRequest,
     ) -> tonic::Result<page_api::GetDbSizeResponse> {
         self.retry
-            .with(async || {
+            .with(async |_| {
                 // Relation metadata is only available on shard 0.
                 let mut client = self.shards.load_full().get_zero().client().await?;
                 client.get_db_size(req).await
@@ -166,8 +166,9 @@ impl PageserverClient {
             .await
     }
 
-    /// Fetches pages. The `request_id` must be unique across all in-flight requests. Automatically
-    /// splits requests that straddle shard boundaries, and assembles the responses.
+    /// Fetches pages. The `request_id` must be unique across all in-flight requests, and the
+    /// `attempt` must be 0 (incremented on retry). Automatically splits requests that straddle
+    /// shard boundaries, and assembles the responses.
     ///
     /// Unlike `page_api::Client`, this automatically converts `status_code` into `tonic::Status`
     /// errors. All responses will have `GetPageStatusCode::Ok`.
@@ -187,6 +188,10 @@ impl PageserverClient {
         if req.block_numbers.is_empty() {
             return Err(tonic::Status::invalid_argument("no block number"));
         }
+        // The request attempt must be 0. The client will increment it internally.
+        if req.request_id.attempt != 0 {
+            return Err(tonic::Status::invalid_argument("request attempt must be 0"));
+        }
 
         // The shards may change while we're fetching pages. We execute the request using a stable
         // view of the shards (especially important for requests that span shards), but retry the
@@ -197,7 +202,11 @@ impl PageserverClient {
         // TODO: the gRPC server and client doesn't yet properly support shard splits. Revisit this
         // once we figure out how to handle these.
         self.retry
-            .with(async || Self::get_page_with_shards(req.clone(), &self.shards.load_full()).await)
+            .with(async |attempt| {
+                let mut req = req.clone();
+                req.request_id.attempt = attempt as u32;
+                Self::get_page_with_shards(req, &self.shards.load_full()).await
+            })
             .await
     }
 
@@ -209,7 +218,7 @@ impl PageserverClient {
     ) -> tonic::Result<page_api::GetPageResponse> {
         // Fast path: request is for a single shard.
         if let Some(shard_id) =
-            GetPageSplitter::is_single_shard(&req, shards.count, shards.stripe_size)
+            GetPageSplitter::for_single_shard(&req, shards.count, shards.stripe_size)
         {
             return Self::get_page_with_shard(req, shards.get(shard_id)?).await;
         }
@@ -229,7 +238,7 @@ impl PageserverClient {
             splitter.add_response(shard_id, shard_response)?;
         }
 
-        splitter.assemble_response()
+        splitter.get_response()
     }
 
     /// Fetches pages on the given shard. Does not retry internally.
@@ -237,9 +246,8 @@ impl PageserverClient {
         req: page_api::GetPageRequest,
         shard: &Shard,
     ) -> tonic::Result<page_api::GetPageResponse> {
-        let expected = req.block_numbers.len();
         let stream = shard.stream(req.request_class.is_bulk()).await;
-        let resp = stream.send(req).await?;
+        let resp = stream.send(req.clone()).await?;
 
         // Convert per-request errors into a tonic::Status.
         if resp.status_code != page_api::GetPageStatusCode::Ok {
@@ -249,11 +257,27 @@ impl PageserverClient {
             ));
         }
 
-        // Check that we received the expected number of pages.
-        let actual = resp.page_images.len();
-        if expected != actual {
+        // Check that we received the expected pages.
+        if req.rel != resp.rel {
             return Err(tonic::Status::internal(format!(
-                "expected {expected} pages, got {actual}",
+                "shard {} returned wrong relation, expected {} got {}",
+                shard.id, req.rel, resp.rel
+            )));
+        }
+        if !req
+            .block_numbers
+            .iter()
+            .copied()
+            .eq(resp.pages.iter().map(|p| p.block_number))
+        {
+            return Err(tonic::Status::internal(format!(
+                "shard {} returned wrong pages, expected {:?} got {:?}",
+                shard.id,
+                req.block_numbers,
+                resp.pages
+                    .iter()
+                    .map(|page| page.block_number)
+                    .collect::<Vec<_>>()
             )));
         }
 
@@ -267,7 +291,7 @@ impl PageserverClient {
         req: page_api::GetRelSizeRequest,
     ) -> tonic::Result<page_api::GetRelSizeResponse> {
         self.retry
-            .with(async || {
+            .with(async |_| {
                 // Relation metadata is only available on shard 0.
                 let mut client = self.shards.load_full().get_zero().client().await?;
                 client.get_rel_size(req).await
@@ -282,7 +306,7 @@ impl PageserverClient {
         req: page_api::GetSlruSegmentRequest,
     ) -> tonic::Result<page_api::GetSlruSegmentResponse> {
         self.retry
-            .with(async || {
+            .with(async |_| {
                 // SLRU segments are only available on shard 0.
                 let mut client = self.shards.load_full().get_zero().client().await?;
                 client.get_slru_segment(req).await
@@ -426,6 +450,8 @@ impl Shards {
 ///   * Bulk client pool: unbounded.
 ///     * Bulk stream pool: MAX_BULK_STREAMS and MAX_BULK_STREAM_QUEUE_DEPTH.
 struct Shard {
+    /// The shard ID.
+    id: ShardIndex,
     /// Unary gRPC client pool.
     client_pool: Arc<ClientPool>,
     /// GetPage stream pool.
@@ -491,6 +517,7 @@ impl Shard {
         );
 
         Ok(Self {
+            id: shard_id,
             client_pool,
             stream_pool,
             bulk_stream_pool,

pageserver/client_grpc/src/pool.rs

@@ -591,6 +591,10 @@ impl StreamPool {
 
         // Track caller response channels by request ID. If the task returns early, these response
         // channels will be dropped and the waiting callers will receive an error.
+        //
+        // NB: this will leak entries if the server doesn't respond to a request (by request ID).
+        // It shouldn't happen, and if it does it will often hold onto queue depth quota anyway and
+        // block further use. But we could consider reaping closed channels after some time.
         let mut callers = HashMap::new();
 
         // Process requests and responses.
@@ -695,6 +699,15 @@ impl Drop for StreamGuard {
 
         // Release the queue depth reservation on drop. This can prematurely decrement it if dropped
         // before the response is received, but that's okay.
+        //
+        // TODO: actually, it's probably not okay. Queue depth release should be moved into the
+        // stream task, such that it continues to account for the queue depth slot until the server
+        // responds. Otherwise, if a slow request times out and keeps blocking the stream, the
+        // server will keep waiting on it and we can pile on subsequent requests (including the
+        // timeout retry) in the same stream and get blocked. But we may also want to avoid blocking
+        // requests on e.g. LSN waits and layer downloads, instead returning early to free up the
+        // stream. Or just scale out streams with a queue depth of 1 to sidestep all head-of-line
+        // blocking. TBD.
         let mut streams = pool.streams.lock().unwrap();
         let entry = streams.get_mut(&self.id).expect("unknown stream");
         assert!(entry.idle_since.is_none(), "active stream marked idle");

pageserver/client_grpc/src/retry.rs

@@ -23,14 +23,14 @@ impl Retry {
     /// If true, log successful requests. For debugging.
     const LOG_SUCCESS: bool = false;
 
-    /// Runs the given async closure with timeouts and retries (exponential backoff). Logs errors,
-    /// using the current tracing span for context.
+    /// Runs the given async closure with timeouts and retries (exponential backoff), passing the
+    /// attempt number starting at 0. Logs errors, using the current tracing span for context.
     ///
     /// Only certain gRPC status codes are retried, see [`Self::should_retry`]. For default
     /// timeouts, see [`Self::REQUEST_TIMEOUT`] and [`Self::TOTAL_TIMEOUT`].
     pub async fn with<T, F, O>(&self, mut f: F) -> tonic::Result<T>
     where
-        F: FnMut() -> O,
+        F: FnMut(usize) -> O, // takes attempt number, starting at 0
         O: Future<Output = tonic::Result<T>>,
     {
         let started = Instant::now();
@@ -47,7 +47,7 @@ impl Retry {
                 }
 
                 let request_started = Instant::now();
-                tokio::time::timeout(Self::REQUEST_TIMEOUT, f())
+                tokio::time::timeout(Self::REQUEST_TIMEOUT, f(retries))
                     .await
                     .map_err(|_| {
                         tonic::Status::deadline_exceeded(format!(

pageserver/client_grpc/src/split.rs

@@ -5,27 +5,24 @@ use bytes::Bytes;
 use pageserver_api::key::rel_block_to_key;
 use pageserver_api::shard::{ShardStripeSize, key_to_shard_number};
 use pageserver_page_api as page_api;
-use utils::shard::{ShardCount, ShardIndex};
+use utils::shard::{ShardCount, ShardIndex, ShardNumber};
 
 /// Splits GetPageRequests that straddle shard boundaries and assembles the responses.
 /// TODO: add tests for this.
 pub struct GetPageSplitter {
-    /// The original request ID. Used for all shard requests.
-    request_id: page_api::RequestID,
     /// Split requests by shard index.
     requests: HashMap<ShardIndex, page_api::GetPageRequest>,
-    /// Maps the offset in `GetPageRequest::block_numbers` to the owning shard. Used to assemble
-    /// the response pages in the same order as the original request.
+    /// The response being assembled. Preallocated with empty pages, to be filled in.
+    response: page_api::GetPageResponse,
+    /// Maps the offset in `request.block_numbers` and `response.pages` to the owning shard. Used
+    /// to assemble the response pages in the same order as the original request.
     block_shards: Vec<ShardIndex>,
-    /// Page responses by shard index. Will be assembled into a single response.
-    responses: HashMap<ShardIndex, Vec<Bytes>>,
 }
 
 impl GetPageSplitter {
     /// Checks if the given request only touches a single shard, and returns the shard ID. This is
     /// the common case, so we check first in order to avoid unnecessary allocations and overhead.
-    /// The caller must ensure that the request has at least one block number, or this will panic.
-    pub fn is_single_shard(
+    pub fn for_single_shard(
         req: &page_api::GetPageRequest,
         count: ShardCount,
         stripe_size: ShardStripeSize,
@@ -35,8 +32,12 @@ impl GetPageSplitter {
             return Some(ShardIndex::unsharded());
         }
 
-        // Find the base shard index for the first page, and compare with the rest.
-        let key = rel_block_to_key(req.rel, *req.block_numbers.first().expect("no pages"));
+        // 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 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);
 
         req.block_numbers
@@ -57,19 +58,19 @@ impl GetPageSplitter {
     ) -> Self {
         // The caller should make sure we don't split requests unnecessarily.
         debug_assert!(
-            Self::is_single_shard(&req, count, stripe_size).is_none(),
+            Self::for_single_shard(&req, count, stripe_size).is_none(),
             "unnecessary request split"
         );
 
         // 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());
-        for blkno in req.block_numbers {
+        for &blkno in &req.block_numbers {
             let key = rel_block_to_key(req.rel, blkno);
             let shard_number = key_to_shard_number(count, stripe_size, &key);
             let shard_id = ShardIndex::new(shard_number, count);
 
-            let shard_req = requests
+            requests
                 .entry(shard_id)
                 .or_insert_with(|| page_api::GetPageRequest {
                     request_id: req.request_id,
@@ -77,20 +78,39 @@ impl GetPageSplitter {
                     rel: req.rel,
                     read_lsn: req.read_lsn,
                     block_numbers: Vec::new(),
-                });
-            shard_req.block_numbers.push(blkno);
+                })
+                .block_numbers
+                .push(blkno);
             block_shards.push(shard_id);
         }
 
-        Self {
+        // Construct a response to be populated by shard responses. Preallocate empty page slots
+        // with the expected block numbers.
+        let response = page_api::GetPageResponse {
             request_id: req.request_id,
-            responses: HashMap::with_capacity(requests.len()),
+            status_code: page_api::GetPageStatusCode::Ok,
+            reason: None,
+            rel: req.rel,
+            pages: req
+                .block_numbers
+                .into_iter()
+                .map(|block_number| {
+                    page_api::Page {
+                        block_number,
+                        image: Bytes::new(), // empty page slot to be filled in
+                    }
+                })
+                .collect(),
+        };
+
+        Self {
             requests,
+            response,
             block_shards,
         }
     }
 
-    /// Drains the per-shard requests, moving them out of the hashmap to avoid extra allocations.
+    /// Drains the per-shard requests, moving them out of the splitter to avoid extra allocations.
     pub fn drain_requests(
         &mut self,
     ) -> impl Iterator<Item = (ShardIndex, page_api::GetPageRequest)> {
@@ -108,72 +128,82 @@ impl GetPageSplitter {
         // The caller should already have converted status codes into tonic::Status.
         if response.status_code != page_api::GetPageStatusCode::Ok {
             return Err(tonic::Status::internal(format!(
-                "unexpected non-OK response for shard {shard_id}: {:?}",
-                response.status_code
+                "unexpected non-OK response for shard {shard_id}: {} {}",
+                response.status_code,
+                response.reason.unwrap_or_default()
             )));
         }
 
-        // The stream pool ensures the response matches the request ID.
-        if response.request_id != self.request_id {
+        if response.request_id != self.response.request_id {
             return Err(tonic::Status::internal(format!(
                 "response ID mismatch for shard {shard_id}: expected {}, got {}",
-                self.request_id, response.request_id
+                self.response.request_id, response.request_id
             )));
         }
 
-        // We only dispatch one request per shard.
-        if self.responses.contains_key(&shard_id) {
+        // Place the shard response pages into the assembled response, in request order.
+        let mut pages = response.pages.into_iter();
+
+        for (i, &s) in self.block_shards.iter().enumerate() {
+            if shard_id != s {
+                continue;
+            }
+
+            let Some(slot) = self.response.pages.get_mut(i) else {
+                return Err(tonic::Status::internal(format!(
+                    "no block_shards slot {i} for shard {shard_id}"
+                )));
+            };
+            let Some(page) = pages.next() else {
+                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(tonic::Status::internal(format!(
+                    "shard {shard_id} returned wrong page at index {i}, expected {} got {}",
+                    slot.block_number, page.block_number
+                )));
+            }
+            if !slot.image.is_empty() {
+                return Err(tonic::Status::internal(format!(
+                    "shard {shard_id} returned duplicate page {} at index {i}",
+                    slot.block_number
+                )));
+            }
+
+            *slot = page;
+        }
+
+        // Make sure we've consumed all pages from the shard response.
+        if let Some(extra_page) = pages.next() {
             return Err(tonic::Status::internal(format!(
-                "duplicate response for shard {shard_id}"
+                "shard {shard_id} returned extra page: {}",
+                extra_page.block_number
             )));
         }
 
-        // Add the response data to the map.
-        self.responses.insert(shard_id, response.page_images);
-
         Ok(())
     }
 
-    /// Assembles the shard responses into a single response. Responses must be present for all
-    /// relevant shards, and the total number of pages must match the original request.
+    /// Fetches the final, assembled response.
     #[allow(clippy::result_large_err)]
-    pub fn assemble_response(self) -> tonic::Result<page_api::GetPageResponse> {
-        let mut response = page_api::GetPageResponse {
-            request_id: self.request_id,
-            status_code: page_api::GetPageStatusCode::Ok,
-            reason: None,
-            page_images: Vec::with_capacity(self.block_shards.len()),
-        };
-
-        // Set up per-shard page iterators we can pull from.
-        let mut shard_responses = HashMap::with_capacity(self.responses.len());
-        for (shard_id, responses) in self.responses {
-            shard_responses.insert(shard_id, responses.into_iter());
-        }
-
-        // Reassemble the responses in the same order as the original request.
-        for shard_id in &self.block_shards {
-            let page = shard_responses
-                .get_mut(shard_id)
-                .ok_or_else(|| {
-                    tonic::Status::internal(format!("missing response for shard {shard_id}"))
-                })?
-                .next()
-                .ok_or_else(|| {
-                    tonic::Status::internal(format!("missing page from shard {shard_id}"))
-                })?;
-            response.page_images.push(page);
-        }
-
-        // Make sure there are no additional pages.
-        for (shard_id, mut pages) in shard_responses {
-            if pages.next().is_some() {
+    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(tonic::Status::internal(format!(
-                    "extra pages returned from shard {shard_id}"
+                    "missing page {} for shard {}",
+                    page.block_number,
+                    self.block_shards
+                        .get(i)
+                        .map(|s| s.to_string())
+                        .unwrap_or_else(|| "?".to_string())
                 )));
             }
         }
 
-        Ok(response)
+        Ok(self.response)
     }
 }

pageserver/page_api/proto/page_service.proto

@@ -153,7 +153,7 @@ message GetDbSizeResponse {
 message GetPageRequest {
   // A request ID. Will be included in the response. Should be unique for
   // in-flight requests on the stream.
-  uint64 request_id = 1;
+  RequestID request_id = 1;
   // The request class.
   GetPageClass request_class = 2;
   // The LSN to read at.
@@ -177,6 +177,14 @@ message GetPageRequest {
   repeated uint32 block_number = 5;
 }
 
+// A Request ID. Should be unique for in-flight requests on a stream. Included in the response.
+message RequestID {
+  // The base request ID.
+  uint64 id = 1;
+  // The request attempt. Starts at 0, incremented on each retry.
+  uint32 attempt = 2;
+}
+
 // A GetPageRequest class. Primarily intended for observability, but may also be
 // used for prioritization in the future.
 enum GetPageClass {
@@ -199,13 +207,26 @@ enum GetPageClass {
 // the entire batch is ready, so no one can make use of the individual pages.
 message GetPageResponse {
   // The original request's ID.
-  uint64 request_id = 1;
-  // The response status code.
+  RequestID request_id = 1;
+  // The response status code. If not OK, the rel and page fields will be empty.
   GetPageStatusCode status_code = 2;
   // A string describing the status, if any.
   string reason = 3;
-  // The 8KB page images, in the same order as the request. Empty if status_code != OK.
-  repeated bytes page_image = 4;
+  // The relation that the pages belong to.
+  RelTag rel = 4;
+  // The page(s), in the same order as the request.
+  repeated Page page = 5;
+}
+
+// A page.
+//
+// TODO: it would be slightly more efficient (but less convenient) to have separate arrays of block
+// numbers and images, but given the 8KB page size it's probably negligible. Benchmark it anyway.
+message Page {
+  // The page number.
+  uint32 block_number = 1;
+  // The materialized page image, as an 8KB byte vector.
+  bytes image = 2;
 }
 
 // A GetPageResponse status code.

pageserver/page_api/src/client.rs

@@ -1,4 +1,5 @@
 use anyhow::Context as _;
+use futures::future::ready;
 use futures::{Stream, StreamExt as _, TryStreamExt as _};
 use tokio::io::AsyncRead;
 use tokio_util::io::StreamReader;
@@ -110,7 +111,7 @@ impl Client {
     ) -> tonic::Result<impl Stream<Item = tonic::Result<GetPageResponse>> + Send + 'static> {
         let reqs = reqs.map(proto::GetPageRequest::from);
         let resps = self.inner.get_pages(reqs).await?.into_inner();
-        Ok(resps.map_ok(GetPageResponse::from))
+        Ok(resps.and_then(|resp| ready(GetPageResponse::try_from(resp).map_err(|err| err.into()))))
     }
 
     /// Returns the size of a relation, as # of blocks.

pageserver/page_api/src/model.rs

@@ -356,7 +356,10 @@ impl TryFrom<proto::GetPageRequest> for GetPageRequest {
             return Err(ProtocolError::Missing("block_number"));
         }
         Ok(Self {
-            request_id: pb.request_id,
+            request_id: pb
+                .request_id
+                .ok_or(ProtocolError::Missing("request_id"))?
+                .into(),
             request_class: pb.request_class.into(),
             read_lsn: pb
                 .read_lsn
@@ -371,7 +374,7 @@ impl TryFrom<proto::GetPageRequest> for GetPageRequest {
 impl From<GetPageRequest> for proto::GetPageRequest {
     fn from(request: GetPageRequest) -> Self {
         Self {
-            request_id: request.request_id,
+            request_id: Some(request.request_id.into()),
             request_class: request.request_class.into(),
             read_lsn: Some(request.read_lsn.into()),
             rel: Some(request.rel.into()),
@@ -380,8 +383,51 @@ impl From<GetPageRequest> for proto::GetPageRequest {
     }
 }
 
-/// A GetPage request ID.
-pub type RequestID = u64;
+/// A GetPage request ID and retry attempt. Should be unique for in-flight requests on a stream.
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, PartialOrd, Ord)]
+pub struct RequestID {
+    /// The base request ID.
+    pub id: u64,
+    // The request attempt. Starts at 0, incremented on each retry.
+    pub attempt: u32,
+}
+
+impl RequestID {
+    /// Creates a new RequestID with the given ID and an initial attempt of 0.
+    pub fn new(id: u64) -> Self {
+        Self { id, attempt: 0 }
+    }
+}
+
+impl Display for RequestID {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}.{}", self.id, self.attempt)
+    }
+}
+
+impl From<proto::RequestId> for RequestID {
+    fn from(pb: proto::RequestId) -> Self {
+        Self {
+            id: pb.id,
+            attempt: pb.attempt,
+        }
+    }
+}
+
+impl From<u64> for RequestID {
+    fn from(id: u64) -> Self {
+        Self::new(id)
+    }
+}
+
+impl From<RequestID> for proto::RequestId {
+    fn from(request_id: RequestID) -> Self {
+        Self {
+            id: request_id.id,
+            attempt: request_id.attempt,
+        }
+    }
+}
 
 /// A GetPage request class.
 #[derive(Clone, Copy, Debug, strum_macros::Display)]
@@ -456,32 +502,41 @@ impl From<GetPageClass> for i32 {
 pub struct GetPageResponse {
     /// The original request's ID.
     pub request_id: RequestID,
-    /// The response status code.
+    /// The response status code. If not OK, the `rel` and `pages` fields will be empty.
     pub status_code: GetPageStatusCode,
     /// A string describing the status, if any.
     pub reason: Option<String>,
-    /// The 8KB page images, in the same order as the request. Empty if status != OK.
-    pub page_images: Vec<Bytes>,
+    /// The relation that the pages belong to.
+    pub rel: RelTag,
+    // The page(s), in the same order as the request.
+    pub pages: Vec<Page>,
 }
 
-impl From<proto::GetPageResponse> for GetPageResponse {
-    fn from(pb: proto::GetPageResponse) -> Self {
-        Self {
-            request_id: pb.request_id,
+impl TryFrom<proto::GetPageResponse> for GetPageResponse {
+    type Error = ProtocolError;
+
+    fn try_from(pb: proto::GetPageResponse) -> Result<Self, ProtocolError> {
+        Ok(Self {
+            request_id: pb
+                .request_id
+                .ok_or(ProtocolError::Missing("request_id"))?
+                .into(),
             status_code: pb.status_code.into(),
             reason: Some(pb.reason).filter(|r| !r.is_empty()),
-            page_images: pb.page_image,
-        }
+            rel: pb.rel.ok_or(ProtocolError::Missing("rel"))?.try_into()?,
+            pages: pb.page.into_iter().map(Page::from).collect(),
+        })
     }
 }
 
 impl From<GetPageResponse> for proto::GetPageResponse {
     fn from(response: GetPageResponse) -> Self {
         Self {
-            request_id: response.request_id,
+            request_id: Some(response.request_id.into()),
             status_code: response.status_code.into(),
             reason: response.reason.unwrap_or_default(),
-            page_image: response.page_images,
+            rel: Some(response.rel.into()),
+            page: response.pages.into_iter().map(proto::Page::from).collect(),
         }
     }
 }
@@ -514,11 +569,39 @@ impl GetPageResponse {
             request_id,
             status_code,
             reason: Some(status.message().to_string()),
-            page_images: Vec::new(),
+            rel: RelTag::default(),
+            pages: Vec::new(),
         })
     }
 }
 
+// A page.
+#[derive(Clone, Debug)]
+pub struct Page {
+    /// The page number.
+    pub block_number: u32,
+    /// The materialized page image, as an 8KB byte vector.
+    pub image: Bytes,
+}
+
+impl From<proto::Page> for Page {
+    fn from(pb: proto::Page) -> Self {
+        Self {
+            block_number: pb.block_number,
+            image: pb.image,
+        }
+    }
+}
+
+impl From<Page> for proto::Page {
+    fn from(page: Page) -> Self {
+        Self {
+            block_number: page.block_number,
+            image: page.image,
+        }
+    }
+}
+
 /// A GetPage response status code.
 ///
 /// These are effectively equivalent to gRPC statuses. However, we use a bidirectional stream

pageserver/pagebench/src/cmd/getpage_latest_lsn.rs

@@ -674,7 +674,7 @@ impl Client for GrpcClient {
         blks: Vec<u32>,
     ) -> anyhow::Result<()> {
         let req = page_api::GetPageRequest {
-            request_id: req_id,
+            request_id: req_id.into(),
             request_class: page_api::GetPageClass::Normal,
             read_lsn: page_api::ReadLsn {
                 request_lsn: req_lsn,
@@ -694,7 +694,10 @@ impl Client for GrpcClient {
             "unexpected status code: {}",
             resp.status_code,
         );
-        Ok((resp.request_id, resp.page_images))
+        Ok((
+            resp.request_id.id,
+            resp.pages.into_iter().map(|p| p.image).collect(),
+        ))
     }
 }
 
@@ -740,7 +743,7 @@ impl Client for RichGrpcClient {
         blks: Vec<u32>,
     ) -> anyhow::Result<()> {
         let req = page_api::GetPageRequest {
-            request_id: req_id,
+            request_id: req_id.into(),
             request_class: page_api::GetPageClass::Normal,
             read_lsn: page_api::ReadLsn {
                 request_lsn: req_lsn,
@@ -761,6 +764,9 @@ impl Client for RichGrpcClient {
 
     async fn recv_get_page(&mut self) -> anyhow::Result<(u64, Vec<Bytes>)> {
         let resp = self.requests.next().await.unwrap()?;
-        Ok((resp.request_id, resp.page_images))
+        Ok((
+            resp.request_id.id,
+            resp.pages.into_iter().map(|p| p.image).collect(),
+        ))
     }
 }

pageserver/src/basebackup.rs

@@ -17,6 +17,7 @@ use anyhow::{Context, anyhow};
 use async_compression::tokio::write::GzipEncoder;
 use bytes::{BufMut, Bytes, BytesMut};
 use fail::fail_point;
+use neon_failpoint as fail;
 use pageserver_api::key::{Key, rel_block_to_key};
 use pageserver_api::reltag::{RelTag, SlruKind};
 use postgres_ffi::pg_constants::{PG_HBA, PGDATA_SPECIAL_FILES};

pageserver/src/bin/pageserver.rs

@@ -68,7 +68,7 @@ const FEATURES: &[&str] = &[
 fn version() -> String {
     format!(
         "{GIT_VERSION} failpoints: {}, features: {:?}",
-        fail::has_failpoints(),
+        neon_failpoint::has_failpoints(),
         FEATURES,
     )
 }
@@ -84,7 +84,7 @@ fn main() -> anyhow::Result<()> {
     }
 
     // Initialize up failpoints support
-    let scenario = failpoint_support::init();
+    failpoint_support::init().unwrap();
 
     let workdir = arg_matches
         .get_one::<String>("workdir")
@@ -221,7 +221,6 @@ fn main() -> anyhow::Result<()> {
 
     start_pageserver(launch_ts, conf, ignored, otel_guard).context("Failed to start pageserver")?;
 
-    scenario.teardown();
     Ok(())
 }
 
@@ -366,16 +365,9 @@ fn start_pageserver(
 
     // If any failpoints were set from FAILPOINTS environment variable,
     // print them to the log for debugging purposes
-    let failpoints = fail::list();
-    if !failpoints.is_empty() {
-        info!(
-            "started with failpoints: {}",
-            failpoints
-                .iter()
-                .map(|(name, actions)| format!("{name}={actions}"))
-                .collect::<Vec<String>>()
-                .join(";")
-        )
+    let failpoints = neon_failpoint::list();
+    for (name, actions) in failpoints {
+        info!("starting with failpoint: {name} {actions}");
     }
 
     // Create and lock PID file. This ensures that there cannot be more than one

pageserver/src/consumption_metrics/disk_cache.rs

@@ -6,6 +6,8 @@ use camino::{Utf8Path, Utf8PathBuf};
 use super::{NewMetricsRoot, NewRawMetric, RawMetric};
 use crate::consumption_metrics::NewMetricsRefRoot;
 
+use neon_failpoint as fail;
+
 pub(super) fn read_metrics_from_serde_value(
     json_value: serde_json::Value,
 ) -> anyhow::Result<Vec<NewRawMetric>> {
@@ -129,7 +131,7 @@ pub(super) async fn flush_metrics_to_disk(
             tempfile.flush()?;
             tempfile.as_file().sync_all()?;
 
-            fail::fail_point!("before-persist-last-metrics-collected");
+            fail::fail_point_sync!("before-persist-last-metrics-collected");
 
             drop(tempfile.persist(&*path).map_err(|e| e.error)?);
 

pageserver/src/deletion_queue/deleter.rs

@@ -8,6 +8,7 @@
 
 use std::time::Duration;
 
+use neon_failpoint as fail;
 use remote_storage::{GenericRemoteStorage, RemotePath, TimeoutOrCancel};
 use tokio_util::sync::CancellationToken;
 use tracing::{info, warn};

pageserver/src/http/openapi_spec.yml

@@ -271,18 +271,6 @@ paths:
             application/json:
               schema:
                 $ref: "#/components/schemas/LsnLease"
-        "500":
-          description: Generic operation error
-          content:
-            application/json:
-              schema:
-                $ref: "#/components/schemas/Error"
-        "503":
-          description: Temporarily unavailable, please retry.
-          content:
-            application/json:
-              schema:
-                $ref: "#/components/schemas/ServiceUnavailableError"
 
   /v1/tenant/{tenant_id}/timeline/{timeline_id}/do_gc:
     parameters:

pageserver/src/http/routes.rs

@@ -28,6 +28,7 @@ use http_utils::{RequestExt, RouterBuilder};
 use humantime::format_rfc3339;
 use hyper::{Body, Request, Response, StatusCode, Uri, header};
 use metrics::launch_timestamp::LaunchTimestamp;
+use neon_failpoint as fail;
 use pageserver_api::models::virtual_file::IoMode;
 use pageserver_api::models::{
     DetachBehavior, DownloadRemoteLayersTaskSpawnRequest, IngestAuxFilesRequest,
@@ -3975,7 +3976,7 @@ pub fn make_router(
         .get("/profile/cpu", |r| request_span(r, profile_cpu_handler))
         .get("/profile/heap", |r| request_span(r, profile_heap_handler))
         .get("/v1/status", |r| api_handler(r, status_handler))
-        .put("/v1/failpoints", |r| {
+        .post("/v1/failpoints", |r| {
             testing_api_handler("manage failpoints", r, failpoints_handler)
         })
         .post("/v1/reload_auth_validation_keys", |r| {

pageserver/src/page_service.rs

@@ -19,6 +19,7 @@ use futures::future::BoxFuture;
 use futures::{FutureExt, Stream};
 use itertools::Itertools;
 use jsonwebtoken::TokenData;
+use neon_failpoint as fail;
 use once_cell::sync::OnceCell;
 use pageserver_api::config::{
     GetVectoredConcurrentIo, PageServicePipeliningConfig, PageServicePipeliningConfigPipelined,
@@ -335,18 +336,21 @@ async fn page_service_conn_main(
 
     let default_timeout_ms = 10 * 60 * 1000; // 10 minutes by default
     let socket_timeout_ms = (|| {
-        fail::fail_point!("simulated-bad-compute-connection", |avg_timeout_ms| {
-            // Exponential distribution for simulating
-            // poor network conditions, expect about avg_timeout_ms to be around 15
-            // in tests
-            if let Some(avg_timeout_ms) = avg_timeout_ms {
-                let avg = avg_timeout_ms.parse::<i64>().unwrap() as f32;
-                let u = rand::random::<f32>();
-                ((1.0 - u).ln() / (-avg)) as u64
-            } else {
-                default_timeout_ms
+        fail::fail_point_sync!(
+            "simulated-bad-compute-connection",
+            |avg_timeout_ms: Option<String>| {
+                // Exponential distribution for simulating
+                // poor network conditions, expect about avg_timeout_ms to be around 15
+                // in tests
+                if let Some(avg_timeout_ms) = avg_timeout_ms {
+                    let avg = avg_timeout_ms.parse::<i64>().unwrap() as f32;
+                    let u = rand::random::<f32>();
+                    ((1.0 - u).ln() / (-avg)) as u64
+                } else {
+                    default_timeout_ms
+                }
             }
-        });
+        );
         default_timeout_ms
     })();
 
@@ -3043,7 +3047,7 @@ where
         _pgb: &mut PostgresBackend<IO>,
         sm: &FeStartupPacket,
     ) -> Result<(), QueryError> {
-        fail::fail_point!("ps::connection-start::startup-packet");
+        fail::fail_point_sync!("ps::connection-start::startup-packet");
 
         if let FeStartupPacket::StartupMessage { params, .. } = sm {
             if let Some(app_name) = params.get("application_name") {
@@ -3338,9 +3342,12 @@ impl GrpcPageServiceHandler {
     }
 
     /// Generates a PagestreamRequest header from a ReadLsn and request ID.
-    fn make_hdr(read_lsn: page_api::ReadLsn, req_id: u64) -> PagestreamRequest {
+    fn make_hdr(
+        read_lsn: page_api::ReadLsn,
+        req_id: Option<page_api::RequestID>,
+    ) -> PagestreamRequest {
         PagestreamRequest {
-            reqid: req_id,
+            reqid: req_id.map(|r| r.id).unwrap_or_default(),
             request_lsn: read_lsn.request_lsn,
             not_modified_since: read_lsn
                 .not_modified_since_lsn
@@ -3450,7 +3457,7 @@ impl GrpcPageServiceHandler {
 
             batch.push(BatchedGetPageRequest {
                 req: PagestreamGetPageRequest {
-                    hdr: Self::make_hdr(req.read_lsn, req.request_id),
+                    hdr: Self::make_hdr(req.read_lsn, Some(req.request_id)),
                     rel: req.rel,
                     blkno,
                 },
@@ -3480,12 +3487,16 @@ impl GrpcPageServiceHandler {
             request_id: req.request_id,
             status_code: page_api::GetPageStatusCode::Ok,
             reason: None,
-            page_images: Vec::with_capacity(results.len()),
+            rel: req.rel,
+            pages: Vec::with_capacity(results.len()),
         };
 
         for result in results {
             match result {
-                Ok((PagestreamBeMessage::GetPage(r), _, _)) => resp.page_images.push(r.page),
+                Ok((PagestreamBeMessage::GetPage(r), _, _)) => resp.pages.push(page_api::Page {
+                    block_number: r.req.blkno,
+                    image: r.page,
+                }),
                 Ok((resp, _, _)) => {
                     return Err(tonic::Status::internal(format!(
                         "unexpected response: {resp:?}"
@@ -3528,7 +3539,7 @@ impl proto::PageService for GrpcPageServiceHandler {
         span_record!(rel=%req.rel, lsn=%req.read_lsn);
 
         let req = PagestreamExistsRequest {
-            hdr: Self::make_hdr(req.read_lsn, 0),
+            hdr: Self::make_hdr(req.read_lsn, None),
             rel: req.rel,
         };
 
@@ -3678,7 +3689,7 @@ impl proto::PageService for GrpcPageServiceHandler {
         span_record!(db_oid=%req.db_oid, lsn=%req.read_lsn);
 
         let req = PagestreamDbSizeRequest {
-            hdr: Self::make_hdr(req.read_lsn, 0),
+            hdr: Self::make_hdr(req.read_lsn, None),
             dbnode: req.db_oid,
         };
 
@@ -3728,7 +3739,7 @@ impl proto::PageService for GrpcPageServiceHandler {
                 .await?
                 .downgrade();
             while let Some(req) = reqs.message().await? {
-                let req_id = req.request_id;
+                let req_id = req.request_id.map(page_api::RequestID::from).unwrap_or_default();
                 let result = Self::get_page(&ctx, &timeline, req, io_concurrency.clone())
                     .instrument(span.clone()) // propagate request span
                     .await;
@@ -3767,7 +3778,7 @@ impl proto::PageService for GrpcPageServiceHandler {
         span_record!(rel=%req.rel, lsn=%req.read_lsn);
 
         let req = PagestreamNblocksRequest {
-            hdr: Self::make_hdr(req.read_lsn, 0),
+            hdr: Self::make_hdr(req.read_lsn, None),
             rel: req.rel,
         };
 
@@ -3800,7 +3811,7 @@ impl proto::PageService for GrpcPageServiceHandler {
         span_record!(kind=%req.kind, segno=%req.segno, lsn=%req.read_lsn);
 
         let req = PagestreamGetSlruSegmentRequest {
-            hdr: Self::make_hdr(req.read_lsn, 0),
+            hdr: Self::make_hdr(req.read_lsn, None),
             kind: req.kind as u8,
             segno: req.segno,
         };

pageserver/src/pgdatadir_mapping.rs

@@ -14,6 +14,7 @@ use crate::{PERF_TRACE_TARGET, ensure_walingest};
 use anyhow::Context;
 use bytes::{Buf, Bytes, BytesMut};
 use enum_map::Enum;
+use neon_failpoint as fail;
 use pageserver_api::key::{
     AUX_FILES_KEY, CHECKPOINT_KEY, CONTROLFILE_KEY, CompactKey, DBDIR_KEY, Key, RelDirExists,
     TWOPHASEDIR_KEY, dbdir_key_range, rel_block_to_key, rel_dir_to_key, rel_key_range,

pageserver/src/tenant.rs

@@ -30,6 +30,7 @@ use enumset::EnumSet;
 use futures::StreamExt;
 use futures::stream::FuturesUnordered;
 use itertools::Itertools as _;
+use neon_failpoint as fail;
 use once_cell::sync::Lazy;
 pub use pageserver_api::models::TenantState;
 use pageserver_api::models::{self, RelSizeMigration};
@@ -9571,7 +9572,7 @@ mod tests {
         writer.finish_write(Lsn(0x30));
         drop(writer);
 
-        fail::cfg(
+        neon_failpoint::configure_failpoint(
             "flush-layer-before-update-remote-consistent-lsn",
             "return()",
         )

pageserver/src/tenant/mgr.rs

@@ -12,6 +12,7 @@ use anyhow::Context;
 use camino::{Utf8DirEntry, Utf8Path, Utf8PathBuf};
 use futures::StreamExt;
 use itertools::Itertools;
+use neon_failpoint as fail;
 use pageserver_api::key::Key;
 use pageserver_api::models::{DetachBehavior, LocationConfigMode};
 use pageserver_api::shard::{

pageserver/src/tenant/remote_timeline_client.rs

@@ -194,6 +194,7 @@ pub(crate) use download::{
 };
 use index::GcCompactionState;
 pub(crate) use index::LayerFileMetadata;
+use neon_failpoint as fail;
 use pageserver_api::models::{RelSizeMigration, TimelineArchivalState, TimelineVisibilityState};
 use pageserver_api::shard::{ShardIndex, TenantShardId};
 use regex::Regex;

pageserver/src/tenant/remote_timeline_client/download.rs

@@ -11,6 +11,7 @@ use std::time::SystemTime;
 
 use anyhow::{Context, anyhow};
 use camino::{Utf8Path, Utf8PathBuf};
+use neon_failpoint as fail;
 use pageserver_api::shard::TenantShardId;
 use remote_storage::{
     DownloadError, DownloadKind, DownloadOpts, GenericRemoteStorage, ListingMode, RemotePath,

pageserver/src/tenant/remote_timeline_client/upload.rs

@@ -8,6 +8,7 @@ use anyhow::{Context, bail};
 use bytes::Bytes;
 use camino::Utf8Path;
 use fail::fail_point;
+use neon_failpoint as fail;
 use pageserver_api::shard::TenantShardId;
 use remote_storage::{GenericRemoteStorage, RemotePath, TimeTravelError};
 use tokio::fs::{self, File};

pageserver/src/tenant/timeline.rs

@@ -39,6 +39,7 @@ use layer_manager::{
     LayerManagerLockHolder, LayerManagerReadGuard, LayerManagerWriteGuard, LockedLayerManager,
     Shutdown,
 };
+use neon_failpoint as fail;
 
 use once_cell::sync::Lazy;
 use pageserver_api::config::tenant_conf_defaults::DEFAULT_PITR_INTERVAL;
@@ -5183,7 +5184,9 @@ impl Timeline {
             *self.applied_gc_cutoff_lsn.read(),
         );
 
-        fail_point!("checkpoint-before-saving-metadata", |x| bail!(
+        neon_failpoint::fail_point_sync!("checkpoint-before-saving-metadata", |x: Option<
+            String,
+        >| bail!(
             "{}",
             x.unwrap()
         ));

pageserver/src/tenant/timeline/compaction.rs

@@ -26,6 +26,7 @@ use enumset::EnumSet;
 use fail::fail_point;
 use futures::FutureExt;
 use itertools::Itertools;
+use neon_failpoint as fail;
 use once_cell::sync::Lazy;
 use pageserver_api::config::tenant_conf_defaults::DEFAULT_CHECKPOINT_DISTANCE;
 use pageserver_api::key::{KEY_SIZE, Key};

pageserver/src/tenant/timeline/delete.rs

@@ -2,6 +2,7 @@ use std::ops::{Deref, DerefMut};
 use std::sync::Arc;
 
 use anyhow::Context;
+use neon_failpoint as fail;
 use pageserver_api::models::TimelineState;
 use pageserver_api::shard::TenantShardId;
 use remote_storage::DownloadError;

pageserver/src/tenant/timeline/detach_ancestor.rs

@@ -4,6 +4,7 @@ use std::sync::Arc;
 use anyhow::Context;
 use bytes::Bytes;
 use http_utils::error::ApiError;
+use neon_failpoint as fail;
 use pageserver_api::key::Key;
 use pageserver_api::keyspace::KeySpace;
 use pageserver_api::models::DetachBehavior;
@@ -1113,7 +1114,7 @@ pub(super) async fn detach_and_reparent(
     // others will fail as if those timelines had been stopped for whatever reason.
     #[cfg(feature = "testing")]
     let failpoint_sem = || -> Option<Arc<Semaphore>> {
-        fail::fail_point!("timeline-detach-ancestor::allow_one_reparented", |_| Some(
+        fail::fail_point_sync!("timeline-detach-ancestor::allow_one_reparented", |_| Some(
             Arc::new(Semaphore::new(1))
         ));
         None

pageserver/src/tenant/timeline/uninit.rs

@@ -5,6 +5,7 @@ use std::sync::Arc;
 
 use anyhow::Context;
 use camino::Utf8PathBuf;
+use neon_failpoint as fail;
 use tracing::{error, info, info_span};
 use utils::fs_ext;
 use utils::id::TimelineId;

pageserver/src/tenant/timeline/walreceiver/walreceiver_connection.rs

@@ -11,6 +11,7 @@ use bytes::BytesMut;
 use chrono::{NaiveDateTime, Utc};
 use fail::fail_point;
 use futures::StreamExt;
+use neon_failpoint as fail;
 use postgres_backend::is_expected_io_error;
 use postgres_connection::PgConnectionConfig;
 use postgres_ffi::WAL_SEGMENT_SIZE;

pgxn/neon/libpagestore.c

@@ -1410,7 +1410,7 @@ pg_init_libpagestore(void)
 							"sharding stripe size",
 							NULL,
 							&stripe_size,
-							32768, 1, INT_MAX,
+							2048, 1, INT_MAX,
 							PGC_SIGHUP,
 							GUC_UNIT_BLOCKS,
 							NULL, NULL, NULL);

pgxn/neon/walproposer.h

@@ -376,6 +376,18 @@ typedef struct PageserverFeedback
 	uint32		shard_number;
 } PageserverFeedback;
 
+/* BEGIN_HADRON */
+typedef struct WalRateLimiter
+{
+	/* If the value is 1, PG backends will hit backpressure. */
+	pg_atomic_uint32 should_limit;
+	/* The number of bytes sent in the current second. */
+	uint64		sent_bytes;
+	/* The last recorded time in microsecond. */
+	TimestampTz last_recorded_time_us;
+} WalRateLimiter;
+/* END_HADRON */
+
 typedef struct WalproposerShmemState
 {
 	pg_atomic_uint64 propEpochStartLsn;
@@ -395,6 +407,11 @@ typedef struct WalproposerShmemState
 
 	/* aggregated feedback with min LSNs across shards */
 	PageserverFeedback min_ps_feedback;
+
+	/* BEGIN_HADRON */
+	/* The WAL rate limiter */
+	WalRateLimiter wal_rate_limiter;
+	/* END_HADRON */
 } WalproposerShmemState;
 
 /*

pgxn/neon/walproposer_pg.c

@@ -66,6 +66,9 @@ int			wal_acceptor_reconnect_timeout = 1000;
 int			wal_acceptor_connection_timeout = 10000;
 int			safekeeper_proto_version = 3;
 char	   *safekeeper_conninfo_options = "";
+/* BEGIN_HADRON */
+int         databricks_max_wal_mb_per_second = -1;
+/* END_HADRON */
 
 /* Set to true in the walproposer bgw. */
 static bool am_walproposer;
@@ -252,6 +255,18 @@ nwp_register_gucs(void)
 							PGC_POSTMASTER,
 							0,
 							NULL, NULL, NULL);
+
+    /* BEGIN_HADRON */
+    DefineCustomIntVariable(
+                            "databricks.max_wal_mb_per_second",
+                            "The maximum WAL MB per second allowed. If breached, sending WAL hit the backpressure. Setting to -1 disables the limit.",
+                            NULL,
+                            &databricks_max_wal_mb_per_second,
+                            -1, -1, INT_MAX,
+                            PGC_SUSET,
+                            GUC_UNIT_MB,
+                            NULL, NULL, NULL);
+    /* END_HADRON */
 }
 
 
@@ -393,6 +408,7 @@ assign_neon_safekeepers(const char *newval, void *extra)
 static uint64
 backpressure_lag_impl(void)
 {
+	struct WalproposerShmemState* state = NULL;
 	if (max_replication_apply_lag > 0 || max_replication_flush_lag > 0 || max_replication_write_lag > 0)
 	{
 		XLogRecPtr	writePtr;
@@ -426,6 +442,18 @@ backpressure_lag_impl(void)
 			return (myFlushLsn - applyPtr - max_replication_apply_lag * MB);
 		}
 	}
+
+	/* BEGIN_HADRON */
+	if (databricks_max_wal_mb_per_second == -1) {
+		return 0;
+	}
+
+	state = GetWalpropShmemState();
+	if (state != NULL && pg_atomic_read_u32(&state->wal_rate_limiter.should_limit) == 1)
+	{
+		return 1;
+	}
+	/* END_HADRON */
 	return 0;
 }
 
@@ -472,6 +500,9 @@ WalproposerShmemInit(void)
 		pg_atomic_init_u64(&walprop_shared->mineLastElectedTerm, 0);
 		pg_atomic_init_u64(&walprop_shared->backpressureThrottlingTime, 0);
 		pg_atomic_init_u64(&walprop_shared->currentClusterSize, 0);
+		/* BEGIN_HADRON */
+		pg_atomic_init_u32(&walprop_shared->wal_rate_limiter.should_limit, 0);
+		/* END_HADRON */
 	}
 	LWLockRelease(AddinShmemInitLock);
 
@@ -487,6 +518,9 @@ WalproposerShmemInit_SyncSafekeeper(void)
 	pg_atomic_init_u64(&walprop_shared->propEpochStartLsn, 0);
 	pg_atomic_init_u64(&walprop_shared->mineLastElectedTerm, 0);
 	pg_atomic_init_u64(&walprop_shared->backpressureThrottlingTime, 0);
+	/* BEGIN_HADRON */
+	pg_atomic_init_u32(&walprop_shared->wal_rate_limiter.should_limit, 0);
+	/* END_HADRON */
 }
 
 #define BACK_PRESSURE_DELAY 10000L // 0.01 sec
@@ -521,7 +555,6 @@ backpressure_throttling_impl(void)
 	if (lag == 0)
 		return retry;
 
-
 	old_status = get_ps_display(&len);
 	new_status = (char *) palloc(len + 64 + 1);
 	memcpy(new_status, old_status, len);
@@ -1458,6 +1491,8 @@ XLogBroadcastWalProposer(WalProposer *wp)
 {
 	XLogRecPtr	startptr;
 	XLogRecPtr	endptr;
+	struct WalproposerShmemState *state = NULL;
+	TimestampTz now = 0;
 
 	/* Start from the last sent position */
 	startptr = sentPtr;
@@ -1502,13 +1537,36 @@ XLogBroadcastWalProposer(WalProposer *wp)
 	 * that arbitrary LSN is eventually reported as written, flushed and
 	 * applied, so that it can measure the elapsed time.
 	 */
-	LagTrackerWrite(endptr, GetCurrentTimestamp());
+	now = GetCurrentTimestamp();
+	LagTrackerWrite(endptr, now);
 
 	/* Do we have any work to do? */
 	Assert(startptr <= endptr);
 	if (endptr <= startptr)
 		return;
 
+	/* BEGIN_HADRON */
+	state = GetWalpropShmemState();
+	if (databricks_max_wal_mb_per_second != -1 && state != NULL)
+	{
+		uint64 max_wal_bytes = (uint64) databricks_max_wal_mb_per_second * 1024 * 1024;
+		struct WalRateLimiter *limiter = &state->wal_rate_limiter;
+
+		if (now - limiter->last_recorded_time_us > USECS_PER_SEC)
+		{
+			/* Reset the rate limiter */
+			limiter->last_recorded_time_us = now;
+			limiter->sent_bytes = 0;
+			pg_atomic_exchange_u32(&limiter->should_limit, 0);
+		}
+		limiter->sent_bytes += (endptr - startptr);
+		if (limiter->sent_bytes > max_wal_bytes)
+		{
+			pg_atomic_exchange_u32(&limiter->should_limit, 1);
+		}
+	}
+	/* END_HADRON */
+
 	WalProposerBroadcast(wp, startptr, endptr);
 	sentPtr = endptr;
 

safekeeper/Cargo.toml

@@ -6,9 +6,7 @@ license.workspace = true
 
 [features]
 default = []
-# Enables test-only APIs, incuding failpoints. In particular, enables the `fail_point!` macro,
-# which adds some runtime cost to run tests on outage conditions
-testing = ["fail/failpoints"]
+testing = ["neon_failpoint/testing"]
 benchmarking = []
 
 [dependencies]
@@ -21,7 +19,7 @@ camino-tempfile.workspace = true
 chrono.workspace = true
 clap = { workspace = true, features = ["derive"] }
 crc32c.workspace = true
-fail.workspace = true
+neon_failpoint.workspace = true
 hex.workspace = true
 humantime.workspace = true
 http.workspace = true

safekeeper/src/bin/safekeeper.rs

@@ -65,7 +65,7 @@ const FEATURES: &[&str] = &[
 fn version() -> String {
     format!(
         "{GIT_VERSION} failpoints: {}, features: {:?}",
-        fail::has_failpoints(),
+        neon_failpoint::has_failpoints(),
         FEATURES,
     )
 }

safekeeper/src/http/routes.rs

@@ -717,7 +717,7 @@ pub fn make_router(
         .get("/profile/cpu", |r| request_span(r, profile_cpu_handler))
         .get("/profile/heap", |r| request_span(r, profile_heap_handler))
         .get("/v1/status", |r| request_span(r, status_handler))
-        .put("/v1/failpoints", |r| {
+        .post("/v1/failpoints", |r| {
             request_span(r, move |r| async {
                 check_permission(&r, None)?;
                 let cancel = CancellationToken::new();

safekeeper/src/send_wal.rs

@@ -872,14 +872,15 @@ impl<IO: AsyncRead + AsyncWrite + Unpin> WalSender<'_, IO> {
     async fn wait_wal(&mut self) -> Result<(), CopyStreamHandlerEnd> {
         loop {
             self.end_pos = self.end_watch.get();
-            let have_something_to_send = (|| {
-                fail::fail_point!(
+            let have_something_to_send = async {
+                neon_failpoint::fail_point!(
                     "sk-pause-send",
                     self.appname.as_deref() != Some("pageserver"),
                     |_| { false }
                 );
                 self.end_pos > self.start_pos
-            })();
+            }
+            .await;
 
             if have_something_to_send {
                 trace!("got end_pos {:?}, streaming", self.end_pos);
@@ -931,14 +932,15 @@ impl<IO: AsyncRead + AsyncWrite + Unpin> WalSender<'_, IO> {
     /// - Err in case of error -- only if 1) term changed while fetching in recovery
     ///   mode 2) watch channel closed, which must never happen.
     async fn wait_for_lsn(&mut self) -> anyhow::Result<Option<Lsn>> {
-        let fp = (|| {
-            fail::fail_point!(
+        let fp = async {
+            neon_failpoint::fail_point!(
                 "sk-pause-send",
                 self.appname.as_deref() != Some("pageserver"),
                 |_| { true }
             );
             false
-        })();
+        }
+        .await;
         if fp {
             tokio::time::sleep(POLL_STATE_TIMEOUT).await;
             return Ok(None);

safekeeper/src/wal_backup.rs

@@ -657,7 +657,7 @@ pub async fn delete_timeline(
 
     pausable_failpoint!("sk-delete-timeline-remote-pause");
 
-    fail::fail_point!("sk-delete-timeline-remote", |_| {
+    neon_failpoint::fail_point!("sk-delete-timeline-remote", |_| {
         Err(anyhow::anyhow!("failpoint: sk-delete-timeline-remote"))
     });
 

safekeeper/src/wal_storage.rs

@@ -301,7 +301,7 @@ impl PhysicalStorage {
                 format!("Failed to open tmp wal file {:?}", &tmp_path)
             })?;
 
-            fail::fail_point!("sk-zero-segment", |_| {
+            neon_failpoint::fail_point!("sk-zero-segment", |_| {
                 info!("sk-zero-segment failpoint hit");
                 Err(anyhow::anyhow!("failpoint: sk-zero-segment"))
             });

storage_controller/Cargo.toml

@@ -22,7 +22,7 @@ clap.workspace = true
 clashmap.workspace = true
 compute_api.workspace = true
 cron.workspace = true
-fail.workspace = true
+neon_failpoint.workspace = true
 futures.workspace = true
 governor.workspace = true
 hex.workspace = true

storage_controller/src/service.rs

@@ -25,6 +25,7 @@ use futures::stream::FuturesUnordered;
 use http_utils::error::ApiError;
 use hyper::Uri;
 use itertools::Itertools;
+use neon_failpoint as fail;
 use pageserver_api::config::PostHogConfig;
 use pageserver_api::controller_api::{
     AvailabilityZone, MetadataHealthRecord, MetadataHealthUpdateRequest, NodeAvailability,
@@ -6026,7 +6027,7 @@ impl Service {
         tenant_id: TenantId,
         split_req: TenantShardSplitRequest,
     ) -> Result<ShardSplitAction, ApiError> {
-        fail::fail_point!("shard-split-validation", |_| Err(ApiError::BadRequest(
+        fail::fail_point_sync!("shard-split-validation", |_| Err(ApiError::BadRequest(
             anyhow::anyhow!("failpoint")
         )));
 

test_runner/fixtures/endpoint/http.py

@@ -159,16 +159,59 @@ class EndpointHttpClient(requests.Session):
         res.raise_for_status()
         return res.json()
 
-    def configure_failpoints(self, *args: tuple[str, str]) -> None:
-        body: list[dict[str, str]] = []
+    def configure_failpoints(
+        self, *args: tuple[str, str] | list[dict[str, str | dict[str, str]]]
+    ) -> None:
+        """Configure failpoints for testing purposes.
 
-        for fp in args:
-            body.append(
-                {
-                    "name": fp[0],
-                    "action": fp[1],
+        Args:
+            *args: Can be one of:
+                - Variable number of (name, actions) tuples
+                - Single list of dicts with keys: name, actions, and optionally context_matchers
+
+        Examples:
+            # Basic failpoints
+            client.configure_failpoints(("test_fp", "return(error)"))
+            client.configure_failpoints(("fp1", "return"), ("fp2", "sleep(1000)"))
+
+            # Probability-based failpoint
+            client.configure_failpoints(("test_fp", "50%return(error)"))
+
+            # Context-based failpoint
+            client.configure_failpoints([{
+                "name": "test_fp",
+                "actions": "return(error)",
+                "context_matchers": {"tenant_id": ".*test.*"}
+            }])
+        """
+
+        request_body: list[dict[str, Any]] = []
+
+        if (
+            len(args) == 1
+            and isinstance(args[0], list)
+            and args[0]
+            and isinstance(args[0][0], dict)
+        ):
+            # Handle list of dicts (context-based failpoints)
+            failpoint_configs = args[0]
+            for config in failpoint_configs:
+                server_config: dict[str, Any] = {
+                    "name": config["name"],
+                    "actions": config["actions"],
                 }
-            )
+                if "context_matchers" in config:
+                    server_config["context_matchers"] = config["context_matchers"]
+                request_body.append(server_config)
+        else:
+            # Handle tuples (basic failpoints)
+            for fp in args:
+                request_body.append(
+                    {
+                        "name": fp[0],
+                        "actions": fp[1],
+                    }
+                )
 
-        res = self.post(f"http://localhost:{self.internal_port}/failpoints", json=body)
+        res = self.post(f"http://localhost:{self.internal_port}/failpoints", json=request_body)
         res.raise_for_status()

test_runner/fixtures/neon_fixtures.py

@@ -2614,22 +2614,68 @@ class NeonStorageController(MetricsGetter, LogUtils):
         )
         return res.json()
 
-    def configure_failpoints(self, config_strings: tuple[str, str] | list[tuple[str, str]]):
-        if isinstance(config_strings, tuple):
-            pairs = [config_strings]
-        else:
-            pairs = config_strings
+    def configure_failpoints(
+        self,
+        config_strings: tuple[str, str]
+        | list[tuple[str, str]]
+        | list[dict[str, str | dict[str, str]]],
+    ):
+        """
+        Configure failpoints for testing purposes.
 
-        log.info(f"Requesting config failpoints: {repr(pairs)}")
+        Args:
+            config_strings: Can be one of:
+                - Single tuple of (name, actions)
+                - List of tuples [(name, actions), ...]
+                - List of dicts with keys: name, actions, and optionally context_matchers
+
+        Examples:
+            # Basic failpoint
+            client.configure_failpoints(("test_fp", "return(error)"))
+
+            # Multiple basic failpoints
+            client.configure_failpoints([("fp1", "return"), ("fp2", "sleep(1000)")])
+
+            # Probability-based failpoint
+            client.configure_failpoints(("test_fp", "50%return(error)"))
+
+            # Context-based failpoint
+            client.configure_failpoints([{
+                "name": "test_fp",
+                "actions": "return(error)",
+                "context_matchers": {"tenant_id": ".*test.*"}
+            }])
+        """
+        # Handle single tuple case
+        if isinstance(config_strings, tuple):
+            config_strings = [config_strings]
+
+        # Convert to server format
+        body: list[dict[str, str | dict[str, str]]] = []
+        for config in config_strings:
+            if isinstance(config, tuple):
+                # Simple (name, actions) tuple
+                body.append({"name": config[0], "actions": config[1]})
+            elif isinstance(config, dict):
+                # Dict with name, actions, and optional context_matchers
+                server_config: dict[str, str | dict[str, str]] = {
+                    "name": config["name"],
+                    "actions": config["actions"],
+                }
+                if "context_matchers" in config:
+                    server_config["context_matchers"] = config["context_matchers"]
+                body.append(server_config)
+            else:
+                raise ValueError(f"Invalid config format: {config}")
 
         res = self.request(
             "PUT",
-            f"{self.api}/debug/v1/failpoints",
-            json=[{"name": name, "actions": actions} for name, actions in pairs],
+            f"{self.api_root()}/debug/v1/failpoints",
+            json=body,
             headers=self.headers(TokenScope.ADMIN),
         )
-        log.info(f"Got failpoints request response code {res.status_code}")
-        res.raise_for_status()
+        if res.status_code != 200:
+            self.raise_api_exception(res)
 
     def get_tenants_placement(self) -> defaultdict[str, dict[str, Any]]:
         """

test_runner/fixtures/pageserver/http.py

@@ -309,25 +309,64 @@ class PageserverHttpClient(requests.Session, MetricsGetter):
     def check_status(self):
         self.get(f"http://localhost:{self.port}/v1/status").raise_for_status()
 
-    def configure_failpoints(self, config_strings: tuple[str, str] | list[tuple[str, str]]):
+    def configure_failpoints(
+        self,
+        config_strings: tuple[str, str]
+        | list[tuple[str, str]]
+        | list[dict[str, str | dict[str, str]]],
+    ):
+        """
+        Configure failpoints for testing purposes.
+
+        Args:
+            config_strings: Can be one of:
+                - Single tuple of (name, actions)
+                - List of tuples [(name, actions), ...]
+                - List of dicts with keys: name, actions, and optionally context_matchers
+
+        Examples:
+            # Basic failpoint
+            client.configure_failpoints(("test_fp", "return(error)"))
+
+            # Multiple basic failpoints
+            client.configure_failpoints([("fp1", "return"), ("fp2", "sleep(1000)")])
+
+            # Probability-based failpoint
+            client.configure_failpoints(("test_fp", "50%return(error)"))
+
+            # Context-based failpoint
+            client.configure_failpoints([{
+                "name": "test_fp",
+                "actions": "return(error)",
+                "context_matchers": {"tenant_id": ".*test.*"}
+            }])
+        """
         self.is_testing_enabled_or_skip()
 
+        # Handle single tuple case
         if isinstance(config_strings, tuple):
-            pairs = [config_strings]
-        else:
-            pairs = config_strings
+            config_strings = [config_strings]
 
-        log.info(f"Requesting config failpoints: {repr(pairs)}")
+        # Convert to server format
+        body: list[dict[str, str | dict[str, str]]] = []
+        for config in config_strings:
+            if isinstance(config, tuple):
+                # Simple (name, actions) tuple
+                body.append({"name": config[0], "actions": config[1]})
+            elif isinstance(config, dict):
+                # Dict with name, actions, and optional context_matchers
+                server_config = {"name": config["name"], "actions": config["actions"]}
+                if "context_matchers" in config:
+                    server_config["context_matchers"] = config["context_matchers"]
+                body.append(server_config)
+            else:
+                raise ValueError(f"Invalid config format: {config}")
 
-        res = self.put(
-            f"http://localhost:{self.port}/v1/failpoints",
-            json=[{"name": name, "actions": actions} for name, actions in pairs],
-        )
-        log.info(f"Got failpoints request response code {res.status_code}")
-        self.verbose_error(res)
-        res_json = res.json()
-        assert res_json is None
-        return res_json
+        res = self.post(f"{self.base_url}/v1/failpoints", json=body)
+        if res.status_code != 200:
+            raise PageserverApiException(
+                f"Failed to configure failpoints: {res.text}", res.status_code
+            )
 
     def reload_auth_validation_keys(self):
         res = self.post(f"http://localhost:{self.port}/v1/reload_auth_validation_keys")

test_runner/fixtures/safekeeper/http.py

@@ -8,7 +8,6 @@ import pytest
 import requests
 
 from fixtures.common_types import Lsn, TenantId, TenantTimelineId, TimelineId
-from fixtures.log_helper import log
 from fixtures.metrics import Metrics, MetricsGetter, parse_metrics
 from fixtures.utils import EnhancedJSONEncoder, wait_until
 
@@ -155,25 +154,62 @@ class SafekeeperHttpClient(requests.Session, MetricsGetter):
         if not self.is_testing_enabled:
             pytest.skip("safekeeper was built without 'testing' feature")
 
-    def configure_failpoints(self, config_strings: tuple[str, str] | list[tuple[str, str]]):
+    def configure_failpoints(
+        self,
+        config_strings: tuple[str, str]
+        | list[tuple[str, str]]
+        | list[dict[str, str | dict[str, str]]],
+    ):
+        """
+        Configure failpoints for testing purposes.
+
+        Args:
+            config_strings: Can be one of:
+                - Single tuple of (name, actions)
+                - List of tuples [(name, actions), ...]
+                - List of dicts with keys: name, actions, and optionally context_matchers
+
+        Examples:
+            # Basic failpoint
+            client.configure_failpoints(("test_fp", "return(error)"))
+
+            # Multiple basic failpoints
+            client.configure_failpoints([("fp1", "return"), ("fp2", "sleep(1000)")])
+
+            # Probability-based failpoint
+            client.configure_failpoints(("test_fp", "50%return(error)"))
+
+            # Context-based failpoint
+            client.configure_failpoints([{
+                "name": "test_fp",
+                "actions": "return(error)",
+                "context_matchers": {"tenant_id": ".*test.*"}
+            }])
+        """
         self.is_testing_enabled_or_skip()
 
+        # Handle single tuple case
         if isinstance(config_strings, tuple):
-            pairs = [config_strings]
-        else:
-            pairs = config_strings
+            config_strings = [config_strings]
 
-        log.info(f"Requesting config failpoints: {repr(pairs)}")
+        # Convert to server format
+        body: list[dict[str, str | dict[str, str]]] = []
+        for config in config_strings:
+            if isinstance(config, tuple):
+                # Simple (name, actions) tuple
+                body.append({"name": config[0], "actions": config[1]})
+            elif isinstance(config, dict):
+                # Dict with name, actions, and optional context_matchers
+                server_config = {"name": config["name"], "actions": config["actions"]}
+                if "context_matchers" in config:
+                    server_config["context_matchers"] = config["context_matchers"]
+                body.append(server_config)
+            else:
+                raise ValueError(f"Invalid config format: {config}")
 
-        res = self.put(
-            f"http://localhost:{self.port}/v1/failpoints",
-            json=[{"name": name, "actions": actions} for name, actions in pairs],
-        )
-        log.info(f"Got failpoints request response code {res.status_code}")
-        res.raise_for_status()
-        res_json = res.json()
-        assert res_json is None
-        return res_json
+        res = self.post(f"http://localhost:{self.port}/v1/failpoints", json=body)
+        if res.status_code != 200:
+            raise RuntimeError(f"Failed to configure failpoints: {res.text}")
 
     def tenant_delete_force(self, tenant_id: TenantId) -> dict[Any, Any]:
         res = self.delete(f"http://localhost:{self.port}/v1/tenant/{tenant_id}")

test_runner/regress/test_lfc_prewarm.py

@@ -40,7 +40,7 @@ def prom_parse(client: EndpointHttpClient) -> dict[str, float]:
 
 def offload_lfc(method: PrewarmMethod, client: EndpointHttpClient, cur: Cursor) -> Any:
     if method == PrewarmMethod.POSTGRES:
-        cur.execute("select get_local_cache_state()")
+        cur.execute("select neon.get_local_cache_state()")
         return cur.fetchall()[0][0]
 
     if method == PrewarmMethod.AUTOPREWARM:
@@ -72,7 +72,7 @@ def prewarm_endpoint(
     elif method == PrewarmMethod.COMPUTE_CTL:
         client.prewarm_lfc()
     elif method == PrewarmMethod.POSTGRES:
-        cur.execute("select prewarm_local_cache(%s)", (lfc_state,))
+        cur.execute("select neon.prewarm_local_cache(%s)", (lfc_state,))
 
 
 def check_prewarmed(
@@ -116,7 +116,7 @@ def test_lfc_prewarm(neon_simple_env: NeonEnv, method: PrewarmMethod):
 
     pg_conn = endpoint.connect()
     pg_cur = pg_conn.cursor()
-    pg_cur.execute("create extension neon")
+    pg_cur.execute("create schema neon; create extension neon with schema neon")
     pg_cur.execute("create database lfc")
 
     lfc_conn = endpoint.connect(dbname="lfc")
@@ -142,10 +142,12 @@ def test_lfc_prewarm(neon_simple_env: NeonEnv, method: PrewarmMethod):
     lfc_cur = lfc_conn.cursor()
     prewarm_endpoint(method, client, pg_cur, lfc_state)
 
-    pg_cur.execute("select lfc_value from neon_lfc_stats where lfc_key='file_cache_used_pages'")
+    pg_cur.execute(
+        "select lfc_value from neon.neon_lfc_stats where lfc_key='file_cache_used_pages'"
+    )
     lfc_used_pages = pg_cur.fetchall()[0][0]
     log.info(f"Used LFC size: {lfc_used_pages}")
-    pg_cur.execute("select * from get_prewarm_info()")
+    pg_cur.execute("select * from neon.get_prewarm_info()")
     total, prewarmed, skipped, _ = pg_cur.fetchall()[0]
     log.info(f"Prewarm info: {total=} {prewarmed=} {skipped=}")
     progress = (prewarmed + skipped) * 100 // total
@@ -186,7 +188,7 @@ def test_lfc_prewarm_under_workload(neon_simple_env: NeonEnv, method: PrewarmMet
 
     pg_conn = endpoint.connect()
     pg_cur = pg_conn.cursor()
-    pg_cur.execute("create extension neon")
+    pg_cur.execute("create schema neon; create extension neon with schema neon")
     pg_cur.execute("CREATE DATABASE lfc")
 
     lfc_conn = endpoint.connect(dbname="lfc")

test_runner/regress/test_replica_promotes.py

@@ -60,7 +60,7 @@ def test_replica_promote(neon_simple_env: NeonEnv, method: PromoteMethod):
 
     with primary.connect() as primary_conn:
         primary_cur = primary_conn.cursor()
-        primary_cur.execute("create extension neon")
+        primary_cur.execute("create schema neon;create extension neon with schema neon")
         primary_cur.execute(
             "create table t(pk bigint GENERATED ALWAYS AS IDENTITY, payload integer)"
         )
@@ -172,7 +172,7 @@ def test_replica_promote_handler_disconnects(neon_simple_env: NeonEnv):
     secondary: Endpoint = env.endpoints.new_replica_start(origin=primary, endpoint_id="secondary")
 
     with primary.connect() as conn, conn.cursor() as cur:
-        cur.execute("create extension neon")
+        cur.execute("create schema neon;create extension neon with schema neon")
         cur.execute("create table t(pk bigint GENERATED ALWAYS AS IDENTITY, payload integer)")
         cur.execute("INSERT INTO t(payload) SELECT generate_series(1, 100)")
         cur.execute("show neon.safekeepers")