feat(proof-gen): request admission limits, per-chain fill budget, single-flight block fills

proof-gen-api-server/bin/server.rs

@@ -157,8 +157,12 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
             }],
             max_batch_size: args.max_batch_size,
             max_batch_span: args.max_batch_span,
+            admission: proof_gen_api_server::config::AdmissionConfig::default(),
         }
     };
+    // MAX_IN_FLIGHT_REQUESTS / MAX_IN_FLIGHT_PER_CHAIN / REQUEST_TIMEOUT_SECS override YAML.
+    let mut config = config;
+    config.admission.apply_env_overrides();
 
     let server = Server::new(config).await?;
     server.run().await?;

proof-gen-api-server/config.example.yaml

@@ -49,6 +49,10 @@ chains:
     # attestation load on a fast chain multiplies this cache by the same factor. On a chain with
     # both a wide window and dense blocks that is enough to exhaust the process.
     # cache:
+    #   # Simultaneous source-block fetches for this chain's merkle cache, across requests and
+    #   # the backfill worker. Caps RPC amplification under a burst of cold requests.
+    #   max_concurrent_block_fills: 16
+    #
     #   # Retention window in source blocks. Omit to derive it as described above.
     #   # Pin it to size the cache independently of attestation cadence.
     #   merkle_retention_blocks: 1000
@@ -84,3 +88,12 @@ max_batch_size: 10
 # Maximum allowed block span (highest − lowest block) in a single batch request.
 # Prevents small batches from forcing proof generation over extremely large ranges.
 max_batch_span: 1000
+
+# Request admission (all optional; omit the block for the defaults shown).
+# `max_batch_size` bounds the work inside one batch; these bound how many proof requests run
+# at once. Excess requests get an immediate 503 with Retry-After, a request past its deadline
+# a 504. /livez, /readyz, /api/v1/health and /metrics are never limited.
+# admission:
+#   max_in_flight_requests: 64     # across all chains (env: MAX_IN_FLIGHT_REQUESTS)
+#   max_in_flight_per_chain: 32    # per chain_key    (env: MAX_IN_FLIGHT_PER_CHAIN)
+#   request_timeout_secs: 120      # end to end       (env: REQUEST_TIMEOUT_SECS)

proof-gen-api-server/src/config.rs

@@ -64,8 +64,19 @@ pub struct ChainCacheConfig {
     /// be served -- roughly `max_entries * attestation_interval * checkpoint_interval` blocks
     /// -- in exchange for bounding a cache that otherwise only ever grows. Opt in knowingly.
     pub checkpoint_cache_max_entries: Option<usize>,
+    /// Upper bound on simultaneous source-block fetches for this chain's merkle cache, across
+    /// every request and the backfill worker together. Without it R concurrent cold requests
+    /// meant R block fetches in flight plus the backfill's own; this caps the RPC amplification
+    /// per chain.
+    pub max_concurrent_block_fills: NonZeroUsize,
 }
 
+/// Default [`ChainCacheConfig::max_concurrent_block_fills`].
+pub const DEFAULT_MAX_CONCURRENT_BLOCK_FILLS: NonZeroUsize = match NonZeroUsize::new(16) {
+    Some(n) => n,
+    None => unreachable!(),
+};
+
 impl Default for ChainCacheConfig {
     fn default() -> Self {
         Self {
@@ -74,10 +85,90 @@ impl Default for ChainCacheConfig {
             block_cache_capacity: DEFAULT_BLOCK_CACHE_CAPACITY,
             merkle_backfill_enabled: true,
             checkpoint_cache_max_entries: None,
+            max_concurrent_block_fills: DEFAULT_MAX_CONCURRENT_BLOCK_FILLS,
         }
     }
 }
 
+/// Process-wide request admission. Bounds what the HTTP layer lets in so a burst of cold
+/// requests degrades into fast `503`s instead of unbounded RPC fan-out, memory growth and
+/// requests that never finish.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct AdmissionConfig {
+    /// Proof requests (single, by-tx, batch) in flight across all chains. Health, readiness
+    /// and metrics are never counted or refused.
+    pub max_in_flight_requests: NonZeroUsize,
+    /// Proof requests in flight per chain, so one chain's storm cannot starve the others.
+    pub max_in_flight_per_chain: NonZeroUsize,
+    /// Deadline for one proof request end to end. Exceeding it answers `504`.
+    pub request_timeout: std::time::Duration,
+}
+
+impl Default for AdmissionConfig {
+    fn default() -> Self {
+        Self {
+            max_in_flight_requests: NonZeroUsize::new(64).expect("non-zero"),
+            max_in_flight_per_chain: NonZeroUsize::new(32).expect("non-zero"),
+            request_timeout: std::time::Duration::from_secs(120),
+        }
+    }
+}
+
+impl AdmissionConfig {
+    /// Apply `MAX_IN_FLIGHT_REQUESTS`, `MAX_IN_FLIGHT_PER_CHAIN` and `REQUEST_TIMEOUT_SECS`
+    /// from the environment when set and valid; anything else is left untouched.
+    pub fn apply_env_overrides(&mut self) {
+        if let Some(n) = std::env::var("MAX_IN_FLIGHT_REQUESTS")
+            .ok()
+            .and_then(|raw| raw.parse::<NonZeroUsize>().ok())
+        {
+            self.max_in_flight_requests = n;
+        }
+        if let Some(n) = std::env::var("MAX_IN_FLIGHT_PER_CHAIN")
+            .ok()
+            .and_then(|raw| raw.parse::<NonZeroUsize>().ok())
+        {
+            self.max_in_flight_per_chain = n;
+        }
+        if let Some(secs) = std::env::var("REQUEST_TIMEOUT_SECS")
+            .ok()
+            .and_then(|raw| raw.parse::<u64>().ok())
+            .filter(|secs| *secs > 0)
+        {
+            self.request_timeout = std::time::Duration::from_secs(secs);
+        }
+    }
+}
+
+/// YAML `admission:` block; every field optional.
+#[derive(Debug, Default, Deserialize)]
+pub struct AdmissionFile {
+    #[serde(default)]
+    pub max_in_flight_requests: Option<NonZeroUsize>,
+    #[serde(default)]
+    pub max_in_flight_per_chain: Option<NonZeroUsize>,
+    #[serde(default)]
+    pub request_timeout_secs: Option<u64>,
+}
+
+fn resolve_admission(file: AdmissionFile) -> Result<AdmissionConfig> {
+    if file.request_timeout_secs == Some(0) {
+        bail!("`admission.request_timeout_secs` must be greater than 0; omit it for the default");
+    }
+    let defaults = AdmissionConfig::default();
+    Ok(AdmissionConfig {
+        max_in_flight_requests: file
+            .max_in_flight_requests
+            .unwrap_or(defaults.max_in_flight_requests),
+        max_in_flight_per_chain: file
+            .max_in_flight_per_chain
+            .unwrap_or(defaults.max_in_flight_per_chain),
+        request_timeout: file
+            .request_timeout_secs
+            .map_or(defaults.request_timeout, std::time::Duration::from_secs),
+    })
+}
+
 /// One source chain (EVM) served by this process, keyed on Creditcoin3.
 #[derive(Debug, Clone)]
 pub struct ChainConfig {
@@ -117,6 +208,7 @@ pub struct Config {
     pub chains: Vec<ChainConfig>,
     pub max_batch_size: NonZeroUsize,
     pub max_batch_span: u64,
+    pub admission: AdmissionConfig,
 }
 
 impl Config {
@@ -138,6 +230,7 @@ impl Config {
             }],
             max_batch_size: DEFAULT_MAX_BATCH_SIZE,
             max_batch_span: DEFAULT_MAX_BATCH_SPAN,
+            admission: AdmissionConfig::default(),
         }
     }
 
@@ -174,6 +267,9 @@ pub struct ConfigFile {
     pub max_batch_size: NonZeroUsize,
     #[serde(default = "default_max_batch_span")]
     pub max_batch_span: u64,
+    /// Optional request admission limits; omit the block for the defaults.
+    #[serde(default)]
+    pub admission: AdmissionFile,
 }
 
 #[derive(Debug, Deserialize)]
@@ -224,6 +320,8 @@ pub struct ChainCacheConfigFile {
     pub merkle_backfill_enabled: Option<bool>,
     #[serde(default)]
     pub checkpoint_cache_max_entries: Option<usize>,
+    #[serde(default)]
+    pub max_concurrent_block_fills: Option<NonZeroUsize>,
 }
 
 /// Deserialize a byte size written either as a number or as a human-readable string.
@@ -331,6 +429,7 @@ impl ConfigFile {
             chains,
             max_batch_size: self.max_batch_size,
             max_batch_span: self.max_batch_span,
+            admission: resolve_admission(self.admission)?,
         })
     }
 }
@@ -376,6 +475,9 @@ fn resolve_cache_config(chain_key: u64, file: ChainCacheConfigFile) -> Result<Ch
             .merkle_backfill_enabled
             .unwrap_or(defaults.merkle_backfill_enabled),
         checkpoint_cache_max_entries: file.checkpoint_cache_max_entries,
+        max_concurrent_block_fills: file
+            .max_concurrent_block_fills
+            .unwrap_or(defaults.max_concurrent_block_fills),
     })
 }
 

proof-gen-api-server/src/lib.rs

@@ -369,7 +369,12 @@ impl Server {
         ContinuityService::spawn_merkle_backfill(service.clone());
 
         let allowed: std::collections::HashSet<u64> = self.config.chain_keys();
-        let app = build_app(service.clone(), allowed, self.prom_metrics.clone());
+        let app = networking::build_app_with_admission(
+            service.clone(),
+            allowed,
+            self.prom_metrics.clone(),
+            self.config.admission.clone(),
+        );
         let (http_shutdown_tx, http_shutdown_rx) = channel::<()>();
 
         let bind_host = &self.config.bind_host;

proof-gen-api-server/src/networking/admission.rs

@@ -0,0 +1,154 @@
+//! Request admission: a process-wide and a per-chain cap on proof requests in flight, plus an
+//! end-to-end deadline per request.
+//!
+//! `max_batch_size` bounds the work inside one batch, not how many requests run at once. With
+//! R simultaneous cold requests the old router let roughly `R × (20 + max_batch_size)`
+//! block-level operations run, each with its own RPC retries, and no request ever timed out.
+//! Refusing the excess up front with a `503` is cheaper for everyone: the caller retries with
+//! back-off, and the requests that were admitted actually finish.
+//!
+//! Health, readiness and metrics are exempt so probes and recovery keep working under load.
+
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::time::Duration;
+
+use axum::{
+    extract::Request,
+    http::{header, StatusCode},
+    middleware::Next,
+    response::{IntoResponse, Response},
+    Extension,
+};
+use serde_json::json;
+use tokio::sync::{OwnedSemaphorePermit, Semaphore};
+
+use crate::config::AdmissionConfig;
+use crate::networking::middleware::extract_chain_key_from_path;
+use crate::prom::{labels::Rejection, ProofGenMetrics};
+
+pub struct Admission {
+    global: Arc<Semaphore>,
+    per_chain: HashMap<u64, Arc<Semaphore>>,
+    request_timeout: Duration,
+    metrics: Arc<ProofGenMetrics>,
+}
+
+impl Admission {
+    pub fn new(
+        config: &AdmissionConfig,
+        chain_keys: impl IntoIterator<Item = u64>,
+        metrics: Arc<ProofGenMetrics>,
+    ) -> Self {
+        Self {
+            global: Arc::new(Semaphore::new(config.max_in_flight_requests.get())),
+            per_chain: chain_keys
+                .into_iter()
+                .map(|k| {
+                    (
+                        k,
+                        Arc::new(Semaphore::new(config.max_in_flight_per_chain.get())),
+                    )
+                })
+                .collect(),
+            request_timeout: config.request_timeout,
+            metrics,
+        }
+    }
+
+    /// Try to admit one proof request for `chain_key`. `Err` names the limit that was hit.
+    fn try_admit(&self, chain_key: u64) -> Result<Permits, Rejection> {
+        let global = self
+            .global
+            .clone()
+            .try_acquire_owned()
+            .map_err(|_| Rejection::Overloaded)?;
+        let chain = match self.per_chain.get(&chain_key) {
+            Some(sem) => Some(
+                sem.clone()
+                    .try_acquire_owned()
+                    .map_err(|_| Rejection::ChainOverloaded)?,
+            ),
+            // Unknown chain: the chain-key validator answers 400 further in; only the global
+            // permit applies.
+            None => None,
+        };
+        Ok(Permits {
+            _global: global,
+            _chain: chain,
+        })
+    }
+}
+
+struct Permits {
+    _global: OwnedSemaphorePermit,
+    _chain: Option<OwnedSemaphorePermit>,
+}
+
+/// Only proof endpoints are admission-controlled; everything else passes untouched.
+fn is_proof_request(request: &Request) -> bool {
+    request.uri().path().starts_with("/api/v1/proof")
+}
+
+pub async fn admission_middleware(
+    Extension(admission): Extension<Arc<Admission>>,
+    request: Request,
+    next: Next,
+) -> Response {
+    if !is_proof_request(&request) {
+        return next.run(request).await;
+    }
+    let chain_key = extract_chain_key_from_path(request.uri()).unwrap_or(u64::MAX);
+    let _permits = match admission.try_admit(chain_key) {
+        Ok(permits) => permits,
+        Err(reason) => {
+            admission.metrics.request_rejected(reason);
+            tracing::warn!(
+                chain_key,
+                ?reason,
+                "🚦 proof request refused by admission control"
+            );
+            return rejected(
+                StatusCode::SERVICE_UNAVAILABLE,
+                match reason {
+                    Rejection::ChainOverloaded => "ChainOverloaded",
+                    _ => "Overloaded",
+                },
+                "too many proof requests in flight; retry with back-off",
+            );
+        }
+    };
+
+    admission.metrics.request_admitted();
+    let outcome = tokio::time::timeout(admission.request_timeout, next.run(request)).await;
+    admission.metrics.request_finished();
+    match outcome {
+        Ok(response) => response,
+        Err(_elapsed) => {
+            admission.metrics.request_rejected(Rejection::Timeout);
+            tracing::warn!(
+                chain_key,
+                timeout = ?admission.request_timeout,
+                "⏱️ proof request exceeded its deadline"
+            );
+            rejected(
+                StatusCode::GATEWAY_TIMEOUT,
+                "RequestTimeout",
+                "proof request exceeded the server's deadline",
+            )
+        }
+    }
+}
+
+fn rejected(status: StatusCode, code: &str, message: &str) -> Response {
+    (
+        status,
+        [(header::RETRY_AFTER, "1")],
+        axum::Json(json!({
+            "code": code,
+            "message": message,
+            "retriable": true,
+        })),
+    )
+        .into_response()
+}