Update settlement contract balances in background task

crates/driver/src/infra/tokens.rs

@@ -3,10 +3,13 @@ use {
         domain::eth,
         infra::{blockchain, Ethereum},
     },
+    ethrpc::current_block::{self, CurrentBlockStream},
+    futures::StreamExt,
     std::{
         collections::{HashMap, HashSet},
         sync::{Arc, RwLock},
     },
+    tracing::Instrument,
 };
 
 #[derive(Clone, Debug)]
@@ -22,10 +25,16 @@ pub struct Fetcher(Arc<Inner>);
 
 impl Fetcher {
     pub fn new(eth: Ethereum) -> Self {
-        Self(Arc::new(Inner {
+        let block_stream = eth.current_block().clone();
+        let inner = Arc::new(Inner {
             eth,
             cache: RwLock::new(HashMap::new()),
-        }))
+        });
+        tokio::task::spawn(
+            update_task(block_stream, Arc::downgrade(&inner))
+                .instrument(tracing::info_span!("token_fetcher")),
+        );
+        Self(inner)
     }
 
     /// Returns the `Metadata` for the given tokens. Note that the result will
@@ -39,6 +48,62 @@ impl Fetcher {
     }
 }
 
+/// Runs a single cache update cycle whenever a new block arrives until the
+/// fetcher is dropped.
+async fn update_task(blocks: CurrentBlockStream, inner: std::sync::Weak<Inner>) {
+    let mut stream = current_block::into_stream(blocks);
+    while stream.next().await.is_some() {
+        let inner = match inner.upgrade() {
+            Some(inner) => inner,
+            // Fetcher was dropped, stop update task.
+            None => break,
+        };
+        if let Err(err) = update_balances(inner).await {
+            tracing::warn!(?err, "error updating token cache");
+        }
+    }
+}
+
+/// Updates the settlement contract's balance for every cached token.
+async fn update_balances(inner: Arc<Inner>) -> Result<(), blockchain::Error> {
+    let settlement = inner.eth.contracts().settlement().address().into();
+    let futures = {
+        let cache = inner.cache.read().unwrap();
+        let tokens = cache.keys().cloned().collect::<Vec<_>>();
+        tracing::debug!(
+            tokens = tokens.len(),
+            "updating settlement contract balances"
+        );
+        tokens.into_iter().map(|token| {
+            let erc20 = inner.eth.erc20(token);
+            async move {
+                Ok::<(eth::TokenAddress, eth::TokenAmount), blockchain::Error>((
+                    token,
+                    erc20.balance(settlement).await?,
+                ))
+            }
+        })
+    };
+
+    // Don't hold on to the lock while fetching balances to allow concurrent
+    // updates. This may lead to new entries arriving in the meantime, however
+    // their balances should already be up-to-date.
+    let mut balances = futures::future::try_join_all(futures)
+        .await?
+        .into_iter()
+        .collect::<HashMap<_, _>>();
+
+    let mut cache = inner.cache.write().unwrap();
+    for (key, entry) in cache.iter_mut() {
+        if let Some(balance) = balances.remove(key) {
+            entry.balance = balance;
+        } else {
+            tracing::info!(?key, "key without balance update");
+        }
+    }
+    Ok(())
+}
+
 /// Provides metadata of tokens.
 struct Inner {
     eth: Ethereum,

crates/e2e/src/setup/colocation.rs

@@ -11,7 +11,7 @@ pub async fn start_solver(weth: H160) -> Url {
         r#"
 weth = "{weth:?}"
 base-tokens = []
-max-hops = 0
+max-hops = 1
 max-partial-attempts = 5
 risk-parameters = [0,0,0,0]
         "#,

crates/e2e/tests/e2e/colocation_buffers.rs

@@ -0,0 +1,108 @@
+use {
+    e2e::{setup::*, tx},
+    ethcontract::prelude::U256,
+    model::{
+        order::{OrderCreation, OrderKind},
+        signature::EcdsaSigningScheme,
+    },
+    secp256k1::SecretKey,
+    shared::ethrpc::Web3,
+    web3::signing::SecretKeyRef,
+};
+
+#[tokio::test]
+#[ignore]
+async fn local_node_buffers() {
+    run_test(onchain_settlement_without_liquidity).await;
+}
+
+async fn onchain_settlement_without_liquidity(web3: Web3) {
+    let mut onchain = OnchainComponents::deploy(web3).await;
+
+    let [solver] = onchain.make_solvers(to_wei(1)).await;
+    let [trader] = onchain.make_accounts(to_wei(1)).await;
+    let [token_a, token_b] = onchain
+        .deploy_tokens_with_weth_uni_v2_pools(to_wei(1_000), to_wei(1_000))
+        .await;
+
+    // Fund trader, settlement accounts, and pool creation
+    token_a.mint(trader.address(), to_wei(100)).await;
+    token_b
+        .mint(onchain.contracts().gp_settlement.address(), to_wei(5))
+        .await;
+    token_a.mint(solver.address(), to_wei(1000)).await;
+    token_b.mint(solver.address(), to_wei(1000)).await;
+
+    // Approve GPv2 for trading
+    tx!(
+        trader.account(),
+        token_a.approve(onchain.contracts().allowance, to_wei(100))
+    );
+
+    // Start system
+    let solver_endpoint = colocation::start_solver(onchain.contracts().weth.address()).await;
+    colocation::start_driver(onchain.contracts(), &solver_endpoint, &solver);
+
+    let services = Services::new(onchain.contracts()).await;
+    services.start_autopilot(vec![
+        "--enable-colocation=true".to_string(),
+        format!(
+            "--trusted-tokens={weth:#x},{token_a:#x},{token_b:#x}",
+            weth = onchain.contracts().weth.address(),
+            token_a = token_a.address(),
+            token_b = token_b.address()
+        ),
+        "--drivers=test_solver|http://localhost:11088/test_solver".to_string(),
+    ]);
+    services.start_api(vec![]).await;
+
+    // Place Order
+    let order = OrderCreation {
+        sell_token: token_a.address(),
+        sell_amount: to_wei(9),
+        fee_amount: to_wei(1),
+        buy_token: token_b.address(),
+        buy_amount: to_wei(5),
+        valid_to: model::time::now_in_epoch_seconds() + 300,
+        kind: OrderKind::Buy,
+        ..Default::default()
+    }
+    .sign(
+        EcdsaSigningScheme::Eip712,
+        &onchain.contracts().domain_separator,
+        SecretKeyRef::from(&SecretKey::from_slice(trader.private_key()).unwrap()),
+    );
+    services.create_order(&order).await.unwrap();
+
+    tracing::info!("waiting for first trade");
+    let trade_happened =
+        || async { token_b.balance_of(trader.address()).call().await.unwrap() == order.buy_amount };
+    wait_for_condition(TIMEOUT, trade_happened).await.unwrap();
+
+    // Check that settlement buffers were traded.
+    let settlement_contract_balance = token_b
+        .balance_of(onchain.contracts().gp_settlement.address())
+        .call()
+        .await
+        .unwrap();
+    // Check that internal buffers were used
+    assert!(settlement_contract_balance == 0.into());
+
+    // Same order can trade again with external liquidity
+    let order = OrderCreation {
+        valid_to: model::time::now_in_epoch_seconds() + 301,
+        ..order
+    }
+    .sign(
+        EcdsaSigningScheme::Eip712,
+        &onchain.contracts().domain_separator,
+        SecretKeyRef::from(&SecretKey::from_slice(trader.private_key()).unwrap()),
+    );
+    services.create_order(&order).await.unwrap();
+
+    tracing::info!("waiting for second trade");
+    let trade_happened = || async {
+        token_b.balance_of(trader.address()).call().await.unwrap() == order.buy_amount * 2
+    };
+    wait_for_condition(TIMEOUT, trade_happened).await.unwrap();
+}

crates/e2e/tests/e2e/main.rs

@@ -6,6 +6,7 @@
 // Each of the following modules contains tests.
 mod app_data;
 mod app_data_signer;
+mod colocation_buffers;
 mod colocation_ethflow;
 mod colocation_hooks;
 mod colocation_partial_fill;