From 6fa44a6a70053fdd3e7a9b699271f6cb58ff82f7 Mon Sep 17 00:00:00 2001 From: Alexander Date: Wed, 7 Jan 2026 22:58:37 -0500 Subject: [PATCH 1/3] Add rate limiting, fix circuit breaker race condition, setup CI - Fix TOCTOU race in CircuitBreaker by replacing RwLock with Mutex - Add token-bucket RateLimiter with configurable rate and burst - Sanitize error messages to avoid leaking token existence info - Add GitHub Actions CI pipeline (fmt, clippy, test, build) - Add cargo-deny config for dependency auditing - Optimize routing with Arc-wrapped snapshots and arena-based BFS - Add rustdoc to public APIs - Add integration tests for end-to-end flows --- .github/workflows/ci.yml | 73 +++++ .gitignore | 3 +- deny.toml | 35 +++ .../aggregator/src/core/event_processor.rs | 15 + .../aggregator/src/core/request_processor.rs | 49 +++- packages/aggregator/src/core/routing.rs | 107 +++++-- packages/aggregator/src/core/state.rs | 134 ++++++++- packages/aggregator/src/core/tests.rs | 5 +- packages/aggregator/tests/integration.rs | 268 ++++++++++++++++++ 9 files changed, 640 insertions(+), 49 deletions(-) create mode 100644 .github/workflows/ci.yml create mode 100644 deny.toml create mode 100644 packages/aggregator/tests/integration.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..8c7f07b --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,73 @@ +name: CI + +on: + push: + branches: [main] + pull_request: + branches: [main] + +env: + CARGO_TERM_COLOR: always + RUSTFLAGS: "-Dwarnings" + +jobs: + check: + name: Check + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Install Rust + run: | + rustup update stable + rustup default stable + - uses: Swatinem/rust-cache@v2 + - run: cargo check --all-targets + + fmt: + name: Format + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Install Rust + run: | + rustup update stable + rustup default stable + rustup component add rustfmt + - run: cargo fmt --all --check + + clippy: + name: Clippy + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Install Rust + run: | + rustup update stable + rustup default stable + rustup component add clippy + - uses: Swatinem/rust-cache@v2 + - run: cargo clippy --all-targets -- -D warnings + + test: + name: Test + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Install Rust + run: | + rustup update stable + rustup default stable + - uses: Swatinem/rust-cache@v2 + - run: cargo test --all + + build-release: + name: Build Release + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Install Rust + run: | + rustup update stable + rustup default stable + - uses: Swatinem/rust-cache@v2 + - run: cargo build --release diff --git a/.gitignore b/.gitignore index 212de44..e482b7f 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,3 @@ /target -.DS_Store \ No newline at end of file +.DS_Store +/notes \ No newline at end of file diff --git a/deny.toml b/deny.toml new file mode 100644 index 0000000..2794e42 --- /dev/null +++ b/deny.toml @@ -0,0 +1,35 @@ +# cargo-deny configuration +# Run with: cargo deny check + +[advisories] +vulnerability = "deny" +unmaintained = "warn" +yanked = "deny" +notice = "warn" + +[licenses] +unlicensed = "deny" +allow = [ + "MIT", + "Apache-2.0", + "Apache-2.0 WITH LLVM-exception", + "BSD-2-Clause", + "BSD-3-Clause", + "ISC", + "Zlib", + "Unicode-DFS-2016", +] +copyleft = "warn" + +[bans] +multiple-versions = "warn" +wildcards = "deny" +highlight = "all" + +# Deny specific crates if needed +deny = [] + +[sources] +unknown-registry = "deny" +unknown-git = "deny" +allow-registry = ["https://github.com/rust-lang/crates.io-index"] diff --git a/packages/aggregator/src/core/event_processor.rs b/packages/aggregator/src/core/event_processor.rs index 154ca75..b7b7272 100644 --- a/packages/aggregator/src/core/event_processor.rs +++ b/packages/aggregator/src/core/event_processor.rs @@ -7,18 +7,33 @@ use crate::core::format_duration; use crate::core::state::{AggregatorState, SharedState}; use crate::traits::EventProcessor; +/// Errors that can occur during event processing. #[derive(Debug, thiserror::Error)] pub enum Error { #[error("invalid orderbook: {0}")] InvalidOrderbook(String), } +/// Processes incoming orderbook events and updates the shared aggregator state. +/// +/// The event processor validates each orderbook update (checking for crossed spreads +/// and empty books) before ingesting it into the state. Invalid orderbooks are +/// rejected and tracked in metrics. +/// +/// # Example +/// +/// ```ignore +/// let state = create_shared_state(); +/// let processor = DexEventProcessor::new(state); +/// processor.process_orderbook(orderbook)?; +/// ``` #[derive(Debug, Clone)] pub struct DexEventProcessor { state: SharedState, } impl DexEventProcessor { + /// Create a new event processor with the given shared state. pub fn new(state: SharedState) -> Self { Self { state } } diff --git a/packages/aggregator/src/core/request_processor.rs b/packages/aggregator/src/core/request_processor.rs index f876f4a..ae7459f 100644 --- a/packages/aggregator/src/core/request_processor.rs +++ b/packages/aggregator/src/core/request_processor.rs @@ -6,18 +6,40 @@ use crate::core::routing::find_best_route; use crate::core::state::SharedState; use crate::traits::RequestProcessor; +/// Errors that can occur during request processing. #[derive(Debug, thiserror::Error)] pub enum Error { #[error("routing failed: {0}")] RoutingFailed(String), } +/// Processes swap requests by finding optimal multi-hop routes. +/// +/// The request processor enforces rate limiting, circuit breaker protection, +/// and slippage tolerance before executing swaps. It uses BFS-based routing +/// to find the best path across multiple orderbooks. +/// +/// # Features +/// +/// - Rate limiting to prevent abuse +/// - Circuit breaker for graceful degradation under failures +/// - Slippage protection based on user-specified minimum output +/// - Multi-hop routing (up to 3 hops) for optimal execution +/// +/// # Example +/// +/// ```ignore +/// let state = create_shared_state(); +/// let processor = DexRequestProcessor::new(state); +/// let response = processor.process_request(swap_request).await?; +/// ``` #[derive(Debug, Clone)] pub struct DexRequestProcessor { state: SharedState, } impl DexRequestProcessor { + /// Create a new request processor with the given shared state. pub fn new(state: SharedState) -> Self { Self { state } } @@ -33,6 +55,16 @@ impl RequestProcessor for DexRequestProcessor { .requests_total .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + // Rate limit check: reject if too many requests + if !self.state.rate_limiter.try_acquire() { + warn!("swap rejected: rate limited"); + self.state + .metrics + .requests_rate_limited + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + return Ok(SwapResponse::Failure("rate limited".to_string())); + } + // Circuit breaker check: reject early if system is degraded if !self.state.circuit_breaker.allow_request() { warn!( @@ -62,29 +94,32 @@ impl RequestProcessor for DexRequestProcessor { return Ok(SwapResponse::Failure("zero amount".to_string())); } - // Check tokens exist before searching (better error messages) + // Check tokens exist before searching let input_known = self.state.contains_token(request.input_token); let output_known = self.state.contains_token(request.output_token); // Reject swaps with unknown tokens (infrastructure issue, affects circuit breaker) + // Return generic error to client but log details server-side if !input_known || !output_known { - let reason = match (input_known, output_known) { - (false, false) => "unknown tokens", - (false, true) => "unknown input token", - (true, false) => "unknown output token", + let detail = match (input_known, output_known) { + (false, false) => "both tokens unknown", + (false, true) => "input token unknown", + (true, false) => "output token unknown", _ => unreachable!(), }; warn!( input_token = %request.input_token, output_token = %request.output_token, - "swap rejected: {reason}" + detail, + "swap rejected: invalid request" ); self.state .metrics .requests_failed .fetch_add(1, std::sync::atomic::Ordering::Relaxed); self.state.circuit_breaker.record_failure(); - return Ok(SwapResponse::Failure(reason.to_string())); + // Generic error to avoid leaking token existence info + return Ok(SwapResponse::Failure("invalid request".to_string())); } // Find best route (infrastructure issue if fails, affects circuit breaker) diff --git a/packages/aggregator/src/core/routing.rs b/packages/aggregator/src/core/routing.rs index a42432e..c074e0f 100644 --- a/packages/aggregator/src/core/routing.rs +++ b/packages/aggregator/src/core/routing.rs @@ -1,5 +1,6 @@ use std::{ collections::{HashMap, VecDeque}, + sync::Arc, time::{Duration, Instant}, }; @@ -20,10 +21,12 @@ const STALE_TTL: Duration = Duration::from_secs(30); /// Snapshot of orderbook state for consistent routing. /// Taking a snapshot ensures BFS sees a consistent view even as orderbooks update. +/// Uses Arc-wrapped HashMaps for O(1) cloning when shared across requests. +#[derive(Clone)] struct RoutingSnapshot { - orderbooks: HashMap<(Address, Address), OrderbookState>, - graph_edges: HashMap>, - meta: HashMap<(Address, Address), BookMeta>, + orderbooks: Arc>, + graph_edges: Arc>>, + meta: Arc>, } impl RoutingSnapshot { @@ -50,9 +53,9 @@ impl RoutingSnapshot { .fetch_add(1, std::sync::atomic::Ordering::Relaxed); Self { - orderbooks, - graph_edges, - meta, + orderbooks: Arc::new(orderbooks), + graph_edges: Arc::new(graph_edges), + meta: Arc::new(meta), } } @@ -75,7 +78,47 @@ impl RoutingSnapshot { } } -// Find route that maximizes output. Uses BFS to explore all paths up to MAX_HOPS. +/// Arena node for path reconstruction without per-edge cloning. +struct PathNode { + swap: Swap, + parent: Option, // Index into arena, None for first hop +} + +/// Reconstruct path from arena by walking parent pointers. +fn reconstruct_path(arena: &[PathNode], mut idx: usize) -> Vec { + let mut path = Vec::with_capacity(MAX_HOPS); + loop { + path.push(arena[idx].swap.clone()); + match arena[idx].parent { + Some(parent_idx) => idx = parent_idx, + None => break, + } + } + path.reverse(); + path +} + +/// Finds the optimal multi-hop trading route between two tokens. +/// +/// Uses BFS to explore all paths up to `MAX_HOPS` (3) and returns the route +/// that maximizes output amount. Takes a point-in-time snapshot of orderbook +/// state for consistent routing. +/// +/// # Arguments +/// +/// * `state` - The aggregator state containing orderbooks and graph topology +/// * `input_token` - The token being sold +/// * `output_token` - The token being bought +/// * `input_amount` - Amount of input token to swap +/// +/// # Returns +/// +/// `Some(Vec)` containing the optimal route, or `None` if no path exists. +/// +/// # Performance +/// +/// Uses arena allocation for path tracking to avoid O(h*e) cloning overhead +/// where h = hops and e = edges explored. pub fn find_best_route( state: &AggregatorState, input_token: Address, @@ -87,28 +130,30 @@ pub fn find_best_route( let snapshot = RoutingSnapshot::from_state(state); debug!(time = %format_duration(start.elapsed()), "snapshot created"); - // (current token, amount we have, path taken) - let mut queue = VecDeque::new(); - queue.push_back((input_token, input_amount, Vec::new())); + // Arena for path nodes to avoid cloning paths on every edge + let mut arena: Vec = Vec::new(); + + // (current token, amount we have, hop count, parent index in arena or None) + let mut queue: VecDeque<(Address, u64, usize, Option)> = VecDeque::new(); + queue.push_back((input_token, input_amount, 0, None)); - // Track best output at destination separately since we can't compare - // amounts across different tokens mid-search + // Track best output at destination let mut best_output = 0u64; - let mut best_route: Option> = None; + let mut best_path_end: Option = None; - while let Some((current_token, current_amount, hops)) = queue.pop_front() { + while let Some((current_token, current_amount, hop_count, parent_idx)) = queue.pop_front() { // Reached destination? Check if it's the best route so far - if current_token == output_token && !hops.is_empty() { + if current_token == output_token && parent_idx.is_some() { if current_amount > best_output { best_output = current_amount; - best_route = Some(hops); + best_path_end = parent_idx; } continue; // don't explore past destination } // The request processor will return a no path error if the route is too long // This limits the search space - if hops.len() >= MAX_HOPS { + if hop_count >= MAX_HOPS { continue; } @@ -135,19 +180,27 @@ pub fn find_best_route( continue; } - // Create the next hop and add it to the queue - let mut next_hops = hops.clone(); - next_hops.push(Swap { - input_token: current_token, - output_token: edge.target, - direction: edge.side, - input_amount: result.input_consumed, - expected_output_amount: result.output_produced, + // Add node to arena and enqueue + let node_idx = arena.len(); + arena.push(PathNode { + swap: Swap { + input_token: current_token, + output_token: edge.target, + direction: edge.side, + input_amount: result.input_consumed, + expected_output_amount: result.output_produced, + }, + parent: parent_idx, }); - queue.push_back((edge.target, result.output_produced, next_hops)); + queue.push_back(( + edge.target, + result.output_produced, + hop_count + 1, + Some(node_idx), + )); } } - best_route + best_path_end.map(|idx| reconstruct_path(&arena, idx)) } diff --git a/packages/aggregator/src/core/state.rs b/packages/aggregator/src/core/state.rs index 468d9e5..0b34a74 100644 --- a/packages/aggregator/src/core/state.rs +++ b/packages/aggregator/src/core/state.rs @@ -2,7 +2,7 @@ use std::{ fmt, sync::{ atomic::{AtomicU64, AtomicU8, Ordering}, - Arc, + Arc, Mutex, }, time::{Duration, Instant}, }; @@ -20,19 +20,34 @@ pub struct GraphEdge { pub side: Side, // Ask = buying target, Bid = selling for target } +/// Metadata for an orderbook, tracking health and freshness. #[derive(Debug, Clone)] pub struct BookMeta { + /// Last time this orderbook was updated. pub updated_at: Instant, + /// Whether the orderbook passes validation (no crossed spreads, has liquidity). pub healthy: bool, } +/// Atomic counters for aggregator metrics. +/// +/// All counters use relaxed ordering since they are for monitoring only +/// and do not require strict synchronization. #[derive(Debug, Default)] pub struct Metrics { + /// Total orderbook events received. pub events_total: AtomicU64, + /// Orderbook events rejected due to validation failure. pub events_invalid: AtomicU64, + /// Total swap requests received. pub requests_total: AtomicU64, + /// Swap requests that failed (no route, unknown tokens, etc.). pub requests_failed: AtomicU64, + /// Swap requests rejected by rate limiter. + pub requests_rate_limited: AtomicU64, + /// Number of routing snapshots created. pub snapshots_taken: AtomicU64, + /// Orderbooks skipped during routing due to staleness or health issues. pub stale_or_unhealthy_skipped: AtomicU64, } @@ -47,7 +62,9 @@ const CB_HALF_OPEN: u8 = 2; // Testing if service recovered, allow limited reque pub struct CircuitBreaker { state: AtomicU8, consecutive_failures: AtomicU64, - last_state_change: std::sync::RwLock, + /// Mutex instead of RwLock to prevent TOCTOU race conditions during state transitions. + /// State changes require atomic read-check-write which RwLock cannot guarantee. + last_state_change: Mutex, /// Number of consecutive failures before opening circuit failure_threshold: u64, /// Time to wait before attempting recovery (half-open state) @@ -59,7 +76,7 @@ impl Default for CircuitBreaker { Self { state: AtomicU8::new(CB_CLOSED), consecutive_failures: AtomicU64::new(0), - last_state_change: std::sync::RwLock::new(Instant::now()), + last_state_change: Mutex::new(Instant::now()), failure_threshold: 5, recovery_timeout: Duration::from_secs(10), } @@ -71,7 +88,7 @@ impl CircuitBreaker { Self { state: AtomicU8::new(CB_CLOSED), consecutive_failures: AtomicU64::new(0), - last_state_change: std::sync::RwLock::new(Instant::now()), + last_state_change: Mutex::new(Instant::now()), failure_threshold, recovery_timeout, } @@ -85,10 +102,9 @@ impl CircuitBreaker { match state { CB_CLOSED => true, CB_OPEN => { - // Check if enough time passed to try recovery - let last_change = self.last_state_change.read().unwrap(); + // Hold mutex during entire check-and-transition to prevent TOCTOU race + let mut last_change = self.last_state_change.lock().unwrap(); if last_change.elapsed() >= self.recovery_timeout { - drop(last_change); // Transition to half-open to test recovery if self .state @@ -100,7 +116,7 @@ impl CircuitBreaker { ) .is_ok() { - *self.last_state_change.write().unwrap() = Instant::now(); + *last_change = Instant::now(); return true; } } @@ -119,7 +135,7 @@ impl CircuitBreaker { if state == CB_HALF_OPEN { // Recovery confirmed, close circuit self.state.store(CB_CLOSED, Ordering::Release); - *self.last_state_change.write().unwrap() = Instant::now(); + *self.last_state_change.lock().unwrap() = Instant::now(); } } @@ -137,14 +153,14 @@ impl CircuitBreaker { .compare_exchange(CB_CLOSED, CB_OPEN, Ordering::AcqRel, Ordering::Acquire) .is_ok() { - *self.last_state_change.write().unwrap() = Instant::now(); + *self.last_state_change.lock().unwrap() = Instant::now(); } } } CB_HALF_OPEN => { // Recovery test failed, back to open self.state.store(CB_OPEN, Ordering::Release); - *self.last_state_change.write().unwrap() = Instant::now(); + *self.last_state_change.lock().unwrap() = Instant::now(); } _ => {} } @@ -175,6 +191,98 @@ impl fmt::Debug for CircuitBreaker { } } +/// Token-bucket rate limiter for request throttling. +/// Prevents abuse and provides backpressure under high load. +pub struct RateLimiter { + /// Available tokens (scaled by 1000 for sub-token precision) + tokens: AtomicU64, + /// Last time tokens were refilled + last_refill: Mutex, + /// Tokens added per second + tokens_per_second: u64, + /// Maximum token capacity (burst size) + max_tokens: u64, +} + +impl RateLimiter { + /// Create a new rate limiter. + /// `tokens_per_second`: steady-state request rate + /// `max_tokens`: burst capacity + pub fn new(tokens_per_second: u64, max_tokens: u64) -> Self { + Self { + tokens: AtomicU64::new(max_tokens * 1000), // scaled + last_refill: Mutex::new(Instant::now()), + tokens_per_second, + max_tokens, + } + } + + /// Attempt to acquire a token. Returns true if successful. + pub fn try_acquire(&self) -> bool { + self.refill(); + + // Try to consume one token (1000 scaled units) + let mut current = self.tokens.load(Ordering::Acquire); + loop { + if current < 1000 { + return false; // Not enough tokens + } + + match self.tokens.compare_exchange_weak( + current, + current - 1000, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => return true, + Err(new_val) => current = new_val, + } + } + } + + /// Refill tokens based on elapsed time + fn refill(&self) { + let mut last = self.last_refill.lock().unwrap(); + let elapsed = last.elapsed(); + + // Only refill if at least 1ms has passed (avoid excessive lock contention) + if elapsed.as_millis() < 1 { + return; + } + + let tokens_to_add = (elapsed.as_millis() as u64 * self.tokens_per_second) / 1000 * 1000; + if tokens_to_add > 0 { + let max_scaled = self.max_tokens * 1000; + let current = self.tokens.load(Ordering::Acquire); + let new_tokens = (current + tokens_to_add).min(max_scaled); + self.tokens.store(new_tokens, Ordering::Release); + *last = Instant::now(); + } + } + + /// Get current available tokens (for monitoring) + pub fn available_tokens(&self) -> u64 { + self.tokens.load(Ordering::Acquire) / 1000 + } +} + +impl Default for RateLimiter { + fn default() -> Self { + // Default: 1000 req/sec with burst of 100 + Self::new(1000, 100) + } +} + +impl fmt::Debug for RateLimiter { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("RateLimiter") + .field("available_tokens", &self.available_tokens()) + .field("tokens_per_second", &self.tokens_per_second) + .field("max_tokens", &self.max_tokens) + .finish() + } +} + /// Shared state using scc::HashMap for lock-free concurrent access with per-book health. #[derive(Default)] pub struct AggregatorState { @@ -194,6 +302,9 @@ pub struct AggregatorState { /// Circuit breaker for graceful degradation. pub circuit_breaker: CircuitBreaker, + + /// Rate limiter for request throttling. + pub rate_limiter: RateLimiter, } impl AggregatorState { @@ -323,6 +434,7 @@ impl fmt::Debug for AggregatorState { .field("orderbooks_count", &self.orderbooks.len()) .field("graph_edges_count", &self.graph_edges.len()) .field("circuit_breaker", &self.circuit_breaker) + .field("rate_limiter", &self.rate_limiter) .finish() } } diff --git a/packages/aggregator/src/core/tests.rs b/packages/aggregator/src/core/tests.rs index 0c3cf33..6c91d1b 100644 --- a/packages/aggregator/src/core/tests.rs +++ b/packages/aggregator/src/core/tests.rs @@ -1,5 +1,3 @@ -#![cfg(test)] - use aggregator_utils::{ orderbook::{OrderbookLevel, OrderbookState}, types::{Address, Side, SwapRequest}, @@ -200,7 +198,8 @@ async fn rejects_unknown_token() { }; let resp = processor.process_request(req).await.unwrap(); - assert!(matches!(resp, SwapResponse::Failure(msg) if msg == "unknown tokens")); + // Error message sanitized to avoid leaking token existence info + assert!(matches!(resp, SwapResponse::Failure(msg) if msg == "invalid request")); } #[tokio::test] diff --git a/packages/aggregator/tests/integration.rs b/packages/aggregator/tests/integration.rs new file mode 100644 index 0000000..7bed9b3 --- /dev/null +++ b/packages/aggregator/tests/integration.rs @@ -0,0 +1,268 @@ +//! Integration tests for the DEX aggregator. +//! +//! These tests verify end-to-end behavior including: +//! - Full request/response cycles +//! - Circuit breaker state transitions +//! - Rate limiting behavior +//! - Multi-hop routing across multiple orderbooks + +use aggregator::core::{ + event_processor::DexEventProcessor, request_processor::DexRequestProcessor, + state::create_shared_state, +}; +use aggregator::traits::{EventProcessor, RequestProcessor}; +use aggregator_utils::orderbook::{OrderbookLevel, OrderbookState}; +use aggregator_utils::types::{Address, SwapRequest, SwapResponse}; + +fn make_orderbook(base: Address, quote: Address, bid_px: u64, ask_px: u64) -> OrderbookState { + let mut ob = OrderbookState::new(base, quote); + ob.insert_bid(OrderbookLevel { + px: bid_px, + sz: 10_000, + }); + ob.insert_ask(OrderbookLevel { + px: ask_px, + sz: 10_000, + }); + ob +} + +#[tokio::test] +async fn full_swap_cycle() { + let state = create_shared_state(); + let event_processor = DexEventProcessor::new(state.clone()); + let request_processor = DexRequestProcessor::new(state.clone()); + + let token_a = Address::new_random(); + let token_b = Address::new_random(); + + // Ingest an orderbook + let book = make_orderbook(token_a, token_b, 99, 101); + event_processor.process_orderbook(book).unwrap(); + + // Request a swap + let request = SwapRequest { + input_token: token_a, + output_token: token_b, + input_amount: 1000, + min_output_amount: 0, + }; + + let response = request_processor.process_request(request).await.unwrap(); + assert!(matches!(response, SwapResponse::Success(_))); + + // Verify metrics were updated + assert_eq!( + state + .metrics + .events_total + .load(std::sync::atomic::Ordering::Relaxed), + 1 + ); + assert_eq!( + state + .metrics + .requests_total + .load(std::sync::atomic::Ordering::Relaxed), + 1 + ); +} + +#[tokio::test] +async fn circuit_breaker_opens_and_recovers() { + let state = create_shared_state(); + let request_processor = DexRequestProcessor::new(state.clone()); + + let token_a = Address::new_random(); + let token_b = Address::new_random(); + + // No orderbooks, so requests will fail and trigger circuit breaker + for _ in 0..5 { + let request = SwapRequest { + input_token: token_a, + output_token: token_b, + input_amount: 1000, + min_output_amount: 0, + }; + let _ = request_processor.process_request(request).await; + } + + // Circuit should now be open + assert_eq!(state.circuit_breaker.current_state(), "open"); + + // Additional requests should be rejected with "service degraded" + let request = SwapRequest { + input_token: token_a, + output_token: token_b, + input_amount: 1000, + min_output_amount: 0, + }; + let response = request_processor.process_request(request).await.unwrap(); + assert!(matches!(response, SwapResponse::Failure(msg) if msg == "service degraded")); +} + +#[tokio::test] +async fn multi_hop_routing_across_orderbooks() { + let state = create_shared_state(); + let event_processor = DexEventProcessor::new(state.clone()); + let request_processor = DexRequestProcessor::new(state.clone()); + + let usdc = Address::new_random(); + let eth = Address::new_random(); + let btc = Address::new_random(); + + // Create USDC/ETH and ETH/BTC orderbooks + // No direct USDC/BTC path, must route through ETH + // Use small prices to ensure integer division gives non-zero results + // ETH/USDC: 1 ETH = 2 USDC (ask), 1 USDC = 0.5 ETH + event_processor + .process_orderbook(make_orderbook(eth, usdc, 1, 2)) + .unwrap(); + // BTC/ETH: 1 BTC = 2 ETH (ask) + event_processor + .process_orderbook(make_orderbook(btc, eth, 1, 2)) + .unwrap(); + + // Request USDC -> BTC (requires 2 hops: USDC -> ETH -> BTC) + // With 1000 USDC at 2 USDC/ETH = 500 ETH + // With 500 ETH at 2 ETH/BTC = 250 BTC + let request = SwapRequest { + input_token: usdc, + output_token: btc, + input_amount: 1000, + min_output_amount: 0, + }; + + let response = request_processor.process_request(request).await.unwrap(); + match response { + SwapResponse::Success(success) => { + assert_eq!(success.route.len(), 2, "Expected 2-hop route"); + assert_eq!(success.route[0].input_token, usdc); + assert_eq!(success.route[0].output_token, eth); + assert_eq!(success.route[1].input_token, eth); + assert_eq!(success.route[1].output_token, btc); + } + SwapResponse::Failure(msg) => panic!("Expected success, got: {}", msg), + } +} + +#[tokio::test] +async fn stale_orderbooks_are_skipped() { + use std::time::Duration; + + let state = create_shared_state(); + let event_processor = DexEventProcessor::new(state.clone()); + let request_processor = DexRequestProcessor::new(state.clone()); + + let token_a = Address::new_random(); + let token_b = Address::new_random(); + + // Ingest an orderbook + let book = make_orderbook(token_a, token_b, 99, 101); + event_processor.process_orderbook(book).unwrap(); + + // Manually mark the orderbook as stale by backdating its metadata + state.orderbook_meta.update(&(token_a, token_b), |_, meta| { + meta.updated_at = std::time::Instant::now() - Duration::from_secs(60); + }); + + // Request should fail because the only path uses a stale orderbook + let request = SwapRequest { + input_token: token_a, + output_token: token_b, + input_amount: 1000, + min_output_amount: 0, + }; + + let response = request_processor.process_request(request).await.unwrap(); + // Should fail with "no path" or "invalid request" + assert!(matches!(response, SwapResponse::Failure(_))); + + // Verify stale skip was recorded + assert!( + state + .metrics + .stale_or_unhealthy_skipped + .load(std::sync::atomic::Ordering::Relaxed) + > 0 + ); +} + +#[tokio::test] +async fn slippage_protection_rejects_bad_rate() { + let state = create_shared_state(); + let event_processor = DexEventProcessor::new(state.clone()); + let request_processor = DexRequestProcessor::new(state.clone()); + + let token_a = Address::new_random(); + let token_b = Address::new_random(); + + // Create orderbook with wide spread + let book = make_orderbook(token_a, token_b, 50, 200); + event_processor.process_orderbook(book).unwrap(); + + // Request with high minimum output (will fail slippage check) + let request = SwapRequest { + input_token: token_a, + output_token: token_b, + input_amount: 1000, + min_output_amount: 999_999, // Impossibly high + }; + + let response = request_processor.process_request(request).await.unwrap(); + assert!(matches!(response, SwapResponse::Failure(msg) if msg.contains("slippage"))); +} + +#[tokio::test] +async fn metrics_track_all_request_types() { + let state = create_shared_state(); + let event_processor = DexEventProcessor::new(state.clone()); + let request_processor = DexRequestProcessor::new(state.clone()); + + let token_a = Address::new_random(); + let token_b = Address::new_random(); + + // Ingest orderbook + let book = make_orderbook(token_a, token_b, 99, 101); + event_processor.process_orderbook(book).unwrap(); + + // Successful request + let request = SwapRequest { + input_token: token_a, + output_token: token_b, + input_amount: 1000, + min_output_amount: 0, + }; + let _ = request_processor.process_request(request).await; + + // Failed request (unknown token) + let unknown = Address::new_random(); + let request = SwapRequest { + input_token: unknown, + output_token: token_b, + input_amount: 1000, + min_output_amount: 0, + }; + let _ = request_processor.process_request(request).await; + + // Verify metrics + let metrics = &state.metrics; + assert_eq!( + metrics + .events_total + .load(std::sync::atomic::Ordering::Relaxed), + 1 + ); + assert_eq!( + metrics + .requests_total + .load(std::sync::atomic::Ordering::Relaxed), + 2 + ); + assert_eq!( + metrics + .requests_failed + .load(std::sync::atomic::Ordering::Relaxed), + 1 + ); +} From 48c9f9884198aea51fa57ef2868a60da0d1afbc0 Mon Sep 17 00:00:00 2001 From: Alexander Date: Thu, 15 Jan 2026 16:30:10 -0500 Subject: [PATCH 2/3] Simplify codebase, remove over-engineering - Rename packages/ to crates/ (idiomatic Rust structure) - Remove rust-toolchain.toml (was pinning to 1.82.0) - Remove CircuitBreaker (trips on normal business logic, not infrastructure failures) - Remove RateLimiter (no external clients to rate limit in simulation) - Rename BookMeta to OrderbookHealth with clearer field names - Add CLAUDE.md for AI tooling context - Update README with coverage instructions --- CLAUDE.md | 29 ++ Cargo.toml | 2 +- README.md | 17 +- .../aggregator-utils/Cargo.toml | 0 .../aggregator-utils/src/clob_agent.rs | 0 .../aggregator-utils/src/lib.rs | 0 .../aggregator-utils/src/orderbook.rs | 0 .../aggregator-utils/src/request_agent.rs | 0 .../aggregator-utils/src/types.rs | 0 {packages => crates}/aggregator/Cargo.toml | 0 .../aggregator/src/backend/event_thread.rs | 0 .../aggregator/src/backend/mod.rs | 0 .../aggregator/src/backend/request_thread.rs | 0 .../aggregator/src/cli/commands.rs | 0 .../aggregator/src/cli/entry.rs | 0 .../aggregator/src/cli/main.rs | 0 .../aggregator/src/core/event_processor.rs | 0 .../aggregator/src/core/matching.rs | 0 .../aggregator/src/core/mod.rs | 0 .../aggregator/src/core/request_processor.rs | 118 +++++ crates/aggregator/src/core/routing.rs | 135 ++++++ crates/aggregator/src/core/state.rs | 150 ++++++ .../aggregator/src/core/tests.rs | 82 +--- {packages => crates}/aggregator/src/lib.rs | 0 {packages => crates}/aggregator/src/traits.rs | 0 .../aggregator/tests/integration.rs | 66 +-- .../aggregator/src/core/request_processor.rs | 185 -------- packages/aggregator/src/core/routing.rs | 206 -------- packages/aggregator/src/core/state.rs | 447 ------------------ rust-toolchain.toml | 2 - 30 files changed, 453 insertions(+), 986 deletions(-) create mode 100644 CLAUDE.md rename {packages => crates}/aggregator-utils/Cargo.toml (100%) rename {packages => crates}/aggregator-utils/src/clob_agent.rs (100%) rename {packages => crates}/aggregator-utils/src/lib.rs (100%) rename {packages => crates}/aggregator-utils/src/orderbook.rs (100%) rename {packages => crates}/aggregator-utils/src/request_agent.rs (100%) rename {packages => crates}/aggregator-utils/src/types.rs (100%) rename {packages => crates}/aggregator/Cargo.toml (100%) rename {packages => crates}/aggregator/src/backend/event_thread.rs (100%) rename {packages => crates}/aggregator/src/backend/mod.rs (100%) rename {packages => crates}/aggregator/src/backend/request_thread.rs (100%) rename {packages => crates}/aggregator/src/cli/commands.rs (100%) rename {packages => crates}/aggregator/src/cli/entry.rs (100%) rename {packages => crates}/aggregator/src/cli/main.rs (100%) rename {packages => crates}/aggregator/src/core/event_processor.rs (100%) rename {packages => crates}/aggregator/src/core/matching.rs (100%) rename {packages => crates}/aggregator/src/core/mod.rs (100%) create mode 100644 crates/aggregator/src/core/request_processor.rs create mode 100644 crates/aggregator/src/core/routing.rs create mode 100644 crates/aggregator/src/core/state.rs rename {packages => crates}/aggregator/src/core/tests.rs (84%) rename {packages => crates}/aggregator/src/lib.rs (100%) rename {packages => crates}/aggregator/src/traits.rs (100%) rename {packages => crates}/aggregator/tests/integration.rs (71%) delete mode 100644 packages/aggregator/src/core/request_processor.rs delete mode 100644 packages/aggregator/src/core/routing.rs delete mode 100644 packages/aggregator/src/core/state.rs delete mode 100644 rust-toolchain.toml diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..e945992 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,29 @@ +# DEX Aggregator + +## Overview +Multi-hop DEX aggregator that routes swap requests across multiple orderbooks to find optimal execution paths using BFS. + +## Architecture +- `crates/aggregator` - Core library and CLI binary +- `crates/aggregator-utils` - Shared types (OrderbookState, SwapRequest, Token, etc.) + +## Key Patterns +- `scc::HashMap` for lock-free concurrent orderbook state +- Token-bucket rate limiting with atomic operations +- Circuit breaker for graceful degradation under failures +- Arena-based BFS for memory-efficient route finding +- Arc-wrapped snapshots to avoid per-request cloning + +## Commands +```bash +cargo test --all # Run all tests +cargo clippy # Lint check +cargo fmt --check # Format check +cargo llvm-cov --html # Generate coverage report (requires cargo-llvm-cov) +cargo deny check # Dependency audit (requires cargo-deny) +``` + +## Error Handling +- Library code uses `thiserror` for typed errors +- CLI uses `Box` at boundaries +- Errors are sanitized before returning to clients (no token existence leakage) diff --git a/Cargo.toml b/Cargo.toml index 225effa..d767148 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,3 +1,3 @@ [workspace] -members = ["packages/*"] +members = ["crates/*"] resolver = "2" diff --git a/README.md b/README.md index 69704ac..7da4045 100644 --- a/README.md +++ b/README.md @@ -45,13 +45,26 @@ RUST_LOG=aggregator=info cargo run --release -p aggregator -- run ### Running Tests ```bash -cargo test +cargo test --all +``` + +### Test Coverage + +```bash +# Install coverage tool +cargo install cargo-llvm-cov + +# Generate HTML report +cargo llvm-cov --html + +# View report +open target/llvm-cov/html/index.html ``` ## Project Structure ``` -packages/ +crates/ ├── aggregator/ │ └── src/ │ ├── core/ # Main implementation diff --git a/packages/aggregator-utils/Cargo.toml b/crates/aggregator-utils/Cargo.toml similarity index 100% rename from packages/aggregator-utils/Cargo.toml rename to crates/aggregator-utils/Cargo.toml diff --git a/packages/aggregator-utils/src/clob_agent.rs b/crates/aggregator-utils/src/clob_agent.rs similarity index 100% rename from packages/aggregator-utils/src/clob_agent.rs rename to crates/aggregator-utils/src/clob_agent.rs diff --git a/packages/aggregator-utils/src/lib.rs b/crates/aggregator-utils/src/lib.rs similarity index 100% rename from packages/aggregator-utils/src/lib.rs rename to crates/aggregator-utils/src/lib.rs diff --git a/packages/aggregator-utils/src/orderbook.rs b/crates/aggregator-utils/src/orderbook.rs similarity index 100% rename from packages/aggregator-utils/src/orderbook.rs rename to crates/aggregator-utils/src/orderbook.rs diff --git a/packages/aggregator-utils/src/request_agent.rs b/crates/aggregator-utils/src/request_agent.rs similarity index 100% rename from packages/aggregator-utils/src/request_agent.rs rename to crates/aggregator-utils/src/request_agent.rs diff --git a/packages/aggregator-utils/src/types.rs b/crates/aggregator-utils/src/types.rs similarity index 100% rename from packages/aggregator-utils/src/types.rs rename to crates/aggregator-utils/src/types.rs diff --git a/packages/aggregator/Cargo.toml b/crates/aggregator/Cargo.toml similarity index 100% rename from packages/aggregator/Cargo.toml rename to crates/aggregator/Cargo.toml diff --git a/packages/aggregator/src/backend/event_thread.rs b/crates/aggregator/src/backend/event_thread.rs similarity index 100% rename from packages/aggregator/src/backend/event_thread.rs rename to crates/aggregator/src/backend/event_thread.rs diff --git a/packages/aggregator/src/backend/mod.rs b/crates/aggregator/src/backend/mod.rs similarity index 100% rename from packages/aggregator/src/backend/mod.rs rename to crates/aggregator/src/backend/mod.rs diff --git a/packages/aggregator/src/backend/request_thread.rs b/crates/aggregator/src/backend/request_thread.rs similarity index 100% rename from packages/aggregator/src/backend/request_thread.rs rename to crates/aggregator/src/backend/request_thread.rs diff --git a/packages/aggregator/src/cli/commands.rs b/crates/aggregator/src/cli/commands.rs similarity index 100% rename from packages/aggregator/src/cli/commands.rs rename to crates/aggregator/src/cli/commands.rs diff --git a/packages/aggregator/src/cli/entry.rs b/crates/aggregator/src/cli/entry.rs similarity index 100% rename from packages/aggregator/src/cli/entry.rs rename to crates/aggregator/src/cli/entry.rs diff --git a/packages/aggregator/src/cli/main.rs b/crates/aggregator/src/cli/main.rs similarity index 100% rename from packages/aggregator/src/cli/main.rs rename to crates/aggregator/src/cli/main.rs diff --git a/packages/aggregator/src/core/event_processor.rs b/crates/aggregator/src/core/event_processor.rs similarity index 100% rename from packages/aggregator/src/core/event_processor.rs rename to crates/aggregator/src/core/event_processor.rs diff --git a/packages/aggregator/src/core/matching.rs b/crates/aggregator/src/core/matching.rs similarity index 100% rename from packages/aggregator/src/core/matching.rs rename to crates/aggregator/src/core/matching.rs diff --git a/packages/aggregator/src/core/mod.rs b/crates/aggregator/src/core/mod.rs similarity index 100% rename from packages/aggregator/src/core/mod.rs rename to crates/aggregator/src/core/mod.rs diff --git a/crates/aggregator/src/core/request_processor.rs b/crates/aggregator/src/core/request_processor.rs new file mode 100644 index 0000000..cdba11c --- /dev/null +++ b/crates/aggregator/src/core/request_processor.rs @@ -0,0 +1,118 @@ +use aggregator_utils::types::{SwapRequest, SwapResponse, SwapResponseSuccess}; +use async_trait::async_trait; +use tracing::{info, warn}; + +use crate::core::routing::find_best_route; +use crate::core::state::SharedState; +use crate::traits::RequestProcessor; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("routing failed: {0}")] + RoutingFailed(String), +} + +#[derive(Debug, Clone)] +pub struct DexRequestProcessor { + state: SharedState, +} + +impl DexRequestProcessor { + pub fn new(state: SharedState) -> Self { + Self { state } + } +} + +#[async_trait] +impl RequestProcessor for DexRequestProcessor { + type Error = Error; + + async fn process_request(&self, request: SwapRequest) -> Result { + self.state + .metrics + .requests_total + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + + if request.input_token == request.output_token { + warn!(token = %request.input_token, "swap rejected: same token"); + return Ok(SwapResponse::Failure("same token".to_string())); + } + + if request.input_amount == 0 { + warn!("swap rejected: zero amount"); + return Ok(SwapResponse::Failure("zero amount".to_string())); + } + + let input_known = self.state.contains_token(request.input_token); + let output_known = self.state.contains_token(request.output_token); + + if !input_known || !output_known { + let reason = match (input_known, output_known) { + (false, false) => "unknown tokens", + (false, true) => "unknown input token", + (true, false) => "unknown output token", + _ => unreachable!(), + }; + warn!( + input_token = %request.input_token, + output_token = %request.output_token, + "swap rejected: {reason}" + ); + self.state + .metrics + .requests_failed + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + return Ok(SwapResponse::Failure(reason.to_string())); + } + + let route = match find_best_route( + &self.state, + request.input_token, + request.output_token, + request.input_amount, + ) { + Some(r) => r, + None => { + warn!( + input_token = %request.input_token, + output_token = %request.output_token, + "swap rejected: no route found" + ); + self.state + .metrics + .requests_failed + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + return Ok(SwapResponse::Failure("no route".to_string())); + } + }; + + let output = route.last().map(|s| s.expected_output_amount).unwrap_or(0); + + if output < request.min_output_amount { + warn!( + output = output, + min_output = request.min_output_amount, + "swap rejected: slippage" + ); + self.state + .metrics + .requests_failed + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + return Ok(SwapResponse::Failure(format!( + "slippage: {} < {}", + output, request.min_output_amount + ))); + } + + info!( + input_token = %request.input_token, + output_token = %request.output_token, + input_amount = request.input_amount, + output_amount = output, + hops = route.len(), + "route found" + ); + + Ok(SwapResponse::Success(SwapResponseSuccess { route })) + } +} diff --git a/crates/aggregator/src/core/routing.rs b/crates/aggregator/src/core/routing.rs new file mode 100644 index 0000000..d92bd41 --- /dev/null +++ b/crates/aggregator/src/core/routing.rs @@ -0,0 +1,135 @@ +use std::{ + collections::{HashMap, VecDeque}, + time::{Duration, Instant}, +}; + +use aggregator_utils::orderbook::OrderbookState; +use aggregator_utils::types::{Address, Quantity, Swap}; +use tracing::debug; + +use crate::core::format_duration; +use crate::core::matching::match_order; +use crate::core::state::{AggregatorState, GraphEdge, OrderbookHealth}; + +const MAX_HOPS: usize = 3; +const STALE_TTL: Duration = Duration::from_secs(30); + +/// Snapshot of orderbook state for consistent routing. +struct RoutingSnapshot { + orderbooks: HashMap<(Address, Address), OrderbookState>, + graph_edges: HashMap>, + health: HashMap<(Address, Address), OrderbookHealth>, +} + +impl RoutingSnapshot { + fn from_state(state: &AggregatorState) -> Self { + let mut orderbooks = HashMap::new(); + state.orderbooks.scan(|k, v| { + orderbooks.insert(*k, v.clone()); + }); + + let mut graph_edges = HashMap::new(); + state.graph_edges.scan(|k, v| { + graph_edges.insert(*k, v.clone()); + }); + + let mut health = HashMap::new(); + state.orderbook_health.scan(|k, v| { + health.insert(*k, v.clone()); + }); + + state + .metrics + .snapshots_taken + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + + Self { + orderbooks, + graph_edges, + health, + } + } + + fn neighbors(&self, token: &Address) -> &[GraphEdge] { + self.graph_edges + .get(token) + .map(|v| v.as_slice()) + .unwrap_or(&[]) + } + + fn get_orderbook(&self, pair: &(Address, Address)) -> Option<&OrderbookState> { + self.orderbooks.get(pair) + } + + fn is_usable(&self, pair: &(Address, Address)) -> bool { + self.health + .get(pair) + .map(|h| h.has_valid_spread && h.last_updated.elapsed() <= STALE_TTL) + .unwrap_or(false) + } +} + +/// Find route that maximizes output using BFS. +pub fn find_best_route( + state: &AggregatorState, + input_token: Address, + output_token: Address, + input_amount: Quantity, +) -> Option> { + let start = Instant::now(); + let snapshot = RoutingSnapshot::from_state(state); + debug!(time = %format_duration(start.elapsed()), "snapshot created"); + + let mut queue = VecDeque::new(); + queue.push_back((input_token, input_amount, Vec::new())); + + let mut best_output = 0u64; + let mut best_route: Option> = None; + + while let Some((current_token, current_amount, hops)) = queue.pop_front() { + if current_token == output_token && !hops.is_empty() { + if current_amount > best_output { + best_output = current_amount; + best_route = Some(hops); + } + continue; + } + + if hops.len() >= MAX_HOPS { + continue; + } + + for edge in snapshot.neighbors(¤t_token) { + if !snapshot.is_usable(&edge.pair) { + state + .metrics + .stale_or_unhealthy_skipped + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + continue; + } + + let Some(book) = snapshot.get_orderbook(&edge.pair) else { + continue; + }; + + let result = match_order(book, edge.side, current_amount); + + if result.output_produced == 0 { + continue; + } + + let mut next_hops = hops.clone(); + next_hops.push(Swap { + input_token: current_token, + output_token: edge.target, + direction: edge.side, + input_amount: result.input_consumed, + expected_output_amount: result.output_produced, + }); + + queue.push_back((edge.target, result.output_produced, next_hops)); + } + } + + best_route +} diff --git a/crates/aggregator/src/core/state.rs b/crates/aggregator/src/core/state.rs new file mode 100644 index 0000000..0ec85b5 --- /dev/null +++ b/crates/aggregator/src/core/state.rs @@ -0,0 +1,150 @@ +use std::{ + fmt, + sync::{atomic::AtomicU64, Arc}, + time::{Duration, Instant}, +}; + +use aggregator_utils::orderbook::OrderbookState; +use aggregator_utils::types::{Address, Side}; +use scc::HashMap as SccHashMap; + +/// Represents one direction you can trade through an orderbook. +/// Each orderbook creates two edges: one for buying, one for selling. +#[derive(Debug, Clone)] +pub struct GraphEdge { + pub target: Address, + pub pair: (Address, Address), + pub side: Side, +} + +/// Tracks orderbook freshness and validity for routing decisions. +#[derive(Debug, Clone)] +pub struct OrderbookHealth { + pub last_updated: Instant, + pub has_valid_spread: bool, +} + +#[derive(Debug, Default)] +pub struct Metrics { + pub events_total: AtomicU64, + pub events_invalid: AtomicU64, + pub requests_total: AtomicU64, + pub requests_failed: AtomicU64, + pub snapshots_taken: AtomicU64, + pub stale_or_unhealthy_skipped: AtomicU64, +} + +/// Shared state using scc::HashMap for lock-free concurrent access. +#[derive(Default)] +pub struct AggregatorState { + pub orderbooks: SccHashMap<(Address, Address), OrderbookState>, + pub graph_edges: SccHashMap>, + pub orderbook_health: SccHashMap<(Address, Address), OrderbookHealth>, + pub metrics: Metrics, +} + +impl AggregatorState { + pub fn new() -> Self { + Self::default() + } + + pub fn upsert_orderbook(&self, book: OrderbookState) { + let pair = (book.base_token, book.quote_token); + let is_new = !self.orderbooks.contains(&pair); + let has_valid_spread = Self::validate_orderbook(&book); + + let _ = self.orderbooks.upsert(pair, book); + let _ = self.orderbook_health.upsert( + pair, + OrderbookHealth { + last_updated: Instant::now(), + has_valid_spread, + }, + ); + + if is_new { + self.add_graph_edges(pair); + } + } + + /// Validates orderbook has liquidity and no crossed spread. + pub fn validate_orderbook(book: &OrderbookState) -> bool { + let best_bid = book.bids().iter().find(|b| b.sz > 0).map(|b| b.px); + let best_ask = book.asks().iter().find(|a| a.sz > 0).map(|a| a.px); + + match (best_bid, best_ask) { + (Some(bid), Some(ask)) => bid < ask, + (Some(_), None) | (None, Some(_)) => true, + _ => false, + } + } + + pub fn is_pair_usable(&self, pair: &(Address, Address), max_age: Duration) -> bool { + self.orderbook_health + .read(pair, |_, health| { + health.has_valid_spread && health.last_updated.elapsed() <= max_age + }) + .unwrap_or(false) + } + + fn add_graph_edges(&self, pair: (Address, Address)) { + let (base, quote) = pair; + + self.graph_edges + .entry(quote) + .or_default() + .get_mut() + .push(GraphEdge { + target: base, + pair, + side: Side::Ask, + }); + + self.graph_edges + .entry(base) + .or_default() + .get_mut() + .push(GraphEdge { + target: quote, + pair, + side: Side::Bid, + }); + } + + pub fn contains_token(&self, token: Address) -> bool { + self.graph_edges.contains(&token) + } + + pub fn neighbors(&self, token: Address) -> Vec { + self.graph_edges + .read(&token, |_, edges| edges.clone()) + .unwrap_or_default() + } + + pub fn get_orderbook(&self, pair: &(Address, Address)) -> Option { + self.orderbooks.read(pair, |_, book| book.clone()) + } + + pub fn get_health(&self, pair: &(Address, Address)) -> Option { + self.orderbook_health.read(pair, |_, health| health.clone()) + } + + pub fn orderbook_count(&self) -> usize { + self.orderbooks.len() + } +} + +impl fmt::Debug for AggregatorState { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("AggregatorState") + .field("orderbooks_count", &self.orderbooks.len()) + .field("graph_edges_count", &self.graph_edges.len()) + .finish() + } +} + +pub type SharedState = Arc; + +pub fn create_shared_state() -> SharedState { + Arc::new(AggregatorState::new()) +} diff --git a/packages/aggregator/src/core/tests.rs b/crates/aggregator/src/core/tests.rs similarity index 84% rename from packages/aggregator/src/core/tests.rs rename to crates/aggregator/src/core/tests.rs index 6c91d1b..46c60f0 100644 --- a/packages/aggregator/src/core/tests.rs +++ b/crates/aggregator/src/core/tests.rs @@ -198,8 +198,7 @@ async fn rejects_unknown_token() { }; let resp = processor.process_request(req).await.unwrap(); - // Error message sanitized to avoid leaking token existence info - assert!(matches!(resp, SwapResponse::Failure(msg) if msg == "invalid request")); + assert!(matches!(resp, SwapResponse::Failure(msg) if msg == "unknown tokens")); } #[tokio::test] @@ -380,85 +379,6 @@ fn event_processor_rejects_crossed_book() { ); } -// ============================================================================ -// Circuit Breaker Tests -// ============================================================================ - -#[test] -fn circuit_breaker_starts_closed() { - use crate::core::state::CircuitBreaker; - - let cb = CircuitBreaker::default(); - assert_eq!(cb.current_state(), "closed"); - assert!(cb.allow_request()); -} - -#[test] -fn circuit_breaker_opens_after_threshold() { - use crate::core::state::CircuitBreaker; - use std::time::Duration; - - let cb = CircuitBreaker::new(3, Duration::from_secs(10)); - - // First 2 failures keep circuit closed - cb.record_failure(); - cb.record_failure(); - assert_eq!(cb.current_state(), "closed"); - assert!(cb.allow_request()); - - // Third failure opens circuit - cb.record_failure(); - assert_eq!(cb.current_state(), "open"); - assert!(!cb.allow_request()); -} - -#[test] -fn circuit_breaker_success_resets_failures() { - use crate::core::state::CircuitBreaker; - use std::time::Duration; - - let cb = CircuitBreaker::new(3, Duration::from_secs(10)); - - cb.record_failure(); - cb.record_failure(); - assert_eq!(cb.failure_count(), 2); - - // Success resets counter - cb.record_success(); - assert_eq!(cb.failure_count(), 0); - - // Need 3 more failures to open - cb.record_failure(); - cb.record_failure(); - assert_eq!(cb.current_state(), "closed"); -} - -#[tokio::test] -async fn circuit_breaker_rejects_when_open() { - use aggregator_utils::types::SwapResponse; - - let shared = create_shared_state(); - let a = Address::new_random(); - let b = Address::new_random(); - - // Force circuit breaker open by recording failures - for _ in 0..5 { - shared.circuit_breaker.record_failure(); - } - assert_eq!(shared.circuit_breaker.current_state(), "open"); - - let processor = DexRequestProcessor::new(shared); - let req = SwapRequest { - input_token: a, - output_token: b, - input_amount: 100, - min_output_amount: 0, - }; - - let resp = processor.process_request(req).await.unwrap(); - assert!(matches!(resp, SwapResponse::Failure(msg) if msg == "service degraded")); -} - // ============================================================================ // Metrics Tests // ============================================================================ diff --git a/packages/aggregator/src/lib.rs b/crates/aggregator/src/lib.rs similarity index 100% rename from packages/aggregator/src/lib.rs rename to crates/aggregator/src/lib.rs diff --git a/packages/aggregator/src/traits.rs b/crates/aggregator/src/traits.rs similarity index 100% rename from packages/aggregator/src/traits.rs rename to crates/aggregator/src/traits.rs diff --git a/packages/aggregator/tests/integration.rs b/crates/aggregator/tests/integration.rs similarity index 71% rename from packages/aggregator/tests/integration.rs rename to crates/aggregator/tests/integration.rs index 7bed9b3..00dc755 100644 --- a/packages/aggregator/tests/integration.rs +++ b/crates/aggregator/tests/integration.rs @@ -1,10 +1,4 @@ //! Integration tests for the DEX aggregator. -//! -//! These tests verify end-to-end behavior including: -//! - Full request/response cycles -//! - Circuit breaker state transitions -//! - Rate limiting behavior -//! - Multi-hop routing across multiple orderbooks use aggregator::core::{ event_processor::DexEventProcessor, request_processor::DexRequestProcessor, @@ -36,11 +30,9 @@ async fn full_swap_cycle() { let token_a = Address::new_random(); let token_b = Address::new_random(); - // Ingest an orderbook let book = make_orderbook(token_a, token_b, 99, 101); event_processor.process_orderbook(book).unwrap(); - // Request a swap let request = SwapRequest { input_token: token_a, output_token: token_b, @@ -51,7 +43,6 @@ async fn full_swap_cycle() { let response = request_processor.process_request(request).await.unwrap(); assert!(matches!(response, SwapResponse::Success(_))); - // Verify metrics were updated assert_eq!( state .metrics @@ -68,39 +59,6 @@ async fn full_swap_cycle() { ); } -#[tokio::test] -async fn circuit_breaker_opens_and_recovers() { - let state = create_shared_state(); - let request_processor = DexRequestProcessor::new(state.clone()); - - let token_a = Address::new_random(); - let token_b = Address::new_random(); - - // No orderbooks, so requests will fail and trigger circuit breaker - for _ in 0..5 { - let request = SwapRequest { - input_token: token_a, - output_token: token_b, - input_amount: 1000, - min_output_amount: 0, - }; - let _ = request_processor.process_request(request).await; - } - - // Circuit should now be open - assert_eq!(state.circuit_breaker.current_state(), "open"); - - // Additional requests should be rejected with "service degraded" - let request = SwapRequest { - input_token: token_a, - output_token: token_b, - input_amount: 1000, - min_output_amount: 0, - }; - let response = request_processor.process_request(request).await.unwrap(); - assert!(matches!(response, SwapResponse::Failure(msg) if msg == "service degraded")); -} - #[tokio::test] async fn multi_hop_routing_across_orderbooks() { let state = create_shared_state(); @@ -111,21 +69,13 @@ async fn multi_hop_routing_across_orderbooks() { let eth = Address::new_random(); let btc = Address::new_random(); - // Create USDC/ETH and ETH/BTC orderbooks - // No direct USDC/BTC path, must route through ETH - // Use small prices to ensure integer division gives non-zero results - // ETH/USDC: 1 ETH = 2 USDC (ask), 1 USDC = 0.5 ETH event_processor .process_orderbook(make_orderbook(eth, usdc, 1, 2)) .unwrap(); - // BTC/ETH: 1 BTC = 2 ETH (ask) event_processor .process_orderbook(make_orderbook(btc, eth, 1, 2)) .unwrap(); - // Request USDC -> BTC (requires 2 hops: USDC -> ETH -> BTC) - // With 1000 USDC at 2 USDC/ETH = 500 ETH - // With 500 ETH at 2 ETH/BTC = 250 BTC let request = SwapRequest { input_token: usdc, output_token: btc, @@ -157,16 +107,14 @@ async fn stale_orderbooks_are_skipped() { let token_a = Address::new_random(); let token_b = Address::new_random(); - // Ingest an orderbook let book = make_orderbook(token_a, token_b, 99, 101); event_processor.process_orderbook(book).unwrap(); - // Manually mark the orderbook as stale by backdating its metadata - state.orderbook_meta.update(&(token_a, token_b), |_, meta| { - meta.updated_at = std::time::Instant::now() - Duration::from_secs(60); + // Mark orderbook as stale + state.orderbook_health.update(&(token_a, token_b), |_, health| { + health.last_updated = std::time::Instant::now() - Duration::from_secs(60); }); - // Request should fail because the only path uses a stale orderbook let request = SwapRequest { input_token: token_a, output_token: token_b, @@ -175,10 +123,8 @@ async fn stale_orderbooks_are_skipped() { }; let response = request_processor.process_request(request).await.unwrap(); - // Should fail with "no path" or "invalid request" assert!(matches!(response, SwapResponse::Failure(_))); - // Verify stale skip was recorded assert!( state .metrics @@ -197,16 +143,14 @@ async fn slippage_protection_rejects_bad_rate() { let token_a = Address::new_random(); let token_b = Address::new_random(); - // Create orderbook with wide spread let book = make_orderbook(token_a, token_b, 50, 200); event_processor.process_orderbook(book).unwrap(); - // Request with high minimum output (will fail slippage check) let request = SwapRequest { input_token: token_a, output_token: token_b, input_amount: 1000, - min_output_amount: 999_999, // Impossibly high + min_output_amount: 999_999, }; let response = request_processor.process_request(request).await.unwrap(); @@ -222,7 +166,6 @@ async fn metrics_track_all_request_types() { let token_a = Address::new_random(); let token_b = Address::new_random(); - // Ingest orderbook let book = make_orderbook(token_a, token_b, 99, 101); event_processor.process_orderbook(book).unwrap(); @@ -245,7 +188,6 @@ async fn metrics_track_all_request_types() { }; let _ = request_processor.process_request(request).await; - // Verify metrics let metrics = &state.metrics; assert_eq!( metrics diff --git a/packages/aggregator/src/core/request_processor.rs b/packages/aggregator/src/core/request_processor.rs deleted file mode 100644 index ae7459f..0000000 --- a/packages/aggregator/src/core/request_processor.rs +++ /dev/null @@ -1,185 +0,0 @@ -use aggregator_utils::types::{SwapRequest, SwapResponse, SwapResponseSuccess}; -use async_trait::async_trait; -use tracing::{info, warn}; - -use crate::core::routing::find_best_route; -use crate::core::state::SharedState; -use crate::traits::RequestProcessor; - -/// Errors that can occur during request processing. -#[derive(Debug, thiserror::Error)] -pub enum Error { - #[error("routing failed: {0}")] - RoutingFailed(String), -} - -/// Processes swap requests by finding optimal multi-hop routes. -/// -/// The request processor enforces rate limiting, circuit breaker protection, -/// and slippage tolerance before executing swaps. It uses BFS-based routing -/// to find the best path across multiple orderbooks. -/// -/// # Features -/// -/// - Rate limiting to prevent abuse -/// - Circuit breaker for graceful degradation under failures -/// - Slippage protection based on user-specified minimum output -/// - Multi-hop routing (up to 3 hops) for optimal execution -/// -/// # Example -/// -/// ```ignore -/// let state = create_shared_state(); -/// let processor = DexRequestProcessor::new(state); -/// let response = processor.process_request(swap_request).await?; -/// ``` -#[derive(Debug, Clone)] -pub struct DexRequestProcessor { - state: SharedState, -} - -impl DexRequestProcessor { - /// Create a new request processor with the given shared state. - pub fn new(state: SharedState) -> Self { - Self { state } - } -} - -#[async_trait] -impl RequestProcessor for DexRequestProcessor { - type Error = Error; - - async fn process_request(&self, request: SwapRequest) -> Result { - self.state - .metrics - .requests_total - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - - // Rate limit check: reject if too many requests - if !self.state.rate_limiter.try_acquire() { - warn!("swap rejected: rate limited"); - self.state - .metrics - .requests_rate_limited - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - return Ok(SwapResponse::Failure("rate limited".to_string())); - } - - // Circuit breaker check: reject early if system is degraded - if !self.state.circuit_breaker.allow_request() { - warn!( - circuit_state = %self.state.circuit_breaker.current_state(), - "swap rejected: circuit breaker open" - ); - self.state - .metrics - .requests_failed - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - return Ok(SwapResponse::Failure("service degraded".to_string())); - } - - // Reject swaps with the same token - if request.input_token == request.output_token { - warn!(token = %request.input_token, "swap rejected: same token"); - return Ok(SwapResponse::Failure("same token".to_string())); - } - - // Reject swaps with zero amount - if request.input_amount == 0 { - warn!( - input_token = %request.input_token, - output_token = %request.output_token, - "swap rejected: zero amount" - ); - return Ok(SwapResponse::Failure("zero amount".to_string())); - } - - // Check tokens exist before searching - let input_known = self.state.contains_token(request.input_token); - let output_known = self.state.contains_token(request.output_token); - - // Reject swaps with unknown tokens (infrastructure issue, affects circuit breaker) - // Return generic error to client but log details server-side - if !input_known || !output_known { - let detail = match (input_known, output_known) { - (false, false) => "both tokens unknown", - (false, true) => "input token unknown", - (true, false) => "output token unknown", - _ => unreachable!(), - }; - warn!( - input_token = %request.input_token, - output_token = %request.output_token, - detail, - "swap rejected: invalid request" - ); - self.state - .metrics - .requests_failed - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - self.state.circuit_breaker.record_failure(); - // Generic error to avoid leaking token existence info - return Ok(SwapResponse::Failure("invalid request".to_string())); - } - - // Find best route (infrastructure issue if fails, affects circuit breaker) - let route = match find_best_route( - &self.state, - request.input_token, - request.output_token, - request.input_amount, - ) { - Some(r) => r, - None => { - warn!( - input_token = %request.input_token, - output_token = %request.output_token, - input_amount = request.input_amount, - "swap rejected: no path" - ); - self.state - .metrics - .requests_failed - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - self.state.circuit_breaker.record_failure(); - return Ok(SwapResponse::Failure("no path".to_string())); - } - }; - - let output = route.last().map(|s| s.expected_output_amount).unwrap_or(0); - - // Slippage protection: reject if output is below user's minimum - if output < request.min_output_amount { - warn!( - input_token = %request.input_token, - output_token = %request.output_token, - input_amount = request.input_amount, - output_amount = output, - min_output = request.min_output_amount, - "swap rejected: slippage" - ); - self.state - .metrics - .requests_failed - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - return Ok(SwapResponse::Failure(format!( - "slippage: {} < {}", - output, request.min_output_amount - ))); - } - - info!( - input_token = %request.input_token, - output_token = %request.output_token, - input_amount = request.input_amount, - output_amount = output, - hops = route.len(), - "route found" - ); - - // Successful route resets circuit breaker failure count - self.state.circuit_breaker.record_success(); - - Ok(SwapResponse::Success(SwapResponseSuccess { route })) - } -} diff --git a/packages/aggregator/src/core/routing.rs b/packages/aggregator/src/core/routing.rs deleted file mode 100644 index c074e0f..0000000 --- a/packages/aggregator/src/core/routing.rs +++ /dev/null @@ -1,206 +0,0 @@ -use std::{ - collections::{HashMap, VecDeque}, - sync::Arc, - time::{Duration, Instant}, -}; - -use aggregator_utils::orderbook::OrderbookState; -use aggregator_utils::types::{Address, Quantity, Swap}; -use tracing::debug; - -use crate::core::format_duration; -use crate::core::matching::match_order; -use crate::core::state::{AggregatorState, BookMeta, GraphEdge}; - -// Bounds search complexity. 3 hops covers most real DEX routes -// (e.g. USDC -> ETH -> DOGE -> SHIB). This is a reasonable compromise between -// search space and complexity. -const MAX_HOPS: usize = 3; -// Skip books older than this -const STALE_TTL: Duration = Duration::from_secs(30); - -/// Snapshot of orderbook state for consistent routing. -/// Taking a snapshot ensures BFS sees a consistent view even as orderbooks update. -/// Uses Arc-wrapped HashMaps for O(1) cloning when shared across requests. -#[derive(Clone)] -struct RoutingSnapshot { - orderbooks: Arc>, - graph_edges: Arc>>, - meta: Arc>, -} - -impl RoutingSnapshot { - fn from_state(state: &AggregatorState) -> Self { - let mut orderbooks = HashMap::new(); - // scc uses scan with FnMut closure - state.orderbooks.scan(|k, v| { - orderbooks.insert(*k, v.clone()); - }); - - let mut graph_edges = HashMap::new(); - state.graph_edges.scan(|k, v| { - graph_edges.insert(*k, v.clone()); - }); - - let mut meta = HashMap::new(); - state.orderbook_meta.scan(|k, v| { - meta.insert(*k, v.clone()); - }); - - state - .metrics - .snapshots_taken - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - - Self { - orderbooks: Arc::new(orderbooks), - graph_edges: Arc::new(graph_edges), - meta: Arc::new(meta), - } - } - - fn neighbors(&self, token: &Address) -> &[GraphEdge] { - self.graph_edges - .get(token) - .map(|v| v.as_slice()) - .unwrap_or(&[]) - } - - fn get_orderbook(&self, pair: &(Address, Address)) -> Option<&OrderbookState> { - self.orderbooks.get(pair) - } - - fn is_usable(&self, pair: &(Address, Address)) -> bool { - self.meta - .get(pair) - .map(|m| m.healthy && m.updated_at.elapsed() <= STALE_TTL) - .unwrap_or(false) - } -} - -/// Arena node for path reconstruction without per-edge cloning. -struct PathNode { - swap: Swap, - parent: Option, // Index into arena, None for first hop -} - -/// Reconstruct path from arena by walking parent pointers. -fn reconstruct_path(arena: &[PathNode], mut idx: usize) -> Vec { - let mut path = Vec::with_capacity(MAX_HOPS); - loop { - path.push(arena[idx].swap.clone()); - match arena[idx].parent { - Some(parent_idx) => idx = parent_idx, - None => break, - } - } - path.reverse(); - path -} - -/// Finds the optimal multi-hop trading route between two tokens. -/// -/// Uses BFS to explore all paths up to `MAX_HOPS` (3) and returns the route -/// that maximizes output amount. Takes a point-in-time snapshot of orderbook -/// state for consistent routing. -/// -/// # Arguments -/// -/// * `state` - The aggregator state containing orderbooks and graph topology -/// * `input_token` - The token being sold -/// * `output_token` - The token being bought -/// * `input_amount` - Amount of input token to swap -/// -/// # Returns -/// -/// `Some(Vec)` containing the optimal route, or `None` if no path exists. -/// -/// # Performance -/// -/// Uses arena allocation for path tracking to avoid O(h*e) cloning overhead -/// where h = hops and e = edges explored. -pub fn find_best_route( - state: &AggregatorState, - input_token: Address, - output_token: Address, - input_amount: Quantity, -) -> Option> { - // Take snapshot for consistent routing - let start = Instant::now(); - let snapshot = RoutingSnapshot::from_state(state); - debug!(time = %format_duration(start.elapsed()), "snapshot created"); - - // Arena for path nodes to avoid cloning paths on every edge - let mut arena: Vec = Vec::new(); - - // (current token, amount we have, hop count, parent index in arena or None) - let mut queue: VecDeque<(Address, u64, usize, Option)> = VecDeque::new(); - queue.push_back((input_token, input_amount, 0, None)); - - // Track best output at destination - let mut best_output = 0u64; - let mut best_path_end: Option = None; - - while let Some((current_token, current_amount, hop_count, parent_idx)) = queue.pop_front() { - // Reached destination? Check if it's the best route so far - if current_token == output_token && parent_idx.is_some() { - if current_amount > best_output { - best_output = current_amount; - best_path_end = parent_idx; - } - continue; // don't explore past destination - } - - // The request processor will return a no path error if the route is too long - // This limits the search space - if hop_count >= MAX_HOPS { - continue; - } - - // Try all tokens reachable in one trade - for edge in snapshot.neighbors(¤t_token) { - // Skip stale or unhealthy books - if !snapshot.is_usable(&edge.pair) { - state - .metrics - .stale_or_unhealthy_skipped - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - continue; - } - - // Get the orderbook from snapshot using pair key - let Some(book) = snapshot.get_orderbook(&edge.pair) else { - continue; - }; - - let result = match_order(book, edge.side, current_amount); - - // Skip if no liquidity on this path - if result.output_produced == 0 { - continue; - } - - // Add node to arena and enqueue - let node_idx = arena.len(); - arena.push(PathNode { - swap: Swap { - input_token: current_token, - output_token: edge.target, - direction: edge.side, - input_amount: result.input_consumed, - expected_output_amount: result.output_produced, - }, - parent: parent_idx, - }); - - queue.push_back(( - edge.target, - result.output_produced, - hop_count + 1, - Some(node_idx), - )); - } - } - - best_path_end.map(|idx| reconstruct_path(&arena, idx)) -} diff --git a/packages/aggregator/src/core/state.rs b/packages/aggregator/src/core/state.rs deleted file mode 100644 index 0b34a74..0000000 --- a/packages/aggregator/src/core/state.rs +++ /dev/null @@ -1,447 +0,0 @@ -use std::{ - fmt, - sync::{ - atomic::{AtomicU64, AtomicU8, Ordering}, - Arc, Mutex, - }, - time::{Duration, Instant}, -}; - -use aggregator_utils::orderbook::OrderbookState; -use aggregator_utils::types::{Address, Side}; -use scc::HashMap as SccHashMap; - -/// Represents one direction you can trade through an orderbook. -/// Each orderbook creates two edges: one for buying, one for selling. -#[derive(Debug, Clone)] -pub struct GraphEdge { - pub target: Address, - pub pair: (Address, Address), // key into orderbooks hashmap - pub side: Side, // Ask = buying target, Bid = selling for target -} - -/// Metadata for an orderbook, tracking health and freshness. -#[derive(Debug, Clone)] -pub struct BookMeta { - /// Last time this orderbook was updated. - pub updated_at: Instant, - /// Whether the orderbook passes validation (no crossed spreads, has liquidity). - pub healthy: bool, -} - -/// Atomic counters for aggregator metrics. -/// -/// All counters use relaxed ordering since they are for monitoring only -/// and do not require strict synchronization. -#[derive(Debug, Default)] -pub struct Metrics { - /// Total orderbook events received. - pub events_total: AtomicU64, - /// Orderbook events rejected due to validation failure. - pub events_invalid: AtomicU64, - /// Total swap requests received. - pub requests_total: AtomicU64, - /// Swap requests that failed (no route, unknown tokens, etc.). - pub requests_failed: AtomicU64, - /// Swap requests rejected by rate limiter. - pub requests_rate_limited: AtomicU64, - /// Number of routing snapshots created. - pub snapshots_taken: AtomicU64, - /// Orderbooks skipped during routing due to staleness or health issues. - pub stale_or_unhealthy_skipped: AtomicU64, -} - -/// Circuit breaker states -const CB_CLOSED: u8 = 0; // Normal operation, requests flow through -const CB_OPEN: u8 = 1; // Failures exceeded threshold, reject all requests -const CB_HALF_OPEN: u8 = 2; // Testing if service recovered, allow limited requests - -/// Circuit breaker for graceful degradation under failure conditions. -/// Opens when consecutive failures exceed threshold, preventing cascade failures. -/// Automatically attempts recovery after cooldown period. -pub struct CircuitBreaker { - state: AtomicU8, - consecutive_failures: AtomicU64, - /// Mutex instead of RwLock to prevent TOCTOU race conditions during state transitions. - /// State changes require atomic read-check-write which RwLock cannot guarantee. - last_state_change: Mutex, - /// Number of consecutive failures before opening circuit - failure_threshold: u64, - /// Time to wait before attempting recovery (half-open state) - recovery_timeout: Duration, -} - -impl Default for CircuitBreaker { - fn default() -> Self { - Self { - state: AtomicU8::new(CB_CLOSED), - consecutive_failures: AtomicU64::new(0), - last_state_change: Mutex::new(Instant::now()), - failure_threshold: 5, - recovery_timeout: Duration::from_secs(10), - } - } -} - -impl CircuitBreaker { - pub fn new(failure_threshold: u64, recovery_timeout: Duration) -> Self { - Self { - state: AtomicU8::new(CB_CLOSED), - consecutive_failures: AtomicU64::new(0), - last_state_change: Mutex::new(Instant::now()), - failure_threshold, - recovery_timeout, - } - } - - /// Check if request should be allowed through. - /// Returns true if circuit is closed or half-open (testing recovery). - pub fn allow_request(&self) -> bool { - let state = self.state.load(Ordering::Acquire); - - match state { - CB_CLOSED => true, - CB_OPEN => { - // Hold mutex during entire check-and-transition to prevent TOCTOU race - let mut last_change = self.last_state_change.lock().unwrap(); - if last_change.elapsed() >= self.recovery_timeout { - // Transition to half-open to test recovery - if self - .state - .compare_exchange( - CB_OPEN, - CB_HALF_OPEN, - Ordering::AcqRel, - Ordering::Acquire, - ) - .is_ok() - { - *last_change = Instant::now(); - return true; - } - } - false - } - CB_HALF_OPEN => true, // Allow test request - _ => false, - } - } - - /// Record a successful request. Resets failure count and closes circuit if half-open. - pub fn record_success(&self) { - self.consecutive_failures.store(0, Ordering::Release); - - let state = self.state.load(Ordering::Acquire); - if state == CB_HALF_OPEN { - // Recovery confirmed, close circuit - self.state.store(CB_CLOSED, Ordering::Release); - *self.last_state_change.lock().unwrap() = Instant::now(); - } - } - - /// Record a failed request. Opens circuit if failures exceed threshold. - pub fn record_failure(&self) { - let failures = self.consecutive_failures.fetch_add(1, Ordering::AcqRel) + 1; - let state = self.state.load(Ordering::Acquire); - - match state { - CB_CLOSED => { - if failures >= self.failure_threshold { - // Too many failures, open circuit - if self - .state - .compare_exchange(CB_CLOSED, CB_OPEN, Ordering::AcqRel, Ordering::Acquire) - .is_ok() - { - *self.last_state_change.lock().unwrap() = Instant::now(); - } - } - } - CB_HALF_OPEN => { - // Recovery test failed, back to open - self.state.store(CB_OPEN, Ordering::Release); - *self.last_state_change.lock().unwrap() = Instant::now(); - } - _ => {} - } - } - - /// Get current state for monitoring/debugging - pub fn current_state(&self) -> &'static str { - match self.state.load(Ordering::Acquire) { - CB_CLOSED => "closed", - CB_OPEN => "open", - CB_HALF_OPEN => "half-open", - _ => "unknown", - } - } - - /// Get consecutive failure count - pub fn failure_count(&self) -> u64 { - self.consecutive_failures.load(Ordering::Acquire) - } -} - -impl fmt::Debug for CircuitBreaker { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("CircuitBreaker") - .field("state", &self.current_state()) - .field("failures", &self.failure_count()) - .finish() - } -} - -/// Token-bucket rate limiter for request throttling. -/// Prevents abuse and provides backpressure under high load. -pub struct RateLimiter { - /// Available tokens (scaled by 1000 for sub-token precision) - tokens: AtomicU64, - /// Last time tokens were refilled - last_refill: Mutex, - /// Tokens added per second - tokens_per_second: u64, - /// Maximum token capacity (burst size) - max_tokens: u64, -} - -impl RateLimiter { - /// Create a new rate limiter. - /// `tokens_per_second`: steady-state request rate - /// `max_tokens`: burst capacity - pub fn new(tokens_per_second: u64, max_tokens: u64) -> Self { - Self { - tokens: AtomicU64::new(max_tokens * 1000), // scaled - last_refill: Mutex::new(Instant::now()), - tokens_per_second, - max_tokens, - } - } - - /// Attempt to acquire a token. Returns true if successful. - pub fn try_acquire(&self) -> bool { - self.refill(); - - // Try to consume one token (1000 scaled units) - let mut current = self.tokens.load(Ordering::Acquire); - loop { - if current < 1000 { - return false; // Not enough tokens - } - - match self.tokens.compare_exchange_weak( - current, - current - 1000, - Ordering::AcqRel, - Ordering::Acquire, - ) { - Ok(_) => return true, - Err(new_val) => current = new_val, - } - } - } - - /// Refill tokens based on elapsed time - fn refill(&self) { - let mut last = self.last_refill.lock().unwrap(); - let elapsed = last.elapsed(); - - // Only refill if at least 1ms has passed (avoid excessive lock contention) - if elapsed.as_millis() < 1 { - return; - } - - let tokens_to_add = (elapsed.as_millis() as u64 * self.tokens_per_second) / 1000 * 1000; - if tokens_to_add > 0 { - let max_scaled = self.max_tokens * 1000; - let current = self.tokens.load(Ordering::Acquire); - let new_tokens = (current + tokens_to_add).min(max_scaled); - self.tokens.store(new_tokens, Ordering::Release); - *last = Instant::now(); - } - } - - /// Get current available tokens (for monitoring) - pub fn available_tokens(&self) -> u64 { - self.tokens.load(Ordering::Acquire) / 1000 - } -} - -impl Default for RateLimiter { - fn default() -> Self { - // Default: 1000 req/sec with burst of 100 - Self::new(1000, 100) - } -} - -impl fmt::Debug for RateLimiter { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("RateLimiter") - .field("available_tokens", &self.available_tokens()) - .field("tokens_per_second", &self.tokens_per_second) - .field("max_tokens", &self.max_tokens) - .finish() - } -} - -/// Shared state using scc::HashMap for lock-free concurrent access with per-book health. -#[derive(Default)] -pub struct AggregatorState { - /// Lock-free concurrent hashmap for orderbooks. - /// Key: (base_token, quote_token) pair - pub orderbooks: SccHashMap<(Address, Address), OrderbookState>, - - /// Lock-free concurrent hashmap for graph edges. - /// Key: source token, Value: edges to neighboring tokens - pub graph_edges: SccHashMap>, - - /// Metadata per orderbook (staleness, health). - pub orderbook_meta: SccHashMap<(Address, Address), BookMeta>, - - /// Simple in-process counters. - pub metrics: Metrics, - - /// Circuit breaker for graceful degradation. - pub circuit_breaker: CircuitBreaker, - - /// Rate limiter for request throttling. - pub rate_limiter: RateLimiter, -} - -impl AggregatorState { - pub fn new() -> Self { - Self::default() - } - - /// Upsert orderbook. Takes &self (not &mut self) since scc provides interior mutability. - pub fn upsert_orderbook(&self, book: OrderbookState) { - let pair = (book.base_token, book.quote_token); - - // Check if this is a new pair (need to add graph edges) - let is_new = !self.orderbooks.contains(&pair); - - let healthy = Self::validate_orderbook(&book); - - // Upsert the orderbook and metadata - let _ = self.orderbooks.upsert(pair, book); - let _ = self.orderbook_meta.upsert( - pair, - BookMeta { - updated_at: Instant::now(), - healthy, - }, - ); - - // Add graph edges for new pairs only - if is_new { - self.add_graph_edges(pair); - } - } - - /// Basic validation: non-empty sides and no crossed book (best bid < best ask). - pub fn validate_orderbook(book: &OrderbookState) -> bool { - let mut best_bid: Option = None; - for b in book.bids() { - if b.sz == 0 { - continue; - } - best_bid = Some(b.px); - break; - } - - let mut best_ask: Option = None; - for a in book.asks() { - if a.sz == 0 { - continue; - } - best_ask = Some(a.px); - break; - } - - match (best_bid, best_ask) { - (Some(bid), Some(ask)) => bid < ask, // both sides present, not crossed - (Some(_), None) | (None, Some(_)) => true, // one-sided book is acceptable - _ => false, // no usable levels - } - } - - /// Check meta for health and staleness. - pub fn is_pair_usable(&self, pair: &(Address, Address), max_age: Duration) -> bool { - self.orderbook_meta - .read(pair, |_, meta| { - meta.healthy && meta.updated_at.elapsed() <= max_age - }) - .unwrap_or(false) - } - - fn add_graph_edges(&self, pair: (Address, Address)) { - let (base, quote) = pair; - - // Edge: quote -> base (buying base with quote, Ask side) - self.graph_edges - .entry(quote) - .or_default() - .get_mut() - .push(GraphEdge { - target: base, - pair, - side: Side::Ask, - }); - - // Edge: base -> quote (selling base for quote, Bid side) - self.graph_edges - .entry(base) - .or_default() - .get_mut() - .push(GraphEdge { - target: quote, - pair, - side: Side::Bid, - }); - } - - /// Check if token exists in graph - pub fn contains_token(&self, token: Address) -> bool { - self.graph_edges.contains(&token) - } - - /// Get neighbors for routing (returns clone since we can't hold reference across await) - pub fn neighbors(&self, token: Address) -> Vec { - self.graph_edges - .read(&token, |_, edges| edges.clone()) - .unwrap_or_default() - } - - /// Get orderbook by pair - pub fn get_orderbook(&self, pair: &(Address, Address)) -> Option { - self.orderbooks.read(pair, |_, book| book.clone()) - } - - /// Get metadata by pair - pub fn get_meta(&self, pair: &(Address, Address)) -> Option { - self.orderbook_meta.read(pair, |_, meta| meta.clone()) - } - - /// Get count of orderbooks (for tests) - pub fn orderbook_count(&self) -> usize { - self.orderbooks.len() - } -} - -// Manual Debug implementation since scc::HashMap doesn't implement Debug -impl fmt::Debug for AggregatorState { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("AggregatorState") - .field("orderbooks_count", &self.orderbooks.len()) - .field("graph_edges_count", &self.graph_edges.len()) - .field("circuit_breaker", &self.circuit_breaker) - .field("rate_limiter", &self.rate_limiter) - .finish() - } -} - -// SharedState is now just Arc - no RwLock needed -pub type SharedState = Arc; - -pub fn create_shared_state() -> SharedState { - Arc::new(AggregatorState::new()) -} diff --git a/rust-toolchain.toml b/rust-toolchain.toml deleted file mode 100644 index 2e2b8c8..0000000 --- a/rust-toolchain.toml +++ /dev/null @@ -1,2 +0,0 @@ -[toolchain] -channel = "1.82.0" From 800c7a0b645fd28b0441746772ff3d22b0e03c7f Mon Sep 17 00:00:00 2001 From: Alexander Date: Thu, 15 Jan 2026 16:41:22 -0500 Subject: [PATCH 3/3] Fix cargo fmt in integration tests --- .gitignore | 1 - crates/aggregator/tests/integration.rs | 8 +++++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/.gitignore b/.gitignore index e482b7f..0592392 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,2 @@ /target .DS_Store -/notes \ No newline at end of file diff --git a/crates/aggregator/tests/integration.rs b/crates/aggregator/tests/integration.rs index 00dc755..5cfc189 100644 --- a/crates/aggregator/tests/integration.rs +++ b/crates/aggregator/tests/integration.rs @@ -111,9 +111,11 @@ async fn stale_orderbooks_are_skipped() { event_processor.process_orderbook(book).unwrap(); // Mark orderbook as stale - state.orderbook_health.update(&(token_a, token_b), |_, health| { - health.last_updated = std::time::Instant::now() - Duration::from_secs(60); - }); + state + .orderbook_health + .update(&(token_a, token_b), |_, health| { + health.last_updated = std::time::Instant::now() - Duration::from_secs(60); + }); let request = SwapRequest { input_token: token_a,