Add expensive op extension

.changes/added/3181.md

@@ -0,0 +1 @@
+Add expensive op extension for full block queries on GraphQL

bin/fuel-core/src/cli/run.rs

@@ -734,6 +734,8 @@ impl Command {
                 required_fuel_block_height_timeout: graphql
                     .required_fuel_block_height_timeout
                     .into(),
+                concurrent_full_block_requests: graphql.concurrent_full_block_requests,
+                full_block_request_timeout: graphql.full_block_request_timeout.into(),
             },
             combined_db_config,
             snapshot_reader,

bin/fuel-core/src/cli/run/graphql.rs

@@ -105,6 +105,12 @@ pub struct GraphQLArgs {
 
     #[clap(flatten)]
     pub costs: QueryCosts,
+
+    #[clap(long = "concurrent-full-block-requests", default_value = "10", env)]
+    pub concurrent_full_block_requests: usize,
+
+    #[clap(long = "full-block-request-timeout", default_value = "3s", env)]
+    pub full_block_request_timeout: humantime::Duration,
 }
 
 /// Costs for individual graphql queries.

crates/fuel-core/src/graphql_api.rs

@@ -59,6 +59,9 @@ pub struct ServiceConfig {
     pub assemble_tx_estimate_predicates_limit: usize,
     /// Configurable cost parameters to limit graphql queries complexity
     pub costs: Costs,
+    /// Configurable parameters to limit number of concurrent requests to the full block API
+    pub concurrent_full_block_requests: usize,
+    pub full_block_request_timeout: Duration,
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]

crates/fuel-core/src/graphql_api/api_service.rs

@@ -1,7 +1,10 @@
 use crate::{
     fuel_core_graphql_api::{
         Config,
-        extensions::unify_response,
+        extensions::{
+            expensive_op_guard::ExpensiveOpGuardFactory,
+            unify_response,
+        },
         ports::{
             BlockProducerPort,
             ChainStateProvider as ChainStateProviderTrait,
@@ -297,6 +300,11 @@ where
         .extension(MetricsExtension::new(
             config.config.query_log_threshold_time,
         ))
+        .extension(ExpensiveOpGuardFactory::new(
+            Arc::new(["block".to_string(), "blocks".to_string()]),
+            config.config.concurrent_full_block_requests,
+            config.config.full_block_request_timeout,
+        ))
         .data(config)
         .data(combined_read_database)
         .data(txpool)

crates/fuel-core/src/graphql_api/extensions.rs

@@ -13,6 +13,8 @@ pub(crate) mod metrics;
 pub(crate) mod required_fuel_block_height;
 pub(crate) mod validation;
 
+pub(crate) mod expensive_op_guard;
+
 // In the case of a successful query, we return the information below on
 // the `response.extensions` level.
 // But in the case of the error, `async_graphql` returns information from extensions

crates/fuel-core/src/graphql_api/extensions/expensive_op_guard.rs

@@ -0,0 +1,220 @@
+use async_graphql::{
+    Positioned,
+    Response,
+    ServerError,
+    ServerResult,
+    Variables,
+    extensions::{
+        Extension,
+        ExtensionContext,
+        ExtensionFactory,
+        NextExecute,
+        NextParseQuery,
+    },
+    parser::types::{
+        ExecutableDocument,
+        FragmentDefinition,
+        Selection,
+        SelectionSet,
+    },
+};
+use std::{
+    sync::{
+        Arc,
+        OnceLock,
+    },
+    time::Duration,
+};
+use tokio::sync::Semaphore;
+
+pub struct ExpensiveOpGuardFactory {
+    expensive_root_field_names: Arc<[String]>,
+    semaphore: Arc<Semaphore>,
+    timeout: Duration,
+}
+
+impl ExpensiveOpGuardFactory {
+    pub fn new(
+        expensive_root_field_names: Arc<[String]>,
+        max_in_flight: usize,
+        timeout: Duration,
+    ) -> Self {
+        Self {
+            expensive_root_field_names,
+            semaphore: Arc::new(Semaphore::new(max_in_flight)),
+            timeout,
+        }
+    }
+}
+
+impl ExtensionFactory for ExpensiveOpGuardFactory {
+    fn create(&self) -> Arc<dyn Extension> {
+        Arc::new(ExpensiveOpGuard {
+            expensive_root_field_names: self.expensive_root_field_names.clone(),
+            semaphore: self.semaphore.clone(),
+            timeout: self.timeout,
+            expensive_root_field_count: OnceLock::new(),
+        })
+    }
+}
+
+pub struct ExpensiveOpGuard {
+    expensive_root_field_names: Arc<[String]>,
+    semaphore: Arc<Semaphore>,
+    timeout: Duration,
+    expensive_root_field_count: OnceLock<usize>,
+}
+
+#[async_trait::async_trait]
+impl Extension for ExpensiveOpGuard {
+    async fn parse_query(
+        &self,
+        _ctx: &ExtensionContext<'_>,
+        _query: &str,
+        _variables: &Variables,
+        next: NextParseQuery<'_>,
+    ) -> ServerResult<ExecutableDocument> {
+        let doc = next.run(_ctx, _query, _variables).await?;
+        let expensive_root_field_count =
+            expensive_root_field_count(&doc, &self.expensive_root_field_names);
+        let _ = self
+            .expensive_root_field_count
+            .set(expensive_root_field_count);
+        Ok(doc)
+    }
+
+    async fn execute(
+        &self,
+        ctx: &ExtensionContext<'_>,
+        operation_name: Option<&str>,
+        next: NextExecute<'_>,
+    ) -> Response {
+        let expensive_root_field_count =
+            self.expensive_root_field_count.get().copied().unwrap_or(0);
+        let is_expensive = expensive_root_field_count > 0;
+
+        tracing::debug!(
+            "Executing operation: {:?}, expensive root field count: {:?}, expected root fields: {:?}, expensive: {:?}, timeout: {:?}, semaphore_size: {:?}",
+            operation_name,
+            expensive_root_field_count,
+            self.expensive_root_field_names,
+            is_expensive,
+            self.timeout,
+            self.semaphore.available_permits(),
+        );
+
+        if is_expensive {
+            self.expensive_execution(
+                ctx,
+                operation_name,
+                next,
+                expensive_root_field_count,
+            )
+            .await
+        } else {
+            next.run(ctx, operation_name).await
+        }
+    }
+}
+impl ExpensiveOpGuard {
+    async fn expensive_execution(
+        &self,
+        ctx: &ExtensionContext<'_>,
+        operation_name: Option<&str>,
+        next: NextExecute<'_>,
+        expensive_root_field_count: usize,
+    ) -> Response {
+        let permit_count = u32::try_from(expensive_root_field_count).unwrap_or(u32::MAX);
+
+        // Concurrency gate (bulkhead)
+        let permit = match self.semaphore.clone().try_acquire_many_owned(permit_count) {
+            Ok(p) => p,
+            Err(_) => {
+                let mut resp = Response::new(async_graphql::Value::Null);
+                resp.errors.push(ServerError::new(
+                    "Rate limit exceeded for this operation",
+                    None,
+                ));
+                return resp;
+            }
+        };
+
+        // Time bound (avoid request pile-ups)
+        let fut = next.run(ctx, operation_name);
+        let starting_time = tokio::time::Instant::now();
+        let out = tokio::time::timeout(self.timeout, fut).await;
+        tracing::warn!(
+            "finished executing in {:?}ns, success: {:?}",
+            starting_time.elapsed(),
+            out.is_ok(),
+        );
+
+        drop(permit);
+
+        match out {
+            Ok(resp) => resp,
+            Err(_) => {
+                let mut resp = Response::new(async_graphql::Value::Null);
+                resp.errors
+                    .push(ServerError::new("Operation timed out", None));
+                resp
+            }
+        }
+    }
+}
+
+fn expensive_root_field_count(
+    doc: &ExecutableDocument,
+    expensive_field_names: &[String],
+) -> usize {
+    doc.operations
+        .iter()
+        .map(|(_, op)| {
+            count_expensive_root_fields(
+                &op.node.selection_set,
+                &doc.fragments,
+                expensive_field_names,
+            )
+        })
+        .sum()
+}
+
+fn count_expensive_root_fields(
+    selection_set: &Positioned<SelectionSet>,
+    fragments: &std::collections::HashMap<
+        async_graphql::Name,
+        Positioned<FragmentDefinition>,
+    >,
+    expensive_field_names: &[String],
+) -> usize {
+    selection_set
+        .node
+        .items
+        .iter()
+        .map(|selection| match &selection.node {
+            Selection::Field(field) => {
+                let field_name = field.node.name.node.as_str();
+                if expensive_field_names.iter().any(|name| name == field_name) {
+                    1
+                } else {
+                    0
+                }
+            }
+            Selection::FragmentSpread(spread) => fragments
+                .get(&spread.node.fragment_name.node)
+                .map(|fragment| {
+                    count_expensive_root_fields(
+                        &fragment.node.selection_set,
+                        fragments,
+                        expensive_field_names,
+                    )
+                })
+                .unwrap_or(0),
+            Selection::InlineFragment(inline) => count_expensive_root_fields(
+                &inline.node.selection_set,
+                fragments,
+                expensive_field_names,
+            ),
+        })
+        .sum()
+}

crates/fuel-core/src/service/config.rs

@@ -244,6 +244,8 @@ impl Config {
                 costs: Default::default(),
                 required_fuel_block_height_tolerance: 10,
                 required_fuel_block_height_timeout: Duration::from_secs(30),
+                concurrent_full_block_requests: 10,
+                full_block_request_timeout: Duration::from_secs(3),
             },
             combined_db_config,
             continue_on_error: false,

tests/Cargo.toml

@@ -34,12 +34,14 @@ not_leader_lock = []
 
 [dependencies]
 anyhow = { workspace = true }
+async-trait = { workspace = true }
 aws-config = { workspace = true, features = [
     "behavior-version-latest",
 ], optional = true }
 aws-sdk-kms = { workspace = true, optional = true }
 aws-sdk-s3 = { workspace = true }
 clap = { workspace = true }
+cynic = { workspace = true }
 
 fuel-core = { path = "../crates/fuel-core", default-features = false, features = ["p2p", "relayer", "wasm-executor", "test-helpers", "backup"] }
 fuel-core-benches = { path = "../benches" }