Skip to content
This repository has been archived by the owner on Dec 26, 2024. It is now read-only.

Commit

Permalink
refactor(network): db executor communicates through channels (#2120)
Browse files Browse the repository at this point in the history 
* refactor(network): remove Data from db executor

* refactor(network): fix cr comments

* refactor(network): db executor communicates through channels (#2121)

* refactor(network): db executor communicates through channels

* refactor(network): fix cr comments

* refactor(network): db executor is removed from NetworkManager and runs in a separate task (#2122)

* refactor(network): db executor is removed from NetworkManager and runs in a separate task

* test(network): return network manager tests (#2124)
  • Loading branch information
@ShahakShama
ShahakShama authored Jun 25, 2024
1 parent 3675c53 commit 89854f5
Show file tree
Hide file tree
Showing 8 changed files with 405 additions and 526 deletions.
1 change: 0 additions & 1 deletion Cargo.lock
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 0 additions & 1 deletion crates/papyrus_network/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,6 @@ thiserror.workspace = true
tokio = { workspace = true, features = ["full", "sync"] }
tracing.workspace = true
unsigned-varint = { workspace = true, features = ["std"] }
async-trait.workspace = true
validator = { workspace = true, features = ["derive"] }

# Binaries dependencies
Expand Down
266 changes: 111 additions & 155 deletions crates/papyrus_network/src/db_executor/mod.rs
View file
Original file line number Diff line number Diff line change
@@ -1,75 +1,32 @@
use std::vec;

use async_trait::async_trait;
use bytes::BufMut;
use futures::channel::mpsc::Sender;
use futures::future::{pending, poll_fn};
#[cfg(test)]
use mockall::automock;
use papyrus_protobuf::protobuf;
use futures::channel::mpsc::SendError;
use futures::{Sink, SinkExt, Stream, StreamExt};
use papyrus_protobuf::converters::ProtobufConversionError;
use papyrus_protobuf::sync::{
BlockHashOrNumber,
ContractDiff,
DataOrFin,
DeclaredClass,
DeprecatedDeclaredClass,
HeaderQuery,
Query,
SignedBlockHeader,
StateDiffChunk,
StateDiffQuery,
};
use papyrus_storage::header::HeaderStorageReader;
use papyrus_storage::state::StateStorageReader;
use papyrus_storage::{db, StorageReader, StorageTxn};
use prost::Message;
use starknet_api::block::BlockNumber;
use starknet_api::state::ThinStateDiff;
use tracing::error;

use crate::DataType;

#[cfg(test)]
mod test;

mod utils;

#[derive(thiserror::Error, Debug)]
#[error("Failed to encode data")]
pub struct DataEncodingError;

#[cfg_attr(test, derive(Debug, PartialEq, Eq))]
#[derive(Clone)]
pub enum Data {
BlockHeaderAndSignature(DataOrFin<SignedBlockHeader>),
StateDiffChunk(DataOrFin<StateDiffChunk>),
}

impl Default for Data {
fn default() -> Self {
// TODO: consider this default data type.
Self::BlockHeaderAndSignature(DataOrFin(None))
}
}

impl Data {
pub fn encode<B>(self, buf: &mut B) -> Result<(), DataEncodingError>
where
B: BufMut,
{
match self {
Data::BlockHeaderAndSignature(maybe_signed_block_header) => {
let block_headers_response =
protobuf::BlockHeadersResponse::from(maybe_signed_block_header);
block_headers_response.encode(buf).map_err(|_| DataEncodingError)
}
Data::StateDiffChunk(maybe_state_diff_chunk) => {
let state_diffs_response =
protobuf::StateDiffsResponse::from(maybe_state_diff_chunk);
state_diffs_response.encode(buf).map_err(|_| DataEncodingError)
}
}
}
}

#[derive(thiserror::Error, Debug)]
pub enum DBExecutorError {
#[error(transparent)]
Expand Down Expand Up @@ -102,44 +59,66 @@ impl DBExecutorError {
}

/// A DBExecutor receives inbound queries and returns their corresponding data.
#[async_trait]
pub trait DBExecutorTrait {
/// Send a query to be executed in the DBExecutor. The query will be run concurrently with the
/// calling code and the result will be over the given channel.
fn register_query(
&mut self,
query: Query,
data_type: impl FetchBlockDataFromDb + Send + 'static,
sender: Sender<Data>,
);

/// Polls incoming queries.
// TODO(shahak): Consume self.
async fn run(&mut self);
}

pub struct DBExecutor {
pub struct DBExecutor<HeaderQueryReceiver, StateDiffQueryReceiver> {
storage_reader: StorageReader,
header_queries_receiver: HeaderQueryReceiver,
state_diff_queries_receiver: StateDiffQueryReceiver,
}

impl DBExecutor {
pub fn new(storage_reader: StorageReader) -> Self {
Self { storage_reader }
impl<HeaderQueryReceiver, StateDiffQueryReceiver, HeaderResponsesSender, StateDiffResponsesSender>
DBExecutor<HeaderQueryReceiver, StateDiffQueryReceiver>
where
HeaderQueryReceiver: Stream<Item = (Result<HeaderQuery, ProtobufConversionError>, HeaderResponsesSender)>
+ Unpin,
HeaderResponsesSender:
Sink<DataOrFin<SignedBlockHeader>, Error = SendError> + Unpin + Send + 'static,
StateDiffQueryReceiver: Stream<Item = (Result<StateDiffQuery, ProtobufConversionError>, StateDiffResponsesSender)>
+ Unpin,
StateDiffResponsesSender:
Sink<DataOrFin<StateDiffChunk>, Error = SendError> + Unpin + Send + 'static,
{
pub async fn run(mut self) {
loop {
tokio::select! {
result = self.header_queries_receiver.next() => {
let (query_result, response_sender) = result.expect(
"Header queries sender was unexpectedly dropped."
);
// TODO(shahak): Report if query_result is Err.
if let Ok(query) = query_result {
self.register_query(query.0, response_sender);
}
}
result = self.state_diff_queries_receiver.next() => {
let (query_result, response_sender) = result.expect(
"State diff queries sender was unexpectedly dropped."
);
// TODO(shahak): Report if query_result is Err.
if let Ok(query) = query_result {
self.register_query(query.0, response_sender);
}
}
};
}
}
}

#[async_trait]
impl DBExecutorTrait for DBExecutor {
fn register_query(
&mut self,
query: Query,
data_type: impl FetchBlockDataFromDb + Send + 'static,
sender: Sender<Data>,
) {
pub fn new(
storage_reader: StorageReader,
header_queries_receiver: HeaderQueryReceiver,
state_diff_queries_receiver: StateDiffQueryReceiver,
) -> Self {
Self { storage_reader, header_queries_receiver, state_diff_queries_receiver }
}

fn register_query<Data, Sender>(&self, query: Query, sender: Sender)
where
Data: FetchBlockDataFromDb + Send + 'static,
Sender: Sink<DataOrFin<Data>> + Unpin + Send + 'static,
DBExecutorError: From<<Sender as Sink<DataOrFin<Data>>>::Error>,
{
let storage_reader_clone = self.storage_reader.clone();
tokio::task::spawn(async move {
let result =
send_data_for_query(storage_reader_clone, query.clone(), data_type, sender).await;
let result = send_data_for_query(storage_reader_clone, query.clone(), sender).await;
if let Err(error) = result {
if error.should_log_in_error_level() {
error!("Running inbound query {query:?} failed on {error:?}");
Expand All @@ -150,77 +129,51 @@ impl DBExecutorTrait for DBExecutor {
}
});
}

async fn run(&mut self) {
// TODO(shahak): Parse incoming queries once we receive them through channel instead of
// through function.
pending::<()>().await
}
}

#[cfg_attr(test, automock)]
// we need to tell clippy to ignore the "needless" lifetime warning because it's not true.
// we do need the lifetime for the automock, following clippy's suggestion will break the code.
#[allow(clippy::needless_lifetimes)]
pub trait FetchBlockDataFromDb {
fn fetch_block_data_from_db<'a>(
&self,
pub trait FetchBlockDataFromDb: Sized {
fn fetch_block_data_from_db(
block_number: BlockNumber,
txn: &StorageTxn<'a, db::RO>,
) -> Result<Vec<Data>, DBExecutorError>;

fn fin(&self) -> Data;
txn: &StorageTxn<'_, db::RO>,
) -> Result<Vec<Self>, DBExecutorError>;
}

impl FetchBlockDataFromDb for DataType {
impl FetchBlockDataFromDb for SignedBlockHeader {
fn fetch_block_data_from_db(
&self,
block_number: BlockNumber,
txn: &StorageTxn<'_, db::RO>,
) -> Result<Vec<Data>, DBExecutorError> {
match self {
DataType::SignedBlockHeader => {
let mut header =
txn.get_block_header(block_number)?.ok_or(DBExecutorError::BlockNotFound {
block_hash_or_number: BlockHashOrNumber::Number(block_number),
})?;
// TODO(shahak) Remove this once central sync fills the state_diff_length field.
if header.state_diff_length.is_none() {
header.state_diff_length = Some(
txn.get_state_diff(block_number)?
.ok_or(DBExecutorError::BlockNotFound {
block_hash_or_number: BlockHashOrNumber::Number(block_number),
})?
.len(),
);
}
let signature = txn
.get_block_signature(block_number)?
.ok_or(DBExecutorError::SignatureNotFound { block_number })?;
Ok(vec![Data::BlockHeaderAndSignature(DataOrFin(Some(SignedBlockHeader {
block_header: header,
signatures: vec![signature],
})))])
}
DataType::StateDiff => {
let thin_state_diff =
txn.get_state_diff(block_number)?.ok_or(DBExecutorError::BlockNotFound {
) -> Result<Vec<Self>, DBExecutorError> {
let mut header =
txn.get_block_header(block_number)?.ok_or(DBExecutorError::BlockNotFound {
block_hash_or_number: BlockHashOrNumber::Number(block_number),
})?;
// TODO(shahak) Remove this once central sync fills the state_diff_length field.
if header.state_diff_length.is_none() {
header.state_diff_length = Some(
txn.get_state_diff(block_number)?
.ok_or(DBExecutorError::BlockNotFound {
block_hash_or_number: BlockHashOrNumber::Number(block_number),
})?;
let vec_data = split_thin_state_diff(thin_state_diff)
.into_iter()
.map(|state_diff_chunk| Data::StateDiffChunk(DataOrFin(Some(state_diff_chunk))))
.collect();
Ok(vec_data)
}
})?
.len(),
);
}
let signature = txn
.get_block_signature(block_number)?
.ok_or(DBExecutorError::SignatureNotFound { block_number })?;
Ok(vec![SignedBlockHeader { block_header: header, signatures: vec![signature] }])
}
}

fn fin(&self) -> Data {
match self {
DataType::SignedBlockHeader => Data::BlockHeaderAndSignature(DataOrFin(None)),
DataType::StateDiff => Data::StateDiffChunk(DataOrFin(None)),
}
impl FetchBlockDataFromDb for StateDiffChunk {
fn fetch_block_data_from_db(
block_number: BlockNumber,
txn: &StorageTxn<'_, db::RO>,
) -> Result<Vec<Self>, DBExecutorError> {
let thin_state_diff =
txn.get_state_diff(block_number)?.ok_or(DBExecutorError::BlockNotFound {
block_hash_or_number: BlockHashOrNumber::Number(block_number),
})?;
Ok(split_thin_state_diff(thin_state_diff))
}
}

Expand Down Expand Up @@ -265,27 +218,32 @@ pub fn split_thin_state_diff(thin_state_diff: ThinStateDiff) -> Vec<StateDiffChu
state_diff_chunks
}

async fn send_data_for_query(
async fn send_data_for_query<Data, Sender>(
storage_reader: StorageReader,
query: Query,
data_type: impl FetchBlockDataFromDb + Send + 'static,
mut sender: Sender<Data>,
) -> Result<(), DBExecutorError> {
let fin = data_type.fin();
mut sender: Sender,
) -> Result<(), DBExecutorError>
where
Data: FetchBlockDataFromDb + Send + 'static,
Sender: Sink<DataOrFin<Data>> + Unpin + Send + 'static,
DBExecutorError: From<<Sender as Sink<DataOrFin<Data>>>::Error>,
{
// If this function fails, we still want to send fin before failing.
let result =
send_data_without_fin_for_query(&storage_reader, query, data_type, &mut sender).await;
poll_fn(|cx| sender.poll_ready(cx)).await?;
sender.start_send(fin)?;
let result = send_data_without_fin_for_query(&storage_reader, query, &mut sender).await;
sender.feed(DataOrFin(None)).await?;
result
}

async fn send_data_without_fin_for_query(
async fn send_data_without_fin_for_query<Data, Sender>(
storage_reader: &StorageReader,
query: Query,
data_type: impl FetchBlockDataFromDb + Send + 'static,
sender: &mut Sender<Data>,
) -> Result<(), DBExecutorError> {
sender: &mut Sender,
) -> Result<(), DBExecutorError>
where
Data: FetchBlockDataFromDb + Send + 'static,
Sender: Sink<DataOrFin<Data>> + Unpin + Send + 'static,
DBExecutorError: From<<Sender as Sink<DataOrFin<Data>>>::Error>,
{
let txn = storage_reader.begin_ro_txn()?;
let start_block_number = match query.start_block {
BlockHashOrNumber::Number(BlockNumber(num)) => num,
Expand All @@ -300,12 +258,10 @@ async fn send_data_without_fin_for_query(
for block_counter in 0..query.limit {
let block_number =
BlockNumber(utils::calculate_block_number(&query, start_block_number, block_counter)?);
let data_vec = data_type.fetch_block_data_from_db(block_number, &txn)?;
// Using poll_fn because Sender::poll_ready is not a future
poll_fn(|cx| sender.poll_ready(cx)).await?;
let data_vec = Data::fetch_block_data_from_db(block_number, &txn)?;
for data in data_vec {
// TODO: consider implement retry mechanism.
sender.start_send(data)?;
sender.feed(DataOrFin(Some(data))).await?;
}
}
Ok(())
Expand Down
Loading

0 comments on commit 89854f5

Please sign in to comment.