kabu_core_router/
swap_router_component.rs

1use eyre::{Result, eyre};
2use reth_tasks::TaskExecutor;
3use std::sync::Arc;
4use tokio::sync::{RwLock, broadcast};
5use tracing::{debug, error, info};
6
7use kabu_core_components::Component;
8use kabu_types_entities::{AccountNonceAndBalanceState, TxSigners};
9use kabu_types_events::{MessageSwapCompose, SwapComposeData, SwapComposeMessage, TxComposeData};
10use reth_node_types::NodePrimitives;
11use revm::DatabaseRef;
12
13/// Swap router component that handles routing of swap requests
14pub struct SwapRouterComponent<DB: Send + Sync + Clone + 'static, LDT: NodePrimitives + 'static = reth_ethereum_primitives::EthPrimitives> {
15    signers: Arc<RwLock<TxSigners<LDT>>>,
16    account_nonce_balance: Arc<RwLock<AccountNonceAndBalanceState>>,
17    swap_compose_rx: broadcast::Receiver<MessageSwapCompose<DB, LDT>>,
18    swap_compose_tx: broadcast::Sender<MessageSwapCompose<DB, LDT>>,
19}
20
21impl<DB, LDT> SwapRouterComponent<DB, LDT>
22where
23    DB: DatabaseRef + Send + Sync + Clone + 'static,
24    LDT: NodePrimitives + 'static,
25{
26    pub fn new(
27        signers: Arc<RwLock<TxSigners<LDT>>>,
28        account_nonce_balance: Arc<RwLock<AccountNonceAndBalanceState>>,
29        swap_compose_tx: broadcast::Sender<MessageSwapCompose<DB, LDT>>,
30    ) -> Self {
31        let swap_compose_rx = swap_compose_tx.subscribe();
32        Self { signers, account_nonce_balance, swap_compose_rx, swap_compose_tx }
33    }
34
35    async fn run(mut self) -> Result<()> {
36        info!("Starting swap router component");
37
38        loop {
39            match self.swap_compose_rx.recv().await {
40                Ok(msg) => {
41                    if let SwapComposeMessage::Prepare(swap_compose) = msg.inner() {
42                        self.handle_encode(swap_compose.clone()).await?;
43                    }
44                }
45                Err(e) => {
46                    error!("Swap compose channel error: {}", e);
47                    return Err(eyre!("Swap compose channel closed"));
48                }
49            }
50        }
51    }
52
53    async fn handle_encode(&self, route_request: SwapComposeData<DB, LDT>) -> Result<()> {
54        debug!("Router handling encode request: {}", route_request.swap);
55
56        let signer = match route_request.tx_compose.eoa {
57            Some(eoa) => self.signers.read().await.get_signer_by_address(&eoa)?,
58            None => self.signers.read().await.get_random_signer().ok_or(eyre!("NO_SIGNER"))?,
59        };
60
61        let nonce = self.account_nonce_balance.read().await.get_account(&signer.address()).unwrap().get_nonce();
62
63        let eth_balance = self.account_nonce_balance.read().await.get_account(&signer.address()).unwrap().get_eth_balance();
64
65        if route_request.tx_compose.next_block_base_fee == 0 {
66            error!("Block base fee is not set");
67            return Err(eyre!("NO_BLOCK_GAS_FEE"));
68        }
69
70        let gas = (route_request.swap.pre_estimate_gas()) * 2;
71
72        let estimate_request = SwapComposeData::<DB, LDT> {
73            tx_compose: TxComposeData::<LDT> { signer: Some(signer), nonce, eth_balance, gas, ..route_request.tx_compose },
74            ..route_request
75        };
76
77        let estimate_request = MessageSwapCompose::estimate(estimate_request);
78
79        match self.swap_compose_tx.send(estimate_request) {
80            Err(e) => {
81                error!("Failed to send estimate request: {}", e);
82                Err(eyre!("ERROR_SENDING_REQUEST: {}", e))
83            }
84            Ok(count) => {
85                debug!("Sent estimate request to {} receivers", count);
86                Ok(())
87            }
88        }
89    }
90}
91
92impl<DB, LDT> Component for SwapRouterComponent<DB, LDT>
93where
94    DB: DatabaseRef + Send + Sync + Clone + 'static,
95    LDT: NodePrimitives + 'static,
96{
97    fn spawn(self, executor: TaskExecutor) -> Result<()> {
98        executor.spawn_critical(self.name(), async move {
99            if let Err(e) = self.run().await {
100                tracing::error!("Swap router component failed: {}", e);
101            }
102        });
103        Ok(())
104    }
105    fn name(&self) -> &'static str {
106        "SwapRouterComponent"
107    }
108}
109
110/// Builder for SwapRouterComponent
111pub struct SwapRouterComponentBuilder {
112    // Builder can be extended with configuration options
113}
114
115impl SwapRouterComponentBuilder {
116    pub fn new() -> Self {
117        Self {}
118    }
119
120    pub fn build<DB, LDT>(
121        self,
122        signers: Arc<RwLock<TxSigners<LDT>>>,
123        account_nonce_balance: Arc<RwLock<AccountNonceAndBalanceState>>,
124        swap_compose_tx: broadcast::Sender<MessageSwapCompose<DB, LDT>>,
125    ) -> SwapRouterComponent<DB, LDT>
126    where
127        DB: DatabaseRef + Send + Sync + Clone + 'static,
128        LDT: NodePrimitives + 'static,
129    {
130        SwapRouterComponent::new(signers, account_nonce_balance, swap_compose_tx)
131    }
132}
133
134impl Default for SwapRouterComponentBuilder {
135    fn default() -> Self {
136        Self::new()
137    }
138}