driver: filter orders based on user balances (#1687)

Progress on #1672.

Filters orders based on trader balances. If the traders don't have
enough balance to cover the order, it will be removed before being sent
to the solver.

TODO after merging:
- Configure ethflow address in infra.
- Extend the test.

### Test Plan

Added automated test.

---------

Co-authored-by: Valentin <[email protected]>
Co-authored-by: Martin Beckmann <[email protected]>
Co-authored-by: Martin Beckmann <[email protected]>

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

@@ -1,8 +1,19 @@
 use {
-    crate::domain::{
-        competition::{self, solution},
-        eth,
+    super::order,
+    crate::{
+        domain::{
+            competition::{self, solution},
+            eth::{self},
+        },
+        infra::{
+            blockchain,
+            observe::{self},
+            Ethereum,
+        },
     },
+    futures::future::join_all,
+    itertools::Itertools,
+    primitive_types::U256,
     std::collections::HashMap,
     thiserror::Error,
 };
@@ -23,26 +34,53 @@ pub struct Auction {
 }
 
 impl Auction {
-    pub fn new(
+    pub async fn new(
         id: Option<Id>,
-        mut orders: Vec<competition::Order>,
+        orders: Vec<competition::Order>,
         tokens: impl Iterator<Item = Token>,
-        gas_price: eth::GasPrice,
         deadline: Deadline,
-        weth: eth::WethAddress,
-    ) -> Result<Self, InvalidTokens> {
+        eth: &Ethereum,
+    ) -> Result<Self, Error> {
         let tokens = Tokens(tokens.map(|token| (token.address, token)).collect());
 
         // Ensure that tokens are included for each order.
+        let weth = eth.contracts().weth_address();
         if !orders.iter().all(|order| {
             tokens.0.contains_key(&order.buy.token.wrap(weth))
                 && tokens.0.contains_key(&order.sell.token.wrap(weth))
         }) {
-            return Err(InvalidTokens);
+            return Err(Error::InvalidTokens);
         }
 
-        // Sort orders such that most likely to be fulfilled come first.
-        orders.sort_by_key(|order| {
+        Ok(Self {
+            id,
+            orders,
+            tokens,
+            gas_price: eth.gas_price().await?,
+            deadline,
+        })
+    }
+
+    /// [`None`] if this auction applies to a quote. See
+    /// [`crate::domain::quote`].
+    pub fn id(&self) -> Option<Id> {
+        self.id
+    }
+
+    /// The orders for the auction.
+    pub fn orders(&self) -> &[competition::Order] {
+        &self.orders
+    }
+
+    /// Prioritize the orders such that those which are more likely to be
+    /// fulfilled come before less likely orders. Filter out orders which
+    /// the trader doesn't have enough balance to pay for.
+    ///
+    /// Prioritization is skipped during quoting. It's only used during
+    /// competition.
+    pub async fn prioritize(mut self, eth: &Ethereum) -> Self {
+        // Sort orders so that most likely to be fulfilled come first.
+        self.orders.sort_by_key(|order| {
             // Market orders are preferred over limit orders, as the expectation is that
             // they should be immediately fulfillable. Liquidity orders come last, as they
             // are the most niche and rarely used.
@@ -55,31 +93,112 @@ impl Auction {
                 class,
                 // If the orders are of the same kind, then sort by likelihood of fulfillment
                 // based on token prices.
-                order.likelihood(&tokens),
+                order.likelihood(&self.tokens),
             ))
         });
 
-        // TODO Filter out orders based on user balance
+        // Fetch balances of each token for each trader.
+        // Has to be separate closure due to compiler bug.
+        let f = |order: &competition::Order| -> (order::Trader, eth::TokenAddress) {
+            (order.trader(), order.sell.token)
+        };
+        let tokens_by_trader = self.orders.iter().map(f).unique();
+        let mut balances: HashMap<
+            (order::Trader, eth::TokenAddress),
+            Result<eth::TokenAmount, crate::infra::blockchain::Error>,
+        > = join_all(tokens_by_trader.map(|(trader, token)| async move {
+            let contract = eth.erc20(token);
+            let balance = contract.balance(trader.into()).await;
+            ((trader, token), balance)
+        }))
+        .await
+        .into_iter()
+        .collect();
+
+        self.orders.retain(|order| {
+            // TODO: We should use balance fetching that takes interactions into account
+            // from `crates/shared/src/account_balances/simulation.rs` instead of hardcoding
+            // an Ethflow exception. https://github.com/cowprotocol/services/issues/1595
+            if Some(order.signature.signer.0) == eth.contracts().ethflow_address().map(|a| a.0) {
+                return true;
+            }
+
+            let remaining_balance = match balances
+                .get_mut(&(order.trader(), order.sell.token))
+                .unwrap()
+            {
+                Ok(balance) => &mut balance.0,
+                Err(err) => {
+                    let reason = observe::OrderExcludedFromAuctionReason::CouldNotFetchBalance(err);
+                    observe::order_excluded_from_auction(order, reason);
+                    return false;
+                }
+            };
 
-        Ok(Self {
-            id,
-            orders,
-            tokens,
-            gas_price,
-            deadline,
-        })
-    }
+            fn max_fill(order: &competition::Order) -> anyhow::Result<U256> {
+                use {
+                    anyhow::Context,
+                    shared::remaining_amounts::{Order as RemainingOrder, Remaining},
+                };
+
+                let remaining = Remaining::from_order(&RemainingOrder {
+                    kind: match order.side {
+                        order::Side::Buy => model::order::OrderKind::Buy,
+                        order::Side::Sell => model::order::OrderKind::Sell,
+                    },
+                    buy_amount: order.buy.amount.0,
+                    sell_amount: order.sell.amount.0,
+                    fee_amount: order.fee.user.0,
+                    executed_amount: match order.partial {
+                        order::Partial::Yes { executed } => executed.0,
+                        order::Partial::No => 0.into(),
+                    },
+                    partially_fillable: match order.partial {
+                        order::Partial::Yes { .. } => true,
+                        order::Partial::No => false,
+                    },
+                })
+                .context("Remaining::from_order")?;
+                let sell = remaining
+                    .remaining(order.sell.amount.0)
+                    .context("remaining_sell")?;
+                let fee = remaining
+                    .remaining(order.fee.user.0)
+                    .context("remaining_fee")?;
+                sell.checked_add(fee).context("add sell and fee")
+            }
+
+            let max_fill = match max_fill(order) {
+                Ok(balance) => balance,
+                Err(err) => {
+                    let reason =
+                        observe::OrderExcludedFromAuctionReason::CouldNotCalculateRemainingAmount(
+                            &err,
+                        );
+                    observe::order_excluded_from_auction(order, reason);
+                    return false;
+                }
+            };
 
-    /// [`None`] if this auction applies to a quote. See
-    /// [`crate::domain::quote`].
-    pub fn id(&self) -> Option<Id> {
-        self.id
-    }
+            let used_balance = match order.is_partial() {
+                true => {
+                    if *remaining_balance == 0.into() {
+                        return false;
+                    }
+                    max_fill.min(*remaining_balance)
+                }
+                false => {
+                    if *remaining_balance < max_fill {
+                        return false;
+                    }
+                    max_fill
+                }
+            };
+            *remaining_balance -= used_balance;
+            true
+        });
 
-    /// The orders for the auction. The orders are sorted such that those which
-    /// are more likely to be fulfilled come before less likely orders.
-    pub fn orders(&self) -> &[competition::Order] {
-        &self.orders
+        self
     }
 
     /// The tokens used in the auction.
@@ -224,10 +343,14 @@ pub struct DeadlineExceeded;
 #[error("invalid auction id")]
 pub struct InvalidId;
 
-#[derive(Debug, Error)]
-#[error("invalid auction tokens")]
-pub struct InvalidTokens;
-
 #[derive(Debug, Error)]
 #[error("price cannot be zero")]
 pub struct InvalidPrice;
+
+#[derive(Debug, Error)]
+pub enum Error {
+    #[error("invalid auction tokens")]
+    InvalidTokens,
+    #[error("blockchain error: {0:?}")]
+    Blockchain(#[from] blockchain::Error),
+}

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

@@ -103,6 +103,10 @@ impl Order {
         }
     }
 
+    pub fn trader(&self) -> Trader {
+        Trader(self.signature.signer)
+    }
+
     pub fn is_partial(&self) -> bool {
         matches!(self.partial, Partial::Yes { .. })
     }
@@ -285,6 +289,16 @@ pub enum BuyTokenBalance {
     Internal,
 }
 
+/// The address which placed the order.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub struct Trader(eth::Address);
+
+impl From<Trader> for eth::Address {
+    fn from(value: Trader) -> Self {
+        value.0
+    }
+}
+
 /// A just-in-time order. JIT orders are added at solving time by the solver to
 /// generate a more optimal solution for the auction. Very similar to a regular
 /// [`Order`].

crates/driver/src/domain/quote.rs

@@ -75,12 +75,10 @@ impl Order {
         tokens: &infra::tokens::Fetcher,
     ) -> Result<Quote, Error> {
         let liquidity = liquidity.fetch(&self.liquidity_pairs()).await;
-        let gas_price = eth.gas_price().await?;
         let timeout = self.deadline.timeout()?;
-        let fake_auction = self
-            .fake_auction(gas_price, eth.contracts().weth_address(), tokens)
-            .await;
-        let solutions = solver.solve(&fake_auction, &liquidity, timeout).await?;
+        let solutions = solver
+            .solve(&self.fake_auction(eth, tokens).await?, &liquidity, timeout)
+            .await?;
         Quote::new(
             eth,
             self,
@@ -95,10 +93,9 @@ impl Order {
 
     async fn fake_auction(
         &self,
-        gas_price: eth::GasPrice,
-        weth: eth::WethAddress,
+        eth: &Ethereum,
         tokens: &infra::tokens::Fetcher,
-    ) -> competition::Auction {
+    ) -> Result<competition::Auction, Error> {
         let tokens = tokens.get(&[self.buy().token, self.sell().token]).await;
 
         let buy_token_metadata = tokens.get(&self.buy().token);
@@ -146,11 +143,14 @@ impl Order {
                 },
             ]
             .into_iter(),
-            gas_price.effective().into(),
             Default::default(),
-            weth,
+            eth,
         )
-        .unwrap()
+        .await
+        .map_err(|err| match err {
+            auction::Error::InvalidTokens => panic!("fake auction with invalid tokens"),
+            auction::Error::Blockchain(e) => e.into(),
+        })
     }
 
     /// The asset being bought, or [`eth::U256::one`] if this is a sell, to

crates/driver/src/infra/api/error.rs

@@ -84,7 +84,7 @@ impl From<api::routes::AuctionError> for (hyper::StatusCode, axum::Json<Error>)
             api::routes::AuctionError::InvalidAuctionId => Kind::InvalidAuctionId,
             api::routes::AuctionError::MissingSurplusFee => Kind::MissingSurplusFee,
             api::routes::AuctionError::InvalidTokens => Kind::InvalidTokens,
-            api::routes::AuctionError::GasPrice(_) => Kind::Unknown,
+            api::routes::AuctionError::Blockchain(_) => Kind::Unknown,
         };
         error.into()
     }

crates/driver/src/infra/api/routes/solve/dto/auction.rs

@@ -131,10 +131,10 @@ impl Auction {
                     trusted: token.trusted,
                 }
             }),
-            eth.gas_price().await.map_err(Error::GasPrice)?,
             self.deadline.into(),
-            eth.contracts().weth_address(),
+            eth,
         )
+        .await
         .map_err(Into::into)
     }
 }
@@ -147,8 +147,8 @@ pub enum Error {
     MissingSurplusFee,
     #[error("invalid tokens in auction")]
     InvalidTokens,
-    #[error("error getting gas price")]
-    GasPrice(#[source] crate::infra::blockchain::Error),
+    #[error("blockchain error: {0:?}")]
+    Blockchain(#[source] crate::infra::blockchain::Error),
 }
 
 impl From<auction::InvalidId> for Error {
@@ -157,9 +157,12 @@ impl From<auction::InvalidId> for Error {
     }
 }
 
-impl From<auction::InvalidTokens> for Error {
-    fn from(_value: auction::InvalidTokens) -> Self {
-        Self::InvalidTokens
+impl From<auction::Error> for Error {
+    fn from(value: auction::Error) -> Self {
+        match value {
+            auction::Error::InvalidTokens => Self::InvalidTokens,
+            auction::Error::Blockchain(err) => Self::Blockchain(err),
+        }
     }
 }
 

crates/driver/src/infra/api/routes/solve/mod.rs

@@ -27,6 +27,7 @@ async fn route(
             .tap_err(|err| {
                 observe::invalid_dto(err, "auction");
             })?;
+        let auction = auction.prioritize(state.eth()).await;
         observe::auction(&auction);
         let competition = state.competition();
         let result = competition.solve(&auction).await;