diff --git a/Cargo.lock b/Cargo.lock index 13d9156f..b11a3ad9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2931,6 +2931,7 @@ dependencies = [ "tikv-jemallocator", "time", "tokio", + "tokio-util", "tracing", "tracing-subscriber 0.3.22", "uuid", @@ -5467,7 +5468,7 @@ dependencies = [ [[package]] name = "rbuilder-primitives" version = "0.1.0" -source = "git+https://github.com/flashbots/rbuilder?rev=95323c0c6b1ac742a5716aba1942f3e06b8b5241#95323c0c6b1ac742a5716aba1942f3e06b8b5241" +source = "git+https://github.com/flashbots/rbuilder?rev=e3b49692d5b4353c62abe828245a44c390f7bec2#e3b49692d5b4353c62abe828245a44c390f7bec2" dependencies = [ "ahash", "alloy-consensus", @@ -5499,7 +5500,6 @@ dependencies = [ "ssz_types", "thiserror 1.0.69", "time", - "tracing", "tree_hash 0.8.0", "tree_hash_derive 0.8.0", "typenum", @@ -5509,7 +5509,7 @@ dependencies = [ [[package]] name = "rbuilder-utils" version = "0.1.0" -source = "git+https://github.com/flashbots/rbuilder?rev=95323c0c6b1ac742a5716aba1942f3e06b8b5241#95323c0c6b1ac742a5716aba1942f3e06b8b5241" +source = "git+https://github.com/flashbots/rbuilder?rev=e3b49692d5b4353c62abe828245a44c390f7bec2#e3b49692d5b4353c62abe828245a44c390f7bec2" dependencies = [ "alloy-primitives 1.5.0", "clickhouse", diff --git a/Cargo.toml b/Cargo.toml index 41093d70..50368032 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,8 +20,8 @@ revm-primitives = { version = "21.0.2", default-features = false } revm-interpreter = { version = "29.0.1", default-features = false } # rbuilder -rbuilder-primitives = { git = "https://github.com/flashbots/rbuilder", rev = "95323c0c6b1ac742a5716aba1942f3e06b8b5241" } -rbuilder-utils = { git = "https://github.com/flashbots/rbuilder", rev = "95323c0c6b1ac742a5716aba1942f3e06b8b5241" , features = [ +rbuilder-primitives = { git = "https://github.com/flashbots/rbuilder", rev = "e3b49692d5b4353c62abe828245a44c390f7bec2" } +rbuilder-utils = { git = "https://github.com/flashbots/rbuilder", rev = "e3b49692d5b4353c62abe828245a44c390f7bec2", features = [ "test-utils" ] } @@ -34,6 +34,7 @@ tokio = { version = "1", default-features = false, features = [ "macros" ] } futures = { version = "0.3" } +tokio-util = "0.7.12" # allocator tikv-jemallocator = { version = "0.6", optional = true } diff --git a/src/cli.rs b/src/cli.rs index 0e3e72dd..5c61eae4 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -98,6 +98,24 @@ pub struct ClickhouseArgs { default_value_t = MAX_DISK_BACKUP_SIZE_BYTES )] pub backup_disk_max_size_bytes: u64, + + /// Send timeout in milliseconds for ClickHouse HTTP requests. Defaults to 2_000. + #[arg( + long = "indexer.clickhouse.send-timeout-ms", + env = "CLICKHOUSE_SEND_TIMEOUT_MS", + id = "CLICKHOUSE_SEND_TIMEOUT_MS", + default_value_t = 2_000 + )] + pub send_timeout_ms: u64, + + /// End-to-end timeout in milliseconds for ClickHouse HTTP requests. Defaults to 3_000. + #[arg( + long = "indexer.clickhouse.end-timeout-ms", + env = "CLICKHOUSE_END_TIMEOUT_MS", + id = "CLICKHOUSE_END_TIMEOUT_MS", + default_value_t = 3_000 + )] + pub end_timeout_ms: u64, } /// Arguments required to setup file-based parquet indexing. @@ -248,6 +266,16 @@ pub struct OrderflowIngressArgs { #[clap(long = "http.enable-gzip", default_value_t = false)] pub gzip_enabled: bool, + /// Maximum local ClickHouse backup disk size in MB above which user RPC (e.g. eth_sendBundle) + /// is rejected with disk full. Defaults to 1024 MB (1 GiB). + #[clap( + long = "disk-max-size-to-accept-user-rpc-mb", + default_value_t = 1024, + env = "DISK_MAX_SIZE_TO_ACCEPT_USER_RPC", + id = "DISK_MAX_SIZE_TO_ACCEPT_USER_RPC" + )] + pub disk_max_size_to_accept_user_rpc_mb: u64, + /// The interval in seconds to update the peer list from BuilderHub. #[clap( long = "peer.update-interval-s", @@ -316,6 +344,7 @@ impl Default for OrderflowIngressArgs { score_bucket_s: 4, log_json: false, gzip_enabled: false, + disk_max_size_to_accept_user_rpc_mb: 1024, tcp_small_clients: NonZero::new(4).expect("non-zero"), tcp_big_clients: 0, io_threads: 4, diff --git a/src/indexer/click/mod.rs b/src/indexer/click/mod.rs index 6ef2b4b3..4e6a55db 100644 --- a/src/indexer/click/mod.rs +++ b/src/indexer/click/mod.rs @@ -1,6 +1,14 @@ //! Indexing functionality powered by Clickhouse. -use std::{fmt::Debug, time::Duration}; +use std::{ + fmt::Debug, + marker::PhantomData, + sync::{ + atomic::{AtomicU64, Ordering}, + LazyLock, + }, + time::Duration, +}; use crate::{ cli::ClickhouseArgs, @@ -19,6 +27,7 @@ use rbuilder_utils::{ }, tasks::TaskExecutor, }; +use tokio::task::JoinHandle; mod models; @@ -32,9 +41,59 @@ fn config_from_clickhouse_args(args: &ClickhouseArgs, validation: bool) -> Click } } -struct MetricsWrapper; +/// little global (puaj) info to easily get the current clickhouse disk size. +#[derive(Default)] +pub(crate) struct ClickhouseLocalBackupDiskSize { + bundles_size: AtomicU64, + bundle_receipts_size: AtomicU64, +} + +impl ClickhouseLocalBackupDiskSize { + pub(crate) fn set_bundles_size(&self, size: u64) { + self.bundles_size.store(size, Ordering::Relaxed); + } + pub(crate) fn set_bundle_receipts_size(&self, size: u64) { + self.bundle_receipts_size.store(size, Ordering::Relaxed); + } + pub(crate) fn disk_size(&self) -> u64 { + self.bundles_size.load(Ordering::Relaxed) + + self.bundle_receipts_size.load(Ordering::Relaxed) + } +} + +/// We store here the current disk size of the backup database to avoid querying the metrics since +/// that would include a string map access. +pub(crate) static CLICKHOUSE_LOCAL_BACKUP_DISK_SIZE: LazyLock = + LazyLock::new(ClickhouseLocalBackupDiskSize::default); + +/// Callback invoked when disk backup size is set. Implement this trait to observe size updates. +pub(crate) trait DiskBackupSizeCallback: Send + Sync { + fn on_disk_backup_size(size_bytes: u64); +} + +struct UpdateBundleSizeCallback; + +impl DiskBackupSizeCallback for UpdateBundleSizeCallback { + fn on_disk_backup_size(size_bytes: u64) { + CLICKHOUSE_LOCAL_BACKUP_DISK_SIZE.set_bundles_size(size_bytes); + } +} + +struct UpdateBundleReceiptsSizeCallback; + +impl DiskBackupSizeCallback for UpdateBundleReceiptsSizeCallback { + fn on_disk_backup_size(size_bytes: u64) { + CLICKHOUSE_LOCAL_BACKUP_DISK_SIZE.set_bundle_receipts_size(size_bytes); + } +} + +struct MetricsWrapper(PhantomData) +where + F: DiskBackupSizeCallback; -impl rbuilder_utils::clickhouse::backup::metrics::Metrics for MetricsWrapper { +impl rbuilder_utils::clickhouse::backup::metrics::Metrics + for MetricsWrapper +{ fn increment_write_failures(err: String) { CLICKHOUSE_METRICS.write_failures(err).inc(); } @@ -58,6 +117,7 @@ impl rbuilder_utils::clickhouse::backup::metrics::Metrics for MetricsWrapper { } fn set_disk_backup_size(size_bytes: u64, batches: usize, order: &'static str) { + F::on_disk_backup_size(size_bytes); CLICKHOUSE_METRICS.backup_size_bytes(order, "disk").set(size_bytes); CLICKHOUSE_METRICS.backup_size_batches(order, "disk").set(batches); } @@ -104,7 +164,7 @@ impl ClickhouseIndexer { receivers: OrderReceivers, task_executor: TaskExecutor, validation: bool, - ) { + ) -> Vec> { let client = config_from_clickhouse_args(&args, validation).into(); tracing::info!("Running with clickhouse indexer"); @@ -116,7 +176,13 @@ impl ClickhouseIndexer { ) .expect("could not create disk backup"); - spawn_clickhouse_inserter_and_backup::( + let send_timeout = Duration::from_millis(args.send_timeout_ms); + let end_timeout = Duration::from_millis(args.end_timeout_ms); + let bundle_inserter_join_handle = spawn_clickhouse_inserter_and_backup::< + SystemBundle, + BundleRow, + MetricsWrapper, + >( &client, receivers.bundle_rx, &task_executor, @@ -124,10 +190,16 @@ impl ClickhouseIndexer { builder_name.clone(), disk_backup.clone(), args.backup_memory_max_size_bytes, + send_timeout, + end_timeout, TARGET_INDEXER, ); - spawn_clickhouse_inserter_and_backup::( + let bundle_receipt_inserter_join_handle = spawn_clickhouse_inserter_and_backup::< + BundleReceipt, + BundleReceiptRow, + MetricsWrapper, + >( &client, receivers.bundle_receipt_rx, &task_executor, @@ -135,8 +207,11 @@ impl ClickhouseIndexer { builder_name.clone(), disk_backup.clone(), args.backup_memory_max_size_bytes, + send_timeout, + end_timeout, TARGET_INDEXER, ); + vec![bundle_inserter_join_handle, bundle_receipt_inserter_join_handle] } } @@ -145,17 +220,30 @@ pub(crate) mod tests { use std::{borrow::Cow, collections::BTreeMap, fs, time::Duration}; use crate::{ - cli::ClickhouseArgs, + cli::{ClickhouseArgs, IndexerArgs, OrderflowIngressArgs}, + consts::FLASHBOTS_SIGNATURE_HEADER, indexer::{ click::{ models::{BundleReceiptRow, BundleRow}, ClickhouseClientConfig, ClickhouseIndexer, }, tests::{bundle_receipt_example, system_bundle_example}, - OrderSenders, BUNDLE_RECEIPTS_TABLE_NAME, BUNDLE_TABLE_NAME, TARGET_INDEXER, + Indexer, OrderSenders, BUNDLE_RECEIPTS_TABLE_NAME, BUNDLE_TABLE_NAME, TARGET_INDEXER, + }, + jsonrpc::{JsonRpcError, JsonRpcResponse, JsonRpcResponseTy, JSONRPC_VERSION_2}, + utils::{ + testutils::{random_raw_bundle_with_tx_count_and_input_size, Random}, + wait_for_critical_tasks, SHUTDOWN_TIMEOUT, }, }; + use alloy_primitives::keccak256; + use alloy_signer::Signer; + use alloy_signer_local::PrivateKeySigner; + use axum::{http::StatusCode, routing::post, Json}; use clickhouse::{error::Result as ClickhouseResult, Client as ClickhouseClient}; + use msg_socket::RepSocket; + use msg_transport::tcp::Tcp; + use rbuilder_primitives::serialize::RawBundle; use rbuilder_utils::{ clickhouse::{ backup::{metrics::NullMetrics, Backup, DiskBackup, DiskBackupConfig, FailedCommit}, @@ -164,6 +252,8 @@ pub(crate) mod tests { }, tasks::TaskManager, }; + use serde_json::json; + use std::net::SocketAddr; use testcontainers::{ core::{ error::Result as TestcontainersResult, wait::HttpWaitStrategy, ContainerPort, WaitFor, @@ -171,7 +261,8 @@ pub(crate) mod tests { runners::AsyncRunner as _, ContainerAsync, Image, }; - use tokio::{runtime::Handle, sync::mpsc}; + use tokio::{net::TcpListener, runtime::Handle, sync::mpsc}; + use tokio_util::sync::CancellationToken; // Uncomment to enable logging during tests. // use tracing::level_filters::LevelFilter; @@ -239,6 +330,7 @@ pub(crate) mod tests { } } + ///Only for testing purposes. impl From for ClickhouseArgs { fn from(config: ClickhouseClientConfig) -> Self { Self { @@ -251,6 +343,8 @@ pub(crate) mod tests { backup_memory_max_size_bytes: 1024 * 1024 * 10, // 10MiB backup_disk_database_path: default_disk_backup_database_path(), backup_disk_max_size_bytes: 1024 * 1024 * 100, // 100MiB + send_timeout_ms: 2_000, + end_timeout_ms: 3_000, } } } @@ -367,7 +461,7 @@ pub(crate) mod tests { let (senders, receivers) = OrderSenders::new(); let validation = false; - ClickhouseIndexer::run( + let indexer_join_handles = ClickhouseIndexer::run( config.into(), builder_name.clone(), receivers, @@ -379,6 +473,8 @@ pub(crate) mod tests { let system_bundle = system_bundle_example(); let system_bundle_row = (system_bundle.clone(), builder_name.clone()).into(); senders.bundle_tx.send(system_bundle.clone()).await.unwrap(); + drop(senders); + wait_for_critical_tasks(indexer_join_handles, SHUTDOWN_TIMEOUT).await; // Wait a bit for bundle to be actually processed before shutting down. tokio::time::sleep(Duration::from_secs(1)).await; @@ -531,4 +627,188 @@ pub(crate) mod tests { drop(image); } } + + /// Integration test: when backup DB exceeds disk_max_size_to_accept_user_rpc_mb, user RPC + /// returns DiskFull; after fixing ClickHouse and draining backup, RPC accepts again. + /// A little long func.. consider improving this. + #[tokio::test(flavor = "multi_thread")] + async fn disk_full_rpc_rejects_then_accepts_after_drain() { + const BACKUP_DISK_MAX_SIZE_BYTES: u64 = 5; + const BUNDLE_TX_COUNT: usize = 10; + const BUNDLE_TX_INPUT_SIZE: usize = 1024; + const DISK_MAX_SIZE_TO_ACCEPT_USER_RPC_MB: u64 = 1; + // We need to fill DISK_MAX_SIZE_TO_ACCEPT_USER_RPC_MB with bundles. + // 2 is to play it safe. + const BUNDLE_COUNT_TO_FILL_DISK: usize = 2 * + (DISK_MAX_SIZE_TO_ACCEPT_USER_RPC_MB * 1024 * 1024) as usize / + (BUNDLE_TX_COUNT * BUNDLE_TX_INPUT_SIZE); + + const FLOWPROXY_START_DELAY_MS: Duration = Duration::from_millis(800); + // Assume 100ms per bundle to clickhouse (it's a LOT) + const DRAIN_TIMEOUT: Duration = + Duration::from_millis(100 * BUNDLE_COUNT_TO_FILL_DISK as u64); + + let mut rng = rand::rng(); + let task_manager = TaskManager::new(tokio::runtime::Handle::current()); + let task_executor = task_manager.executor(); + + // 1. Start ClickHouse without tables so inserts fail and go to backup. + let (image, client, config) = create_test_clickhouse_client(false).await.unwrap(); + + let temp_dir = tempfile::tempdir().unwrap(); + let backup_path = temp_dir.path().join("clickhouse-backup.db"); + + let mut clickhouse_args: ClickhouseArgs = config.clone().into(); + clickhouse_args.backup_disk_database_path = backup_path.to_string_lossy().to_string(); + clickhouse_args.backup_disk_max_size_bytes = BACKUP_DISK_MAX_SIZE_BYTES * 1024 * 1024; + + // 2. Mock builder so forwarder requests complete. + let builder_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let builder_port = builder_listener.local_addr().unwrap().port(); + let builder_url = format!("http://127.0.0.1:{builder_port}"); + let app = + axum::Router::new().route("/", post(|| async { (StatusCode::OK, Json(json!({}))) })); + tokio::spawn(async move { axum::serve(builder_listener, app).await.unwrap() }); + + let mut args = OrderflowIngressArgs::default().gzip_enabled().disable_builder_hub(); + args.peer_update_interval_s = 5; + args.indexing = IndexerArgs { clickhouse: Some(clickhouse_args), parquet: None }; + args.disk_max_size_to_accept_user_rpc_mb = DISK_MAX_SIZE_TO_ACCEPT_USER_RPC_MB; // 1 MiB threshold + args.builder_url = Some(builder_url.clone()); + + let user_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let user_port = user_listener.local_addr().unwrap().port(); + let user_url = format!("http://127.0.0.1:{user_port}"); + + // dummy_system_listener will not really be used but run_with_listeners needs one. + let mut dummy_system_listener = RepSocket::new(Tcp::default()); + let dummy_system_addr: SocketAddr = "127.0.0.1:0".parse().unwrap(); + dummy_system_listener.bind(dummy_system_addr).await.expect("to bind system listener"); + + let (indexer_handle, _indexer_join_handles) = + Indexer::run(args.indexing.clone(), args.builder_name.clone(), task_executor.clone()); + let cancellation_token = CancellationToken::new(); + + tokio::spawn(async move { + crate::run_with_listeners( + args, + user_listener, + dummy_system_listener, + None, + task_executor, + indexer_handle, + cancellation_token, + ) + .await + .unwrap(); + }); + + tokio::time::sleep(FLOWPROXY_START_DELAY_MS).await; + + let reqwest_client = reqwest::Client::default(); + let signer = PrivateKeySigner::random(); + + // 3. Phase 1: send bundles until backup exceeds threshold and we get DiskFull. + // Delay between requests so the indexer can commit batches (which fail and go to backup). + let mut got_disk_full = false; + for _ in 0..BUNDLE_COUNT_TO_FILL_DISK { + let bundle = random_raw_bundle_with_tx_count_and_input_size( + &mut rng, + BUNDLE_TX_COUNT, + Some(BUNDLE_TX_INPUT_SIZE), + ); + let response = send_bundle_req(&reqwest_client, &user_url, &signer, &bundle).await; + if is_disk_full(response).await { + got_disk_full = true; + break; + } + tokio::time::sleep(Duration::from_millis(30)).await; + } + assert!(got_disk_full, "expected RPC to eventually return DiskFull"); + + // 4. Phase 2: create tables so backup can drain. + create_clickhouse_bundles_table(&client).await.unwrap(); + create_clickhouse_bundle_receipts_table(&client).await.unwrap(); + + // 5. Poll until RPC accepts again (backup drained). + let poll_interval = Duration::from_millis(500); + let mut elapsed = Duration::ZERO; + while elapsed < DRAIN_TIMEOUT { + tokio::time::sleep(poll_interval).await; + elapsed += poll_interval; + let bundle = RawBundle::random(&mut rng); + let response = send_bundle_req(&reqwest_client, &user_url, &signer, &bundle).await; + if response.status().is_success() { + let body = response.bytes().await.unwrap(); + let parsed: JsonRpcResponse = + serde_json::from_slice(body.as_ref()).expect("valid json"); + if matches!(parsed.result_or_error, JsonRpcResponseTy::Result(_)) { + break; + } + } + } + assert!(elapsed < DRAIN_TIMEOUT, "RPC did not accept again within {:?}", DRAIN_TIMEOUT); + + // 6. Phase 3: one more successful eth_sendBundle. + let bundle = RawBundle::random(&mut rng); + let response = send_bundle_req(&reqwest_client, &user_url, &signer, &bundle).await; + assert!( + response.status().is_success(), + "expected success after drain, got {}", + response.text().await.unwrap_or_default() + ); + let body = response.bytes().await.unwrap(); + let parsed: JsonRpcResponse = + serde_json::from_slice(body.as_ref()).unwrap(); + assert!( + matches!(parsed.result_or_error, JsonRpcResponseTy::Result(_)), + "expected result, got {:?}", + parsed.result_or_error + ); + + drop(image); + } + + async fn is_disk_full(response: reqwest::Response) -> bool { + let status = response.status(); + let body = response.bytes().await.unwrap(); + if !status.is_success() { + let parsed: JsonRpcResponse<()> = match serde_json::from_slice(body.as_ref()) { + Ok(p) => p, + Err(_) => return false, + }; + matches!( + parsed.result_or_error, + JsonRpcResponseTy::Error { code: -32603, message: JsonRpcError::DiskFull } + ) + } else { + false + } + } + + async fn send_bundle_req( + client: &reqwest::Client, + url: &str, + signer: &PrivateKeySigner, + bundle: &RawBundle, + ) -> reqwest::Response { + let body = json!({ + "id": 0, + "jsonrpc": JSONRPC_VERSION_2, + "method": "eth_sendBundle", + "params": [bundle] + }); + let body_bytes = serde_json::to_vec(&body).unwrap(); + let sighash = format!("{:?}", keccak256(&body_bytes)); + let sig = signer.sign_message(sighash.as_bytes()).await.unwrap(); + let signature_header = format!("{:?}:{}", signer.address(), sig); + client + .post(url) + .header(reqwest::header::CONTENT_TYPE, "application/json") + .header(FLASHBOTS_SIGNATURE_HEADER, signature_header) + .body(body_bytes) + .send() + .await + .unwrap() + } } diff --git a/src/indexer/click/models.rs b/src/indexer/click/models.rs index b25da3c1..b60e8f13 100644 --- a/src/indexer/click/models.rs +++ b/src/indexer/click/models.rs @@ -259,7 +259,7 @@ impl From<(SystemBundle, String)> for BundleRow { .unwrap_or_default(), // Decoded bundles always have a uuid. internal_uuid: decoded.uuid, - replacement_uuid: decoded.replacement_data.clone().map(|r| r.key.id), + replacement_uuid: decoded.replacement_data.map(|r| r.key.id), replacement_nonce: bundle.raw_bundle.metadata.replacement_nonce, signer_address: Some(bundle.metadata.signer), builder_name, diff --git a/src/indexer/mod.rs b/src/indexer/mod.rs index b95132af..0fea37e5 100644 --- a/src/indexer/mod.rs +++ b/src/indexer/mod.rs @@ -4,7 +4,8 @@ use std::fmt::Debug; use rbuilder_utils::tasks::TaskExecutor; -use tokio::sync::mpsc; +use tokio::{sync::mpsc, task::JoinHandle}; +use tracing::error; use crate::{ cli::IndexerArgs, @@ -77,33 +78,38 @@ impl OrderSenders { pub struct Indexer; impl Indexer { + /// Returns the IndexerHandle to send data and a vector of join handles to the indexer tasks so + /// we can wait for them to finish on shutdown. pub fn run( args: IndexerArgs, builder_name: String, task_executor: TaskExecutor, - ) -> IndexerHandle { + ) -> (IndexerHandle, Vec>) { let (senders, receivers) = OrderSenders::new(); match (args.clickhouse, args.parquet) { (None, None) => { MockIndexer.run(receivers, task_executor); - IndexerHandle::new(senders) + (IndexerHandle::new(senders), vec![]) } (Some(clickhouse), None) => { let validation = false; - ClickhouseIndexer::run( + let join_handles = ClickhouseIndexer::run( clickhouse, builder_name, receivers, task_executor, validation, ); - IndexerHandle::new(senders) + (IndexerHandle::new(senders), join_handles) } (None, Some(parquet)) => { ParquetIndexer::run(parquet, builder_name, receivers, task_executor) .expect("failed to start parquet indexer"); - IndexerHandle::new(senders) + error!( + "Parquet indexer does not support proper shutdown, returning empty join handles" + ); + (IndexerHandle::new(senders), vec![]) } (Some(_), Some(_)) => { unreachable!("Cannot specify both clickhouse and parquet indexer"); diff --git a/src/ingress/mod.rs b/src/ingress/mod.rs index 84fbb932..5fdb60d9 100644 --- a/src/ingress/mod.rs +++ b/src/ingress/mod.rs @@ -7,7 +7,7 @@ use crate::{ }, entity::{Entity, EntityBuilderStats, EntityData, EntityRequest, EntityScores, SpamThresholds}, forwarder::IngressForwarders, - indexer::{IndexerHandle, OrderIndexer as _}, + indexer::{click::CLICKHOUSE_LOCAL_BACKUP_DISK_SIZE, IndexerHandle, OrderIndexer as _}, jsonrpc::{JsonRpcError, JsonRpcRequest, JsonRpcResponse}, metrics::{IngressMetrics, SYSTEM_METRICS}, primitives::{ @@ -49,6 +49,7 @@ use std::{ time::{Duration, Instant}, }; use time::UtcDateTime; +use tokio_util::sync::CancellationToken; use tracing::*; pub mod error; @@ -75,6 +76,8 @@ pub struct OrderflowIngress { pub local_builder_url: Option, pub builder_ready_endpoint: Option, pub indexer_handle: IndexerHandle, + /// Maximum local ClickHouse backup disk size in bytes above which user RPC is rejected. + pub disk_max_size_to_accept_user_rpc: u64, // Metrics pub(crate) user_metrics: IngressMetrics, @@ -129,7 +132,7 @@ impl OrderflowIngress { /// Perform maintenance task for internal orderflow ingress state. #[tracing::instrument(skip_all, name = "ingress_maintanance")] - pub async fn maintenance(&self) { + pub fn maintenance(&self) { let len_before = self.entities.len(); tracing::info!(entries = len_before, "starting state maintenance"); @@ -141,6 +144,10 @@ impl OrderflowIngress { tracing::info!(entries = len_after, num_removed, "finished state maintenance"); } + fn clickhouse_backup_disk_size_is_ok(&self) -> bool { + CLICKHOUSE_LOCAL_BACKUP_DISK_SIZE.disk_size() <= self.disk_max_size_to_accept_user_rpc + } + #[tracing::instrument(skip_all, name = "ingress", fields( handler = "user", @@ -152,6 +159,9 @@ impl OrderflowIngress { headers: HeaderMap, body: axum::body::Bytes, ) -> JsonRpcResponse { + if !ingress.clickhouse_backup_disk_size_is_ok() { + return JsonRpcResponse::error(Value::Null, JsonRpcError::DiskFull); + } let received_at = UtcInstant::now(); let body = match maybe_decompress(ingress.gzip_enabled, &headers, body) { @@ -285,6 +295,13 @@ impl OrderflowIngress { /// returns 200 if the local builder is not configured. #[tracing::instrument(skip_all, name = "ingress_readyz")] pub async fn ready_handler(State(ingress): State>) -> Response { + if !ingress.clickhouse_backup_disk_size_is_ok() { + return Response::builder() + .status(StatusCode::SERVICE_UNAVAILABLE) + .body(Body::from("clickhouse backup too big")) + .unwrap(); + } + if let Some(ref url) = ingress.builder_ready_endpoint { let client = reqwest::Client::builder() .timeout(Duration::from_secs(DEFAULT_HTTP_TIMEOUT_SECS)) @@ -296,7 +313,7 @@ impl OrderflowIngress { tracing::error!(%url, "error sending readyz request"); return Response::builder() .status(StatusCode::SERVICE_UNAVAILABLE) - .body(Body::from("not ready")) + .body(Body::from("builder not answering readyz request")) .unwrap(); }; @@ -307,7 +324,7 @@ impl OrderflowIngress { tracing::error!(%url, status = %response.status(), "local builder is not ready"); return Response::builder() .status(StatusCode::SERVICE_UNAVAILABLE) - .body(Body::from("not ready")) + .body(Body::from("builder not ready")) .unwrap(); } } @@ -784,16 +801,25 @@ impl IngressSocket { Self { reply_socket: socket, ingress_state, task_executor } } - pub async fn listen(mut self) { - while let Some(req) = self.reply_socket.next().await { - let data = req.msg().clone(); - let state = self.ingress_state.clone(); - self.task_executor.spawn(async { - let response = OrderflowIngress::system_handler(state, data).await; - if let Err(e) = req.respond(bitcode::encode(&response).into()) { - tracing::error!(?e, "failed to respond to request"); + pub async fn listen(mut self, cancellation_token: CancellationToken) { + loop { + tokio::select! { + Some(req) = self.reply_socket.next() => { + let data = req.msg().clone(); + let state = self.ingress_state.clone(); + self.task_executor.spawn(async { + let response = OrderflowIngress::system_handler(state, data).await; + if let Err(e) = req.respond(bitcode::encode(&response).into()) { + tracing::error!(?e, "failed to respond to request"); + } + }); + } + _ = cancellation_token.cancelled() => { + info!("Cancellation token cancelled, stopping ingress socket listener"); + break; } - }); + + } } } } diff --git a/src/jsonrpc.rs b/src/jsonrpc.rs index b74980a5..7d15b734 100644 --- a/src/jsonrpc.rs +++ b/src/jsonrpc.rs @@ -177,6 +177,8 @@ pub enum JsonRpcError { RateLimited, #[error("Internal error")] Internal, + #[error("Disk full")] + DiskFull, #[error("{0}")] Unknown(String), } @@ -192,6 +194,7 @@ impl FromStr for JsonRpcError { "Invalid params" => Self::InvalidParams, "Rate limited" => Self::RateLimited, "Internal error" => Self::Internal, + "Disk full" => Self::DiskFull, s => { if s.starts_with("Method not found: ") { Self::MethodNotFound( @@ -239,9 +242,11 @@ impl JsonRpcError { Self::InvalidRequest => -32600, Self::MethodNotFound(_) => -32601, Self::InvalidParams => -32602, - Self::RateLimited | Self::Internal | Self::Unknown(_) | Self::InvalidSignature => { - -32603 - } + Self::RateLimited | + Self::Internal | + Self::Unknown(_) | + Self::InvalidSignature | + Self::DiskFull => -32603, } } @@ -254,7 +259,7 @@ impl JsonRpcError { Self::InvalidSignature => StatusCode::BAD_REQUEST, Self::MethodNotFound(_) => StatusCode::NOT_FOUND, Self::RateLimited => StatusCode::TOO_MANY_REQUESTS, - Self::Internal | Self::Unknown(_) => StatusCode::INTERNAL_SERVER_ERROR, + Self::Internal | Self::Unknown(_) | Self::DiskFull => StatusCode::INTERNAL_SERVER_ERROR, } } } diff --git a/src/lib.rs b/src/lib.rs index 672901d7..c011b455 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,11 +7,11 @@ use crate::{ client::{default_http_builder, HttpClientPool}, http::spawn_http_forwarder, }, + indexer::IndexerHandle, ingress::IngressSocket, metrics::IngressMetrics, primitives::SystemBundleDecoder, priority::workers::PriorityWorkers, - runner::CliContext, statics::LOCAL_PEER_STORE, }; use alloy_signer_local::PrivateKeySigner; @@ -28,6 +28,7 @@ use forwarder::{IngressForwarders, PeerHandle}; use msg_socket::RepSocket; use msg_transport::tcp::Tcp; use prometric::exporter::ExporterBuilder; +use rbuilder_utils::tasks::TaskExecutor; use reqwest::Url; use std::{ net::SocketAddr, @@ -36,8 +37,9 @@ use std::{ sync::Arc, time::{Duration, Instant}, }; -use tokio::net::TcpListener; -use tracing::level_filters::LevelFilter; +use tokio::{net::TcpListener, select}; +use tokio_util::sync::CancellationToken; +use tracing::{info, level_filters::LevelFilter}; use tracing_subscriber::{layer::SubscriberExt as _, util::SubscriberInitExt, EnvFilter}; pub mod cli; @@ -46,7 +48,7 @@ use cli::OrderflowIngressArgs; pub mod ingress; use ingress::OrderflowIngress; -use crate::{cache::OrderCache, indexer::Indexer}; +use crate::cache::OrderCache; pub mod builderhub; mod cache; @@ -59,13 +61,17 @@ pub mod metrics; pub mod primitives; pub mod priority; pub mod rate_limit; -pub mod runner; pub mod statics; pub mod trace; pub mod utils; pub mod validation; -pub async fn run(args: OrderflowIngressArgs, ctx: CliContext) -> eyre::Result<()> { +pub async fn run( + args: OrderflowIngressArgs, + task_executor: TaskExecutor, + indexer_handle: IndexerHandle, + cancellation_token: CancellationToken, +) -> eyre::Result<()> { fdlimit::raise_fd_limit()?; if let Some(ref metrics_addr) = args.metrics { @@ -74,6 +80,9 @@ pub async fn run(args: OrderflowIngressArgs, ctx: CliContext) -> eyre::Result<() // Set build info metric metrics::BUILD_INFO_METRICS.info(env!("CARGO_PKG_VERSION"), env!("GIT_HASH")).set(1); + metrics::CLICKHOUSE_METRICS + .disk_max_size_to_accept_user_rpc_bytes() + .set(args.disk_max_size_to_accept_user_rpc_mb.saturating_mul(1024 * 1024)); } let user_listener = TcpListener::bind(&args.user_listen_addr).await?; @@ -85,15 +94,28 @@ pub async fn run(args: OrderflowIngressArgs, ctx: CliContext) -> eyre::Result<() } else { None }; - run_with_listeners(args, user_listener, system_listener, builder_listener, ctx).await + run_with_listeners( + args, + user_listener, + system_listener, + builder_listener, + task_executor, + indexer_handle, + cancellation_token, + ) + .await } +/// Cancellation is a little ugly, just added enough to make it drop indexer_handle but it's a mix +/// of cancellation_token + task_executor which also has a shutdown method. pub async fn run_with_listeners( args: OrderflowIngressArgs, user_listener: TcpListener, system_listener: RepSocket, builder_listener: Option, - ctx: CliContext, + task_executor: TaskExecutor, + indexer_handle: IndexerHandle, + cancellation_token: CancellationToken, ) -> eyre::Result<()> { // Initialize tracing. let registry = tracing_subscriber::registry().with( @@ -105,8 +127,6 @@ pub async fn run_with_listeners( let _ = registry.with(tracing_subscriber::fmt::layer()).try_init(); } - let indexer_handle = Indexer::run(args.indexing, args.builder_name, ctx.task_executor.clone()); - let orderflow_signer = match args.orderflow_signer { Some(signer) => { tracing::warn!( @@ -141,10 +161,10 @@ pub async fn run_with_listeners( peer_update_config, builder_hub, peers.clone(), - ctx.task_executor.clone(), + task_executor.clone(), ); - ctx.task_executor + task_executor .spawn_critical("run_update_peers", peer_updater.run(args.peer_update_interval_s)); } else { tracing::warn!("No BuilderHub URL provided, running with local peer store"); @@ -154,14 +174,9 @@ pub async fn run_with_listeners( .register(local_signer, Some(system_listener.local_addr().expect("bound").port())); let peers = peers.clone(); - let peer_updater = PeersUpdater::new( - peer_update_config, - peer_store, - peers.clone(), - ctx.task_executor.clone(), - ); - - ctx.task_executor + let peer_updater = + PeersUpdater::new(peer_update_config, peer_store, peers.clone(), task_executor.clone()); + task_executor .spawn_critical("local_update_peers", peer_updater.run(args.peer_update_interval_s)); } @@ -176,7 +191,7 @@ pub async fn run_with_listeners( builder_url.to_string(), // Use 1 client here, this is still using HTTP/1.1 with internal connection pooling. HttpClientPool::new(NonZero::new(1).unwrap(), || client.clone()), - &ctx.task_executor, + &task_executor, )?; IngressForwarders::new(local_sender, peers, orderflow_signer, workers.clone()) @@ -210,25 +225,35 @@ pub async fn run_with_listeners( local_builder_url: builder_url, builder_ready_endpoint, indexer_handle, + disk_max_size_to_accept_user_rpc: args.disk_max_size_to_accept_user_rpc_mb * 1024 * 1024, user_metrics: IngressMetrics::builder().with_label("handler", "user").build(), system_metrics: IngressMetrics::builder().with_label("handler", "system").build(), }); // Spawn a state maintenance task. - tokio::spawn({ + let cancellation_token_clone = cancellation_token.clone(); + task_executor.spawn({ let ingress = ingress.clone(); async move { loop { - tokio::time::sleep(Duration::from_secs(60)).await; - ingress.maintenance().await; + info!("starting state maintenance!!"); + select! { + _ = cancellation_token_clone.cancelled() => { + info!("Cancellation token cancelled, stopping state maintenance"); + break; + } + _ = tokio::time::sleep(Duration::from_secs(60)) => { + ingress.maintenance(); + } + } } } }); tracing::info!(addr = ?system_listener.local_addr(), "starting system tcp listener"); let ingress_socket = - IngressSocket::new(system_listener, ingress.clone(), ctx.task_executor.clone()); - ctx.task_executor.spawn(ingress_socket.listen()); + IngressSocket::new(system_listener, ingress.clone(), task_executor.clone()); + task_executor.spawn(ingress_socket.listen(cancellation_token)); // Spawn user facing HTTP server for accepting bundles and raw transactions. let user_router = Router::new() diff --git a/src/main.rs b/src/main.rs index 8f5b85bb..a722dd94 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,9 +1,15 @@ +use std::{future::Future, time::Duration}; + use clap::Parser; use flowproxy::{ cli::OrderflowIngressArgs, - runner::{CliContext, CliRunner}, + indexer::Indexer, trace::init_tracing, + utils::{wait_for_critical_tasks, SHUTDOWN_TIMEOUT}, }; +use rbuilder_utils::tasks::{PanickedTaskError, TaskManager}; +use tokio_util::sync::CancellationToken; +use tracing::{error, info}; #[cfg(all(feature = "jemalloc", unix))] type AllocatorInner = tikv_jemallocator::Jemalloc; @@ -33,12 +39,115 @@ fn main() { .enable_all() .build() .expect("failed to create runtime"); + let task_manager = TaskManager::new(tokio_runtime.handle().clone()); + info!("Main task started"); + + // Executes the main task command until it finished or ctrl-c was fired. + // IMPORTANT: flowproxy::run has no nice cancellation and will be stopped being polled abruptly + // so it must not contain any critical tasks that need proper shutdown. + tokio_runtime.block_on(run_with_shutdown(args, task_manager)); + + info!("Main task finished. Shutting down tokio runtime"); + + if let Err(error) = wait_tokio_runtime_shutdown(tokio_runtime, Duration::from_secs(5)) { + error!(?error, "Flow proxy terminated with error"); + } +} + +async fn run_with_shutdown(args: OrderflowIngressArgs, mut task_manager: TaskManager) { + let task_executor = task_manager.executor(); + let (indexer_handle, indexer_join_handles) = + Indexer::run(args.indexing.clone(), args.builder_name.clone(), task_executor.clone()); + let cancellation_token = CancellationToken::new(); + let main_task = flowproxy::run(args, task_executor, indexer_handle, cancellation_token.clone()); + match run_to_completion_or_panic(&mut task_manager, run_until_ctrl_c(main_task)).await { + Ok(()) => { + tracing::warn!(target = "cli", "shutting down gracefully"); + } + Err(err) => { + tracing::error!(?err, target = "cli", "shutting down due to error"); + } + } + // This kills some tasks launched by flowproxy::run that release the last references to the + // indexer_handle. + cancellation_token.cancel(); + // At this point all the rpc was abruptly dropped which dropped the indexer_handle and that will + // allow the indexer core to process all pending data and start shutting down. + wait_for_critical_tasks(indexer_join_handles, SHUTDOWN_TIMEOUT).await; + // We already have a chance to critical tasks to finish by themselves, so we can now call the + // graceful shutdown. + task_manager.graceful_shutdown_with_timeout(SHUTDOWN_TIMEOUT); +} + +fn wait_tokio_runtime_shutdown( + tokio_runtime: tokio::runtime::Runtime, + timeout: Duration, +) -> Result<(), std::sync::mpsc::RecvTimeoutError> { + // `drop(tokio_runtime)` would block the current thread until its pools + // (including blocking pool) are shutdown. Since we want to exit as soon as possible, drop + // it on a separate thread and wait for up to 5 seconds for this operation to + // complete. + let (tx, rx) = std::sync::mpsc::channel(); + std::thread::Builder::new() + .name("tokio-runtime-shutdown".to_string()) + .spawn(move || { + drop(tokio_runtime); + let _ = tx.send(()); + }) + .unwrap(); + + rx.recv_timeout(timeout).inspect_err(|err| { + tracing::debug!(target: "reth::cli", %err, "tokio runtime shutdown timed out"); + }) +} - let runner = CliRunner::from_runtime(tokio_runtime); +/// Runs the future to completion or until: +/// - `ctrl-c` is received. +/// - `SIGTERM` is received (unix only). +async fn run_until_ctrl_c(fut: F) -> Result<(), E> +where + F: Future>, + E: Send + Sync + 'static + From, +{ + let ctrl_c = tokio::signal::ctrl_c(); - let command = |ctx: CliContext| flowproxy::run(args, ctx); + let mut stream = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())?; + let sigterm = stream.recv(); + let sigterm = Box::pin(sigterm); + let ctrl_c = Box::pin(ctrl_c); + let fut = Box::pin(fut); + + tokio::select! { + _ = ctrl_c => { + tracing::info!("Received ctrl-c"); + }, + _ = sigterm => { + tracing::info!("Received SIGTERM"); + }, + res = fut => res?, + } + + Ok(()) +} - if let Err(e) = runner.run_command_until_exit(command) { - eprintln!("Orderflow proxy terminated with error: {e}"); +/// Runs the given future to completion or until a critical task panicked. +/// +/// Returns the error if a task panicked, or the given future returned an error. +async fn run_to_completion_or_panic(tasks: &mut TaskManager, fut: F) -> Result<(), E> +where + F: Future>, + E: Send + Sync + From + 'static, +{ + { + let fut = Box::pin(fut); + tokio::select! { + task_manager_result = tasks => { + if let Err(panicked_error) = task_manager_result { + return Err(panicked_error.into()); + } + }, + res = fut => res?, + } } + Ok(()) } diff --git a/src/metrics.rs b/src/metrics.rs index ac3b553e..d864a4b3 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -189,6 +189,9 @@ pub struct ClickhouseMetrics { /// Errors encountered during Clickhouse disk backup. #[metric(labels = ["order", "error"])] backup_disk_errors: Counter, + /// Configured max disk size (bytes) above which user RPC is rejected. + #[metric(rename = "disk_max_size_to_accept_user_rpc_bytes")] + disk_max_size_to_accept_user_rpc_bytes: Gauge, } #[metrics(scope = "indexer_parquet")] diff --git a/src/runner/mod.rs b/src/runner/mod.rs deleted file mode 100644 index e1c91373..00000000 --- a/src/runner/mod.rs +++ /dev/null @@ -1,135 +0,0 @@ -//! Asynchronous task runner utilities. - -use std::{future::Future, time::Duration}; - -use rbuilder_utils::tasks::{PanickedTaskError, TaskExecutor, TaskManager}; - -#[derive(Debug, Clone)] -pub struct CliContext { - pub task_executor: TaskExecutor, -} - -/// Executes CLI commands. -/// -/// Provides utilities for running a cli command to completion. -#[derive(Debug)] -#[non_exhaustive] -pub struct CliRunner { - tokio_runtime: tokio::runtime::Runtime, -} - -impl CliRunner { - /// Create a new [`CliRunner`] from a provided tokio [`Runtime`](tokio::runtime::Runtime). - pub const fn from_runtime(tokio_runtime: tokio::runtime::Runtime) -> Self { - Self { tokio_runtime } - } -} - -// === impl CliRunner === - -impl CliRunner { - /// Executes the given _async_ command on the tokio runtime until the command future resolves or - /// until the process receives a `SIGINT` or `SIGTERM` signal. - /// - /// Tasks spawned by the command via the [`TaskExecutor`] are shut down and an attempt is made - /// to drive their shutdown to completion after the command has finished. - pub fn run_command_until_exit( - self, - command: impl FnOnce(CliContext) -> F, - ) -> Result<(), E> - where - F: Future>, - E: Send + Sync + From + From + 'static, - { - let tokio_runtime = self.tokio_runtime; - let mut task_manager = TaskManager::new(tokio_runtime.handle().clone()); - let task_executor = task_manager.executor(); - let context = CliContext { task_executor }; - - // Executes the command until it finished or ctrl-c was fired - let command_res = tokio_runtime.block_on(run_to_completion_or_panic( - &mut task_manager, - run_until_ctrl_c(command(context)), - )); - - if command_res.is_err() { - tracing::error!(target: "cli", "shutting down due to error"); - } else { - tracing::debug!(target: "cli", "shutting down gracefully"); - // after the command has finished or exit signal was received we shutdown the task - // manager which fires the shutdown signal to all tasks spawned via the task - // executor and awaiting on tasks spawned with graceful shutdown - task_manager.graceful_shutdown_with_timeout(Duration::from_secs(5)); - } - - // `drop(tokio_runtime)` would block the current thread until its pools - // (including blocking pool) are shutdown. Since we want to exit as soon as possible, drop - // it on a separate thread and wait for up to 5 seconds for this operation to - // complete. - let (tx, rx) = std::sync::mpsc::channel(); - std::thread::Builder::new() - .name("tokio-runtime-shutdown".to_string()) - .spawn(move || { - drop(tokio_runtime); - let _ = tx.send(()); - }) - .unwrap(); - - let _ = rx.recv_timeout(Duration::from_secs(5)).inspect_err(|err| { - tracing::debug!(target: "reth::cli", %err, "tokio runtime shutdown timed out"); - }); - - command_res - } -} - -/// Runs the future to completion or until: -/// - `ctrl-c` is received. -/// - `SIGTERM` is received (unix only). -async fn run_until_ctrl_c(fut: F) -> Result<(), E> -where - F: Future>, - E: Send + Sync + 'static + From, -{ - let ctrl_c = tokio::signal::ctrl_c(); - - let mut stream = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())?; - let sigterm = stream.recv(); - let sigterm = Box::pin(sigterm); - let ctrl_c = Box::pin(ctrl_c); - let fut = Box::pin(fut); - - tokio::select! { - _ = ctrl_c => { - tracing::info!("Received ctrl-c"); - }, - _ = sigterm => { - tracing::info!("Received SIGTERM"); - }, - res = fut => res?, - } - - Ok(()) -} - -/// Runs the given future to completion or until a critical task panicked. -/// -/// Returns the error if a task panicked, or the given future returned an error. -async fn run_to_completion_or_panic(tasks: &mut TaskManager, fut: F) -> Result<(), E> -where - F: Future>, - E: Send + Sync + From + 'static, -{ - { - let fut = Box::pin(fut); - tokio::select! { - task_manager_result = tasks => { - if let Err(panicked_error) = task_manager_result { - return Err(panicked_error.into()); - } - }, - res = fut => res?, - } - } - Ok(()) -} diff --git a/src/utils.rs b/src/utils.rs index 044d79e3..99a4c368 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -1,8 +1,11 @@ use alloy_eips::eip2718::EIP4844_TX_TYPE_ID; use alloy_primitives::Bytes; use alloy_rlp::{Buf as _, Header}; +use futures::{stream::FuturesUnordered, StreamExt}; use std::time::{Duration, Instant}; use time::UtcDateTime; +use tokio::task::JoinHandle; +use tracing::{error, info}; use uuid::Uuid; use crate::{statics::START, validation::MAINNET_CHAIN_ID}; @@ -140,20 +143,26 @@ pub mod testutils { impl Random for TxEip1559 { fn random(rng: &mut R) -> Self { - let max_fee_per_gas = rng.random(); - let max_priority_fee_per_gas = rng.random_range(0..max_fee_per_gas); + let input_len = rng.random_range(0..=1024); + create_tx_eip1559_with_input_size(rng, input_len) + } + } - Self { - chain_id: 1, - nonce: 0, - gas_limit: 100_000, - max_fee_per_gas, - max_priority_fee_per_gas, - to: TxKind::Call(Address::random_with(rng)), - value: U256::random_with(rng), - access_list: Default::default(), - input: Bytes::random(rng), - } + pub fn create_tx_eip1559_with_input_size(rng: &mut R, data_size: usize) -> TxEip1559 { + let max_fee_per_gas = rng.random(); + let max_priority_fee_per_gas = rng.random_range(0..max_fee_per_gas); + let mut bytes = vec![0u8; data_size]; + rng.fill_bytes(&mut bytes); + TxEip1559 { + chain_id: 1, + nonce: 0, + gas_limit: 100_000, + max_fee_per_gas, + max_priority_fee_per_gas, + to: TxKind::Call(Address::random_with(rng)), + value: U256::random_with(rng), + access_list: Default::default(), + input: bytes.into(), } } @@ -244,41 +253,88 @@ pub mod testutils { } } + /// Create a random [`RawBundle`] with a fixed number of EIP-1559 transactions. + /// When `input_size` is `Some(n)`, each transaction's input has exactly `n` bytes (via + /// [`create_tx_eip1559_with_input_size`]). When `None`, each transaction uses random input + /// size (same as [`TxEip1559::random`]). + pub fn random_raw_bundle_with_tx_count_and_input_size( + rng: &mut R, + tx_count: usize, + input_size: Option, + ) -> RawBundle { + let txs = (0..tx_count) + .map(|_| { + let signer = PrivateKeySigner::random(); + let tx = EthereumTypedTransaction::Eip1559(match input_size { + Some(n) => create_tx_eip1559_with_input_size(rng, n), + None => TxEip1559::random(rng), + }); + let sighash = tx.signature_hash(); + let signature = signer.sign_hash_sync(&sighash).unwrap(); + TxEnvelope::new_unhashed(tx, signature).encoded_2718().into() + }) + .collect(); + + RawBundle { + txs, + metadata: RawBundleMetadata { + reverting_tx_hashes: vec![], + dropping_tx_hashes: vec![], + refund_tx_hashes: None, + signing_address: None, + version: Some("v2".to_string()), + block_number: None, + replacement_uuid: None, + refund_identity: None, + uuid: None, + min_timestamp: None, + max_timestamp: None, + replacement_nonce: Some(rng.random()), + refund_percent: Some(rng.random_range(0..100)), + refund_recipient: Some(Address::random_with(rng)), + delayed_refund: None, + bundle_hash: None, + }, + } + } + impl Random for RawBundle { /// Generate a random bundle with transactions of type Eip1559. fn random(rng: &mut R) -> Self { let txs_len = rng.random_range(1..=10); - // We only generate Eip1559 here. - let txs = (0..txs_len) - .map(|_| { - let signer = PrivateKeySigner::random(); - let tx = EthereumTypedTransaction::Eip1559(TxEip1559::random(rng)); - let sighash = tx.signature_hash(); - let signature = signer.sign_hash_sync(&sighash).unwrap(); - TxEnvelope::new_unhashed(tx, signature).encoded_2718().into() - }) - .collect(); + random_raw_bundle_with_tx_count_and_input_size(rng, txs_len, None) + } + } +} - Self { - txs, - metadata: RawBundleMetadata { - reverting_tx_hashes: vec![], - dropping_tx_hashes: vec![], - refund_tx_hashes: None, - signing_address: None, - version: Some("v2".to_string()), - block_number: None, - replacement_uuid: None, - refund_identity: None, - uuid: None, - min_timestamp: None, - max_timestamp: None, - replacement_nonce: Some(rng.random()), - refund_percent: Some(rng.random_range(0..100)), - refund_recipient: Some(Address::random_with(rng)), - delayed_refund: None, - bundle_hash: None, - }, +/// This time out should be enough for the inserter to flush all pending clickhouse data (timeout is +/// clickhouse usually a few secs) and local DB data (disk flush time). +pub const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(20); + +/// Consider move this to rbuilder-utils. +/// Waits for critical_tasks to finish by themselves up to grateful_timeout. +pub async fn wait_for_critical_tasks( + critical_tasks: Vec>, + grateful_timeout: Duration, +) { + let mut critical_tasks: FuturesUnordered<_> = critical_tasks.into_iter().collect(); + let critical_deadline = tokio::time::Instant::now() + grateful_timeout; + loop { + tokio::select! { + biased; + result = critical_tasks.next() => { + match result { + Some(Err(err)) => error!(?err, "Critical task handle await error"), + Some(Ok(())) => {} + None => { + info!("All critical tasks finished ok"); + break; + } + } + } + _ = tokio::time::sleep_until(critical_deadline) => { + error!(pending_task_count = critical_tasks.len(), "Critical tasks shutdown timeout reached"); + break; } } } diff --git a/tests/common/mod.rs b/tests/common/mod.rs index 25d15c01..9ba66292 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -10,18 +10,20 @@ use axum::{extract::State, routing::post, Router}; use flowproxy::{ cli::OrderflowIngressArgs, consts::FLASHBOTS_SIGNATURE_HEADER, + indexer::Indexer, ingress::maybe_decompress, jsonrpc::{JsonRpcError, JsonRpcRequest, JsonRpcResponse, JSONRPC_VERSION_2}, - runner::CliContext, }; use hyper::{header, HeaderMap}; use msg_socket::RepSocket; use msg_transport::tcp::Tcp; use rbuilder_primitives::serialize::RawBundle; +use rbuilder_utils::tasks::TaskManager; use revm_primitives::keccak256; use serde::de::DeserializeOwned; use serde_json::{json, Value}; use tokio::{net::TcpListener, sync::mpsc}; +use tokio_util::sync::CancellationToken; #[cfg(target_os = "linux")] use testcontainers::{ @@ -36,22 +38,27 @@ pub(crate) struct IngressClient { pub(crate) async fn spawn_ingress_with_args( args: OrderflowIngressArgs, + task_manager: &TaskManager, ) -> IngressClient { let user_listener = TcpListener::bind(&args.user_listen_addr).await.unwrap(); let mut system_listener = RepSocket::new(Tcp::default()); system_listener.bind(args.system_listen_addr).await.expect("to bind tcp socket address"); let builder_listener = None; let address = user_listener.local_addr().unwrap(); - - let task_manager = rbuilder_utils::tasks::TaskManager::current(); - + let task_executor = task_manager.executor(); tokio::spawn(async move { + let (indexer_handle, _indexer_join_handles) = + Indexer::run(args.indexing.clone(), args.builder_name.clone(), task_executor.clone()); + let cancellation_token = CancellationToken::new(); + flowproxy::run_with_listeners( args, user_listener, system_listener, builder_listener, - CliContext { task_executor: task_manager.executor() }, + task_executor, + indexer_handle, + cancellation_token, ) .await .unwrap(); @@ -64,11 +71,15 @@ pub(crate) async fn spawn_ingress_with_args( } } -pub(crate) async fn spawn_ingress(builder_url: Option) -> IngressClient { +#[allow(dead_code)] +pub(crate) async fn spawn_ingress( + builder_url: Option, + task_manager: &TaskManager, +) -> IngressClient { let mut args = OrderflowIngressArgs::default().gzip_enabled().disable_builder_hub(); args.peer_update_interval_s = 5; args.builder_url = builder_url; - spawn_ingress_with_args(args).await + spawn_ingress_with_args(args, task_manager).await } impl IngressClient { diff --git a/tests/ingress.rs b/tests/ingress.rs index fcbeb80a..09f4e44c 100644 --- a/tests/ingress.rs +++ b/tests/ingress.rs @@ -34,7 +34,8 @@ mod assert { async fn ingress_http_e2e() { let mut rng = rand::rng(); let mut builder = BuilderReceiver::spawn().await; - let client = spawn_ingress(Some(builder.url())).await; + let task_manager = rbuilder_utils::tasks::TaskManager::current(); + let client = spawn_ingress(Some(builder.url()), &task_manager).await; let empty = json!({}); let response = diff --git a/tests/network.rs b/tests/network.rs index 7e60ebcb..fb6fabd1 100644 --- a/tests/network.rs +++ b/tests/network.rs @@ -14,12 +14,12 @@ use tracing::{debug, info}; async fn network_e2e_bundle_tx_works() { let _ = tracing_subscriber::fmt::try_init(); info!("starting network e2e tcp test"); - + let task_manager = rbuilder_utils::tasks::TaskManager::current(); let mut rng = rand::rng(); let mut builder1 = BuilderReceiver::spawn().await; let mut builder2 = BuilderReceiver::spawn().await; - let client1 = spawn_ingress(Some(builder1.url())).await; - let client2 = spawn_ingress(Some(builder2.url())).await; + let client1 = spawn_ingress(Some(builder1.url()), &task_manager).await; + let client2 = spawn_ingress(Some(builder2.url()), &task_manager).await; // Wait for the proxies to be ready and connected to each other. tokio::time::sleep(Duration::from_secs(10)).await; @@ -104,7 +104,7 @@ mod linux { async fn network_e2e_tls() { let testdata_dir = PathBuf::from(concat!(env!("CARGO_MANIFEST_DIR"), "/tests/testdata")); let cert_dir = testdata_dir.join("certificates"); - + let task_manager = rbuilder_utils::tasks::TaskManager::current(); let signer1 = PrivateKeySigner::random(); let signer2 = PrivateKeySigner::random(); @@ -126,8 +126,8 @@ mod linux { args.builder_url = Some(builder1.url()); args2.builder_url = Some(builder2.url()); - let client1 = spawn_ingress_with_args(args).await; - let _client2 = spawn_ingress_with_args(args2).await; + let client1 = spawn_ingress_with_args(args, &task_manager).await; + let _client2 = spawn_ingress_with_args(args2, &task_manager).await; tokio::time::sleep(Duration::from_secs(1)).await; diff --git a/tests/testdata/openssl.cnf b/tests/testdata/openssl.cnf index a8142f1f..fc30d3cb 100644 --- a/tests/testdata/openssl.cnf +++ b/tests/testdata/openssl.cnf @@ -12,4 +12,5 @@ subjectAltName = @alt_names [alt_names] DNS.1 = localhost -IP.1 = 127.0.0.1 \ No newline at end of file +IP.1 = 127.0.0.1 +IP.2 = ::1