Cap concurrent Electrum checkpoint Merkle proof builds

src/config.rs

@@ -74,6 +74,7 @@ pub struct Config {
     pub utxos_limit: usize,
     pub electrum_txs_limit: usize,
     pub electrum_subscription_limit: usize,
+    pub electrum_checkpoint_proof_concurrency_limit: usize,
     pub electrum_banner: String,
     pub rpc_logging: RpcLogging,
     pub zmq_addr: Option<SocketAddr>,
@@ -292,6 +293,12 @@ impl Config {
                     .help("Maximum number of scripthash subscriptions a single Electrum connection may hold. Every subscription costs a history lookup on each new block, so an unbounded count lets one client impose unbounded recurring work. Re-subscribing to an already-tracked scripthash is always allowed. 0 = unlimited.")
                     .default_value("10000")
                     .takes_value(true)
+            ).arg(
+                Arg::with_name("electrum_checkpoint_proof_concurrency_limit")
+                    .long("electrum-checkpoint-proof-concurrency-limit")
+                    .help("Maximum number of blockchain.block.header(s) checkpoint Merkle proof builds (triggered by a non-zero cp_height) allowed to run at once, process-wide. Each build hashes every header from genesis up to cp_height, so an unbounded count lets concurrent cheap requests pin every CPU core. Requests past the cap fail immediately rather than queueing. 0 = reject all such requests.")
+                    .default_value("2")
+                    .takes_value(true)
             ).arg(
                 Arg::with_name("electrum_banner")
                     .long("electrum-banner")
@@ -589,6 +596,11 @@ impl Config {
             electrum_rpc_max_request_num_bytes,
             electrum_txs_limit: value_t_or_exit!(m, "electrum_txs_limit", usize),
             electrum_subscription_limit: value_t_or_exit!(m, "electrum_subscription_limit", usize),
+            electrum_checkpoint_proof_concurrency_limit: value_t_or_exit!(
+                m,
+                "electrum_checkpoint_proof_concurrency_limit",
+                usize
+            ),
             electrum_banner,
             rpc_logging: {
                 let params = RpcLogging {

src/electrum/server.rs

@@ -183,6 +183,7 @@ struct Connection {
     txs_limit: usize,
     subscription_limit: usize,
     max_request_bytes: usize,
+    checkpoint_proof_concurrency_limit: usize,
     #[cfg(feature = "electrum-discovery")]
     discovery: Option<Arc<DiscoveryManager>>,
     rpc_logging: RpcLogging,
@@ -206,6 +207,7 @@ impl Connection {
         txs_limit: usize,
         subscription_limit: usize,
         max_request_bytes: usize,
+        checkpoint_proof_concurrency_limit: usize,
         #[cfg(feature = "electrum-discovery")] discovery: Option<Arc<DiscoveryManager>>,
         rpc_logging: RpcLogging,
         salt: String,
@@ -221,6 +223,7 @@ impl Connection {
             txs_limit,
             subscription_limit,
             max_request_bytes,
+            checkpoint_proof_concurrency_limit,
             #[cfg(feature = "electrum-discovery")]
             discovery,
             rpc_logging,
@@ -309,7 +312,12 @@ impl Connection {
         if cp_height == 0 {
             return Ok(json!(raw_header_hex));
         }
-        let (branch, root) = get_header_merkle_proof(self.query.chain(), height, cp_height)?;
+        let (branch, root) = get_header_merkle_proof(
+            self.query.chain(),
+            height,
+            cp_height,
+            self.checkpoint_proof_concurrency_limit,
+        )?;
 
         Ok(json!({
             "header": raw_header_hex,
@@ -341,8 +349,12 @@ impl Connection {
             }));
         }
 
-        let (branch, root) =
-            get_header_merkle_proof(self.query.chain(), start_height + (count - 1), cp_height)?;
+        let (branch, root) = get_header_merkle_proof(
+            self.query.chain(),
+            start_height + (count - 1),
+            cp_height,
+            self.checkpoint_proof_concurrency_limit,
+        )?;
 
         Ok(json!({
             "count": headers.len(),
@@ -1238,6 +1250,7 @@ impl RPC {
         let txs_limit = config.electrum_txs_limit;
         let subscription_limit = config.electrum_subscription_limit;
         let max_request_bytes = config.electrum_rpc_max_request_num_bytes;
+        let checkpoint_proof_concurrency_limit = config.electrum_checkpoint_proof_concurrency_limit;
         let conn_max_age = config.electrum_rpc_conn_max_age;
 
         RPC {
@@ -1290,6 +1303,7 @@ impl RPC {
                             txs_limit,
                             subscription_limit,
                             max_request_bytes,
+                            checkpoint_proof_concurrency_limit,
                             #[cfg(feature = "electrum-discovery")]
                             discovery,
                             rpc_logging,

src/util/electrum_merkle.rs

@@ -2,9 +2,40 @@ use crate::chain::{BlockHash, Txid};
 use crate::errors::*;
 use crate::new_index::ChainQuery;
 use bitcoin::hashes::{sha256d::Hash as Sha256dHash, Hash};
+use std::sync::atomic::{AtomicUsize, Ordering};
 
 use electrs_macros::trace;
 
+static INFLIGHT_CHECKPOINT_PROOFS: AtomicUsize = AtomicUsize::new(0);
+
+struct InflightCheckpointProofGuard;
+
+impl InflightCheckpointProofGuard {
+    fn acquire(limit: usize) -> Result<Self> {
+        let mut current = INFLIGHT_CHECKPOINT_PROOFS.load(Ordering::Relaxed);
+        loop {
+            if current >= limit {
+                bail!("too many concurrent checkpoint merkle proof requests, try again later");
+            }
+            match INFLIGHT_CHECKPOINT_PROOFS.compare_exchange_weak(
+                current,
+                current + 1,
+                Ordering::AcqRel,
+                Ordering::Relaxed,
+            ) {
+                Ok(_) => return Ok(Self),
+                Err(observed) => current = observed,
+            }
+        }
+    }
+}
+
+impl Drop for InflightCheckpointProofGuard {
+    fn drop(&mut self) {
+        INFLIGHT_CHECKPOINT_PROOFS.fetch_sub(1, Ordering::AcqRel);
+    }
+}
+
 #[trace]
 pub fn get_tx_merkle_proof(
     chain: &ChainQuery,
@@ -29,6 +60,7 @@ pub fn get_header_merkle_proof(
     chain: &ChainQuery,
     height: usize,
     cp_height: usize,
+    checkpoint_proof_concurrency_limit: usize,
 ) -> Result<(Vec<Sha256dHash>, Sha256dHash)> {
     if cp_height < height {
         bail!("cp_height #{} < height #{}", cp_height, height);
@@ -43,6 +75,8 @@ pub fn get_header_merkle_proof(
         );
     }
 
+    let _guard = InflightCheckpointProofGuard::acquire(checkpoint_proof_concurrency_limit)?;
+
     let heights: Vec<usize> = (0..=cp_height).collect();
     let header_hashes: Vec<BlockHash> = heights
         .into_iter()
@@ -107,3 +141,34 @@ fn create_merkle_branch_and_root(
     }
     (merkle, hashes[0])
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // Single test so the shared static isn't touched by other tests running
+    // concurrently in the same process.
+    #[test]
+    fn checkpoint_proof_guard_caps_concurrency_and_releases_on_drop() {
+        const LIMIT: usize = 2;
+        let mut guards = Vec::new();
+        for _ in 0..LIMIT {
+            guards.push(InflightCheckpointProofGuard::acquire(LIMIT).unwrap());
+        }
+
+        // All permits are taken: the next caller must be rejected outright
+        // instead of blocking or panicking.
+        assert!(InflightCheckpointProofGuard::acquire(LIMIT).is_err());
+
+        // Releasing one permit (drop, e.g. on early return via `?`) must let
+        // the next acquire through.
+        guards.pop();
+        let extra = InflightCheckpointProofGuard::acquire(LIMIT).unwrap();
+
+        // Back at the cap: still no free permits.
+        assert!(InflightCheckpointProofGuard::acquire(LIMIT).is_err());
+
+        drop(extra);
+        drop(guards);
+    }
+}

tests/common.rs

@@ -115,6 +115,7 @@ impl TestRunner {
             utxos_limit: 100,
             electrum_txs_limit: 100,
             electrum_subscription_limit: 10_000,
+            electrum_checkpoint_proof_concurrency_limit: 2,
             electrum_banner: "".into(),
             rpc_logging: RpcLogging::default(),
             zmq_addr: None,