reth_rpc/

reth.rs

use std::{future::Future, sync::Arc};

use alloy_consensus::BlockHeader;
use alloy_eips::BlockId;
use alloy_primitives::{map::AddressMap, U256, U64};
use futures::{Stream, StreamExt};
use jsonrpsee::{core::RpcResult, PendingSubscriptionSink, SubscriptionMessage, SubscriptionSink};
use reth_chain_state::{
    CanonStateNotification, CanonStateSubscriptions, ForkChoiceSubscriptions,
    PersistedBlockSubscriptions,
};
use reth_errors::{RethError, RethResult};
use reth_evm::{database::StateProviderDatabase, execute::Executor, ConfigureEvm};
use reth_execution_types::{Chain, ExecutionOutcome};
use reth_primitives_traits::{NodePrimitives, SealedHeader};
use reth_rpc_api::{RethApiServer, RethJitAction};
use reth_rpc_eth_types::{EthApiError, EthResult};
use reth_storage_api::{
    BlockReader, BlockReaderIdExt, ChangeSetReader, StateProvider, StateProviderFactory,
    TransactionVariant,
};
use reth_tasks::{pool::BlockingTaskGuard, CancelOnDrop, Runtime};
use serde::Serialize;
use serde_json::value::RawValue;
use tokio::sync::oneshot;

/// `reth` API implementation.
///
/// This type provides the functionality for handling `reth` prototype RPC requests.
pub struct RethApi<Provider, EvmConfig> {
    inner: Arc<RethApiInner<Provider, EvmConfig>>,
}

// === impl RethApi ===

impl<Provider, EvmConfig> RethApi<Provider, EvmConfig> {
    /// The provider that can interact with the chain.
    pub fn provider(&self) -> &Provider {
        &self.inner.provider
    }

    /// The evm config.
    pub fn evm_config(&self) -> &EvmConfig {
        &self.inner.evm_config
    }

    /// Create a new instance of the [`RethApi`]
    pub fn new(
        provider: Provider,
        evm_config: EvmConfig,
        blocking_task_guard: BlockingTaskGuard,
        task_spawner: Runtime,
    ) -> Self {
        let inner =
            Arc::new(RethApiInner { provider, evm_config, blocking_task_guard, task_spawner });
        Self { inner }
    }
}

impl<Provider, EvmConfig> RethApi<Provider, EvmConfig>
where
    Provider: BlockReaderIdExt + ChangeSetReader + StateProviderFactory + 'static,
    EvmConfig: Send + Sync + 'static,
{
    /// Executes the future on a new blocking task.
    async fn on_blocking_task<C, F, R>(&self, c: C) -> EthResult<R>
    where
        C: FnOnce(Self) -> F,
        F: Future<Output = EthResult<R>> + Send + 'static,
        R: Send + 'static,
    {
        let (tx, rx) = oneshot::channel();
        let this = self.clone();
        let f = c(this);
        self.inner.task_spawner.spawn_blocking_task(async move {
            let res = f.await;
            let _ = tx.send(res);
        });
        rx.await.map_err(|_| EthApiError::InternalEthError)?
    }

    /// Returns post-block balances for accounts whose balance changed in the block.
    ///
    /// Missing accounts are treated as having zero balance.
    pub async fn balance_changes_in_block(&self, block_id: BlockId) -> EthResult<AddressMap<U256>> {
        self.on_blocking_task(async move |this| this.try_balance_changes_in_block(block_id)).await
    }

    fn try_balance_changes_in_block(&self, block_id: BlockId) -> EthResult<AddressMap<U256>> {
        let Some(block_number) = self.provider().block_number_for_id(block_id)? else {
            return Err(EthApiError::HeaderNotFound(block_id))
        };

        let state = self.provider().state_by_block_id(block_id)?;
        let accounts_before = self.provider().account_block_changeset(block_number)?;
        let hash_map = accounts_before.iter().try_fold(
            AddressMap::default(),
            |mut hash_map, account_before| -> RethResult<_> {
                let current_balance =
                    state.account_balance(&account_before.address)?.unwrap_or_default();
                let prev_balance =
                    account_before.info.as_ref().map(|info| info.balance).unwrap_or_default();
                if current_balance != prev_balance {
                    hash_map.insert(account_before.address, current_balance);
                }
                Ok(hash_map)
            },
        )?;
        Ok(hash_map)
    }
}

impl<N, Provider, EvmConfig> RethApi<Provider, EvmConfig>
where
    N: NodePrimitives,
    Provider: BlockReaderIdExt
        + ChangeSetReader
        + StateProviderFactory
        + BlockReader<Block = N::Block>
        + CanonStateSubscriptions<Primitives = N>
        + 'static,
    EvmConfig: ConfigureEvm<Primitives = N> + 'static,
{
    /// Re-executes one or more consecutive blocks and returns the execution outcome.
    pub async fn block_execution_outcome(
        &self,
        block_id: BlockId,
        count: Option<U64>,
    ) -> EthResult<Option<ExecutionOutcome<N::Receipt>>> {
        const MAX_BLOCK_COUNT: u64 = 128;

        let block_count = count.map(|c| c.to::<u64>()).unwrap_or(1);
        if block_count == 0 || block_count > MAX_BLOCK_COUNT {
            return Err(EthApiError::InvalidParams(format!(
                "block count must be between 1 and {MAX_BLOCK_COUNT}, got {block_count}"
            )))
        }

        let permit = self
            .inner
            .blocking_task_guard
            .clone()
            .acquire_owned()
            .await
            .map_err(|_| EthApiError::InternalEthError)?;
        let guard = CancelOnDrop::default();
        let cancel = guard.clone();
        let outcome = self
            .on_blocking_task(async move |this| {
                let _permit = permit;
                this.try_block_execution_outcome(block_id, block_count, &cancel)
            })
            .await;
        drop(guard);
        outcome
    }

    fn try_block_execution_outcome(
        &self,
        block_id: BlockId,
        block_count: u64,
        cancel: &CancelOnDrop,
    ) -> EthResult<Option<ExecutionOutcome<N::Receipt>>> {
        let Some(start_block) = self.provider().block_number_for_id(block_id)? else {
            return Ok(None)
        };

        if start_block == 0 {
            return Ok(Some(ExecutionOutcome::default()))
        }

        let state_provider = self.provider().history_by_block_number(start_block - 1)?;
        let db = StateProviderDatabase::new((&state_provider).into_evm_state_provider());

        let mut blocks = Vec::with_capacity(block_count as usize);
        for block_number in start_block..start_block + block_count {
            let Some(block) = self
                .provider()
                .recovered_block(block_number.into(), TransactionVariant::WithHash)?
            else {
                if block_number == start_block {
                    return Ok(None)
                }
                break;
            };
            blocks.push(block);
        }

        // stop between blocks once the request is dropped
        let blocks = blocks.iter().take_while(|_| !cancel.is_cancelled());
        let outcome = self.evm_config().executor(db).execute_batch(blocks).map_err(
            |e: reth_evm::execute::BlockExecutionError| {
                EthApiError::Internal(reth_errors::RethError::Other(e.into()))
            },
        )?;
        if cancel.is_cancelled() {
            return Err(EthApiError::InternalEthError)
        }

        Ok(Some(outcome))
    }
}

impl<Provider, EvmConfig> RethApiServer for RethApi<Provider, EvmConfig>
where
    Provider: BlockReaderIdExt
        + ChangeSetReader
        + StateProviderFactory
        + BlockReader<
            Block = <<Provider as CanonStateSubscriptions>::Primitives as NodePrimitives>::Block,
        > + CanonStateSubscriptions
        + ForkChoiceSubscriptions<
            Header = <<Provider as CanonStateSubscriptions>::Primitives as NodePrimitives>::BlockHeader,
        >
        + PersistedBlockSubscriptions
        + 'static,
    EvmConfig: ConfigureEvm<Primitives = <Provider as CanonStateSubscriptions>::Primitives> + 'static,
{
    /// Handler for `reth_getBalanceChangesInBlock`
    async fn reth_get_balance_changes_in_block(
        &self,
        block_id: BlockId,
    ) -> RpcResult<AddressMap<U256>> {
        Ok(Self::balance_changes_in_block(self, block_id).await?)
    }

    /// Handler for `reth_getBlockExecutionOutcome`
    async fn reth_get_block_execution_outcome(
        &self,
        block_id: BlockId,
        count: Option<U64>,
    ) -> RpcResult<Option<Box<RawValue>>> {
        let outcome = Self::block_execution_outcome(self, block_id, count).await?;
        match outcome {
            Some(outcome) => {
                let raw = serde_json::value::to_raw_value(&outcome).map_err(|e| {
                    EthApiError::Internal(reth_errors::RethError::msg(e.to_string()))
                })?;
                Ok(Some(raw))
            }
            None => Ok(None),
        }
    }

    /// Handler for `reth_jit`
    async fn reth_jit(&self, action: RethJitAction) -> RpcResult<()> {
        let Some(jit_backend) = self.evm_config().jit_backend() else {
            return Ok(());
        };

        match action {
            RethJitAction::Enable => jit_backend
                .set_enabled(true)
                .map_err(|err| EthApiError::Internal(RethError::msg(err)))?,
            RethJitAction::Disable => jit_backend
                .set_enabled(false)
                .map_err(|err| EthApiError::Internal(RethError::msg(err)))?,
            RethJitAction::Pause => jit_backend.pause(),
            RethJitAction::Unpause => jit_backend.resume(),
            RethJitAction::Clear => jit_backend.clear(),
        }

        Ok(())
    }

    /// Handler for `reth_subscribeChainNotifications`
    async fn reth_subscribe_chain_notifications(
        &self,
        pending: PendingSubscriptionSink,
    ) -> jsonrpsee::core::SubscriptionResult {
        let sink = pending.accept().await?;
        let stream = self.provider().canonical_state_stream();
        self.inner.task_spawner.spawn_task(pipe_from_stream(sink, stream));

        Ok(())
    }

    /// Handler for `reth_subscribePersistedBlock`
    async fn reth_subscribe_persisted_block(
        &self,
        pending: PendingSubscriptionSink,
    ) -> jsonrpsee::core::SubscriptionResult {
        let sink = pending.accept().await?;
        let stream = self.provider().persisted_block_stream();
        self.inner.task_spawner.spawn_task(pipe_from_stream(sink, stream));

        Ok(())
    }

    /// Handler for `reth_subscribeFinalizedChainNotifications`
    async fn reth_subscribe_finalized_chain_notifications(
        &self,
        pending: PendingSubscriptionSink,
    ) -> jsonrpsee::core::SubscriptionResult {
        let sink = pending.accept().await?;
        let canon_stream = self.provider().canonical_state_stream();
        let finalized_stream = self.provider().finalized_block_stream();
        self.inner.task_spawner.spawn_task(finalized_chain_notifications(
            sink,
            canon_stream,
            finalized_stream,
        ));

        Ok(())
    }
}

/// Pipes all stream items to the subscription sink.
async fn pipe_from_stream<S, T>(sink: SubscriptionSink, mut stream: S)
where
    S: Stream<Item = T> + Unpin,
    T: Serialize,
{
    loop {
        tokio::select! {
            _ = sink.closed() => {
                break
            }
            maybe_item = stream.next() => {
                let Some(item) = maybe_item else {
                    break
                };
                let msg = match SubscriptionMessage::new(sink.method_name(), sink.subscription_id(), &item) {
                    Ok(msg) => msg,
                    Err(err) => {
                        tracing::error!(target: "rpc::reth", %err, "Failed to serialize subscription message");
                        break
                    }
                };
                if sink.send(msg).await.is_err() {
                    break;
                }
            }
        }
    }
}

/// Buffers committed chain notifications and emits them when a new finalized block is received.
async fn finalized_chain_notifications<N>(
    sink: SubscriptionSink,
    mut canon_stream: reth_chain_state::CanonStateNotificationStream<N>,
    mut finalized_stream: reth_chain_state::ForkChoiceStream<SealedHeader<N::BlockHeader>>,
) where
    N: NodePrimitives,
{
    let mut buffered: Vec<CanonStateNotification<N>> = Vec::new();

    loop {
        tokio::select! {
            _ = sink.closed() => {
                break
            }
            maybe_canon = canon_stream.next() => {
                let Some(notification) = maybe_canon else { break };
                match &notification {
                    CanonStateNotification::Commit { .. } => {
                        buffered.push(notification);
                    }
                    CanonStateNotification::Reorg { old, new } => {
                        let first_reverted = old.first().number();
                        buffered.retain_mut(|notification| {
                            let chain = notification.committed();
                            if chain.first().number() >= first_reverted {
                                return false
                            }
                            // Preserve the canonical prefix of a segment crossing the fork.
                            if chain.tip().number() >= first_reverted {
                                let (blocks, mut outcome, mut trie_data) = (*chain).clone().into_inner();
                                outcome.revert_to(first_reverted - 1);
                                trie_data.split_off(&first_reverted);
                                *notification = CanonStateNotification::Commit {
                                    new: Arc::new(Chain::new(
                                        blocks.into_blocks().take_while(|b| b.number() < first_reverted),
                                        outcome,
                                        trie_data,
                                    )),
                                };
                            }
                            true
                        });
                        if !new.is_empty() {
                            buffered.push(CanonStateNotification::Commit { new: new.clone() });
                        }
                    }
                }
            }
            maybe_finalized = finalized_stream.next() => {
                let Some(finalized_header) = maybe_finalized else { break };
                let finalized_num = finalized_header.number();

                let mut committed = Vec::new();
                buffered.retain(|n| {
                    if *n.committed().range().end() <= finalized_num {
                        committed.push(n.clone());
                        false
                    } else {
                        true
                    }
                });

                if committed.is_empty() {
                    continue;
                }

                committed.sort_by_key(|n| *n.committed().range().start());

                let msg = match SubscriptionMessage::new(
                    sink.method_name(),
                    sink.subscription_id(),
                    &committed,
                ) {
                    Ok(msg) => msg,
                    Err(err) => {
                        tracing::error!(target: "rpc::reth", %err, "Failed to serialize finalized chain notification");
                        break
                    }
                };
                if sink.send(msg).await.is_err() {
                    break;
                }
            }
        }
    }
}

impl<Provider, EvmConfig> std::fmt::Debug for RethApi<Provider, EvmConfig> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("RethApi").finish_non_exhaustive()
    }
}

impl<Provider, EvmConfig> Clone for RethApi<Provider, EvmConfig> {
    fn clone(&self) -> Self {
        Self { inner: Arc::clone(&self.inner) }
    }
}

struct RethApiInner<Provider, EvmConfig> {
    /// The provider that can interact with the chain.
    provider: Provider,
    /// The EVM configuration used to create block executors.
    evm_config: EvmConfig,
    /// Guard to restrict the number of concurrent block re-execution requests.
    blocking_task_guard: BlockingTaskGuard,
    /// The type that can spawn tasks which would otherwise block.
    task_spawner: Runtime,
}

#[cfg(test)]
mod tests {
    use super::*;
    use alloy_consensus::{Block, Header};
    use alloy_genesis::{Genesis, GenesisAccount};
    use alloy_primitives::Address;
    use evm2::evm::AccountInfo;
    use reth_chainspec::ChainSpecBuilder;
    use reth_db_common::init::init_genesis;
    use reth_primitives_traits::Block as _;
    use reth_provider::{
        providers::BlockchainProvider, test_utils::create_test_provider_factory_with_chain_spec,
        BlockWriter,
    };
    use reth_trie_common::{bundle_state::BundleState, HashedPostState, KeccakKeyHasher};

    #[test]
    fn balance_changes_ignore_account_existence() {
        let transitions = [
            (1, None, Some(0)),      // Zero-balance creation.
            (2, Some(0), None),      // Zero-balance deletion.
            (3, None, Some(10)),     // Funded creation.
            (4, Some(10), None),     // Funded deletion.
            (5, Some(10), Some(20)), // Balance increase.
            (6, Some(10), Some(0)),  // Existing account emptied.
            (7, Some(10), Some(10)), // Nonce-only change.
        ]
        .map(|(id, before, after)| {
            (
                Address::with_last_byte(id),
                before.map(|balance| AccountInfo {
                    balance: U256::from(balance),
                    nonce: 1,
                    ..Default::default()
                }),
                after.map(|balance| AccountInfo {
                    balance: U256::from(balance),
                    nonce: 2,
                    ..Default::default()
                }),
            )
        });
        let genesis = Genesis::default().extend_accounts(transitions.iter().filter_map(
            |(address, before, _)| {
                before.as_ref().map(|account| {
                    (
                        *address,
                        GenesisAccount::default()
                            .with_balance(account.balance)
                            .with_nonce(Some(account.nonce)),
                    )
                })
            },
        ));
        let factory = create_test_provider_factory_with_chain_spec(Arc::new(
            ChainSpecBuilder::mainnet().genesis(genesis).build(),
        ));
        let genesis_hash = init_genesis(&factory).unwrap();
        let block = Block {
            header: Header { parent_hash: genesis_hash, number: 1, ..Default::default() },
            body: Default::default(),
        };
        let outcome = ExecutionOutcome {
            bundle: BundleState::new(
                transitions.iter().map(|(address, before, after)| {
                    (*address, before.clone(), after.clone(), Default::default())
                }),
                [transitions
                    .iter()
                    .map(|(address, before, _)| (*address, Some(before.clone()), []))],
                [],
            ),
            first_block: 1,
            receipts: vec![vec![]],
            ..Default::default()
        };
        let provider_rw = factory.provider_rw().unwrap();
        provider_rw
            .append_blocks_with_state(
                vec![block.seal_slow().try_recover().unwrap()],
                &outcome,
                HashedPostState::from_bundle_state::<KeccakKeyHasher>(&outcome.bundle.state)
                    .into_sorted(),
            )
            .unwrap();
        provider_rw.commit().unwrap();
        let api = RethApi::new(
            BlockchainProvider::new(factory).unwrap(),
            (),
            BlockingTaskGuard::new(1),
            Runtime::test(),
        );

        assert_eq!(
            api.try_balance_changes_in_block(1.into()).unwrap(),
            AddressMap::from_iter([
                (Address::with_last_byte(3), U256::from(10)),
                (Address::with_last_byte(4), U256::ZERO),
                (Address::with_last_byte(5), U256::from(20)),
                (Address::with_last_byte(6), U256::ZERO),
            ])
        );
    }
}