Skip to content

Commit

Permalink
refactor(pool): consolidate clear state method
Browse files Browse the repository at this point in the history
  • Loading branch information
alex-miao committed Jan 29, 2024
1 parent ac6ca3f commit e7c003e
Show file tree
Hide file tree
Showing 8 changed files with 64 additions and 123 deletions.
17 changes: 4 additions & 13 deletions crates/pool/proto/op_pool/op_pool.proto
Original file line number Diff line number Diff line change
Expand Up @@ -139,9 +139,6 @@ service OpPool {
// Clears the bundler mempool and reputation data of paymasters/accounts/factories/aggregators
rpc DebugClearState (DebugClearStateRequest) returns (DebugClearStateResponse);

// Clears the bundler mempool but not reputation data
rpc DebugClearMempool (DebugClearMempoolRequest) returns (DebugClearMempoolResponse);

// Dumps the current UserOperations mempool
rpc DebugDumpMempool (DebugDumpMempoolRequest) returns (DebugDumpMempoolResponse);

Expand Down Expand Up @@ -293,7 +290,10 @@ message UpdateEntitiesResponse {
}
message UpdateEntitiesSuccess {}

message DebugClearStateRequest {}
message DebugClearStateRequest {
bool clear_mempool = 1;
bool clear_reputation = 2;
}
message DebugClearStateResponse {
oneof result {
DebugClearStateSuccess success = 1;
Expand All @@ -302,15 +302,6 @@ message DebugClearStateResponse {
}
message DebugClearStateSuccess {}

message DebugClearMempoolRequest {}
message DebugClearMempoolResponse {
oneof result {
DebugClearMempoolSuccess success = 1;
MempoolError failure = 2;
}
}
message DebugClearMempoolSuccess {}

message DebugDumpMempoolRequest {
bytes entry_point = 1;
}
Expand Down
10 changes: 2 additions & 8 deletions crates/pool/src/mempool/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -91,14 +91,8 @@ pub trait Mempool: Send + Sync + 'static {

/// Debug methods

/// Clears both the mempool and reputation
fn clear(&self);

// Clears the reputation system
fn clear_reputation(&self);

/// Clears the mempool
fn clear_mempool(&self);
/// Clears the mempool of UOs or reputation of all addresses
fn clear_state(&self, clear_mempool: bool, clear_reputation: bool);

/// Dumps the mempool's reputation tracking
fn dump_reputation(&self) -> Vec<Reputation>;
Expand Down
20 changes: 8 additions & 12 deletions crates/pool/src/mempool/uo_pool.rs
Original file line number Diff line number Diff line change
Expand Up @@ -581,17 +581,13 @@ where
self.state.read().pool.get_operation_by_hash(hash)
}

fn clear(&self) {
self.clear_mempool();
self.clear_reputation();
}

fn clear_reputation(&self) {
self.reputation.clear()
}

fn clear_mempool(&self) {
self.state.write().pool.clear()
fn clear_state(&self, clear_mempool: bool, clear_reputation: bool) {
if clear_mempool {
self.state.write().pool.clear()
}
if clear_reputation {
self.reputation.clear()
}
}

fn dump_reputation(&self) -> Vec<Reputation> {
Expand Down Expand Up @@ -736,7 +732,7 @@ mod tests {
.unwrap();
}
check_ops(pool.best_operations(3, 0).unwrap(), uos);
pool.clear();
pool.clear_state(true, true);
assert_eq!(pool.best_operations(3, 0).unwrap(), vec![]);
}

Expand Down
48 changes: 17 additions & 31 deletions crates/pool/src/server/local.rs
Original file line number Diff line number Diff line change
Expand Up @@ -184,24 +184,22 @@ impl PoolServer for LocalPoolHandle {
}
}

async fn debug_clear_state(&self) -> Result<(), PoolServerError> {
let req = ServerRequestKind::DebugClearState;
async fn debug_clear_state(
&self,
clear_mempool: bool,
clear_reputation: bool,
) -> Result<(), PoolServerError> {
let req = ServerRequestKind::DebugClearState {
clear_mempool,
clear_reputation,
};
let resp = self.send(req).await?;
match resp {
ServerResponse::DebugClearState => Ok(()),
_ => Err(PoolServerError::UnexpectedResponse),
}
}

async fn debug_clear_mempool(&self) -> Result<(), PoolServerError> {
let req = ServerRequestKind::DebugClearMempool;
let resp = self.send(req).await?;
match resp {
ServerResponse::DebugClearMempool => Ok(()),
_ => Err(PoolServerError::UnexpectedResponse),
}
}

async fn debug_dump_mempool(&self, entry_point: Address) -> PoolResult<Vec<PoolOperation>> {
let req = ServerRequestKind::DebugDumpMempool { entry_point };
let resp = self.send(req).await?;
Expand Down Expand Up @@ -371,16 +369,9 @@ where
Ok(())
}

fn debug_clear_state(&self) -> PoolResult<()> {
for mempool in self.mempools.values() {
mempool.clear();
}
Ok(())
}

fn debug_clear_mempool(&self) -> PoolResult<()> {
fn debug_clear_state(&self, clear_mempool: bool, clear_reputation: bool) -> PoolResult<()> {
for mempool in self.mempools.values() {
mempool.clear_mempool();
mempool.clear_state(clear_mempool, clear_reputation);
}
Ok(())
}
Expand Down Expand Up @@ -493,18 +484,12 @@ where
Err(e) => Err(e),
}
},
ServerRequestKind::DebugClearState => {
match self.debug_clear_state() {
ServerRequestKind::DebugClearState { clear_mempool, clear_reputation } => {
match self.debug_clear_state(clear_mempool, clear_reputation) {
Ok(_) => Ok(ServerResponse::DebugClearState),
Err(e) => Err(e),
}
},
ServerRequestKind::DebugClearMempool => {
match self.debug_clear_mempool() {
Ok(_) => Ok(ServerResponse::DebugClearMempool),
Err(e) => Err(e),
}
},
ServerRequestKind::DebugDumpMempool { entry_point } => {
match self.debug_dump_mempool(entry_point) {
Ok(ops) => Ok(ServerResponse::DebugDumpMempool { ops }),
Expand Down Expand Up @@ -592,8 +577,10 @@ enum ServerRequestKind {
entry_point: Address,
entity_updates: Vec<EntityUpdate>,
},
DebugClearState,
DebugClearMempool,
DebugClearState {
clear_mempool: bool,
clear_reputation: bool,
},
DebugDumpMempool {
entry_point: Address,
},
Expand Down Expand Up @@ -632,7 +619,6 @@ enum ServerResponse {
RemoveOps,
UpdateEntities,
DebugClearState,
DebugClearMempool,
DebugDumpMempool {
ops: Vec<PoolOperation>,
},
Expand Down
9 changes: 5 additions & 4 deletions crates/pool/src/server/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -91,10 +91,11 @@ pub trait PoolServer: Send + Sync + 'static {
async fn subscribe_new_heads(&self) -> PoolResult<Pin<Box<dyn Stream<Item = NewHead> + Send>>>;

/// Clear the pool state, used for debug methods
async fn debug_clear_state(&self) -> PoolResult<()>;

/// Clear the mempool state, used for debug methods
async fn debug_clear_mempool(&self) -> PoolResult<()>;
async fn debug_clear_state(
&self,
clear_mempool: bool,
clear_reputation: bool,
) -> PoolResult<()>;

/// Dump all operations in the pool, used for debug methods
async fn debug_dump_mempool(&self, entry_point: Address) -> PoolResult<Vec<PoolOperation>>;
Expand Down
44 changes: 16 additions & 28 deletions crates/pool/src/server/remote/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,14 +33,13 @@ use tonic_health::{
};

use super::protos::{
self, add_op_response, debug_clear_mempool_response, debug_clear_state_response,
debug_dump_mempool_response, debug_dump_reputation_response, debug_set_reputation_response,
get_op_by_hash_response, get_ops_response, get_reputation_status_response,
get_stake_status_response, op_pool_client::OpPoolClient, remove_ops_response,
update_entities_response, AddOpRequest, DebugClearMempoolRequest, DebugClearStateRequest,
DebugDumpMempoolRequest, DebugDumpReputationRequest, DebugSetReputationRequest, GetOpsRequest,
GetReputationStatusRequest, GetStakeStatusRequest, RemoveOpsRequest, SubscribeNewHeadsRequest,
SubscribeNewHeadsResponse, UpdateEntitiesRequest,
self, add_op_response, debug_clear_state_response, debug_dump_mempool_response,
debug_dump_reputation_response, debug_set_reputation_response, get_op_by_hash_response,
get_ops_response, get_reputation_status_response, get_stake_status_response,
op_pool_client::OpPoolClient, remove_ops_response, update_entities_response, AddOpRequest,
DebugClearStateRequest, DebugDumpMempoolRequest, DebugDumpReputationRequest,
DebugSetReputationRequest, GetOpsRequest, GetReputationStatusRequest, GetStakeStatusRequest,
RemoveOpsRequest, SubscribeNewHeadsRequest, SubscribeNewHeadsResponse, UpdateEntitiesRequest,
};
use crate::{
mempool::{PoolOperation, Reputation, StakeStatus},
Expand Down Expand Up @@ -265,11 +264,18 @@ impl PoolServer for RemotePoolClient {
}
}

async fn debug_clear_state(&self) -> PoolResult<()> {
async fn debug_clear_state(
&self,
clear_mempool: bool,
clear_reputation: bool,
) -> PoolResult<()> {
let res = self
.op_pool_client
.clone()
.debug_clear_state(DebugClearStateRequest {})
.debug_clear_state(DebugClearStateRequest {
clear_mempool,
clear_reputation,
})
.await?
.into_inner()
.result;
Expand All @@ -283,24 +289,6 @@ impl PoolServer for RemotePoolClient {
}
}

async fn debug_clear_mempool(&self) -> PoolResult<()> {
let res = self
.op_pool_client
.clone()
.debug_clear_mempool(DebugClearMempoolRequest {})
.await?
.into_inner()
.result;

match res {
Some(debug_clear_mempool_response::Result::Success(_)) => Ok(()),
Some(debug_clear_mempool_response::Result::Failure(f)) => Err(f.try_into()?),
None => Err(PoolServerError::Other(anyhow::anyhow!(
"should have received result from op pool"
)))?,
}
}

async fn debug_dump_mempool(&self, entry_point: Address) -> PoolResult<Vec<PoolOperation>> {
let res = self
.op_pool_client
Expand Down
35 changes: 10 additions & 25 deletions crates/pool/src/server/remote/server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,13 +30,11 @@ use tokio_util::sync::CancellationToken;
use tonic::{transport::Server, Request, Response, Result, Status};

use super::protos::{
add_op_response, debug_clear_mempool_response, debug_clear_state_response,
debug_dump_mempool_response, debug_dump_reputation_response, debug_set_reputation_response,
get_op_by_hash_response, get_ops_response, get_reputation_status_response,
get_stake_status_response,
add_op_response, debug_clear_state_response, debug_dump_mempool_response,
debug_dump_reputation_response, debug_set_reputation_response, get_op_by_hash_response,
get_ops_response, get_reputation_status_response, get_stake_status_response,
op_pool_server::{OpPool, OpPoolServer},
remove_ops_response, update_entities_response, AddOpRequest, AddOpResponse, AddOpSuccess,
DebugClearMempoolRequest, DebugClearMempoolResponse, DebugClearMempoolSuccess,
DebugClearStateRequest, DebugClearStateResponse, DebugClearStateSuccess,
DebugDumpMempoolRequest, DebugDumpMempoolResponse, DebugDumpMempoolSuccess,
DebugDumpReputationRequest, DebugDumpReputationResponse, DebugDumpReputationSuccess,
Expand Down Expand Up @@ -271,9 +269,14 @@ impl OpPool for OpPoolImpl {

async fn debug_clear_state(
&self,
_request: Request<DebugClearStateRequest>,
request: Request<DebugClearStateRequest>,
) -> Result<Response<DebugClearStateResponse>> {
let resp = match self.local_pool.debug_clear_state().await {
let req = request.into_inner();
let resp = match self
.local_pool
.debug_clear_state(req.clear_mempool, req.clear_reputation)
.await
{
Ok(_) => DebugClearStateResponse {
result: Some(debug_clear_state_response::Result::Success(
DebugClearStateSuccess {},
Expand All @@ -287,24 +290,6 @@ impl OpPool for OpPoolImpl {
Ok(Response::new(resp))
}

async fn debug_clear_mempool(
&self,
_request: Request<DebugClearMempoolRequest>,
) -> Result<Response<DebugClearMempoolResponse>> {
let resp = match self.local_pool.debug_clear_mempool().await {
Ok(_) => DebugClearMempoolResponse {
result: Some(debug_clear_mempool_response::Result::Success(
DebugClearMempoolSuccess {},
)),
},
Err(error) => DebugClearMempoolResponse {
result: Some(debug_clear_mempool_response::Result::Failure(error.into())),
},
};

Ok(Response::new(resp))
}

async fn debug_dump_mempool(
&self,
request: Request<DebugDumpMempoolRequest>,
Expand Down
4 changes: 2 additions & 2 deletions crates/rpc/src/debug.rs
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,7 @@ where
async fn bundler_clear_state(&self) -> RpcResult<String> {
let _ = self
.pool
.debug_clear_state()
.debug_clear_state(true, true)
.await
.map_err(|e| rpc_err(INTERNAL_ERROR_CODE, e.to_string()))?;

Expand All @@ -104,7 +104,7 @@ where
async fn bundler_clear_mempool(&self) -> RpcResult<String> {
let _ = self
.pool
.debug_clear_mempool()
.debug_clear_state(true, false)
.await
.map_err(|e| rpc_err(INTERNAL_ERROR_CODE, e.to_string()))?;

Expand Down

0 comments on commit e7c003e

Please sign in to comment.