refactor: use router builder

Signed-off-by: Gustavo Inacio <[email protected]>

Author: gusinacio

Cargo.lock

@@ -29,7 +29,11 @@ alloy = { version = "=0.5.4", features = [
 ], default-features = false }
 clap = "4.4.3"
 lazy_static = "1.4.0"
-axum = { version = "0.7.7", default-features = false }
+axum = { version = "0.7.9", default-features = false, features = [
+  "tokio",
+  "http1",
+  "http2",
+] }
 tokio = "1.40"
 prometheus = "0.13.3"
 anyhow = { version = "1.0.72" }

Cargo.toml

@@ -39,7 +39,8 @@ graphql = { git = "https://github.com/edgeandnode/toolshed", tag = "graphql-v0.3
 tap_core.workspace = true
 uuid.workspace = true
 alloy.workspace = true
-tower_governor = "0.4.0"
+tower_governor = "0.4.3"
+governor = "0.6.0"
 tower-http = { version = "0.6.2", features = [
   "auth",
   "cors",
@@ -55,6 +56,7 @@ cost-model = { git = "https://github.com/graphprotocol/agora", rev = "3ed34ca" }
 bip39.workspace = true
 tower = "0.5.1"
 pin-project = "1.1.7"
+typed-builder = "0.20.0"
 
 [dev-dependencies]
 hex-literal = "0.4.1"

crates/service/Cargo.toml

@@ -9,12 +9,15 @@ use std::time::Duration;
 use sqlx::{postgres::PgPoolOptions, PgPool};
 use tracing::debug;
 
+const DATABASE_TIMEOUT: Duration = Duration::from_secs(30);
+const DATABASE_MAX_CONNECTIONS: u32 = 50;
+
 pub async fn connect(url: &str) -> PgPool {
     debug!("Connecting to database");
 
     PgPoolOptions::new()
-        .max_connections(50)
-        .acquire_timeout(Duration::from_secs(3))
+        .max_connections(DATABASE_MAX_CONNECTIONS)
+        .acquire_timeout(DATABASE_TIMEOUT)
         .connect(url)
         .await
         .expect("Should be able to connect to the database")

crates/service/src/database/mod.rs

@@ -1,8 +1,16 @@
 // Copyright 2023-, Edge & Node, GraphOps, and Semiotic Labs.
 // SPDX-License-Identifier: Apache-2.0
 
+use std::net::SocketAddr;
+
+use axum::{routing::get, serve, Router};
 use lazy_static::lazy_static;
-use prometheus::{register_counter_vec, register_histogram_vec, CounterVec, HistogramVec};
+use prometheus::{
+    register_counter_vec, register_histogram_vec, CounterVec, HistogramVec, TextEncoder,
+};
+use reqwest::StatusCode;
+use tokio::net::TcpListener;
+use tracing::{error, info};
 
 lazy_static! {
     /// Metric registered in global registry for
@@ -37,3 +45,37 @@ lazy_static! {
     .unwrap();
 
 }
+
+pub fn serve_metrics(host_and_port: SocketAddr) {
+    info!(address = %host_and_port, "Serving prometheus metrics");
+
+    tokio::spawn(async move {
+        let router = Router::new().route(
+            "/metrics",
+            get(|| async {
+                let metric_families = prometheus::gather();
+                let encoder = TextEncoder::new();
+
+                match encoder.encode_to_string(&metric_families) {
+                    Ok(s) => (StatusCode::OK, s),
+                    Err(e) => {
+                        error!("Error encoding metrics: {}", e);
+                        (
+                            StatusCode::INTERNAL_SERVER_ERROR,
+                            format!("Error encoding metrics: {}", e),
+                        )
+                    }
+                }
+            }),
+        );
+
+        serve(
+            TcpListener::bind(host_and_port)
+                .await
+                .expect("Failed to bind to metrics port"),
+            router.into_make_service(),
+        )
+        .await
+        .expect("Failed to serve metrics")
+    });
+}

crates/service/src/metrics.rs

@@ -5,19 +5,16 @@ use std::str::FromStr;
 
 use crate::database::cost_model::{self, CostModel};
 use async_graphql::{Context, EmptyMutation, EmptySubscription, Object, Schema, SimpleObject};
-use async_graphql_axum::{GraphQLRequest, GraphQLResponse};
-use axum::extract::State;
 use lazy_static::lazy_static;
 use prometheus::{
     register_counter, register_counter_vec, register_histogram, register_histogram_vec, Counter,
     CounterVec, Histogram, HistogramVec,
 };
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
+use sqlx::PgPool;
 use thegraph_core::DeploymentId;
 
-use crate::service::SubgraphServiceState;
-
 lazy_static! {
     pub static ref COST_MODEL_METRIC: HistogramVec = register_histogram_vec!(
         "indexer_cost_model_seconds",
@@ -126,7 +123,7 @@ impl Query {
         ctx: &Context<'_>,
         deployment_ids: Vec<DeploymentId>,
     ) -> Result<Vec<GraphQlCostModel>, anyhow::Error> {
-        let pool = &ctx.data_unchecked::<SubgraphServiceState>().database;
+        let pool = &ctx.data_unchecked::<PgPool>();
         let cost_models = cost_model::cost_models(pool, &deployment_ids).await?;
         Ok(cost_models.into_iter().map(|m| m.into()).collect())
     }
@@ -136,7 +133,7 @@ impl Query {
         ctx: &Context<'_>,
         deployment_id: DeploymentId,
     ) -> Result<Option<GraphQlCostModel>, anyhow::Error> {
-        let pool = &ctx.data_unchecked::<SubgraphServiceState>().database;
+        let pool = &ctx.data_unchecked::<PgPool>();
         cost_model::cost_model(pool, &deployment_id)
             .await
             .map(|model_opt| model_opt.map(GraphQlCostModel::from))
@@ -145,17 +142,8 @@ impl Query {
 
 pub type CostSchema = Schema<Query, EmptyMutation, EmptySubscription>;
 
-pub async fn build_schema() -> CostSchema {
-    Schema::build(Query, EmptyMutation, EmptySubscription).finish()
-}
-
-pub async fn cost(
-    State(state): State<SubgraphServiceState>,
-    req: GraphQLRequest,
-) -> GraphQLResponse {
-    state
-        .cost_schema
-        .execute(req.into_inner().data(state.clone()))
-        .await
-        .into()
+pub async fn build_schema(data: PgPool) -> CostSchema {
+    Schema::build(Query, EmptyMutation, EmptySubscription)
+        .data(data)
+        .finish()
 }

crates/service/src/routes/cost.rs

@@ -5,18 +5,56 @@ use std::time::Duration;
 use std::{str::FromStr, sync::Arc};
 
 use anyhow::bail;
-use async_graphql::{Context, FieldResult, Object, SimpleObject};
+use async_graphql::{Context, EmptySubscription, FieldResult, Object, Schema, SimpleObject};
 use base64::{engine::general_purpose::STANDARD, Engine};
+use indexer_config::{BlockchainConfig, DipsConfig};
 use indexer_dips::alloy::dyn_abi::Eip712Domain;
 use indexer_dips::SignedCancellationRequest;
 use indexer_dips::{
     alloy::core::primitives::Address, alloy_rlp::Decodable, SignedIndexingAgreementVoucher,
     SubgraphIndexingVoucherMetadata,
 };
+use thegraph_core::attestation::eip712_domain;
 use uuid::Uuid;
 
 use crate::database::dips::AgreementStore;
 
+pub type DipsSchema = Schema<AgreementQuery, AgreementMutation, EmptySubscription>;
+pub type DipsStore = Arc<dyn AgreementStore>;
+
+pub fn build_schema(
+    indexer_address: Address,
+    DipsConfig {
+        allowed_payers,
+        cancellation_time_tolerance,
+    }: &DipsConfig,
+    BlockchainConfig {
+        chain_id,
+        receipts_verifier_address,
+    }: &BlockchainConfig,
+    agreement_store: DipsStore,
+    prices: Vec<Price>,
+) -> DipsSchema {
+    Schema::build(
+        AgreementQuery {},
+        AgreementMutation {
+            expected_payee: indexer_address,
+            allowed_payers: allowed_payers.clone(),
+            domain: eip712_domain(
+                // 42161, // arbitrum
+                *chain_id as u64,
+                *receipts_verifier_address,
+            ),
+            cancel_voucher_time_tolerance: cancellation_time_tolerance
+                .unwrap_or(Duration::from_secs(5)),
+        },
+        EmptySubscription,
+    )
+    .data(agreement_store)
+    .data(prices)
+    .finish()
+}
+
 pub enum NetworkProtocol {
     ArbitrumMainnet,
 }

crates/service/src/routes/dips.rs

@@ -1,9 +1,7 @@
 // Copyright 2023-, Edge & Node, GraphOps, and Semiotic Labs.
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::{
-    error::SubgraphServiceError, middleware::AttestationInput, service::SubgraphServiceState,
-};
+use crate::{error::SubgraphServiceError, middleware::AttestationInput, service::GraphNodeState};
 use axum::{
     extract::{Path, State},
     http::{HeaderValue, Response},
@@ -17,7 +15,7 @@ const GRAPH_ATTESTABLE: &str = "graph-attestable";
 
 pub async fn request_handler(
     Path(deployment): Path<DeploymentId>,
-    State(state): State<SubgraphServiceState>,
+    State(state): State<GraphNodeState>,
     req: String,
 ) -> Result<impl IntoResponse, SubgraphServiceError> {
     trace!("Handling request for deployment `{deployment}`");

crates/service/src/routes/request_handler.rs

@@ -12,7 +12,7 @@ use thegraph_graphql_http::{
     http_client::{ReqwestExt, ResponseError},
 };
 
-use crate::{error::SubgraphServiceError, service::SubgraphServiceState};
+use crate::{error::SubgraphServiceError, service::GraphNodeState};
 
 lazy_static::lazy_static! {
     static ref SUPPORTED_ROOT_FIELDS: HashSet<&'static str> =
@@ -57,7 +57,7 @@ impl IntoRequestParameters for WrappedGraphQLRequest {
 
 // Custom middleware function to process the request before reaching the main handler
 pub async fn status(
-    State(state): State<SubgraphServiceState>,
+    State(state): State<GraphNodeState>,
     request: GraphQLRequest,
 ) -> Result<impl IntoResponse, SubgraphServiceError> {
     let request = request.into_inner();