driver: parallelize unsupported order detection

crates/balance-overrides/src/balance/mod.rs

@@ -891,7 +891,7 @@ mod tests {
 
         let strategy = detector.detect(NATIVE_ETH, user).await.unwrap();
 
-        std::assert_matches!(strategy, Strategy::NativeEth);
+        std::assert_eq!(strategy, Strategy::NativeEth);
 
         // ETH is not an ERC20 token so we can't do an `eth_call` on `balanceOf()`.
         // Additionally `eth_getBalance` does not support state overrides.

crates/driver/Cargo.toml

@@ -93,6 +93,7 @@ app-data = { workspace = true, features = ["test_helpers"] }
 contracts = { workspace = true }
 ethrpc = { workspace = true, features = ["test-util"] }
 maplit = { workspace = true }
+mockall = { workspace = true }
 tempfile = { workspace = true }
 tokio = { workspace = true, features = ["process", "test-util"] }
 

crates/driver/src/domain/competition/mod.rs

@@ -38,7 +38,7 @@ use {
         time::Instant,
     },
     tokio::{sync::mpsc, task},
-    tracing::{Instrument, instrument},
+    tracing::Instrument,
 };
 
 pub mod auction;
@@ -259,6 +259,22 @@ impl Drop for SettleTaskHandle {
     }
 }
 
+struct AbortGuard(Vec<task::AbortHandle>);
+
+impl AbortGuard {
+    fn disarm(&mut self) {
+        self.0.clear();
+    }
+}
+
+impl Drop for AbortGuard {
+    fn drop(&mut self) {
+        for handle in &self.0 {
+            handle.abort();
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct Competition {
     pub solver: Solver,
@@ -330,15 +346,44 @@ impl Competition {
                 Error::MalformedRequest
             })?;
 
-        let auction = self.assemble_auction(&tasks).await;
+        let mut auction = self.assemble_auction(&tasks).await;
 
-        let liquidity = async {
-            match self.solver.liquidity() {
-                solver::Liquidity::Fetch => tasks.liquidity.await,
-                solver::Liquidity::Skip => Arc::new(Vec::new()),
+        // We can run bad token filtering and liquidity fetching in parallel
+        let risk_detector = self.risk_detector.clone();
+        let flashloans_enabled = self.solver.config().flashloans_enabled;
+        let liquidity_mode = self.solver.liquidity();
+        let auction_handle = tokio::spawn(
+            async move {
+                risk_detector
+                    .without_unsupported_orders(&mut auction.orders, flashloans_enabled)
+                    .await;
+                auction
             }
-        }
-        .await;
+            .in_current_span(),
+        );
+        let liquidity_handle = tokio::spawn(
+            async move {
+                match liquidity_mode {
+                    solver::Liquidity::Fetch => tasks.liquidity.await,
+                    solver::Liquidity::Skip => Arc::new(Vec::new()),
+                }
+            }
+            .in_current_span(),
+        );
+        let mut drop_guard = AbortGuard(vec![
+            auction_handle.abort_handle(),
+            liquidity_handle.abort_handle(),
+        ]);
+        let (auction, liquidity) = tokio::join!(auction_handle, liquidity_handle);
+        drop_guard.disarm();
+        let auction = auction.map_err(|err| {
+            tracing::error!(?err, "order filtering task failed");
+            Error::InternalError(err.to_string())
+        })?;
+        let liquidity = liquidity.map_err(|err| {
+            tracing::error!(?err, "liquidity fetch task failed");
+            Error::InternalError(err.to_string())
+        })?;
 
         let elapsed = start.elapsed();
         metrics::get()
@@ -634,14 +679,13 @@ impl Competition {
             tasks.app_data.clone()
         );
 
-        let auction = Self::run_blocking_with_timer("update_orders", move || {
+        Self::run_blocking_with_timer("update_orders", move || {
             // Same as before with sort_orders, we use spawn_blocking() because a lot of CPU
             // bound computations are happening and we want to avoid blocking
             // the runtime.
             Self::update_orders(auction, balances, app_data, cow_amm_orders)
         })
-        .await;
-        self.without_unsupported_orders(auction).await
+        .await
     }
 
     // Oders already need to be sorted from most relevant to least relevant so that
@@ -970,16 +1014,6 @@ impl Competition {
         }
         Ok(())
     }
-
-    #[instrument(skip_all)]
-    async fn without_unsupported_orders(&self, mut auction: Auction) -> Auction {
-        if !self.solver.config().flashloans_enabled {
-            auction.orders.retain(|o| o.app_data.flashloan().is_none());
-        }
-        self.risk_detector
-            .filter_unsupported_orders_in_auction(auction)
-            .await
-    }
 }
 
 const MAX_SOLUTIONS_TO_MERGE: usize = 10;
@@ -1101,4 +1135,6 @@ pub enum Error {
     NoValidOrdersFound,
     #[error("could not parse the request")]
     MalformedRequest,
+    #[error("internal error: {0}")]
+    InternalError(String),
 }

crates/driver/src/domain/competition/risk_detector/bad_tokens/simulation.rs

@@ -5,10 +5,12 @@ use {
         infra::{self, observe::metrics},
     },
     bad_tokens::{TokenQuality, trace_call::TraceCallDetectorRaw},
+    eth_domain_types as eth,
     futures::FutureExt,
     model::interaction::InteractionData,
     request_sharing::BoxRequestSharing,
     std::{
+        ops::Deref,
         sync::Arc,
         time::{Duration, Instant},
     },
@@ -21,6 +23,29 @@ use {
 #[derive(Clone)]
 pub struct Detector(Arc<Inner>);
 
+#[cfg_attr(test, mockall::automock)]
+#[async_trait::async_trait]
+pub trait DetectorApi: Send + Sync {
+    async fn determine_sell_token_quality(&self, order: &Order, now: Instant) -> Quality;
+    fn get_quality(&self, token: &eth::TokenAddress, now: Instant) -> Quality;
+    fn evict_outdated_entries(&self);
+}
+
+#[async_trait::async_trait]
+impl DetectorApi for Detector {
+    async fn determine_sell_token_quality(&self, order: &Order, now: Instant) -> Quality {
+        self.determine_sell_token_quality(order, now).await
+    }
+
+    fn get_quality(&self, token: &eth::TokenAddress, now: Instant) -> Quality {
+        Deref::deref(self).get_quality(token, now)
+    }
+
+    fn evict_outdated_entries(&self) {
+        Deref::deref(self).evict_outdated_entries()
+    }
+}
+
 struct Inner {
     cache: Cache,
     detector: TraceCallDetectorRaw,
@@ -107,7 +132,7 @@ impl Detector {
     }
 }
 
-impl std::ops::Deref for Detector {
+impl Deref for Detector {
     type Target = Cache;
 
     fn deref(&self) -> &Self::Target {