refactor(pool): consolidate clear state method

crates/pool/proto/op_pool/op_pool.proto

@@ -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);
 
@@ -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;
@@ -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;
 }

crates/pool/src/mempool/mod.rs

@@ -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>;

crates/pool/src/mempool/uo_pool.rs

@@ -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> {
@@ -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![]);
     }
 

crates/pool/src/server/local.rs

@@ -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?;
@@ -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(())
     }
@@ -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 }),
@@ -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,
     },
@@ -632,7 +619,6 @@ enum ServerResponse {
     RemoveOps,
     UpdateEntities,
     DebugClearState,
-    DebugClearMempool,
     DebugDumpMempool {
         ops: Vec<PoolOperation>,
     },

crates/pool/src/server/mod.rs

@@ -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>>;

crates/pool/src/server/remote/client.rs

@@ -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},
@@ -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;
@@ -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

crates/pool/src/server/remote/server.rs

@@ -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,
@@ -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 {},
@@ -287,23 +290,23 @@ 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_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,

crates/rpc/src/debug.rs

@@ -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()))?;
 
@@ -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()))?;