diff --git a/Cargo.lock b/Cargo.lock index 34cdc38..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" @@ -3622,7 +3631,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 +3664,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 +3674,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 +3698,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 +3708,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 +3732,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 +3747,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 +3777,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 +3797,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 +3815,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 +3842,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 +3860,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 +3877,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 +3892,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 +3913,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 +3929,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 +3942,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 +3962,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 +3972,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 +3987,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 +4005,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", @@ -4170,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" @@ -4281,7 +4311,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 +5864,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 +6518,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 +7124,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=c4b85c1#c4b85c132fe5fa308be284556d539f012c0760b0" dependencies = [ "anyhow", "base64 0.22.1", @@ -7122,7 +7152,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=c4b85c1#c4b85c132fe5fa308be284556d539f012c0760b0" dependencies = [ "async-trait", "clap", @@ -7143,7 +7173,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=c4b85c1#c4b85c132fe5fa308be284556d539f012c0760b0" dependencies = [ "bytemuck", "flate2", @@ -7159,7 +7189,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=c4b85c1#c4b85c132fe5fa308be284556d539f012c0760b0" dependencies = [ "anyhow", "async-trait", @@ -7283,6 +7313,7 @@ dependencies = [ "substrait", "thiserror 1.0.65", "tokio", + "tokio-graceful-shutdown", "tokio-rusqlite", "tokio-stream", "tokio-util", @@ -7762,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 badd665..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" @@ -61,10 +63,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 = "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 +74,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 = "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 +# 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/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/mod.rs b/src/controller/mod.rs index 6b409dd..991191f 100644 --- a/src/controller/mod.rs +++ b/src/controller/mod.rs @@ -1,5 +1,6 @@ pub mod assignments; pub mod p2p; 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 8853bda..76314ef 100644 --- a/src/controller/p2p.rs +++ b/src/controller/p2p.rs @@ -1,19 +1,23 @@ 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; 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::{ + 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}; @@ -21,18 +25,18 @@ use tracing::{info, instrument, warn, Instrument}; use crate::{ cli::Args, compute_units::{ - self, allocations_checker::{self, AllocationsChecker}, + RateLimitStatus, }, controller::worker::QueryType, logs_storage::LogsStorage, metrics, query::result::{QueryError, QueryResult}, - run_all, storage::layout::DataChunk, util::{timestamp_now_ms, UseOnce}, }; +use super::query_deps::{CuChecker, QueryRunner}; use super::worker::Worker; const WORKER_VERSION: &str = env!("CARGO_PKG_VERSION"); @@ -40,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); @@ -48,31 +55,53 @@ const MAX_PENDING_ASSIGNMENTS: usize = 5; const MAX_LOGS_SIZE: usize = sqd_network_transport::protocol::MAX_LOGS_RESPONSE_SIZE as usize - 100 * 1024; +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. +enum Logged { + Ok { + data_hash: Vec, + uncompressed_size: u64, + last_block: u64, + timings: TimeReport, + }, + Err(query_error::Err), +} + pub struct P2PController { - worker: Arc, + worker: Worker, worker_status: RwLock, assignment_check_interval: Duration, assignment_fetch_timeout: Duration, assignment_fetch_max_delay: Duration, raw_event_stream: UseOnce, - transport_handle: WorkerTransportHandle, + // 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, 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, + 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( - worker: Arc, + worker: Worker, transport_builder: P2PTransportBuilder, args: Args, ) -> Result>> { @@ -90,11 +119,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); @@ -108,7 +133,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, @@ -120,107 +145,103 @@ 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)), }) } +/// 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()) - .for_each_concurrent(CONCURRENT_QUERY_MESSAGES, |(peer_id, query, resp_chan)| { - let this = self.clone(); - tokio::spawn(async move { - this.handle_query(peer_id, query, resp_chan, QueryType::PlainQuery) - .await; - }) - .map(|r| r.unwrap()) - }) + .for_each_concurrent( + CONCURRENT_QUERY_MESSAGES, + |AdmittedQuery { + peer_id, + query, + resp_chan, + retry_after, + }| { + tokio::spawn(self.handle_query( + peer_id, + query, + resp_chan, + retry_after, + QueryType::PlainQuery, + )) + .map(|r| r.unwrap()) + }, + ) .await; info!("Query processing task finished"); } - async fn run_sql_queries_loop(self: Arc, 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()) - .for_each_concurrent(CONCURRENT_QUERY_MESSAGES, |(peer_id, query, resp_chan)| { - let this = self.clone(); - tokio::spawn(async move { - this.handle_query(peer_id, query, resp_chan, QueryType::SqlQuery) - .await; - }) - .map(|r| r.unwrap()) - }) + .for_each_concurrent( + CONCURRENT_QUERY_MESSAGES, + |AdmittedQuery { + peer_id, + query, + resp_chan, + retry_after, + }| { + tokio::spawn(self.handle_query( + peer_id, + query, + resp_chan, + retry_after, + QueryType::SqlQuery, + )) + .map(|r| r.unwrap()) + }, + ) .await; info!("SQL Query processing task finished"); } async fn run_assignments_loop( - &self, + &'static self, cancellation_token: CancellationToken, assignment_check_interval: Duration, ) { @@ -281,7 +302,7 @@ impl + Send + 'static> P2PController + Send + 'static> P2PController + Send + 'static> P2PController { - 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(_)) => { - warn!("Queries queue is full. Dropping query from {peer_id}"); - } - Err(mpsc::error::TrySendError::Closed(_)) => { - break; + let query = match Query::decode(request.as_ref()) { + Ok(query) => 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.admit_query(peer_id, query, resp_chan, &self.queries_tx, "Query") { + break; } } WorkerEvent::SqlQuery { peer_id, - query, + request, resp_chan, } => { - 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(_)) => { - warn!("SQL Queries queue is full. Dropping query from {peer_id}"); - } - Err(mpsc::error::TrySendError::Closed(_)) => { - break; + let query = match Query::decode(request.as_ref()) { + Ok(query) => 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.admit_query( + peer_id, + query, + resp_chan, + &self.sql_queries_tx, + "SQL query", + ) { + break; } } - WorkerEvent::LogsRequest { request, resp_chan } => { + 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(_)) => { @@ -474,12 +512,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"); + tokio::spawn(async move { + 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:?}"); } - } + }); } } } @@ -487,261 +525,199 @@ impl + Send + 'static> P2PController bool { - if !query.verify_signature(peer_id, self.worker_id) { - tracing::warn!("Rejected query with invalid signature from {}", peer_id); - return false; - } - 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 - ); - return false; - } - // TODO: check rate limits here - // TODO: check that query_id has not been used before - true - } - - #[instrument(skip_all, fields(query_id = %query.query_id, peer_id = %peer_id, dataset = %query.dataset))] - async fn handle_query( - &self, + /// 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( + &'static self, peer_id: PeerId, query: Query, - resp_chan: ResponseChannel, - query_type: QueryType, - ) { - let query_id = query.query_id.clone(); - let compression = query.compression(); - - let (mut result, retry_after) = self.process_query(peer_id, &query, query_type).await; - if let Err(e) = &result { - warn!("Query {query_id} by {peer_id} execution failed: {e:?}"); - } - - metrics::query_executed(&result); - - // Cloning is much cheaper than hash computation and we need to keep the result for logging - match self - .send_query_result( - query_id, - result.clone(), - resp_chan, - retry_after, - compression, - ) - .await - { - Ok((compression_duration, signing_duration)) => { - 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:?}"), - } - - if let Some(log) = self.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:?}"); - } + resp_chan: ResponseSender, + queue_tx: &mpsc::Sender, + protocol: &str, + ) -> bool { + if let Err(err) = self.validate_query(&query, peer_id) { + self.spawn_error_response(query.query_id, err, None, resp_chan); + return true; } - } - - /// Returns query result and the time to wait before sending the next query - #[instrument(skip_all)] - async fn process_query( - &self, - 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 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.spawn_error_response( + query.query_id, + query_error::Err::ServerOverloaded(()), + Some(DEFAULT_BACKOFF), + resp_chan, ); + return true; } - } - - 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, - ); + Err(mpsc::error::TrySendError::Closed(())) => return false, }; - - 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); + // 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.spawn_error_response( + query.query_id, + query_error::Err::TooManyRequests(()), + None, + resp_chan, + ); } - 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)); + RateLimitStatus::Paused(retry_after) => { + self.spawn_error_response( + query.query_id, + query_error::Err::TooManyRequests(()), + Some(retry_after), + resp_chan, + ); } - 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); + RateLimitStatus::Spent(retry_after) => { + permit.send(AdmittedQuery { + peer_id, + query, + resp_chan, + retry_after, + }); } } - (result, retry_after) + true } - #[instrument(skip_all)] - async fn send_query_result( + /// 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, - result: QueryResult, - resp_chan: ResponseChannel, + err: query_error::Err, 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(); + resp_chan: ResponseSender, + ) { let mut msg = sqd_messages::QueryResult { query_id, - result: Some(query_result), - retry_after_ms: retry_after.map(|duration| duration.as_millis() as u32), + result: Some(err.into()), + retry_after_ms: retry_after.map(|d| d.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}"); + 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:?}"); + } + } - tracing::trace!("Sending query result"); - // TODO: propagate backpressure from the transport lib - self.transport_handle - .send_query_result(msg, resp_chan) - .map_err(|_| anyhow!("queue full"))?; - - Ok((compression_duration, signing_duration)) + /// 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( + &'static self, + 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)] - async fn generate_log( + fn validate_query(&self, query: &Query, peer_id: PeerId) -> Result<(), query_error::Err> { + if !query.verify_signature(peer_id, self.worker_id) { + 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() + { + warn!( + "Rejected query with invalid timestamp ({}) from {peer_id}", + query.timestamp_ms + ); + return Err(query_error::Err::BadRequest( + "timestamp out of allowed range".to_owned(), + )); + } + // TODO: check that query_id has not been used before + 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, - query_result: &QueryResult, + peer_id: PeerId, 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(), - }; + resp_chan: ResponseSender, + retry_after: Option, + query_type: QueryType, + ) { + let outcome = execute( + &self.worker, + &self.allocations_checker, + peer_id, + &query, + query_type, + ) + .await; + if let Err(e) = &outcome { + warn!( + "Query {} by {peer_id} execution failed: {e:?}", + query.query_id + ); + } + metrics::query_executed(&outcome); - 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 + // Overload is transient, so hint a retry; otherwise keep the admission hint. + let retry_after = if matches!(outcome, Err(QueryError::ServiceOverloaded)) { + Some(DEFAULT_BACKOFF) + } else { + retry_after }; - 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(), - }) + let compression = query.compression(); + let (message, logged) = build_delivery( + &self.keypair, + &query.query_id, + outcome, + retry_after, + compression, + ) + .await; + + // 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:?}"); + } + + 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:?}"); + } } async fn handle_logs_request( &self, request: LogsRequest, - resp_chan: ResponseChannel, + resp_chan: ResponseSender, ) -> Result<()> { let msg = self .logs_storage @@ -757,7 +733,7 @@ impl + Send + 'static> P2PController( + worker: &W, + allocations_checker: &A, + peer_id: PeerId, + query: &Query, + query_type: QueryType, +) -> 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(), + )); + } + } + + 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())); + }; + + 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; + }; + + 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 allocation_chip < 1. { + allocations_checker.refund(peer_id, 1. - allocation_chip); + } + 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 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 +} + +/// 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), + }; + + 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 +1071,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 +1108,224 @@ mod tests { assert_eq!(pending.into_iter().collect::>(), vec![6]); } } + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use sqd_messages::{query_result, 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"; + // 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: QUERY_ID.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 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_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 outcome = execute( + &worker, + &checker, + PeerId::random(), + &query, + QueryType::PlainQuery, + ) + .await; + + assert!(matches!(outcome, Err(QueryError::ServiceOverloaded))); + assert_eq!(checker.net_spent(), 1.0); + assert_eq!(worker.calls(), 1); + } + + // 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_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 outcome = execute( + &worker, + &checker, + PeerId::random(), + &query, + QueryType::PlainQuery, + ) + .await; + + 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 outcome = execute( + &worker, + &checker, + PeerId::random(), + &query, + QueryType::PlainQuery, + ) + .await; + + 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()); + // 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 outcome = execute( + &worker, + &checker, + PeerId::random(), + &query, + QueryType::PlainQuery, + ) + .await; + + assert!(outcome.is_ok()); + assert!( + (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 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_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 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 successful_delivery_logs_a_summary() { + let keypair = Keypair::generate_ed25519(); + let (_message, logged) = + build_delivery(&keypair, QUERY_ID, ok_result(), None, Compression::None).await; + + let query = query_with(Compression::None as i32, Some((100, 200)), CHUNK_ID); + 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); + assert_eq!(summary.uncompressed_data_size, 5); + assert!(!summary.data_hash.is_empty()); + } + other => panic!("expected an Ok summary, got {other:?}"), + } + } + + // Every admitted query is logged; admission already guaranteed it's billable, so there's no + // skip path. + #[tokio::test(flavor = "multi_thread")] + 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 = 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(_)))); + } + 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..253edef --- /dev/null +++ b/src/controller/query_deps.rs @@ -0,0 +1,149 @@ +//! 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, + controller::worker::{QueryType, Worker}, + query::result::QueryResult, + types::dataset::Dataset, +}; + +/// 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 refund(&self, portal_id: PeerId, allocation_chip: f32); +} + +impl CuChecker for AllocationsChecker { + 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() + } + + /// 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; + } + } + + /// 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() + } + } +} diff --git a/src/main.rs b/src/main.rs index d526081..531297e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -24,11 +24,11 @@ #![cfg_attr(test, allow(clippy::all))] use std::borrow::Cow; -use std::sync::Arc; +use std::time::Duration; use anyhow::Result; use clap::Parser; -use tokio_util::sync::CancellationToken; +use tokio_graceful_shutdown::{SubsystemBuilder, SubsystemHandle, Toplevel}; use tracing_subscriber::fmt::format::FmtSpan; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, Layer}; @@ -104,25 +104,8 @@ fn setup_sentry(args: &Args, peer_id: String) -> 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)?; @@ -147,32 +130,35 @@ 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 = Arc::new(Worker::new(state_manager, args.parallel_queries)); + let worker = 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 = 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 (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 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 +170,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/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, } 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); - } -}