opt(torii-core): move off queryqueue for executing tx #2460
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,237 @@ | ||
| use std::collections::VecDeque; | ||
| use std::mem; | ||
|
|
||
| use anyhow::{Context, Result}; | ||
| use dojo_types::schema::{Struct, Ty}; | ||
| use sqlx::query::Query; | ||
| use sqlx::sqlite::SqliteArguments; | ||
| use sqlx::{FromRow, Pool, Sqlite, Transaction}; | ||
| use starknet::core::types::Felt; | ||
| use tokio::sync::broadcast::{Receiver, Sender}; | ||
| use tokio::sync::mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender}; | ||
|
|
||
| use crate::simple_broker::SimpleBroker; | ||
| use crate::types::{ | ||
| Entity as EntityUpdated, Event as EventEmitted, EventMessage as EventMessageUpdated, | ||
| Model as ModelRegistered, | ||
| }; | ||
|
|
||
| #[derive(Debug, Clone)] | ||
| pub enum Argument { | ||
| Null, | ||
| Int(i64), | ||
| Bool(bool), | ||
| String(String), | ||
| FieldElement(Felt), | ||
| } | ||
|
|
||
| #[derive(Debug, Clone)] | ||
| pub enum BrokerMessage { | ||
| ModelRegistered(ModelRegistered), | ||
| EntityUpdated(EntityUpdated), | ||
| EventMessageUpdated(EventMessageUpdated), | ||
| EventEmitted(EventEmitted), | ||
| } | ||
|
|
||
| #[derive(Debug, Clone)] | ||
| pub struct DeleteEntityQuery { | ||
| pub entity_id: String, | ||
| pub event_id: String, | ||
| pub block_timestamp: String, | ||
| pub ty: Ty, | ||
| } | ||
|
|
||
| #[derive(Debug, Clone)] | ||
| pub enum QueryType { | ||
| SetEntity(Ty), | ||
| DeleteEntity(DeleteEntityQuery), | ||
| EventMessage(Ty), | ||
| RegisterModel, | ||
| StoreEvent, | ||
| Execute, | ||
| Other, | ||
| } | ||
|
|
||
| pub struct Executor<'c> { | ||
| pool: Pool<Sqlite>, | ||
| transaction: Transaction<'c, Sqlite>, | ||
| publish_queue: VecDeque<BrokerMessage>, | ||
| rx: UnboundedReceiver<QueryMessage>, | ||
| shutdown_rx: Receiver<()>, | ||
| } | ||
|
|
||
| pub struct QueryMessage { | ||
| pub statement: String, | ||
| pub arguments: Vec<Argument>, | ||
| pub query_type: QueryType, | ||
| } | ||
|
|
||
| impl QueryMessage { | ||
| pub fn new(statement: String, arguments: Vec<Argument>, query_type: QueryType) -> Self { | ||
| Self { statement, arguments, query_type } | ||
| } | ||
|
|
||
| pub fn other(statement: String, arguments: Vec<Argument>) -> Self { | ||
| Self { statement, arguments, query_type: QueryType::Other } | ||
| } | ||
|
|
||
| pub fn execute() -> Self { | ||
| Self { statement: "".to_string(), arguments: vec![], query_type: QueryType::Execute } | ||
| } | ||
| } | ||
|
|
||
| impl<'c> Executor<'c> { | ||
| pub async fn new( | ||
| pool: Pool<Sqlite>, | ||
| shutdown_tx: Sender<()>, | ||
| ) -> Result<(Self, UnboundedSender<QueryMessage>)> { | ||
| let (tx, rx) = unbounded_channel(); | ||
lambda-0x marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| let transaction = pool.begin().await?; | ||
| let publish_queue = VecDeque::new(); | ||
| let shutdown_rx = shutdown_tx.subscribe(); | ||
|
|
||
| Ok((Executor { pool, transaction, publish_queue, rx, shutdown_rx }, tx)) | ||
| } | ||
|
|
||
| pub async fn run(&mut self) -> Result<()> { | ||
| loop { | ||
| tokio::select! { | ||
| _ = self.shutdown_rx.recv() => { | ||
| break Ok(()); | ||
| } | ||
| Some(msg) = self.rx.recv() => { | ||
| let QueryMessage { statement, arguments, query_type } = msg; | ||
| let mut query = sqlx::query(&statement); | ||
|
|
||
| for arg in &arguments { | ||
| query = match arg { | ||
| Argument::Null => query.bind(None::<String>), | ||
| Argument::Int(integer) => query.bind(integer), | ||
| Argument::Bool(bool) => query.bind(bool), | ||
| Argument::String(string) => query.bind(string), | ||
| Argument::FieldElement(felt) => query.bind(format!("{:#x}", felt)), | ||
| } | ||
| } | ||
|
|
||
| self.handle_query_type(query, query_type, &statement, &arguments).await?; | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
There was a problem hiding this comment. Consider improving error handling in the transaction lifecycle. Ohayo, sensei! While the overall implementation is solid, there's room for improvement in error handling, particularly in the transaction lifecycle. Currently, if an error occurs during query execution, the transaction may remain open, potentially leading to database inconsistencies. Consider wrapping the main loop in a pub async fn run(&mut self) -> Result<()> {
loop {
match self.process_messages().await {
Ok(should_break) => {
if should_break {
break;
}
}
Err(e) => {
// Rollback the transaction on error
self.transaction.rollback().await?;
// Start a new transaction for the next iteration
self.transaction = self.pool.begin().await?;
// Log the error or handle it as appropriate
eprintln!("Error processing messages: {:?}", e);
}
}
}
Ok(())
}
async fn process_messages(&mut self) -> Result<bool> {
tokio::select! {
_ = self.shutdown_rx.recv() => {
return Ok(true);
}
Some(msg) = self.rx.recv() => {
let QueryMessage { statement, arguments, query_type } = msg;
// ... (rest of the existing code)
self.handle_query_type(query, query_type, &statement, &arguments).await?;
}
}
Ok(false)
}This approach ensures that the transaction is rolled back in case of errors, maintaining database consistency. |
||
|
|
||
| async fn handle_query_type<'a>( | ||
| &mut self, | ||
| query: Query<'a, Sqlite, SqliteArguments<'a>>, | ||
| query_type: QueryType, | ||
| statement: &str, | ||
| arguments: &[Argument], | ||
| ) -> Result<()> { | ||
| let tx = &mut self.transaction; | ||
|
|
||
| match query_type { | ||
| QueryType::SetEntity(entity) => { | ||
| let row = query.fetch_one(&mut **tx).await.with_context(|| { | ||
| format!("Failed to execute query: {:?}, args: {:?}", statement, arguments) | ||
| })?; | ||
| let mut entity_updated = EntityUpdated::from_row(&row)?; | ||
| entity_updated.updated_model = Some(entity); | ||
| entity_updated.deleted = false; | ||
| let broker_message = BrokerMessage::EntityUpdated(entity_updated); | ||
| self.publish_queue.push_back(broker_message); | ||
| } | ||
| QueryType::DeleteEntity(entity) => { | ||
| let delete_model = query.execute(&mut **tx).await.with_context(|| { | ||
| format!("Failed to execute query: {:?}, args: {:?}", statement, arguments) | ||
| })?; | ||
| if delete_model.rows_affected() == 0 { | ||
| return Ok(()); | ||
| } | ||
|
|
||
| let row = sqlx::query( | ||
| "UPDATE entities SET updated_at=CURRENT_TIMESTAMP, executed_at=?, event_id=? \ | ||
| WHERE id = ? RETURNING *", | ||
| ) | ||
| .bind(entity.block_timestamp) | ||
| .bind(entity.event_id) | ||
| .bind(entity.entity_id) | ||
| .fetch_one(&mut **tx) | ||
| .await?; | ||
| let mut entity_updated = EntityUpdated::from_row(&row)?; | ||
| entity_updated.updated_model = | ||
| Some(Ty::Struct(Struct { name: entity.ty.name(), children: vec![] })); | ||
|
|
||
| let count = sqlx::query_scalar::<_, i64>( | ||
| "SELECT count(*) FROM entity_model WHERE entity_id = ?", | ||
| ) | ||
| .bind(entity_updated.id.clone()) | ||
| .fetch_one(&mut **tx) | ||
| .await?; | ||
|
|
||
| // Delete entity if all of its models are deleted | ||
| if count == 0 { | ||
| sqlx::query("DELETE FROM entities WHERE id = ?") | ||
| .bind(entity_updated.id.clone()) | ||
| .execute(&mut **tx) | ||
| .await?; | ||
| entity_updated.deleted = true; | ||
| } | ||
|
|
||
| let broker_message = BrokerMessage::EntityUpdated(entity_updated); | ||
| self.publish_queue.push_back(broker_message); | ||
| } | ||
| QueryType::RegisterModel => { | ||
| let row = query.fetch_one(&mut **tx).await.with_context(|| { | ||
| format!("Failed to execute query: {:?}, args: {:?}", statement, arguments) | ||
| })?; | ||
| let model_registered = ModelRegistered::from_row(&row)?; | ||
| self.publish_queue.push_back(BrokerMessage::ModelRegistered(model_registered)); | ||
| } | ||
| QueryType::EventMessage(entity) => { | ||
| let row = query.fetch_one(&mut **tx).await.with_context(|| { | ||
| format!("Failed to execute query: {:?}, args: {:?}", statement, arguments) | ||
| })?; | ||
| let mut event_message = EventMessageUpdated::from_row(&row)?; | ||
| event_message.updated_model = Some(entity); | ||
| let broker_message = BrokerMessage::EventMessageUpdated(event_message); | ||
| self.publish_queue.push_back(broker_message); | ||
| } | ||
| QueryType::StoreEvent => { | ||
| let row = query.fetch_one(&mut **tx).await.with_context(|| { | ||
| format!("Failed to execute query: {:?}, args: {:?}", statement, arguments) | ||
| })?; | ||
| let event = EventEmitted::from_row(&row)?; | ||
| self.publish_queue.push_back(BrokerMessage::EventEmitted(event)); | ||
| } | ||
| QueryType::Execute => { | ||
| self.execute().await?; | ||
| } | ||
| QueryType::Other => { | ||
| query.execute(&mut **tx).await.with_context(|| { | ||
| format!("Failed to execute query: {:?}, args: {:?}", statement, arguments) | ||
| })?; | ||
| } | ||
| } | ||
|
|
||
| Ok(()) | ||
| } | ||
|
|
||
| async fn execute(&mut self) -> Result<()> { | ||
| let transaction = mem::replace(&mut self.transaction, self.pool.begin().await?); | ||
| transaction.commit().await?; | ||
|
|
||
| while let Some(message) = self.publish_queue.pop_front() { | ||
| send_broker_message(message); | ||
| } | ||
|
There was a problem hiding this comment. 🛠️ Refactor suggestion Consider asynchronous message publishing for improved performance. Ohayo, sensei! To potentially improve performance, especially under heavy load, consider publishing broker messages asynchronously. This approach can prevent blocking the main execution flow while messages are being published. Here's a suggestion for modifying the loop: use futures::future::join_all;
// ...
let publish_futures: Vec<_> = self.publish_queue
.drain(..)
.map(|message| {
tokio::spawn(async move {
send_broker_message(message);
})
})
.collect();
join_all(publish_futures).await;This modification spawns a new task for each message, allowing them to be published concurrently. The
Larkooo marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
|
||
| Ok(()) | ||
| } | ||
| } | ||
|
|
||
| fn send_broker_message(message: BrokerMessage) { | ||
| match message { | ||
| BrokerMessage::ModelRegistered(model) => SimpleBroker::publish(model), | ||
| BrokerMessage::EntityUpdated(entity) => SimpleBroker::publish(entity), | ||
| BrokerMessage::EventMessageUpdated(event) => SimpleBroker::publish(event), | ||
| BrokerMessage::EventEmitted(event) => SimpleBroker::publish(event), | ||
| } | ||
| } | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Sensei, handle errors from joined tasks in the
JoinSetThe loop
while let Some(_) = set.join_next().await {}ignores the results of the tasks. By not handling theResultfromjoin_next(), any errors occurring in the spawned tasks may be missed. Consider capturing and handling these results to ensure that any task failures are properly addressed.