From 8b88e0a20080d888661831e7517c70b1698bf769 Mon Sep 17 00:00:00 2001 From: Dzmitry Kalabuk Date: Wed, 8 Jul 2026 13:56:24 +0300 Subject: [PATCH 1/9] Add query pipeline test harness Extract the compute-unit accounting (`process_query`) and log building (`generate_log`) into free functions generic over two dependency traits, `QueryRunner` and `CuChecker`, and add mock implementations. This lets the query pipeline be unit-tested without standing up the transport, storage, or a real query engine. Tests cover the current behavior: an overloaded engine is fully refunded, malformed params are rejected before the CU spend, partial chunks refund the unused fraction, and no-allocation queries are not logged. --- src/controller/mod.rs | 1 + src/controller/p2p.rs | 446 ++++++++++++++++++++++++----------- src/controller/query_deps.rs | 153 ++++++++++++ 3 files changed, 467 insertions(+), 133 deletions(-) create mode 100644 src/controller/query_deps.rs diff --git a/src/controller/mod.rs b/src/controller/mod.rs index 6b409dd..6139ed9 100644 --- a/src/controller/mod.rs +++ b/src/controller/mod.rs @@ -1,5 +1,6 @@ pub mod assignments; pub mod p2p; +pub mod query_deps; pub mod polars_target; pub mod sql_request; pub mod worker; diff --git a/src/controller/p2p.rs b/src/controller/p2p.rs index 8853bda..6db95c2 100644 --- a/src/controller/p2p.rs +++ b/src/controller/p2p.rs @@ -21,8 +21,8 @@ use tracing::{info, instrument, warn, Instrument}; use crate::{ cli::Args, compute_units::{ - self, allocations_checker::{self, AllocationsChecker}, + RateLimitStatus, }, controller::worker::QueryType, logs_storage::LogsStorage, @@ -33,6 +33,7 @@ use crate::{ util::{timestamp_now_ms, UseOnce}, }; +use super::query_deps::{CuChecker, QueryRunner}; use super::worker::Worker; const WORKER_VERSION: &str = env!("CARGO_PKG_VERSION"); @@ -519,7 +520,14 @@ impl + Send + 'static> P2PController + Send + 'static> P2PController tracing::error!("Couldn't send query result: {e:?}"), } - if let Some(log) = self.generate_log(&result, query, peer_id).await { + if let Some(log) = generate_log(&result, query, peer_id).await { if log.encoded_len() > MAX_LOGS_SIZE { warn!("Query log is too big: {log:?}"); return; @@ -558,92 +566,6 @@ impl + Send + 'static> P2PController (QueryResult, Option) { - match query.compression { - c if c == sqd_messages::Compression::Gzip as i32 - || c == sqd_messages::Compression::Zstd as i32 - || c == sqd_messages::Compression::None as i32 => {} - _ => { - return ( - Err(QueryError::BadRequest( - "Unsupported compression type".to_owned(), - )), - None, - ); - } - } - - let Some(block_range) = query - .block_range - .map(|sqd_messages::Range { begin, end }| (begin, end)) - else { - return ( - Err(QueryError::BadRequest("block_range is required".to_owned())), - None, - ); - }; - - let mut allocation_chip = 1.0f32; - - if let Ok(chunk) = query.chunk_id.parse::() { - let (begin, end) = block_range; - let active_len = std::cmp::min(chunk.last_block.into(), end) - .saturating_sub(std::cmp::max(chunk.first_block.into(), begin)) - .max(1); - let chunk_len = Into::::into(chunk.last_block) - .saturating_sub(chunk.first_block.into()) - .max(1); - allocation_chip = active_len as f32 / chunk_len as f32; - }; - - // We claim 1. allocation first and refund unused allocation later. It's done to prevent burst overloading with small requests. - let status = match self.allocations_checker.try_spend(peer_id, 1.) { - compute_units::RateLimitStatus::NoAllocation => { - // This error means that we don't have any allocation for particular peer_id at all (e.g. no allocation on contract) - return (Err(QueryError::NoAllocation), None); - } - compute_units::RateLimitStatus::Paused(retry_after) => { - // This error means that we don't have allocation at the moment, but it may be available after retry_after ms - return (Err(QueryError::NoAllocation), Some(retry_after)); - } - status => status, - }; - let mut retry_after = status.retry_after(); - - let result = self - .worker - .run_query( - &query.query, - query.dataset.clone(), - block_range, - &query.chunk_id, - Some(peer_id), - query_type, - ) - .await - .inspect_err(|err| tracing::error!("error processing query: {err}")); - - if let Err(QueryError::ServiceOverloaded) = result { - // Refund everything as we were not able to process request - self.allocations_checker.refund(peer_id, 1.); - retry_after = Some(DEFAULT_BACKOFF); - } else { - if allocation_chip < 1. { - // We refund unused allocation - self.allocations_checker - .refund(peer_id, 1. - allocation_chip); - } - } - (result, retry_after) - } - #[instrument(skip_all)] async fn send_query_result( &self, @@ -695,49 +617,6 @@ impl + Send + 'static> P2PController Option { - use query_executed::Result; - - let result = match query_result { - Ok(result) => Result::Ok(sqd_messages::QueryOkSummary { - uncompressed_data_size: result.data.len() as u64, - data_hash: result.sha3_256().await, - last_block: result.last_block, - }), - Err(QueryError::NoAllocation) => return None, - Err(e) => query_error::Err::from(e).into(), - }; - - let exec_time_report = match query_result { - Ok(result) => Some(TimeReport { - parsing_time_micros: result.time_report.parsing_time.as_micros() as u32, - execution_time_micros: result.time_report.execution_time.as_micros() as u32, - compression_time_micros: result.time_report.compression_time.as_micros() as u32, - signing_time_micros: result.time_report.signing_time.as_micros() as u32, - serialization_time_micros: result.time_report.serialization_time.as_micros() as u32, - }), - Err(_) => None, // TODO: always measure execution time - }; - - Some(QueryExecuted { - client_id: client_id.to_string(), - query: Some(query), - exec_time_micros: exec_time_report - .as_ref() - .map_or(0, |report| report.execution_time_micros), - exec_time_report, - timestamp_ms: timestamp_now_ms(), // TODO: use time of receiving query - result: Some(result), - worker_version: WORKER_VERSION.to_string(), - }) - } - async fn handle_logs_request( &self, request: LogsRequest, @@ -777,6 +656,133 @@ async fn wait_for_assignment_applied( } } +/// Run a query and account for its compute units: claim one up front, refund the unused chunk +/// fraction after. Generic over the engine and rate limiter for mock-based tests (see +/// [`super::query_deps`]). +#[instrument(skip_all)] +async fn process_query( + worker: &W, + allocations_checker: &A, + peer_id: PeerId, + query: &Query, + query_type: QueryType, +) -> (QueryResult, Option) { + match query.compression { + c if c == sqd_messages::Compression::Gzip as i32 + || c == sqd_messages::Compression::Zstd as i32 + || c == sqd_messages::Compression::None as i32 => {} + _ => { + return ( + Err(QueryError::BadRequest( + "Unsupported compression type".to_owned(), + )), + None, + ); + } + } + + let Some(block_range) = query + .block_range + .map(|sqd_messages::Range { begin, end }| (begin, end)) + else { + return ( + Err(QueryError::BadRequest("block_range is required".to_owned())), + None, + ); + }; + + let mut allocation_chip = 1.0f32; + + if let Ok(chunk) = query.chunk_id.parse::() { + let (begin, end) = block_range; + let active_len = std::cmp::min(chunk.last_block.into(), end) + .saturating_sub(std::cmp::max(chunk.first_block.into(), begin)) + .max(1); + let chunk_len = Into::::into(chunk.last_block) + .saturating_sub(chunk.first_block.into()) + .max(1); + allocation_chip = active_len as f32 / chunk_len as f32; + }; + + // We claim 1. allocation first and refund unused allocation later. It's done to prevent burst overloading with small requests. + let status = match allocations_checker.try_spend(peer_id, 1.) { + RateLimitStatus::NoAllocation => { + // This error means that we don't have any allocation for particular peer_id at all (e.g. no allocation on contract) + return (Err(QueryError::NoAllocation), None); + } + RateLimitStatus::Paused(retry_after) => { + // This error means that we don't have allocation at the moment, but it may be available after retry_after ms + return (Err(QueryError::NoAllocation), Some(retry_after)); + } + status => status, + }; + let mut retry_after = status.retry_after(); + + let result = worker + .run_query( + &query.query, + query.dataset.clone(), + block_range, + &query.chunk_id, + Some(peer_id), + query_type, + ) + .await + .inspect_err(|err| tracing::error!("error processing query: {err}")); + + if let Err(QueryError::ServiceOverloaded) = result { + // Refund everything as we were not able to process request + allocations_checker.refund(peer_id, 1.); + retry_after = Some(DEFAULT_BACKOFF); + } else if allocation_chip < 1. { + // We refund unused allocation + allocations_checker.refund(peer_id, 1. - allocation_chip); + } + (result, retry_after) +} + +#[instrument(skip_all)] +async fn generate_log( + query_result: &QueryResult, + query: Query, + client_id: PeerId, +) -> Option { + use query_executed::Result; + + let result = match query_result { + Ok(result) => Result::Ok(sqd_messages::QueryOkSummary { + uncompressed_data_size: result.data.len() as u64, + data_hash: result.sha3_256().await, + last_block: result.last_block, + }), + Err(QueryError::NoAllocation) => return None, + Err(e) => query_error::Err::from(e).into(), + }; + + let exec_time_report = match query_result { + Ok(result) => Some(TimeReport { + parsing_time_micros: result.time_report.parsing_time.as_micros() as u32, + execution_time_micros: result.time_report.execution_time.as_micros() as u32, + compression_time_micros: result.time_report.compression_time.as_micros() as u32, + signing_time_micros: result.time_report.signing_time.as_micros() as u32, + serialization_time_micros: result.time_report.serialization_time.as_micros() as u32, + }), + Err(_) => None, // TODO: always measure execution time + }; + + Some(QueryExecuted { + client_id: client_id.to_string(), + query: Some(query), + exec_time_micros: exec_time_report + .as_ref() + .map_or(0, |report| report.execution_time_micros), + exec_time_report, + timestamp_ms: timestamp_now_ms(), // TODO: use time of receiving query + result: Some(result), + worker_version: WORKER_VERSION.to_string(), + }) +} + #[tracing::instrument(skip_all)] async fn get_worker_status( worker: &Worker, @@ -890,7 +896,7 @@ fn check_peer_id(peer_id: PeerId, filename: PathBuf) { } #[cfg(all(test, feature = "mvcc-chunks"))] -mod tests { +mod assignment_tests { use super::*; #[test] @@ -927,3 +933,177 @@ mod tests { assert_eq!(pending.into_iter().collect::>(), vec![6]); } } + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use sqd_messages::{query_executed, Compression, Range}; + + use crate::{ + controller::query_deps::mocks::{MockChecker, MockWorker}, + query::result::QueryOk, + }; + + use super::*; + + // first_block = 100, last_block = 200 → chunk length 100. + const CHUNK_ID: &str = "0000000000/0000000100-0000000200-abcde"; + + fn query_with(compression: i32, block_range: Option<(u64, u64)>, chunk_id: &str) -> Query { + Query { + query_id: "00000000-0000-0000-0000-000000000000".to_owned(), + dataset: "s3://dataset".to_owned(), + query: "{}".to_owned(), + chunk_id: chunk_id.to_owned(), + block_range: block_range.map(|(begin, end)| Range { begin, end }), + compression, + ..Default::default() + } + } + + fn ok_result() -> QueryResult { + Ok(QueryOk::new( + b"hello".to_vec(), + 1, + 200, + Duration::ZERO, + Duration::ZERO, + Duration::ZERO, + )) + } + + // An overloaded engine is fully refunded, so the query consumes no compute unit. + #[tokio::test(flavor = "multi_thread")] + async fn overloaded_query_consumes_no_cu() { + let checker = MockChecker::new(RateLimitStatus::Spent(None)); + let worker = MockWorker::new(Err(QueryError::ServiceOverloaded)); + let query = query_with(Compression::None as i32, Some((100, 200)), CHUNK_ID); + + let (result, retry) = + process_query(&worker, &checker, PeerId::random(), &query, QueryType::PlainQuery).await; + + assert!(matches!(result, Err(QueryError::ServiceOverloaded))); + assert_eq!(retry, Some(DEFAULT_BACKOFF)); + assert_eq!(checker.net_spent(), 0.0); + assert_eq!(worker.calls(), 1); + } + + // Malformed params are rejected before the CU spend: no retry hint, no charge, engine not run. + #[tokio::test(flavor = "multi_thread")] + async fn malformed_query_is_not_charged() { + let checker = MockChecker::new(RateLimitStatus::Spent(Some(Duration::from_secs(5)))); + let worker = MockWorker::new(ok_result()); + let query = query_with(999, Some((100, 200)), CHUNK_ID); + + let (result, retry) = + process_query(&worker, &checker, PeerId::random(), &query, QueryType::PlainQuery).await; + + assert!(matches!(result, Err(QueryError::BadRequest(_)))); + assert_eq!(retry, None); + assert_eq!(checker.net_spent(), 0.0); + assert_eq!(worker.calls(), 0); + } + + #[tokio::test(flavor = "multi_thread")] + async fn successful_full_chunk_consumes_one_cu() { + let checker = MockChecker::new(RateLimitStatus::Spent(None)); + let worker = MockWorker::new(ok_result()); + let query = query_with(Compression::None as i32, Some((100, 200)), CHUNK_ID); + + let (result, retry) = + process_query(&worker, &checker, PeerId::random(), &query, QueryType::PlainQuery).await; + + assert!(result.is_ok()); + assert_eq!(retry, None); + assert_eq!(checker.net_spent(), 1.0); + } + + #[tokio::test(flavor = "multi_thread")] + async fn partial_chunk_refunds_unused_fraction() { + let checker = MockChecker::new(RateLimitStatus::Spent(None)); + let worker = MockWorker::new(ok_result()); + let query = query_with(Compression::None as i32, Some((100, 150)), CHUNK_ID); + + let (result, _) = + process_query(&worker, &checker, PeerId::random(), &query, QueryType::PlainQuery).await; + + assert!(result.is_ok()); + assert!( + (checker.net_spent() - 0.5).abs() < 1e-6, + "expected 0.5 CU, got {}", + checker.net_spent() + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn no_allocation_is_rejected_without_running() { + let checker = MockChecker::new(RateLimitStatus::NoAllocation); + let worker = MockWorker::new(ok_result()); + let query = query_with(Compression::None as i32, Some((100, 200)), CHUNK_ID); + + let (result, retry) = + process_query(&worker, &checker, PeerId::random(), &query, QueryType::PlainQuery).await; + + assert!(matches!(result, Err(QueryError::NoAllocation))); + assert_eq!(retry, None); + assert_eq!(worker.calls(), 0); + assert_eq!(checker.net_spent(), 0.0); + } + + #[tokio::test(flavor = "multi_thread")] + async fn rate_limited_is_rejected_with_retry() { + let retry = Duration::from_secs(2); + let checker = MockChecker::new(RateLimitStatus::Paused(retry)); + let worker = MockWorker::new(ok_result()); + let query = query_with(Compression::None as i32, Some((100, 200)), CHUNK_ID); + + let (result, got_retry) = + process_query(&worker, &checker, PeerId::random(), &query, QueryType::PlainQuery).await; + + assert!(matches!(result, Err(QueryError::NoAllocation))); + assert_eq!(got_retry, Some(retry)); + assert_eq!(worker.calls(), 0); + } + + #[tokio::test(flavor = "multi_thread")] + async fn no_allocation_is_not_logged() { + let query = query_with(Compression::None as i32, Some((100, 200)), CHUNK_ID); + let log = generate_log(&Err(QueryError::NoAllocation), query, PeerId::random()).await; + assert!(log.is_none()); + } + + #[tokio::test(flavor = "multi_thread")] + async fn log_for_success_summarizes_result() { + let query = query_with(Compression::None as i32, Some((100, 200)), CHUNK_ID); + let log = generate_log(&ok_result(), query, PeerId::random()) + .await + .expect("success should be logged"); + match log.result.expect("log should carry a result") { + query_executed::Result::Ok(summary) => { + assert_eq!(summary.last_block, 200); + assert_eq!(summary.uncompressed_data_size, 5); + assert!(!summary.data_hash.is_empty()); + } + other => panic!("expected an Ok summary, got {other:?}"), + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn log_for_error_records_the_error() { + let query = query_with(Compression::None as i32, Some((100, 200)), CHUNK_ID); + let log = generate_log( + &Err(QueryError::BadRequest("bad".to_owned())), + query, + PeerId::random(), + ) + .await + .expect("errors past the CU bar are logged"); + match log.result.expect("log should carry a result") { + query_executed::Result::Err(e) => { + assert!(matches!(e.err, Some(query_error::Err::BadRequest(_)))); + } + other => panic!("expected an Err result, got {other:?}"), + } + } +} diff --git a/src/controller/query_deps.rs b/src/controller/query_deps.rs new file mode 100644 index 0000000..f47041b --- /dev/null +++ b/src/controller/query_deps.rs @@ -0,0 +1,153 @@ +//! Dependency seams for the query-serving pipeline (`process_query`, `generate_log`) in +//! [`super::p2p`]. Production runs against the real `Worker` and `AllocationsChecker`; the traits +//! let tests inject an engine with a fixed outcome and a rate limiter that records what it spent, +//! with no transport or storage. + +use async_trait::async_trait; +use sqd_network_transport::PeerId; + +use crate::{ + compute_units::{allocations_checker::AllocationsChecker, RateLimitStatus}, + controller::worker::{QueryType, Worker}, + query::result::QueryResult, + types::dataset::Dataset, +}; + +/// Compute-unit rate limiting, behind a trait so tests can set a verdict and record spending. +pub trait CuChecker { + fn try_spend(&self, portal_id: PeerId, allocation_chip: f32) -> RateLimitStatus; + fn refund(&self, portal_id: PeerId, allocation_chip: f32); +} + +impl CuChecker for AllocationsChecker { + fn try_spend(&self, portal_id: PeerId, allocation_chip: f32) -> RateLimitStatus { + AllocationsChecker::try_spend(self, portal_id, allocation_chip) + } + + fn refund(&self, portal_id: PeerId, allocation_chip: f32) { + AllocationsChecker::refund(self, portal_id, allocation_chip) + } +} + +/// The query engine, behind a trait so tests can supply a fixed outcome. +#[async_trait] +pub trait QueryRunner { + async fn run_query( + &self, + query_str: &str, + dataset: Dataset, + block_range: (u64, u64), + chunk_id: &str, + client_id: Option, + query_type: QueryType, + ) -> QueryResult; +} + +#[async_trait] +impl QueryRunner for Worker { + async fn run_query( + &self, + query_str: &str, + dataset: Dataset, + block_range: (u64, u64), + chunk_id: &str, + client_id: Option, + query_type: QueryType, + ) -> QueryResult { + Worker::run_query( + self, + query_str, + dataset, + block_range, + chunk_id, + client_id, + query_type, + ) + .await + } +} + +#[cfg(test)] +pub mod mocks { + use std::sync::atomic::{AtomicU32, Ordering}; + + use async_trait::async_trait; + use parking_lot::Mutex; + use sqd_network_transport::PeerId; + + use crate::{ + compute_units::RateLimitStatus, controller::worker::QueryType, query::result::QueryResult, + types::dataset::Dataset, + }; + + use super::{CuChecker, QueryRunner}; + + /// Returns a fixed verdict and tracks net compute units. Like the real checker, only a `Spent` + /// verdict deducts. + pub struct MockChecker { + verdict: RateLimitStatus, + net_spent: Mutex, + } + + impl MockChecker { + pub fn new(verdict: RateLimitStatus) -> Self { + Self { + verdict, + net_spent: Mutex::new(0.0), + } + } + + /// Net compute units charged to the peer: spent minus refunded. + pub fn net_spent(&self) -> f32 { + *self.net_spent.lock() + } + } + + impl CuChecker for MockChecker { + fn try_spend(&self, _portal_id: PeerId, allocation_chip: f32) -> RateLimitStatus { + if let RateLimitStatus::Spent(_) = self.verdict { + *self.net_spent.lock() += allocation_chip; + } + self.verdict + } + + fn refund(&self, _portal_id: PeerId, allocation_chip: f32) { + *self.net_spent.lock() -= allocation_chip; + } + } + + /// Returns a fixed outcome and counts its calls. + pub struct MockWorker { + outcome: QueryResult, + calls: AtomicU32, + } + + impl MockWorker { + pub fn new(outcome: QueryResult) -> Self { + Self { + outcome, + calls: AtomicU32::new(0), + } + } + + pub fn calls(&self) -> u32 { + self.calls.load(Ordering::SeqCst) + } + } + + #[async_trait] + impl QueryRunner for MockWorker { + async fn run_query( + &self, + _query_str: &str, + _dataset: Dataset, + _block_range: (u64, u64), + _chunk_id: &str, + _client_id: Option, + _query_type: QueryType, + ) -> QueryResult { + self.calls.fetch_add(1, Ordering::SeqCst); + self.outcome.clone() + } + } +} From 4609d937ff74d7b16d46a9119084a6f0dbf48f0f Mon Sep 17 00:00:00 2001 From: Dzmitry Kalabuk Date: Tue, 7 Jul 2026 15:38:46 +0300 Subject: [PATCH 2/9] Migrate worker to stream_server transport Bump sqd-network to the stream_server server behaviour (libp2p-stream instead of request_response): responses are written to the stream directly, so the send_* handle methods are now async. On a full processing queue the worker now returns a signed too_many_requests error instead of dropping the stream, which had left the client with an opaque transport error. Co-Authored-By: Claude Opus 4.8 (1M context) --- Cargo.lock | 57 +++++++++++++++--------------- Cargo.toml | 17 ++++++--- src/controller/p2p.rs | 81 ++++++++++++++++++++++++++----------------- 3 files changed, 89 insertions(+), 66 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 34cdc38..f6fa76f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3622,7 +3622,7 @@ dependencies = [ [[package]] name = "libp2p" version = "0.56.0" -source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=c0ed330#c0ed3308db3a7deda74a688e13b69007942c41ab" +source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=fac3b0c9#fac3b0c97757d6d8d6042f5423d9467a94ea80b0" dependencies = [ "bytes", "either", @@ -3655,7 +3655,7 @@ dependencies = [ [[package]] name = "libp2p-allow-block-list" version = "0.6.0" -source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=c0ed330#c0ed3308db3a7deda74a688e13b69007942c41ab" +source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=fac3b0c9#fac3b0c97757d6d8d6042f5423d9467a94ea80b0" dependencies = [ "libp2p-core", "libp2p-identity", @@ -3665,7 +3665,7 @@ dependencies = [ [[package]] name = "libp2p-autonat" version = "0.15.0" -source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=c0ed330#c0ed3308db3a7deda74a688e13b69007942c41ab" +source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=fac3b0c9#fac3b0c97757d6d8d6042f5423d9467a94ea80b0" dependencies = [ "async-trait", "asynchronous-codec", @@ -3689,7 +3689,7 @@ dependencies = [ [[package]] name = "libp2p-connection-limits" version = "0.6.0" -source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=c0ed330#c0ed3308db3a7deda74a688e13b69007942c41ab" +source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=fac3b0c9#fac3b0c97757d6d8d6042f5423d9467a94ea80b0" dependencies = [ "libp2p-core", "libp2p-identity", @@ -3699,7 +3699,7 @@ dependencies = [ [[package]] name = "libp2p-core" version = "0.43.1" -source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=c0ed330#c0ed3308db3a7deda74a688e13b69007942c41ab" +source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=fac3b0c9#fac3b0c97757d6d8d6042f5423d9467a94ea80b0" dependencies = [ "either", "fnv", @@ -3723,7 +3723,7 @@ dependencies = [ [[package]] name = "libp2p-dns" version = "0.44.0" -source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=c0ed330#c0ed3308db3a7deda74a688e13b69007942c41ab" +source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=fac3b0c9#fac3b0c97757d6d8d6042f5423d9467a94ea80b0" dependencies = [ "async-trait", "futures", @@ -3738,7 +3738,7 @@ dependencies = [ [[package]] name = "libp2p-gossipsub" version = "0.50.0" -source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=c0ed330#c0ed3308db3a7deda74a688e13b69007942c41ab" +source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=fac3b0c9#fac3b0c97757d6d8d6042f5423d9467a94ea80b0" dependencies = [ "async-channel", "asynchronous-codec", @@ -3768,7 +3768,7 @@ dependencies = [ [[package]] name = "libp2p-identify" version = "0.47.0" -source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=c0ed330#c0ed3308db3a7deda74a688e13b69007942c41ab" +source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=fac3b0c9#fac3b0c97757d6d8d6042f5423d9467a94ea80b0" dependencies = [ "asynchronous-codec", "either", @@ -3788,8 +3788,7 @@ dependencies = [ [[package]] name = "libp2p-identity" version = "0.2.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3104e13b51e4711ff5738caa1fb54467c8604c2e94d607e27745bcf709068774" +source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=fac3b0c9#fac3b0c97757d6d8d6042f5423d9467a94ea80b0" dependencies = [ "bs58", "ed25519-dalek", @@ -3807,7 +3806,7 @@ dependencies = [ [[package]] name = "libp2p-kad" version = "0.49.0" -source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=c0ed330#c0ed3308db3a7deda74a688e13b69007942c41ab" +source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=fac3b0c9#fac3b0c97757d6d8d6042f5423d9467a94ea80b0" dependencies = [ "asynchronous-codec", "bytes", @@ -3834,7 +3833,7 @@ dependencies = [ [[package]] name = "libp2p-mdns" version = "0.48.0" -source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=c0ed330#c0ed3308db3a7deda74a688e13b69007942c41ab" +source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=fac3b0c9#fac3b0c97757d6d8d6042f5423d9467a94ea80b0" dependencies = [ "futures", "hickory-proto", @@ -3852,7 +3851,7 @@ dependencies = [ [[package]] name = "libp2p-metrics" version = "0.17.0" -source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=c0ed330#c0ed3308db3a7deda74a688e13b69007942c41ab" +source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=fac3b0c9#fac3b0c97757d6d8d6042f5423d9467a94ea80b0" dependencies = [ "futures", "libp2p-core", @@ -3869,7 +3868,7 @@ dependencies = [ [[package]] name = "libp2p-ping" version = "0.47.0" -source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=c0ed330#c0ed3308db3a7deda74a688e13b69007942c41ab" +source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=fac3b0c9#fac3b0c97757d6d8d6042f5423d9467a94ea80b0" dependencies = [ "futures", "futures-timer", @@ -3884,7 +3883,7 @@ dependencies = [ [[package]] name = "libp2p-quic" version = "0.13.0" -source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=c0ed330#c0ed3308db3a7deda74a688e13b69007942c41ab" +source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=fac3b0c9#fac3b0c97757d6d8d6042f5423d9467a94ea80b0" dependencies = [ "futures", "futures-timer", @@ -3905,7 +3904,7 @@ dependencies = [ [[package]] name = "libp2p-request-response" version = "0.29.0" -source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=c0ed330#c0ed3308db3a7deda74a688e13b69007942c41ab" +source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=fac3b0c9#fac3b0c97757d6d8d6042f5423d9467a94ea80b0" dependencies = [ "async-trait", "futures", @@ -3921,7 +3920,7 @@ dependencies = [ [[package]] name = "libp2p-stream" version = "0.4.0-alpha" -source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=c0ed330#c0ed3308db3a7deda74a688e13b69007942c41ab" +source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=fac3b0c9#fac3b0c97757d6d8d6042f5423d9467a94ea80b0" dependencies = [ "futures", "libp2p-core", @@ -3934,7 +3933,7 @@ dependencies = [ [[package]] name = "libp2p-swarm" version = "0.47.0" -source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=c0ed330#c0ed3308db3a7deda74a688e13b69007942c41ab" +source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=fac3b0c9#fac3b0c97757d6d8d6042f5423d9467a94ea80b0" dependencies = [ "either", "fnv", @@ -3954,7 +3953,7 @@ dependencies = [ [[package]] name = "libp2p-swarm-derive" version = "0.35.1" -source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=c0ed330#c0ed3308db3a7deda74a688e13b69007942c41ab" +source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=fac3b0c9#fac3b0c97757d6d8d6042f5423d9467a94ea80b0" dependencies = [ "heck 0.5.0", "quote", @@ -3964,7 +3963,7 @@ dependencies = [ [[package]] name = "libp2p-tcp" version = "0.44.0" -source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=c0ed330#c0ed3308db3a7deda74a688e13b69007942c41ab" +source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=fac3b0c9#fac3b0c97757d6d8d6042f5423d9467a94ea80b0" dependencies = [ "futures", "futures-timer", @@ -3979,7 +3978,7 @@ dependencies = [ [[package]] name = "libp2p-tls" version = "0.6.2" -source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=c0ed330#c0ed3308db3a7deda74a688e13b69007942c41ab" +source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=fac3b0c9#fac3b0c97757d6d8d6042f5423d9467a94ea80b0" dependencies = [ "futures", "futures-rustls", @@ -3997,7 +3996,7 @@ dependencies = [ [[package]] name = "libp2p-upnp" version = "0.5.1" -source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=c0ed330#c0ed3308db3a7deda74a688e13b69007942c41ab" +source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=fac3b0c9#fac3b0c97757d6d8d6042f5423d9467a94ea80b0" dependencies = [ "futures", "futures-timer", @@ -4281,7 +4280,7 @@ checksum = "defc4c55412d89136f966bbb339008b474350e5e6e78d2714439c386b3137a03" [[package]] name = "multistream-select" version = "0.13.0" -source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=c0ed330#c0ed3308db3a7deda74a688e13b69007942c41ab" +source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=fac3b0c9#fac3b0c97757d6d8d6042f5423d9467a94ea80b0" dependencies = [ "bytes", "futures", @@ -5834,7 +5833,7 @@ dependencies = [ [[package]] name = "quick-protobuf-codec" version = "0.3.1" -source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=c0ed330#c0ed3308db3a7deda74a688e13b69007942c41ab" +source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=fac3b0c9#fac3b0c97757d6d8d6042f5423d9467a94ea80b0" dependencies = [ "asynchronous-codec", "bytes", @@ -6488,7 +6487,7 @@ checksum = "0e819f2bc632f285be6d7cd36e25940d45b2391dd6d9b939e79de557f7014248" [[package]] name = "rw-stream-sink" version = "0.4.0" -source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=c0ed330#c0ed3308db3a7deda74a688e13b69007942c41ab" +source = "git+https://github.com/kalabukdima/rust-libp2p.git?rev=fac3b0c9#fac3b0c97757d6d8d6042f5423d9467a94ea80b0" dependencies = [ "futures", "pin-project", @@ -7094,7 +7093,7 @@ dependencies = [ [[package]] name = "sqd-assignments" version = "0.1.0" -source = "git+https://github.com/subsquid/sqd-network.git?rev=292e67f#292e67f39d392592de609dda6323b1c2e638d271" +source = "git+https://github.com/subsquid/sqd-network.git?rev=541684c#541684c1ad9cd25a6c1423ec699eb8b6dcf65d1e" dependencies = [ "anyhow", "base64 0.22.1", @@ -7122,7 +7121,7 @@ dependencies = [ [[package]] name = "sqd-contract-client" version = "1.3.0" -source = "git+https://github.com/subsquid/sqd-network.git?rev=292e67f#292e67f39d392592de609dda6323b1c2e638d271" +source = "git+https://github.com/subsquid/sqd-network.git?rev=541684c#541684c1ad9cd25a6c1423ec699eb8b6dcf65d1e" dependencies = [ "async-trait", "clap", @@ -7143,7 +7142,7 @@ dependencies = [ [[package]] name = "sqd-messages" version = "2.1.0" -source = "git+https://github.com/subsquid/sqd-network.git?rev=292e67f#292e67f39d392592de609dda6323b1c2e638d271" +source = "git+https://github.com/subsquid/sqd-network.git?rev=541684c#541684c1ad9cd25a6c1423ec699eb8b6dcf65d1e" dependencies = [ "bytemuck", "flate2", @@ -7159,7 +7158,7 @@ dependencies = [ [[package]] name = "sqd-network-transport" version = "3.1.0" -source = "git+https://github.com/subsquid/sqd-network.git?rev=292e67f#292e67f39d392592de609dda6323b1c2e638d271" +source = "git+https://github.com/subsquid/sqd-network.git?rev=541684c#541684c1ad9cd25a6c1423ec699eb8b6dcf65d1e" dependencies = [ "anyhow", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index badd665..804ac74 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -61,10 +61,10 @@ url = "2.5.2" walkdir = "2.5.0" zstd = "0.13" -sqd-assignments = { git = "https://github.com/subsquid/sqd-network.git", rev = "292e67f", features = ["reader"] } -sqd-contract-client = { git = "https://github.com/subsquid/sqd-network.git", rev = "292e67f", version = "1.2.1" } -sqd-messages = { git = "https://github.com/subsquid/sqd-network.git", rev = "292e67f", version = "2.0.2", features = ["bitstring"] } -sqd-network-transport = { git = "https://github.com/subsquid/sqd-network.git", rev = "292e67f", version = "3.0.0", features = ["worker", "metrics"] } +sqd-assignments = { git = "https://github.com/subsquid/sqd-network.git", rev = "541684c", features = ["reader"] } +sqd-contract-client = { git = "https://github.com/subsquid/sqd-network.git", rev = "541684c", version = "1.2.1" } +sqd-messages = { git = "https://github.com/subsquid/sqd-network.git", rev = "541684c", version = "2.0.2", features = ["bitstring"] } +sqd-network-transport = { git = "https://github.com/subsquid/sqd-network.git", rev = "541684c", version = "3.0.0", features = ["worker", "metrics"] } sqd-query = { git = "https://github.com/subsquid/data.git", rev = "b2c59d9", features = ["parquet"] } sqd-polars = { git = "https://github.com/subsquid/data.git", rev = "b2c59d9" } @@ -72,7 +72,14 @@ sqd-polars = { git = "https://github.com/subsquid/data.git", rev = "b2c59d9" } sql_query_plan = {git = "https://github.com/subsquid/qplan.git", rev = "658f88f" } [dev-dependencies] -sqd-assignments = { git = "https://github.com/subsquid/sqd-network.git", rev = "292e67f", features = ["builder"] } +sqd-assignments = { git = "https://github.com/subsquid/sqd-network.git", rev = "541684c", features = ["builder"] } + +# sqd-network-transport depends on the kalabukdima/rust-libp2p fork, whose crates build against their +# in-tree libp2p-identity (path = "identity"). Redirect the crates.io libp2p-identity to the same +# source, otherwise two incompatible copies of libp2p_identity are linked. This mirrors the +# [patch.crates-io] in the sqd-network workspace, which is not inherited by downstream crates. +[patch.crates-io] +libp2p-identity = { git = "https://github.com/kalabukdima/rust-libp2p.git", rev = "fac3b0c9" } [profile.release] debug = true diff --git a/src/controller/p2p.rs b/src/controller/p2p.rs index 6db95c2..2e834f8 100644 --- a/src/controller/p2p.rs +++ b/src/controller/p2p.rs @@ -6,12 +6,12 @@ use camino::Utf8PathBuf as PathBuf; use futures::{FutureExt, Stream, StreamExt}; use parking_lot::RwLock; use sqd_messages::{ - query_error, query_executed, BitString, LogsRequest, ProstMsg, Query, QueryExecuted, QueryLogs, + query_error, query_executed, BitString, LogsRequest, ProstMsg, Query, QueryExecuted, TimeReport, WorkerStatus, }; use sqd_network_transport::{ - protocol, Keypair, P2PTransportBuilder, PeerId, QueueFull, ResponseChannel, WorkerConfig, - WorkerEvent, WorkerTransportHandle, + protocol, Keypair, P2PTransportBuilder, PeerId, ResponseSender, WorkerConfig, WorkerEvent, + WorkerTransportHandle, }; use tokio::{sync::mpsc, time::MissedTickBehavior}; use tokio_stream::wrappers::{IntervalStream, ReceiverStream}; @@ -62,14 +62,12 @@ pub struct P2PController { worker_id: PeerId, keypair: Keypair, assignment_url: String, - queries_tx: mpsc::Sender<(PeerId, Query, ResponseChannel)>, - queries_rx: - UseOnce)>>, - sql_queries_tx: mpsc::Sender<(PeerId, Query, ResponseChannel)>, - sql_queries_rx: - UseOnce)>>, - log_requests_tx: mpsc::Sender<(LogsRequest, ResponseChannel)>, - log_requests_rx: UseOnce)>>, + queries_tx: mpsc::Sender<(PeerId, Query, ResponseSender)>, + queries_rx: UseOnce>, + sql_queries_tx: mpsc::Sender<(PeerId, Query, ResponseSender)>, + sql_queries_rx: UseOnce>, + log_requests_tx: mpsc::Sender<(LogsRequest, ResponseSender)>, + log_requests_rx: UseOnce>, } pub async fn create_p2p_controller( @@ -91,11 +89,7 @@ pub async fn create_p2p_controller( let worker_status = get_worker_status(&worker, allocations_checker.current_epoch()).await; - let mut config = WorkerConfig::default(); - config.status_queue_size = std::env::var("WORKER_STATUS_QUEUE_SIZE") - .ok() - .and_then(|s| s.parse::().ok()) - .unwrap_or(1000); + let config = WorkerConfig::default(); let (event_stream, transport_handle) = transport_builder.build_worker(config).await?; let (queries_tx, queries_rx) = mpsc::channel(QUERIES_POOL_SIZE); @@ -416,7 +410,7 @@ impl + Send + 'static> P2PController, cancellation_token: CancellationToken) { let event_stream = self .raw_event_stream .take() @@ -436,8 +430,9 @@ impl + Send + 'static> P2PController {} - Err(mpsc::error::TrySendError::Full(_)) => { - warn!("Queries queue is full. Dropping query from {peer_id}"); + Err(mpsc::error::TrySendError::Full((_, query, resp_chan))) => { + warn!("Queries queue is full. Rejecting query from {peer_id}"); + self.clone().reject_overloaded(query, resp_chan); } Err(mpsc::error::TrySendError::Closed(_)) => { break; @@ -454,8 +449,9 @@ impl + Send + 'static> P2PController {} - Err(mpsc::error::TrySendError::Full(_)) => { - warn!("SQL Queries queue is full. Dropping query from {peer_id}"); + Err(mpsc::error::TrySendError::Full((_, query, resp_chan))) => { + warn!("SQL Queries queue is full. Rejecting query from {peer_id}"); + self.clone().reject_overloaded(query, resp_chan); } Err(mpsc::error::TrySendError::Closed(_)) => { break; @@ -475,12 +471,12 @@ impl + Send + 'static> P2PController { let status = self.worker_status.read().clone(); - match self.transport_handle.send_status(status, resp_chan) { - Ok(_) => {} - Err(QueueFull) => { - warn!("Couldn't respond with status to {peer_id}: out queue full"); + let this = self.clone(); + tokio::spawn(async move { + if let Err(e) = this.transport_handle.send_status(status, resp_chan).await { + warn!("Couldn't respond with status to {peer_id}: {e:?}"); } - } + }); } } } @@ -488,6 +484,27 @@ impl + Send + 'static> P2PController, query: Query, resp_chan: ResponseSender) { + tokio::spawn(async move { + let mut msg = sqd_messages::QueryResult { + query_id: query.query_id, + result: Some(query_error::Err::TooManyRequests(()).into()), + retry_after_ms: Some(DEFAULT_BACKOFF.as_millis() as u32), + signature: Default::default(), + }; + if let Err(e) = msg.sign(&self.keypair) { + warn!("Couldn't sign too_many_requests response: {e}"); + return; + } + if let Err(e) = self.transport_handle.send_query_result(msg, resp_chan).await { + warn!("Couldn't send too_many_requests response: {e:?}"); + } + }); + } + #[instrument(skip_all)] fn validate_query(&self, query: &Query, peer_id: PeerId) -> bool { if !query.verify_signature(peer_id, self.worker_id) { @@ -514,7 +531,7 @@ impl + Send + 'static> P2PController, + resp_chan: ResponseSender, query_type: QueryType, ) { let query_id = query.query_id.clone(); @@ -571,7 +588,7 @@ impl + Send + 'static> P2PController, + resp_chan: ResponseSender, retry_after: Option, compression: sqd_messages::Compression, ) -> Result<(Duration, Duration)> { @@ -609,10 +626,10 @@ impl + Send + 'static> P2PController + Send + 'static> P2PController, + resp_chan: ResponseSender, ) -> Result<()> { let msg = self .logs_storage @@ -636,7 +653,7 @@ impl + Send + 'static> P2PController Date: Wed, 8 Jul 2026 10:38:25 +0300 Subject: [PATCH 3/9] Decode and encode network protobufs in the worker The transport now hands over raw request bytes and takes raw response bytes. Decode Query/LogsRequest in the event loop (off the swarm loop) and encode responses via ResponseSender::send directly, replacing the removed WorkerTransportHandle::send_* methods. Co-Authored-By: Claude Opus 4.8 (1M context) --- Cargo.lock | 8 +++---- Cargo.toml | 10 ++++---- src/controller/p2p.rs | 53 ++++++++++++++++++++++++++++++++----------- 3 files changed, 49 insertions(+), 22 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f6fa76f..a04f50f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7093,7 +7093,7 @@ dependencies = [ [[package]] name = "sqd-assignments" version = "0.1.0" -source = "git+https://github.com/subsquid/sqd-network.git?rev=541684c#541684c1ad9cd25a6c1423ec699eb8b6dcf65d1e" +source = "git+https://github.com/subsquid/sqd-network.git?rev=c4b85c1#c4b85c132fe5fa308be284556d539f012c0760b0" dependencies = [ "anyhow", "base64 0.22.1", @@ -7121,7 +7121,7 @@ dependencies = [ [[package]] name = "sqd-contract-client" version = "1.3.0" -source = "git+https://github.com/subsquid/sqd-network.git?rev=541684c#541684c1ad9cd25a6c1423ec699eb8b6dcf65d1e" +source = "git+https://github.com/subsquid/sqd-network.git?rev=c4b85c1#c4b85c132fe5fa308be284556d539f012c0760b0" dependencies = [ "async-trait", "clap", @@ -7142,7 +7142,7 @@ dependencies = [ [[package]] name = "sqd-messages" version = "2.1.0" -source = "git+https://github.com/subsquid/sqd-network.git?rev=541684c#541684c1ad9cd25a6c1423ec699eb8b6dcf65d1e" +source = "git+https://github.com/subsquid/sqd-network.git?rev=c4b85c1#c4b85c132fe5fa308be284556d539f012c0760b0" dependencies = [ "bytemuck", "flate2", @@ -7158,7 +7158,7 @@ dependencies = [ [[package]] name = "sqd-network-transport" version = "3.1.0" -source = "git+https://github.com/subsquid/sqd-network.git?rev=541684c#541684c1ad9cd25a6c1423ec699eb8b6dcf65d1e" +source = "git+https://github.com/subsquid/sqd-network.git?rev=c4b85c1#c4b85c132fe5fa308be284556d539f012c0760b0" dependencies = [ "anyhow", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index 804ac74..9dbc478 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -61,10 +61,10 @@ url = "2.5.2" walkdir = "2.5.0" zstd = "0.13" -sqd-assignments = { git = "https://github.com/subsquid/sqd-network.git", rev = "541684c", features = ["reader"] } -sqd-contract-client = { git = "https://github.com/subsquid/sqd-network.git", rev = "541684c", version = "1.2.1" } -sqd-messages = { git = "https://github.com/subsquid/sqd-network.git", rev = "541684c", version = "2.0.2", features = ["bitstring"] } -sqd-network-transport = { git = "https://github.com/subsquid/sqd-network.git", rev = "541684c", version = "3.0.0", features = ["worker", "metrics"] } +sqd-assignments = { git = "https://github.com/subsquid/sqd-network.git", rev = "c4b85c1", features = ["reader"] } +sqd-contract-client = { git = "https://github.com/subsquid/sqd-network.git", rev = "c4b85c1", version = "1.2.1" } +sqd-messages = { git = "https://github.com/subsquid/sqd-network.git", rev = "c4b85c1", version = "2.0.2", features = ["bitstring"] } +sqd-network-transport = { git = "https://github.com/subsquid/sqd-network.git", rev = "c4b85c1", version = "3.0.0", features = ["worker", "metrics"] } sqd-query = { git = "https://github.com/subsquid/data.git", rev = "b2c59d9", features = ["parquet"] } sqd-polars = { git = "https://github.com/subsquid/data.git", rev = "b2c59d9" } @@ -72,7 +72,7 @@ sqd-polars = { git = "https://github.com/subsquid/data.git", rev = "b2c59d9" } sql_query_plan = {git = "https://github.com/subsquid/qplan.git", rev = "658f88f" } [dev-dependencies] -sqd-assignments = { git = "https://github.com/subsquid/sqd-network.git", rev = "541684c", features = ["builder"] } +sqd-assignments = { git = "https://github.com/subsquid/sqd-network.git", rev = "c4b85c1", features = ["builder"] } # sqd-network-transport depends on the kalabukdima/rust-libp2p fork, whose crates build against their # in-tree libp2p-identity (path = "identity"). Redirect the crates.io libp2p-identity to the same diff --git a/src/controller/p2p.rs b/src/controller/p2p.rs index 2e834f8..4f85fa4 100644 --- a/src/controller/p2p.rs +++ b/src/controller/p2p.rs @@ -56,7 +56,9 @@ pub struct P2PController { assignment_fetch_timeout: Duration, assignment_fetch_max_delay: Duration, raw_event_stream: UseOnce, - transport_handle: WorkerTransportHandle, + // Keeps the transport running: it shuts down when the last handle drops. Responses go through + // each request's `ResponseSender`, so this handle is only held, never called. + _transport_handle: WorkerTransportHandle, logs_storage: LogsStorage, allocations_checker: AllocationsChecker, worker_id: PeerId, @@ -103,7 +105,7 @@ pub async fn create_p2p_controller( assignment_fetch_timeout: args.assignment_fetch_timeout, assignment_fetch_max_delay: args.assignment_fetch_max_delay, raw_event_stream: UseOnce::new(event_stream), - transport_handle, + _transport_handle: transport_handle, logs_storage: LogsStorage::new(args.data_dir.join("logs.db").as_str()).await?, allocations_checker, worker_id, @@ -422,9 +424,16 @@ impl + Send + 'static> P2PController { + let query = match Query::decode(request.as_ref()) { + Ok(query) => query, + Err(e) => { + warn!("Failed to decode query from {peer_id}: {e}"); + continue; + } + }; if !self.validate_query(&query, peer_id) { continue; } @@ -441,9 +450,16 @@ impl + Send + 'static> P2PController { + let query = match Query::decode(request.as_ref()) { + Ok(query) => query, + Err(e) => { + warn!("Failed to decode SQL query from {peer_id}: {e}"); + continue; + } + }; if !self.validate_query(&query, peer_id) { continue; } @@ -458,7 +474,18 @@ impl + Send + 'static> P2PController { + WorkerEvent::LogsRequest { + peer_id, + request, + resp_chan, + } => { + let request = match LogsRequest::decode(request.as_ref()) { + Ok(request) => request, + Err(e) => { + warn!("Failed to decode logs request from {peer_id}: {e}"); + continue; + } + }; match self.log_requests_tx.try_send((request, resp_chan)) { Ok(_) => {} Err(mpsc::error::TrySendError::Full(_)) => { @@ -471,9 +498,9 @@ impl + Send + 'static> P2PController { let status = self.worker_status.read().clone(); - let this = self.clone(); tokio::spawn(async move { - if let Err(e) = this.transport_handle.send_status(status, resp_chan).await { + tracing::debug!("Sending worker status to {peer_id}"); + if let Err(e) = resp_chan.send(&status.encode_to_vec()).await { warn!("Couldn't respond with status to {peer_id}: {e:?}"); } }); @@ -485,8 +512,8 @@ impl + Send + 'static> P2PController, query: Query, resp_chan: ResponseSender) { tokio::spawn(async move { let mut msg = sqd_messages::QueryResult { @@ -499,7 +526,7 @@ impl + Send + 'static> P2PController + Send + 'static> P2PController + Send + 'static> P2PController Date: Wed, 8 Jul 2026 13:35:15 +0300 Subject: [PATCH 4/9] Tie query logging to compute-unit spend The log entry a worker stores (and gets rewarded on) was built independently of the response the client received, so an oversized or unsignable result was sent as a `server_error` while still being logged as a successful query. Make the compute-unit spend the single boundary: a query is logged iff it consumed a CU, and the log is a projection of the response that was actually sent. - Move admission (verify, reserve slot, spend one CU) into the event loop; reject pre-admission queries with signed typed errors and a backoff hint instead of silently dropping the stream - Build the wire response and its log from one outcome in `build_delivery`, so an oversized/unsignable result downgrades to `server_error` in both - Charge and log malformed params and engine overload; only unprovable or unbudgeted rejects (bad signature/timestamp, no allocation, rate limit) stay out of the log store, which bounds a log-flood DoS - Bound concurrent reject responses with a semaphore --- src/compute_units/rate_limiter.rs | 10 - src/controller/p2p.rs | 723 ++++++++++++++++++------------ src/controller/query_deps.rs | 24 +- src/query/result.rs | 5 +- 4 files changed, 448 insertions(+), 314 deletions(-) diff --git a/src/compute_units/rate_limiter.rs b/src/compute_units/rate_limiter.rs index 3b1765c..6ae6525 100644 --- a/src/compute_units/rate_limiter.rs +++ b/src/compute_units/rate_limiter.rs @@ -17,16 +17,6 @@ pub enum RateLimitStatus { NoAllocation, } -impl RateLimitStatus { - pub fn retry_after(&self) -> Option { - match self { - RateLimitStatus::Spent(retry_after) => *retry_after, - RateLimitStatus::Paused(retry_after) => Some(*retry_after), - RateLimitStatus::NoAllocation => None, - } - } -} - const MAX_TOKENS: f32 = 3.0f32; struct Bucket { diff --git a/src/controller/p2p.rs b/src/controller/p2p.rs index 4f85fa4..103e503 100644 --- a/src/controller/p2p.rs +++ b/src/controller/p2p.rs @@ -1,7 +1,7 @@ use std::collections::VecDeque; use std::{env, sync::Arc, time::Duration}; -use anyhow::{anyhow, Result}; +use anyhow::Result; use camino::Utf8PathBuf as PathBuf; use futures::{FutureExt, Stream, StreamExt}; use parking_lot::RwLock; @@ -13,7 +13,10 @@ use sqd_network_transport::{ protocol, Keypair, P2PTransportBuilder, PeerId, ResponseSender, WorkerConfig, WorkerEvent, WorkerTransportHandle, }; -use tokio::{sync::mpsc, time::MissedTickBehavior}; +use tokio::{ + sync::{mpsc, Semaphore}, + time::MissedTickBehavior, +}; use tokio_stream::wrappers::{IntervalStream, ReceiverStream}; use tokio_util::sync::CancellationToken; use tracing::{info, instrument, warn, Instrument}; @@ -41,6 +44,9 @@ const LOG_REQUESTS_QUEUE_SIZE: usize = 4; const QUERIES_POOL_SIZE: usize = 16; const CONCURRENT_QUERY_MESSAGES: usize = 32; const DEFAULT_BACKOFF: Duration = Duration::from_secs(1); +/// Caps concurrent reject sends. Past this, the response is dropped — a cheap stream reset — rather +/// than spawning unbounded signing tasks under a flood. +const MAX_CONCURRENT_REJECTS: usize = 64; const LOGS_KEEP_DURATION: Duration = Duration::from_secs(3600 * 2); const LOGS_CLEANUP_INTERVAL: Duration = Duration::from_secs(60); const STATUS_UPDATE_INTERVAL: Duration = Duration::from_secs(60); @@ -49,6 +55,21 @@ const MAX_PENDING_ASSIGNMENTS: usize = 5; const MAX_LOGS_SIZE: usize = sqd_network_transport::protocol::MAX_LOGS_RESPONSE_SIZE as usize - 100 * 1024; +/// The trailing `Option` is the rate-limit backoff hint carried from admission. +type AdmittedQuery = (PeerId, Query, ResponseSender, Option); + +/// What the log records for a served query, built alongside the wire message in [`build_delivery`] +/// so the two can't diverge — a downgraded (oversized or unsignable) result is `Err` in both. +enum Logged { + Ok { + data_hash: Vec, + uncompressed_size: u64, + last_block: u64, + timings: TimeReport, + }, + Err(query_error::Err), +} + pub struct P2PController { worker: Arc, worker_status: RwLock, @@ -56,20 +77,22 @@ pub struct P2PController { assignment_fetch_timeout: Duration, assignment_fetch_max_delay: Duration, raw_event_stream: UseOnce, - // Keeps the transport running: it shuts down when the last handle drops. Responses go through - // each request's `ResponseSender`, so this handle is only held, never called. + // Held to keep the transport running; it stops when the last handle drops. Responses go through + // each request's `ResponseSender`, never this handle. _transport_handle: WorkerTransportHandle, logs_storage: LogsStorage, allocations_checker: AllocationsChecker, worker_id: PeerId, keypair: Keypair, assignment_url: String, - queries_tx: mpsc::Sender<(PeerId, Query, ResponseSender)>, - queries_rx: UseOnce>, - sql_queries_tx: mpsc::Sender<(PeerId, Query, ResponseSender)>, - sql_queries_rx: UseOnce>, + queries_tx: mpsc::Sender, + queries_rx: UseOnce>, + sql_queries_tx: mpsc::Sender, + sql_queries_rx: UseOnce>, log_requests_tx: mpsc::Sender<(LogsRequest, ResponseSender)>, log_requests_rx: UseOnce>, + // Caps concurrent reject sends; see `spawn_error_response`. + reject_semaphore: Arc, } pub async fn create_p2p_controller( @@ -117,6 +140,7 @@ pub async fn create_p2p_controller( sql_queries_rx: UseOnce::new(sql_queries_rx), log_requests_tx, log_requests_rx: UseOnce::new(log_requests_rx), + reject_semaphore: Arc::new(Semaphore::new(MAX_CONCURRENT_REJECTS)), }) } @@ -188,14 +212,23 @@ impl + Send + 'static> P2PController + Send + 'static> P2PController + Send + 'static> P2PController query, Err(e) => { + // Without a decodable query_id there's nothing to bind a signed + // response to. warn!("Failed to decode query from {peer_id}: {e}"); continue; } }; - if !self.validate_query(&query, peer_id) { - continue; - } - match self.queries_tx.try_send((peer_id, query, resp_chan)) { - Ok(_) => {} - Err(mpsc::error::TrySendError::Full((_, query, resp_chan))) => { - warn!("Queries queue is full. Rejecting query from {peer_id}"); - self.clone().reject_overloaded(query, resp_chan); - } - Err(mpsc::error::TrySendError::Closed(_)) => { - break; - } + if !self.admit_query(peer_id, query, resp_chan, &self.queries_tx, "Query") { + break; } } WorkerEvent::SqlQuery { @@ -456,22 +490,15 @@ impl + Send + 'static> P2PController query, Err(e) => { + // Without a decodable query_id there's nothing to bind a signed + // response to. warn!("Failed to decode SQL query from {peer_id}: {e}"); continue; } }; - if !self.validate_query(&query, peer_id) { - continue; - } - match self.sql_queries_tx.try_send((peer_id, query, resp_chan)) { - Ok(_) => {} - Err(mpsc::error::TrySendError::Full((_, query, resp_chan))) => { - warn!("SQL Queries queue is full. Rejecting query from {peer_id}"); - self.clone().reject_overloaded(query, resp_chan); - } - Err(mpsc::error::TrySendError::Closed(_)) => { - break; - } + if !self.admit_query(peer_id, query, resp_chan, &self.sql_queries_tx, "SQL query") + { + break; } } WorkerEvent::LogsRequest { @@ -511,60 +538,144 @@ impl + Send + 'static> P2PController, query: Query, resp_chan: ResponseSender) { - tokio::spawn(async move { - let mut msg = sqd_messages::QueryResult { - query_id: query.query_id, - result: Some(query_error::Err::TooManyRequests(()).into()), - retry_after_ms: Some(DEFAULT_BACKOFF.as_millis() as u32), - signature: Default::default(), - }; - if let Err(e) = msg.sign(&self.keypair) { - warn!("Couldn't sign too_many_requests response: {e}"); - return; + /// The admission bar: reserve a queue slot, then spend one compute unit. Past it a query is + /// billable and always logged; rejected here (queue full, no allocation, rate limited) it gets a + /// typed error but no log. Reserving before spending guarantees a spent unit always enqueues. + /// Returns `false` when the receiver is gone, to stop the event loop. + fn admit_query( + self: &Arc, + peer_id: PeerId, + query: Query, + resp_chan: ResponseSender, + queue_tx: &mpsc::Sender, + protocol: &str, + ) -> bool { + if let Err(err) = self.validate_query(&query, peer_id) { + self.clone() + .spawn_error_response(query.query_id, err, None, resp_chan); + return true; + } + let permit = match queue_tx.try_reserve() { + Ok(permit) => permit, + Err(mpsc::error::TrySendError::Full(())) => { + warn!("{protocol} queue is full. Rejecting query from {peer_id}"); + self.clone().spawn_error_response( + query.query_id, + query_error::Err::TooManyRequests(()), + Some(DEFAULT_BACKOFF), + resp_chan, + ); + return true; } - if let Err(e) = resp_chan.send(&msg.encode_to_vec()).await { - warn!("Couldn't send too_many_requests response: {e:?}"); + Err(mpsc::error::TrySendError::Closed(())) => return false, + }; + // Claim a full unit up front, refund the unused fraction after; otherwise a burst of small + // requests could overload the worker. + match self.allocations_checker.try_spend(peer_id, 1.) { + RateLimitStatus::NoAllocation => { + self.clone().spawn_error_response( + query.query_id, + query_error::Err::TooManyRequests(()), + None, + resp_chan, + ); + } + RateLimitStatus::Paused(retry_after) => { + self.clone().spawn_error_response( + query.query_id, + query_error::Err::TooManyRequests(()), + Some(retry_after), + resp_chan, + ); } + RateLimitStatus::Spent(retry_after) => { + permit.send((peer_id, query, resp_chan, retry_after)); + } + } + true + } + + /// Send a signed error, so a rejected query gets a typed reason instead of an opaque stream + /// reset. Best-effort: on a signing or send failure it drops the response, resetting the stream. + async fn send_error_result( + &self, + query_id: String, + err: query_error::Err, + retry_after: Option, + resp_chan: ResponseSender, + ) { + let mut msg = sqd_messages::QueryResult { + query_id, + result: Some(err.into()), + retry_after_ms: retry_after.map(|d| d.as_millis() as u32), + signature: Default::default(), + }; + if let Err(e) = msg.sign(&self.keypair) { + warn!("Couldn't sign error response: {e}"); + return; + } + if let Err(e) = resp_chan.send(&msg.encode_to_vec()).await { + warn!("Couldn't send error response: {e:?}"); + } + } + + /// Run [`Self::send_error_result`] off the event loop so the write can't stall it. `reject_semaphore` + /// caps in-flight sends; past the limit the response is dropped instead of spawning unbounded tasks. + fn spawn_error_response( + self: Arc, + query_id: String, + err: query_error::Err, + retry_after: Option, + resp_chan: ResponseSender, + ) { + let Ok(permit) = self.reject_semaphore.clone().try_acquire_owned() else { + warn!("Too many pending rejections, dropping response for query {query_id}"); + return; + }; + tokio::spawn(async move { + let _permit = permit; + self.send_error_result(query_id, err, retry_after, resp_chan) + .await; }); } + /// Checks the signature and timestamp freshness; the `Err` is the response to send back. #[instrument(skip_all)] - fn validate_query(&self, query: &Query, peer_id: PeerId) -> bool { + fn validate_query(&self, query: &Query, peer_id: PeerId) -> Result<(), query_error::Err> { if !query.verify_signature(peer_id, self.worker_id) { - tracing::warn!("Rejected query with invalid signature from {}", peer_id); - return false; + warn!("Rejected query with invalid signature from {peer_id}"); + return Err(query_error::Err::BadRequest( + "invalid query signature".to_owned(), + )); } if query.timestamp_ms.abs_diff(timestamp_now_ms()) as u128 > protocol::MAX_TIME_LAG.as_millis() { - tracing::warn!( - "Rejected query with invalid timestamp ({}) from {}", - query.timestamp_ms, - peer_id + warn!( + "Rejected query with invalid timestamp ({}) from {peer_id}", + query.timestamp_ms ); - return false; + return Err(query_error::Err::BadRequest( + "timestamp out of allowed range".to_owned(), + )); } // TODO: check rate limits here // TODO: check that query_id has not been used before - true + Ok(()) } + /// Serve an admitted query. It already spent a compute unit, so it always produces a response + /// and a matching log — both come from the same outcome, so they can't diverge. #[instrument(skip_all, fields(query_id = %query.query_id, peer_id = %peer_id, dataset = %query.dataset))] async fn handle_query( &self, peer_id: PeerId, query: Query, resp_chan: ResponseSender, + retry_after: Option, query_type: QueryType, ) { - let query_id = query.query_id.clone(); - let compression = query.compression(); - - let (mut result, retry_after) = process_query( + let outcome = execute( &*self.worker, &self.allocations_checker, peer_id, @@ -572,93 +683,36 @@ impl + Send + 'static> P2PController { - let _ = result.as_mut().map(|v| { - v.time_report.compression_time = compression_duration; - v.time_report.signing_time = signing_duration; - }); - } - Err(e) => tracing::error!("Couldn't send query result: {e:?}"), - } + let compression = query.compression(); + let (message, logged) = + build_delivery(&self.keypair, &query.query_id, outcome, retry_after, compression).await; - if let Some(log) = generate_log(&result, query, peer_id).await { - if log.encoded_len() > MAX_LOGS_SIZE { - warn!("Query log is too big: {log:?}"); - return; - } - let result = self.logs_storage.save_log(log).await; - if let Err(e) = result { - warn!("Couldn't save query log: {e:?}"); - } + // Send before logging: the unit was spent at admission so the log happens regardless, and + // the SQLite write stays off the response path. + if let Err(e) = resp_chan.send(&message.encode_to_vec()).await { + warn!("Couldn't send query result to {peer_id}: {e:?}"); } - } - #[instrument(skip_all)] - async fn send_query_result( - &self, - query_id: String, - result: QueryResult, - resp_chan: ResponseSender, - retry_after: Option, - compression: sqd_messages::Compression, - ) -> Result<(Duration, Duration)> { - let compression_timer = std::time::Instant::now(); - let query_result = match result { - Ok(result) => { - let data = match compression { - sqd_messages::Compression::None => result.data, - sqd_messages::Compression::Gzip => result.data_gzip().await, - sqd_messages::Compression::Zstd => result.data_zstd().await, - }; - sqd_messages::query_result::Result::Ok(sqd_messages::QueryOk { - data, - last_block: result.last_block, - }) - } - Err(e) => query_error::Err::from(e).into(), - }; - let compression_duration = compression_timer.elapsed(); - let mut msg = sqd_messages::QueryResult { - query_id, - result: Some(query_result), - retry_after_ms: retry_after.map(|duration| duration.as_millis() as u32), - signature: Default::default(), - }; - let signing_timer = std::time::Instant::now(); - let _span = tracing::debug_span!("sign_query_result"); - tokio::task::block_in_place(|| msg.sign(&self.keypair).map_err(|e| anyhow!(e)))?; - drop(_span); - let signing_duration = signing_timer.elapsed(); - - let result_size = msg.encoded_len() as u64; - if result_size > protocol::MAX_QUERY_RESULT_SIZE { - anyhow::bail!("query result size too large: {result_size}"); + let log = build_log(query, peer_id, logged); + if log.encoded_len() > MAX_LOGS_SIZE { + warn!("Query log is too big: {log:?}"); + return; + } + if let Err(e) = self.logs_storage.save_log(log).await { + warn!("Couldn't save query log: {e:?}"); } - - tracing::trace!("Sending query result"); - resp_chan - .send(&msg.encode_to_vec()) - .await - .map_err(|e| anyhow!("couldn't send query result: {e}"))?; - - Ok((compression_duration, signing_duration)) } async fn handle_logs_request( @@ -700,28 +754,25 @@ async fn wait_for_assignment_applied( } } -/// Run a query and account for its compute units: claim one up front, refund the unused chunk -/// fraction after. Generic over the engine and rate limiter for mock-based tests (see -/// [`super::query_deps`]). +/// Run an admitted query and refund the unused chunk fraction. The unit was spent at admission, so +/// malformed params here are still charged and logged. Generic over the engine and rate limiter for +/// mock-based tests (see [`super::query_deps`]). #[instrument(skip_all)] -async fn process_query( +async fn execute( worker: &W, allocations_checker: &A, peer_id: PeerId, query: &Query, query_type: QueryType, -) -> (QueryResult, Option) { +) -> QueryResult { match query.compression { c if c == sqd_messages::Compression::Gzip as i32 || c == sqd_messages::Compression::Zstd as i32 || c == sqd_messages::Compression::None as i32 => {} _ => { - return ( - Err(QueryError::BadRequest( - "Unsupported compression type".to_owned(), - )), - None, - ); + return Err(QueryError::BadRequest( + "Unsupported compression type".to_owned(), + )); } } @@ -729,14 +780,10 @@ async fn process_query( .block_range .map(|sqd_messages::Range { begin, end }| (begin, end)) else { - return ( - Err(QueryError::BadRequest("block_range is required".to_owned())), - None, - ); + return Err(QueryError::BadRequest("block_range is required".to_owned())); }; let mut allocation_chip = 1.0f32; - if let Ok(chunk) = query.chunk_id.parse::() { let (begin, end) = block_range; let active_len = std::cmp::min(chunk.last_block.into(), end) @@ -748,20 +795,6 @@ async fn process_query( allocation_chip = active_len as f32 / chunk_len as f32; }; - // We claim 1. allocation first and refund unused allocation later. It's done to prevent burst overloading with small requests. - let status = match allocations_checker.try_spend(peer_id, 1.) { - RateLimitStatus::NoAllocation => { - // This error means that we don't have any allocation for particular peer_id at all (e.g. no allocation on contract) - return (Err(QueryError::NoAllocation), None); - } - RateLimitStatus::Paused(retry_after) => { - // This error means that we don't have allocation at the moment, but it may be available after retry_after ms - return (Err(QueryError::NoAllocation), Some(retry_after)); - } - status => status, - }; - let mut retry_after = status.retry_after(); - let result = worker .run_query( &query.query, @@ -774,47 +807,144 @@ async fn process_query( .await .inspect_err(|err| tracing::error!("error processing query: {err}")); - if let Err(QueryError::ServiceOverloaded) = result { - // Refund everything as we were not able to process request - allocations_checker.refund(peer_id, 1.); - retry_after = Some(DEFAULT_BACKOFF); - } else if allocation_chip < 1. { - // We refund unused allocation + if allocation_chip < 1. { allocations_checker.refund(peer_id, 1. - allocation_chip); } - (result, retry_after) + result } +/// Build the wire message and the log record from one outcome, so they can't diverge: an oversized +/// or unsignable success downgrades to a signed `server_error` in both. #[instrument(skip_all)] -async fn generate_log( - query_result: &QueryResult, - query: Query, - client_id: PeerId, -) -> Option { - use query_executed::Result; - - let result = match query_result { - Ok(result) => Result::Ok(sqd_messages::QueryOkSummary { - uncompressed_data_size: result.data.len() as u64, - data_hash: result.sha3_256().await, - last_block: result.last_block, - }), - Err(QueryError::NoAllocation) => return None, - Err(e) => query_error::Err::from(e).into(), +async fn build_delivery( + keypair: &Keypair, + query_id: &str, + outcome: QueryResult, + retry_after: Option, + compression: sqd_messages::Compression, +) -> (sqd_messages::QueryResult, Logged) { + let retry_after_ms = retry_after.map(|d| d.as_millis() as u32); + let qok = match outcome { + Ok(qok) => qok, + Err(e) => { + let err = query_error::Err::from(e); + return ( + signed_result(keypair, query_id, err.clone().into(), retry_after_ms), + Logged::Err(err), + ); + } + }; + + let data_hash = qok.sha3_256().await; + let uncompressed_size = qok.data.len() as u64; + let last_block = qok.last_block; + let time_report = qok.time_report.clone(); + + let compression_timer = std::time::Instant::now(); + let data = match compression { + sqd_messages::Compression::None => qok.data, + sqd_messages::Compression::Gzip => qok.data_gzip().await, + sqd_messages::Compression::Zstd => qok.data_zstd().await, + }; + let compression_time = compression_timer.elapsed(); + + let mut msg = sqd_messages::QueryResult { + query_id: query_id.to_owned(), + result: Some(sqd_messages::query_result::Result::Ok(sqd_messages::QueryOk { + data, + last_block, + })), + retry_after_ms, + signature: Default::default(), + }; + let signing_timer = std::time::Instant::now(); + let _span = tracing::debug_span!("sign_query_result").entered(); + let sign_result = tokio::task::block_in_place(|| msg.sign(keypair)); + drop(_span); + let signing_time = signing_timer.elapsed(); + + let too_large = msg.encoded_len() as u64 > protocol::MAX_QUERY_RESULT_SIZE; + if let Err(e) = &sign_result { + // Unreachable given a 36-byte query_id and 32-byte hash, but downgrade defensively so + // client and log still agree. + warn!("Couldn't sign query result for {query_id}: {e}"); + } + if too_large { + warn!("Query result for {query_id} is too large: {} bytes", msg.encoded_len()); + } + if sign_result.is_err() || too_large { + let reason = if sign_result.is_err() { + "failed to sign query result" + } else { + "query result too large" + }; + let err = query_error::Err::ServerError(reason.to_owned()); + // A downgraded result carries no retry hint. + return ( + signed_result(keypair, query_id, err.clone().into(), None), + Logged::Err(err), + ); + } + + let timings = TimeReport { + parsing_time_micros: time_report.parsing_time.as_micros() as u32, + execution_time_micros: time_report.execution_time.as_micros() as u32, + compression_time_micros: compression_time.as_micros() as u32, + signing_time_micros: signing_time.as_micros() as u32, + serialization_time_micros: time_report.serialization_time.as_micros() as u32, + }; + ( + msg, + Logged::Ok { + data_hash, + uncompressed_size, + last_block, + timings, + }, + ) +} + +/// Sign a prepared `QueryResult`. Best-effort: a signing failure is logged and the message goes out +/// unsigned, which the client rejects. +fn signed_result( + keypair: &Keypair, + query_id: &str, + result: sqd_messages::query_result::Result, + retry_after_ms: Option, +) -> sqd_messages::QueryResult { + let mut msg = sqd_messages::QueryResult { + query_id: query_id.to_owned(), + result: Some(result), + retry_after_ms, + signature: Default::default(), }; + if let Err(e) = msg.sign(keypair) { + warn!("Couldn't sign query result for {query_id}: {e}"); + } + msg +} - let exec_time_report = match query_result { - Ok(result) => Some(TimeReport { - parsing_time_micros: result.time_report.parsing_time.as_micros() as u32, - execution_time_micros: result.time_report.execution_time.as_micros() as u32, - compression_time_micros: result.time_report.compression_time.as_micros() as u32, - signing_time_micros: result.time_report.signing_time.as_micros() as u32, - serialization_time_micros: result.time_report.serialization_time.as_micros() as u32, - }), - Err(_) => None, // TODO: always measure execution time +/// Build the log from the delivery projection. Every admitted query is logged — no skip path — and +/// the projection already carries whatever error the client saw. +fn build_log(query: Query, client_id: PeerId, logged: Logged) -> QueryExecuted { + let (result, exec_time_report) = match logged { + Logged::Ok { + data_hash, + uncompressed_size, + last_block, + timings, + } => ( + query_executed::Result::Ok(sqd_messages::QueryOkSummary { + uncompressed_data_size: uncompressed_size, + data_hash, + last_block, + }), + Some(timings), + ), + Logged::Err(err) => (query_executed::Result::from(err), None), }; - Some(QueryExecuted { + QueryExecuted { client_id: client_id.to_string(), query: Some(query), exec_time_micros: exec_time_report @@ -824,7 +954,7 @@ async fn generate_log( timestamp_ms: timestamp_now_ms(), // TODO: use time of receiving query result: Some(result), worker_version: WORKER_VERSION.to_string(), - }) + } } #[tracing::instrument(skip_all)] @@ -982,7 +1112,7 @@ mod assignment_tests { mod tests { use std::time::Duration; - use sqd_messages::{query_executed, Compression, Range}; + use sqd_messages::{query_result, Compression, Range}; use crate::{ controller::query_deps::mocks::{MockChecker, MockWorker}, @@ -993,10 +1123,12 @@ mod tests { // first_block = 100, last_block = 200 → chunk length 100. const CHUNK_ID: &str = "0000000000/0000000100-0000000200-abcde"; + // 36 chars — the UUID length the signature scheme requires. + const QUERY_ID: &str = "00000000-0000-0000-0000-000000000000"; fn query_with(compression: i32, block_range: Option<(u64, u64)>, chunk_id: &str) -> Query { Query { - query_id: "00000000-0000-0000-0000-000000000000".to_owned(), + query_id: QUERY_ID.to_owned(), dataset: "s3://dataset".to_owned(), query: "{}".to_owned(), chunk_id: chunk_id.to_owned(), @@ -1017,112 +1149,130 @@ mod tests { )) } - // An overloaded engine is fully refunded, so the query consumes no compute unit. + // An overloaded engine keeps the compute unit: admission spent it and the served stage doesn't + // refund on overload. #[tokio::test(flavor = "multi_thread")] - async fn overloaded_query_consumes_no_cu() { + async fn overloaded_query_consumes_a_cu() { let checker = MockChecker::new(RateLimitStatus::Spent(None)); + checker.try_spend(1.0); // admission let worker = MockWorker::new(Err(QueryError::ServiceOverloaded)); let query = query_with(Compression::None as i32, Some((100, 200)), CHUNK_ID); - let (result, retry) = - process_query(&worker, &checker, PeerId::random(), &query, QueryType::PlainQuery).await; + let outcome = + execute(&worker, &checker, PeerId::random(), &query, QueryType::PlainQuery).await; - assert!(matches!(result, Err(QueryError::ServiceOverloaded))); - assert_eq!(retry, Some(DEFAULT_BACKOFF)); - assert_eq!(checker.net_spent(), 0.0); + assert!(matches!(outcome, Err(QueryError::ServiceOverloaded))); + assert_eq!(checker.net_spent(), 1.0); assert_eq!(worker.calls(), 1); } - // Malformed params are rejected before the CU spend: no retry hint, no charge, engine not run. + // Malformed params are validated after admission, so they're charged the full unit and never + // reach the engine. #[tokio::test(flavor = "multi_thread")] - async fn malformed_query_is_not_charged() { + async fn malformed_query_is_charged_a_cu() { let checker = MockChecker::new(RateLimitStatus::Spent(Some(Duration::from_secs(5)))); + checker.try_spend(1.0); let worker = MockWorker::new(ok_result()); let query = query_with(999, Some((100, 200)), CHUNK_ID); - let (result, retry) = - process_query(&worker, &checker, PeerId::random(), &query, QueryType::PlainQuery).await; + let outcome = + execute(&worker, &checker, PeerId::random(), &query, QueryType::PlainQuery).await; - assert!(matches!(result, Err(QueryError::BadRequest(_)))); - assert_eq!(retry, None); - assert_eq!(checker.net_spent(), 0.0); + assert!(matches!(outcome, Err(QueryError::BadRequest(_)))); assert_eq!(worker.calls(), 0); + assert_eq!(checker.net_spent(), 1.0); } #[tokio::test(flavor = "multi_thread")] async fn successful_full_chunk_consumes_one_cu() { let checker = MockChecker::new(RateLimitStatus::Spent(None)); + checker.try_spend(1.0); let worker = MockWorker::new(ok_result()); let query = query_with(Compression::None as i32, Some((100, 200)), CHUNK_ID); - let (result, retry) = - process_query(&worker, &checker, PeerId::random(), &query, QueryType::PlainQuery).await; + let outcome = + execute(&worker, &checker, PeerId::random(), &query, QueryType::PlainQuery).await; - assert!(result.is_ok()); - assert_eq!(retry, None); + assert!(outcome.is_ok()); assert_eq!(checker.net_spent(), 1.0); } #[tokio::test(flavor = "multi_thread")] async fn partial_chunk_refunds_unused_fraction() { let checker = MockChecker::new(RateLimitStatus::Spent(None)); + checker.try_spend(1.0); let worker = MockWorker::new(ok_result()); - let query = query_with(Compression::None as i32, Some((100, 150)), CHUNK_ID); + // 30% of [100, 200] — asymmetric, so a flipped fraction would fail. + let query = query_with(Compression::None as i32, Some((100, 130)), CHUNK_ID); - let (result, _) = - process_query(&worker, &checker, PeerId::random(), &query, QueryType::PlainQuery).await; + let outcome = + execute(&worker, &checker, PeerId::random(), &query, QueryType::PlainQuery).await; - assert!(result.is_ok()); + assert!(outcome.is_ok()); assert!( - (checker.net_spent() - 0.5).abs() < 1e-6, - "expected 0.5 CU, got {}", + (checker.net_spent() - 0.3).abs() < 1e-6, + "expected 0.3 CU, got {}", checker.net_spent() ); } + // The admission backoff hint rides along with a BadRequest response. #[tokio::test(flavor = "multi_thread")] - async fn no_allocation_is_rejected_without_running() { - let checker = MockChecker::new(RateLimitStatus::NoAllocation); - let worker = MockWorker::new(ok_result()); - let query = query_with(Compression::None as i32, Some((100, 200)), CHUNK_ID); - - let (result, retry) = - process_query(&worker, &checker, PeerId::random(), &query, QueryType::PlainQuery).await; + async fn bad_request_response_carries_retry_after() { + let keypair = Keypair::generate_ed25519(); + let outcome = Err(QueryError::BadRequest("timestamp out of allowed range".to_owned())); + + let (message, logged) = build_delivery( + &keypair, + QUERY_ID, + outcome, + Some(Duration::from_secs(5)), + Compression::None, + ) + .await; - assert!(matches!(result, Err(QueryError::NoAllocation))); - assert_eq!(retry, None); - assert_eq!(worker.calls(), 0); - assert_eq!(checker.net_spent(), 0.0); + assert_eq!(message.retry_after_ms, Some(5000)); + assert!(matches!( + message.result, + Some(query_result::Result::Err(sqd_messages::QueryError { + err: Some(query_error::Err::BadRequest(_)), + })) + )); + assert!(matches!( + logged, + Logged::Err(query_error::Err::BadRequest(_)) + )); } + // A result that can't be signed (forced here with a non-UUID query_id) downgrades to + // server_error in both the response and the log. #[tokio::test(flavor = "multi_thread")] - async fn rate_limited_is_rejected_with_retry() { - let retry = Duration::from_secs(2); - let checker = MockChecker::new(RateLimitStatus::Paused(retry)); - let worker = MockWorker::new(ok_result()); - let query = query_with(Compression::None as i32, Some((100, 200)), CHUNK_ID); - - let (result, got_retry) = - process_query(&worker, &checker, PeerId::random(), &query, QueryType::PlainQuery).await; - - assert!(matches!(result, Err(QueryError::NoAllocation))); - assert_eq!(got_retry, Some(retry)); - assert_eq!(worker.calls(), 0); + async fn unsignable_result_downgrades_response_and_log_together() { + let keypair = Keypair::generate_ed25519(); + + let (message, logged) = + build_delivery(&keypair, "not-a-uuid", ok_result(), None, Compression::None).await; + + assert!(matches!( + message.result, + Some(query_result::Result::Err(sqd_messages::QueryError { + err: Some(query_error::Err::ServerError(_)), + })) + )); + assert!(matches!( + logged, + Logged::Err(query_error::Err::ServerError(_)) + )); } #[tokio::test(flavor = "multi_thread")] - async fn no_allocation_is_not_logged() { - let query = query_with(Compression::None as i32, Some((100, 200)), CHUNK_ID); - let log = generate_log(&Err(QueryError::NoAllocation), query, PeerId::random()).await; - assert!(log.is_none()); - } + async fn successful_delivery_logs_a_summary() { + let keypair = Keypair::generate_ed25519(); + let (_message, logged) = + build_delivery(&keypair, QUERY_ID, ok_result(), None, Compression::None).await; - #[tokio::test(flavor = "multi_thread")] - async fn log_for_success_summarizes_result() { let query = query_with(Compression::None as i32, Some((100, 200)), CHUNK_ID); - let log = generate_log(&ok_result(), query, PeerId::random()) - .await - .expect("success should be logged"); + let log = build_log(query, PeerId::random(), logged); match log.result.expect("log should carry a result") { query_executed::Result::Ok(summary) => { assert_eq!(summary.last_block, 200); @@ -1133,17 +1283,18 @@ mod tests { } } + // Every admitted query is logged; admission already guaranteed it's billable, so there's no + // skip path. #[tokio::test(flavor = "multi_thread")] - async fn log_for_error_records_the_error() { + async fn admitted_error_is_always_logged() { + let keypair = Keypair::generate_ed25519(); + let outcome = Err(QueryError::BadRequest("bad".to_owned())); + let (_message, logged) = + build_delivery(&keypair, QUERY_ID, outcome, None, Compression::None).await; + let query = query_with(Compression::None as i32, Some((100, 200)), CHUNK_ID); - let log = generate_log( - &Err(QueryError::BadRequest("bad".to_owned())), - query, - PeerId::random(), - ) - .await - .expect("errors past the CU bar are logged"); - match log.result.expect("log should carry a result") { + let log = build_log(query, PeerId::random(), logged); + match log.result.expect("errors past the CU bar are logged") { query_executed::Result::Err(e) => { assert!(matches!(e.err, Some(query_error::Err::BadRequest(_)))); } diff --git a/src/controller/query_deps.rs b/src/controller/query_deps.rs index f47041b..253edef 100644 --- a/src/controller/query_deps.rs +++ b/src/controller/query_deps.rs @@ -1,29 +1,24 @@ -//! Dependency seams for the query-serving pipeline (`process_query`, `generate_log`) in -//! [`super::p2p`]. Production runs against the real `Worker` and `AllocationsChecker`; the traits -//! let tests inject an engine with a fixed outcome and a rate limiter that records what it spent, -//! with no transport or storage. +//! Dependency seams for the query-serving pipeline in [`super::p2p`]. Production runs against the +//! real `Worker` and `AllocationsChecker`; the traits let tests inject an engine with a fixed +//! outcome and a rate limiter that records what it spent, with no transport or storage. use async_trait::async_trait; use sqd_network_transport::PeerId; use crate::{ - compute_units::{allocations_checker::AllocationsChecker, RateLimitStatus}, + compute_units::allocations_checker::AllocationsChecker, controller::worker::{QueryType, Worker}, query::result::QueryResult, types::dataset::Dataset, }; -/// Compute-unit rate limiting, behind a trait so tests can set a verdict and record spending. +/// The refund half of compute-unit accounting. A unit is spent at admission on the concrete +/// `AllocationsChecker`; the served stage only refunds the fraction the query didn't use. pub trait CuChecker { - fn try_spend(&self, portal_id: PeerId, allocation_chip: f32) -> RateLimitStatus; fn refund(&self, portal_id: PeerId, allocation_chip: f32); } impl CuChecker for AllocationsChecker { - fn try_spend(&self, portal_id: PeerId, allocation_chip: f32) -> RateLimitStatus { - AllocationsChecker::try_spend(self, portal_id, allocation_chip) - } - fn refund(&self, portal_id: PeerId, allocation_chip: f32) { AllocationsChecker::refund(self, portal_id, allocation_chip) } @@ -101,16 +96,17 @@ pub mod mocks { pub fn net_spent(&self) -> f32 { *self.net_spent.lock() } - } - impl CuChecker for MockChecker { - fn try_spend(&self, _portal_id: PeerId, allocation_chip: f32) -> RateLimitStatus { + /// Stand in for the admission spend that production does before the served stage. + pub fn try_spend(&self, allocation_chip: f32) -> RateLimitStatus { if let RateLimitStatus::Spent(_) = self.verdict { *self.net_spent.lock() += allocation_chip; } self.verdict } + } + impl CuChecker for MockChecker { fn refund(&self, _portal_id: PeerId, allocation_chip: f32) { *self.net_spent.lock() -= allocation_chip; } diff --git a/src/query/result.rs b/src/query/result.rs index fca0af9..adee1dd 100644 --- a/src/query/result.rs +++ b/src/query/result.rs @@ -8,13 +8,12 @@ use crate::util::hash::sha3_256; pub type QueryResult = std::result::Result; +/// Timings from the query engine. Compression and signing are measured at delivery, not here. #[derive(Debug, Clone, Default)] pub struct WorkerTimeReport { pub parsing_time: Duration, pub execution_time: Duration, pub serialization_time: Duration, - pub compression_time: Duration, - pub signing_time: Duration, } #[derive(Debug, Clone)] @@ -41,8 +40,6 @@ impl QueryOk { parsing_time, execution_time, serialization_time, - compression_time: Duration::from_secs(0), - signing_time: Duration::from_secs(0), }, last_block, } From aae54ff5e71bb229d124f5fff8a24d9f562e6213 Mon Sep 17 00:00:00 2001 From: Dzmitry Kalabuk Date: Wed, 8 Jul 2026 16:30:07 +0300 Subject: [PATCH 5/9] Fix the code style --- src/controller/p2p.rs | 31 +++++++++++++++++++++++++------ 1 file changed, 25 insertions(+), 6 deletions(-) diff --git a/src/controller/p2p.rs b/src/controller/p2p.rs index 103e503..23fbf79 100644 --- a/src/controller/p2p.rs +++ b/src/controller/p2p.rs @@ -55,8 +55,13 @@ const MAX_PENDING_ASSIGNMENTS: usize = 5; const MAX_LOGS_SIZE: usize = sqd_network_transport::protocol::MAX_LOGS_RESPONSE_SIZE as usize - 100 * 1024; -/// The trailing `Option` is the rate-limit backoff hint carried from admission. -type AdmittedQuery = (PeerId, Query, ResponseSender, Option); +struct AdmittedQuery { + peer_id: PeerId, + query: Query, + resp_chan: ResponseSender, + /// Rate-limit backoff hint carried from admission. + retry_after: Option, +} /// What the log records for a served query, built alongside the wire message in [`build_delivery`] /// so the two can't diverge — a downgraded (oversized or unsignable) result is `Err` in both. @@ -214,7 +219,12 @@ impl + Send + 'static> P2PController + Send + 'static> P2PController + Send + 'static> P2PController { - permit.send((peer_id, query, resp_chan, retry_after)); + permit.send(AdmittedQuery { + peer_id, + query, + resp_chan, + retry_after, + }); } } true @@ -659,7 +679,6 @@ impl + Send + 'static> P2PController Date: Thu, 9 Jul 2026 09:18:14 +0300 Subject: [PATCH 6/9] Fix error type on overload --- src/controller/p2p.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/controller/p2p.rs b/src/controller/p2p.rs index 23fbf79..02df04f 100644 --- a/src/controller/p2p.rs +++ b/src/controller/p2p.rs @@ -576,7 +576,7 @@ impl + Send + 'static> P2PController Date: Wed, 15 Jul 2026 10:02:18 +0300 Subject: [PATCH 7/9] Use tokio-graceful-shutdown for task supervision - Replace the run_all! macro and hand-wired cancellation token tree with a subsystem tree that drains gracefully and cancels stragglers after a 5s timeout instead of hanging forever - Leak the controller for &'static access, removing the Arc clone per spawned task - Report a loop that dies on its own as a named subsystem error instead of a silent all-Ok shutdown - Pin the dependency to an exact version; the audited tarball hash is recorded in Cargo.lock Co-Authored-By: Claude Fable 5 --- Cargo.lock | 48 +++++++++++++ Cargo.toml | 2 + src/controller/p2p.rs | 158 +++++++++++++++++------------------------- src/main.rs | 79 +++++++++------------ src/util/mod.rs | 2 - src/util/run_all.rs | 55 --------------- 6 files changed, 145 insertions(+), 199 deletions(-) delete mode 100644 src/util/run_all.rs diff --git a/Cargo.lock b/Cargo.lock index a04f50f..7bcab10 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -544,6 +544,15 @@ version = "0.15.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ae037714f313c1353189ead58ef9eec30a8e8dc101b2622d461418fd59e28a9" +[[package]] +name = "atomic" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89cbf775b137e9b968e67227ef7f775587cde3fd31b0d8599dbd0f598a48340" +dependencies = [ + "bytemuck", +] + [[package]] name = "atomic-waker" version = "1.1.2" @@ -4169,6 +4178,28 @@ dependencies = [ "libc", ] +[[package]] +name = "miette" +version = "7.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f98efec8807c63c752b5bd61f862c165c115b0a35685bdcfd9238c7aeb592b7" +dependencies = [ + "cfg-if", + "miette-derive", + "unicode-width", +] + +[[package]] +name = "miette-derive" +version = "7.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db5b29714e950dbb20d5e6f74f9dcec4edbcc1067bb7f8ed198c097b8c1a818b" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.111", +] + [[package]] name = "mimalloc" version = "0.1.43" @@ -7282,6 +7313,7 @@ dependencies = [ "substrait", "thiserror 1.0.65", "tokio", + "tokio-graceful-shutdown", "tokio-rusqlite", "tokio-stream", "tokio-util", @@ -7761,6 +7793,22 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "tokio-graceful-shutdown" +version = "0.19.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "374fc6755143a79c052be194b45598b75af02aeeb07430acc872afdca7a43652" +dependencies = [ + "atomic", + "bytemuck", + "miette", + "pin-project-lite", + "thiserror 2.0.11", + "tokio", + "tokio-util", + "tracing", +] + [[package]] name = "tokio-macros" version = "2.6.0" diff --git a/Cargo.toml b/Cargo.toml index 9dbc478..4128115 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -50,6 +50,8 @@ sha3 = "0.10.8" substrait = { version = "0.48.0", features = ["serde", "parse"] } thiserror = "1.0.57" tokio = { version = "1.35.1", features = ["full", "tracing", "test-util"] } +# Exact pin; the tarball hash in Cargo.lock guards against tampering +tokio-graceful-shutdown = "=0.19.3" tokio-rusqlite = "0.5.1" tokio-stream = { version = "0.1.14", features = ["sync"] } tokio-util = "0.7.10" diff --git a/src/controller/p2p.rs b/src/controller/p2p.rs index 02df04f..2d6b90f 100644 --- a/src/controller/p2p.rs +++ b/src/controller/p2p.rs @@ -17,6 +17,7 @@ use tokio::{ sync::{mpsc, Semaphore}, time::MissedTickBehavior, }; +use tokio_graceful_shutdown::{SubsystemBuilder, SubsystemHandle}; use tokio_stream::wrappers::{IntervalStream, ReceiverStream}; use tokio_util::sync::CancellationToken; use tracing::{info, instrument, warn, Instrument}; @@ -31,7 +32,6 @@ use crate::{ logs_storage::LogsStorage, metrics, query::result::{QueryError, QueryResult}, - run_all, storage::layout::DataChunk, util::{timestamp_now_ms, UseOnce}, }; @@ -76,7 +76,7 @@ enum Logged { } pub struct P2PController { - worker: Arc, + worker: Worker, worker_status: RwLock, assignment_check_interval: Duration, assignment_fetch_timeout: Duration, @@ -101,7 +101,7 @@ pub struct P2PController { } pub async fn create_p2p_controller( - worker: Arc, + worker: Worker, transport_builder: P2PTransportBuilder, args: Args, ) -> Result>> { @@ -149,71 +149,46 @@ pub async fn create_p2p_controller( }) } +/// Runs a cancellable loop as a subsystem; exiting for any reason other than the requested +/// shutdown is an error. +fn start_loop(s: &SubsystemHandle, name: &'static str, run: F) +where + F: FnOnce(CancellationToken) -> Fut + Send + 'static, + Fut: std::future::Future + Send, +{ + s.start(SubsystemBuilder::new( + name, + async move |sub: &mut SubsystemHandle| { + run(sub.create_cancellation_token()).await; + if sub.is_shutdown_requested() { + Ok(()) + } else { + Err(anyhow::anyhow!("{} exited unexpectedly", sub.name())) + } + }, + )); +} + impl + Send + 'static> P2PController { - pub async fn run(self: Arc, cancellation_token: CancellationToken) { - let this = self.clone(); - let token = cancellation_token.child_token(); - let event_task = tokio::spawn(async move { this.run_event_loop(token).await }); - - let this = self.clone(); - let token = cancellation_token.child_token(); - let queries_task = tokio::spawn(async move { this.run_queries_loop(token).await }); - - let this = self.clone(); - let token = cancellation_token.child_token(); - let sql_queries_task = tokio::spawn(async move { this.run_sql_queries_loop(token).await }); - - let this = self.clone(); - let token = cancellation_token.child_token(); - let assignments_task = tokio::spawn(async move { - this.run_assignments_loop(token, this.assignment_check_interval) - .await + pub fn start_subsystems(&'static self, s: &SubsystemHandle) { + start_loop(s, "transport_events", |t| self.run_event_loop(t)); + start_loop(s, "queries", |t| self.run_queries_loop(t)); + start_loop(s, "sql_queries", |t| self.run_sql_queries_loop(t)); + start_loop(s, "assignments", |t| { + self.run_assignments_loop(t, self.assignment_check_interval) }); - - let this = self.clone(); - let token = cancellation_token.child_token(); - let logs_task: tokio::task::JoinHandle<()> = - tokio::spawn(async move { this.run_logs_loop(token).await }); - - let this = self.clone(); - let token = cancellation_token.child_token(); - let logs_cleanup_task: tokio::task::JoinHandle<()> = tokio::spawn(async move { - this.run_logs_cleanup_loop(token, LOGS_CLEANUP_INTERVAL) - .await + start_loop(s, "logs", |t| self.run_logs_loop(t)); + start_loop(s, "logs_cleanup", |t| { + self.run_logs_cleanup_loop(t, LOGS_CLEANUP_INTERVAL) }); - - let this = self.clone(); - let token = cancellation_token.child_token(); - let status_update_task: tokio::task::JoinHandle<()> = tokio::spawn(async move { - this.run_status_update_loop(token, STATUS_UPDATE_INTERVAL) - .await + start_loop(s, "status_updates", |t| { + self.run_status_update_loop(t, STATUS_UPDATE_INTERVAL) }); - - let this = self.clone(); - let token = cancellation_token.child_token(); - let worker_task: tokio::task::JoinHandle<()> = - tokio::spawn(async move { this.worker.run(token).await }); - - let this = self.clone(); - let token = cancellation_token.child_token(); - let allocations_task: tokio::task::JoinHandle<()> = - tokio::spawn(async move { this.allocations_checker.run(token).await }); - - let _ = run_all!( - cancellation_token, - event_task, - queries_task, - sql_queries_task, - assignments_task, - logs_task, - logs_cleanup_task, - status_update_task, - worker_task, - allocations_task, - ); + start_loop(s, "state_manager", |t| self.worker.run(t)); + start_loop(s, "allocations", |t| self.allocations_checker.run(t)); } - async fn run_queries_loop(self: Arc, cancellation_token: CancellationToken) { + async fn run_queries_loop(&'static self, cancellation_token: CancellationToken) { let queries_rx = self.queries_rx.take().unwrap(); ReceiverStream::new(queries_rx) .take_until(cancellation_token.cancelled_owned()) @@ -225,17 +200,13 @@ impl + Send + 'static> P2PController + Send + 'static> P2PController, cancellation_token: CancellationToken) { + async fn run_sql_queries_loop(&'static self, cancellation_token: CancellationToken) { let sql_queries_rx = self.sql_queries_rx.take().unwrap(); ReceiverStream::new(sql_queries_rx) .take_until(cancellation_token.cancelled_owned()) @@ -255,17 +226,13 @@ impl + Send + 'static> P2PController + Send + 'static> P2PController + Send + 'static> P2PController + Send + 'static> P2PController, cancellation_token: CancellationToken) { + async fn run_event_loop(&'static self, cancellation_token: CancellationToken) { let event_stream = self .raw_event_stream .take() @@ -558,7 +525,7 @@ impl + Send + 'static> P2PController, + &'static self, peer_id: PeerId, query: Query, resp_chan: ResponseSender, @@ -566,15 +533,14 @@ impl + Send + 'static> P2PController bool { if let Err(err) = self.validate_query(&query, peer_id) { - self.clone() - .spawn_error_response(query.query_id, err, None, resp_chan); + self.spawn_error_response(query.query_id, err, None, resp_chan); return true; } let permit = match queue_tx.try_reserve() { Ok(permit) => permit, Err(mpsc::error::TrySendError::Full(())) => { warn!("{protocol} queue is full. Rejecting query from {peer_id}"); - self.clone().spawn_error_response( + self.spawn_error_response( query.query_id, query_error::Err::ServerOverloaded(()), Some(DEFAULT_BACKOFF), @@ -588,7 +554,7 @@ impl + Send + 'static> P2PController { - self.clone().spawn_error_response( + self.spawn_error_response( query.query_id, query_error::Err::TooManyRequests(()), None, @@ -596,7 +562,7 @@ impl + Send + 'static> P2PController { - self.clone().spawn_error_response( + self.spawn_error_response( query.query_id, query_error::Err::TooManyRequests(()), Some(retry_after), @@ -642,7 +608,7 @@ impl + Send + 'static> P2PController, + &'static self, query_id: String, err: query_error::Err, retry_after: Option, @@ -695,7 +661,7 @@ impl + Send + 'static> P2PController Option }) } -fn create_cancellation_token() -> Result { - use tokio::signal::unix::{signal, SignalKind}; - - let token = CancellationToken::new(); - let copy = token.clone(); - let mut sigint = signal(SignalKind::interrupt())?; - let mut sigterm = signal(SignalKind::terminate())?; - tokio::spawn(async move { - tokio::select!( - _ = sigint.recv() => { - copy.cancel(); - }, - _ = sigterm.recv() => { - copy.cancel(); - }, - ); - }); - Ok(token) -} +/// Subsystems that don't finish gracefully within this window get cancelled. +const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5); async fn run(mut args: Args) -> anyhow::Result<()> { setup_tracing(&args)?; @@ -151,28 +134,29 @@ async fn run(mut args: Args) -> anyhow::Result<()> { let _sentry_guard = setup_sentry(&args_clone, peer_id.to_string()); } - let worker = Arc::new(Worker::new(state_manager, args.parallel_queries)); - - let cancellation_token = create_cancellation_token()?; - let controller_fut = async { - let controller = tokio::select! { - controller = create_p2p_controller(worker, transport_builder, args_clone) => Arc::new(controller?), - _ = cancellation_token.cancelled() => return Ok(()), - }; - controller.run(cancellation_token.clone()).await; - anyhow::Ok(()) - }; - - let (controller_result, server_result) = run_all!( - cancellation_token, - controller_fut, - tokio::spawn( - HttpServer::new(peer_id, metrics_registry) - .run(args.prometheus_port, cancellation_token.child_token()) - ), - ); - controller_result?; - server_result??; + let worker = Worker::new(state_manager, args.parallel_queries); + + let controller = create_p2p_controller(worker, transport_builder, args_clone).await?; + // Leaked to give the subsystem tasks `&'static` access; lives until process exit anyway + let controller = &*Box::leak(Box::new(controller)); + + let http_server = HttpServer::new(peer_id, metrics_registry); + let prometheus_port = args.prometheus_port; + + Toplevel::new(async move |s: &mut SubsystemHandle| { + controller.start_subsystems(s); + s.start(SubsystemBuilder::new( + "http_server", + async move |sub: &mut SubsystemHandle| { + http_server + .run(prometheus_port, sub.create_cancellation_token()) + .await + }, + )); + }) + .catch_signals() + .handle_shutdown_requests(SHUTDOWN_TIMEOUT) + .await?; tracing::info!("Shutting down"); Ok(()) @@ -184,8 +168,11 @@ fn main() -> anyhow::Result<()> { init_single_threaded(&args)?; - tokio::runtime::Builder::new_multi_thread() + let runtime = tokio::runtime::Builder::new_multi_thread() .enable_all() - .build()? - .block_on(run(args)) + .build()?; + let result = runtime.block_on(run(args)); + // Don't let stuck spawn_blocking tasks (unabortable) delay the exit + runtime.shutdown_timeout(Duration::from_secs(1)); + result } diff --git a/src/util/mod.rs b/src/util/mod.rs index 627dc7f..825c895 100644 --- a/src/util/mod.rs +++ b/src/util/mod.rs @@ -1,8 +1,6 @@ pub mod hash; pub mod iterator; mod once; -#[macro_use] -pub mod run_all; #[cfg(test)] pub mod tests; pub mod timestamp; diff --git a/src/util/run_all.rs b/src/util/run_all.rs deleted file mode 100644 index 7c9eaa6..0000000 --- a/src/util/run_all.rs +++ /dev/null @@ -1,55 +0,0 @@ -// Runs all futures concurrently. If any of them finishes, cancels the rest. -#[macro_export] -macro_rules! run_all { - ( $cancellation_token:expr, $( $fut:expr ),+ $(,)?) => { - tokio::join!( - $( - async { - let result = $fut.await; - $cancellation_token.cancel(); - result - } - ),* - ) - }; -} - -#[cfg(test)] -mod tests { - use tokio_util::sync::CancellationToken; - - #[tokio::test(start_paused = true)] - async fn test_run_all() { - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::Arc; - use tokio::time::{sleep, Duration}; - - let cancellation_token = CancellationToken::new(); - let cancellation_token_clone = cancellation_token.clone(); - - let counter = Arc::new(AtomicUsize::new(0)); - let fut1 = async { - tokio::select! { - _ = cancellation_token_clone.cancelled() => return, - _ = sleep(Duration::from_secs(1)) => counter.fetch_add(1, Ordering::Relaxed), - }; - }; - let cancellation_token_clone = cancellation_token.clone(); - let fut2 = async { - tokio::select! { - _ = cancellation_token_clone.cancelled() => return, - _ = sleep(Duration::from_secs(2)) => counter.fetch_add(1, Ordering::Relaxed), - }; - }; - let cancellation_token_clone = cancellation_token.clone(); - let fut3 = async { - tokio::select! { - _ = cancellation_token_clone.cancelled() => return, - _ = sleep(Duration::from_secs(3)) => counter.fetch_add(1, Ordering::Relaxed), - }; - }; - - run_all!(cancellation_token, fut1, fut2, fut3); - assert_eq!(counter.load(Ordering::Relaxed), 1); - } -} From f41702efcc588cda061b7f05e5d4160b0e073c7a Mon Sep 17 00:00:00 2001 From: Dzmitry Kalabuk Date: Wed, 15 Jul 2026 10:52:09 +0300 Subject: [PATCH 8/9] Fix formatting --- src/controller/mod.rs | 2 +- src/controller/p2p.rs | 80 +++++++++++++++++++++++++++++++++---------- 2 files changed, 62 insertions(+), 20 deletions(-) diff --git a/src/controller/mod.rs b/src/controller/mod.rs index 6139ed9..991191f 100644 --- a/src/controller/mod.rs +++ b/src/controller/mod.rs @@ -1,6 +1,6 @@ pub mod assignments; pub mod p2p; -pub mod query_deps; pub mod polars_target; +pub mod query_deps; pub mod sql_request; pub mod worker; diff --git a/src/controller/p2p.rs b/src/controller/p2p.rs index 2d6b90f..76314ef 100644 --- a/src/controller/p2p.rs +++ b/src/controller/p2p.rs @@ -478,8 +478,13 @@ impl + Send + 'static> P2PController + Send + 'static> P2PController + Send + 'static> P2PController Date: Wed, 15 Jul 2026 11:06:02 +0300 Subject: [PATCH 9/9] Fix immediate drop of the Sentry guard --- src/main.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/main.rs b/src/main.rs index 017e477..531297e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -130,9 +130,11 @@ async fn run(mut args: Args) -> anyhow::Result<()> { ) .await?; - if args_clone.sentry_is_enabled { - let _sentry_guard = setup_sentry(&args_clone, peer_id.to_string()); - } + let _sentry_guard = if args_clone.sentry_is_enabled { + setup_sentry(&args_clone, peer_id.to_string()) + } else { + None + }; let worker = Worker::new(state_manager, args.parallel_queries);