feat: add functionality to use remote mempool client/server

starkware-libs · Jul 28, 2024 · d2d9fa4 · d2d9fa4
1 parent 9c48900
commit d2d9fa4
Show file tree

Hide file tree

Showing 6 changed files with 60 additions and 9 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/mempool/src/communication.rs b/crates/mempool/src/communication.rs
@@ -1,7 +1,9 @@
+use std::net::IpAddr;
+
 use async_trait::async_trait;
 use starknet_mempool_infra::component_definitions::ComponentRequestHandler;
 use starknet_mempool_infra::component_runner::ComponentStarter;
-use starknet_mempool_infra::component_server::LocalComponentServer;
+use starknet_mempool_infra::component_server::{LocalComponentServer, RemoteComponentServer};
 use starknet_mempool_types::communication::{
     MempoolRequest,
     MempoolRequestAndResponseSender,
@@ -15,6 +17,9 @@ use crate::mempool::Mempool;
 pub type MempoolServer =
     LocalComponentServer<MempoolCommunicationWrapper, MempoolRequest, MempoolResponse>;
 
+pub type RemoteMempoolServer =
+    RemoteComponentServer<MempoolCommunicationWrapper, MempoolRequest, MempoolResponse>;
+
 pub fn create_mempool_server(
     mempool: Mempool,
     rx_mempool: Receiver<MempoolRequestAndResponseSender>,
@@ -23,6 +28,15 @@ pub fn create_mempool_server(
     LocalComponentServer::new(communication_wrapper, rx_mempool)
 }
 
+pub fn create_remote_mempool_server(
+    mempool: Mempool,
+    ip_address: IpAddr,
+    port: u16,
+) -> RemoteMempoolServer {
+    let communication_wrapper = MempoolCommunicationWrapper::new(mempool);
+    RemoteComponentServer::new(communication_wrapper, ip_address, port)
+}
+
 /// Wraps the mempool to enable inbound async communication from other components.
 pub struct MempoolCommunicationWrapper {
     mempool: Mempool,

diff --git a/crates/mempool_types/Cargo.toml b/crates/mempool_types/Cargo.toml
@@ -12,5 +12,6 @@ workspace = true
 async-trait.workspace = true
 starknet_api = { path = "../starknet_api", version = "0.13.0-rc.0"}
 mockall.workspace = true
+serde = { workspace = true, feat = ["derive"] }
 starknet_mempool_infra = { path = "../mempool_infra" }
 thiserror.workspace = true
diff --git a/crates/mempool_types/src/communication.rs b/crates/mempool_types/src/communication.rs
@@ -3,14 +3,20 @@ use std::sync::Arc;
 use async_trait::async_trait;
 use mockall::predicate::*;
 use mockall::*;
-use starknet_mempool_infra::component_client::{ClientError, LocalComponentClient};
+use serde::{Deserialize, Serialize};
+use starknet_mempool_infra::component_client::{
+    ClientError,
+    LocalComponentClient,
+    RemoteComponentClient,
+};
 use starknet_mempool_infra::component_definitions::ComponentRequestAndResponseSender;
 use thiserror::Error;
 
 use crate::errors::MempoolError;
 use crate::mempool_types::{MempoolInput, ThinTransaction};
 
 pub type MempoolClientImpl = LocalComponentClient<MempoolRequest, MempoolResponse>;
+pub type RemoteMempoolClientImpl = RemoteComponentClient<MempoolRequest, MempoolResponse>;
 pub type MempoolResult<T> = Result<T, MempoolError>;
 pub type MempoolClientResult<T> = Result<T, MempoolClientError>;
 pub type MempoolRequestAndResponseSender =
@@ -26,13 +32,13 @@ pub trait MempoolClient: Send + Sync {
     async fn get_txs(&self, n_txs: usize) -> MempoolClientResult<Vec<ThinTransaction>>;
 }
 
-#[derive(Debug)]
+#[derive(Debug, Serialize, Deserialize)]
 pub enum MempoolRequest {
     AddTransaction(MempoolInput),
     GetTransactions(usize),
 }
 
-#[derive(Debug)]
+#[derive(Debug, Serialize, Deserialize)]
 pub enum MempoolResponse {
     AddTransaction(MempoolResult<()>),
     GetTransactions(MempoolResult<Vec<ThinTransaction>>),
@@ -72,3 +78,30 @@ impl MempoolClient for MempoolClientImpl {
         }
     }
 }
+
+#[async_trait]
+impl MempoolClient for RemoteMempoolClientImpl {
+    async fn add_tx(&self, mempool_input: MempoolInput) -> MempoolClientResult<()> {
+        let request = MempoolRequest::AddTransaction(mempool_input);
+        let response = self.send(request).await?;
+        match response {
+            MempoolResponse::AddTransaction(Ok(response)) => Ok(response),
+            MempoolResponse::AddTransaction(Err(response)) => {
+                Err(MempoolClientError::MempoolError(response))
+            }
+            _ => Err(MempoolClientError::ClientError(ClientError::UnexpectedResponse)),
+        }
+    }
+
+    async fn get_txs(&self, n_txs: usize) -> MempoolClientResult<Vec<ThinTransaction>> {
+        let request = MempoolRequest::GetTransactions(n_txs);
+        let response = self.send(request).await?;
+        match response {
+            MempoolResponse::GetTransactions(Ok(response)) => Ok(response),
+            MempoolResponse::GetTransactions(Err(response)) => {
+                Err(MempoolClientError::MempoolError(response))
+            }
+            _ => Err(MempoolClientError::ClientError(ClientError::UnexpectedResponse)),
+        }
+    }
+}
diff --git a/crates/mempool_types/src/errors.rs b/crates/mempool_types/src/errors.rs
@@ -1,8 +1,9 @@
+use serde::{Deserialize, Serialize};
 use starknet_api::core::{ContractAddress, Nonce};
 use starknet_api::transaction::TransactionHash;
 use thiserror::Error;
 
-#[derive(Clone, Debug, Error, PartialEq, Eq)]
+#[derive(Clone, Debug, Error, PartialEq, Eq, Serialize, Deserialize)]
 pub enum MempoolError {
     #[error("Duplicate transaction, sender address: {address}, nonce: {:?}", nonce)]
     DuplicateNonce { address: ContractAddress, nonce: Nonce },

diff --git a/crates/mempool_types/src/mempool_types.rs b/crates/mempool_types/src/mempool_types.rs
@@ -1,30 +1,31 @@
+use serde::{Deserialize, Serialize};
 use starknet_api::core::{ContractAddress, Nonce};
 use starknet_api::transaction::{Tip, TransactionHash};
 
 use crate::errors::MempoolError;
 
-#[derive(Clone, Debug, Default, Eq, PartialEq)]
+#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
 pub struct ThinTransaction {
     pub sender_address: ContractAddress,
     pub tx_hash: TransactionHash,
     pub tip: Tip,
     pub nonce: Nonce,
 }
 
-#[derive(Clone, Copy, Debug, Default, PartialEq)]
+#[derive(Clone, Copy, Debug, Default, PartialEq, Serialize, Deserialize)]
 pub struct AccountState {
     pub nonce: Nonce,
     // TODO: add balance field when needed.
 }
 
-#[derive(Clone, Copy, Debug, Default, PartialEq)]
+#[derive(Clone, Copy, Debug, Default, PartialEq, Serialize, Deserialize)]
 pub struct Account {
     // TODO(Ayelet): Consider removing this field as it is duplicated in ThinTransaction.
     pub sender_address: ContractAddress,
     pub state: AccountState,
 }
 
-#[derive(Clone, Debug, Default, PartialEq)]
+#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
 pub struct MempoolInput {
     pub tx: ThinTransaction,
     pub account: Account,