feat(pool): allow staked senders to have multiple UOs in best_operations

alchemyplatform · Jun 14, 2024 · 1280ade · 1280ade
1 parent b39116b
commit 1280ade
Showing 1 changed file with 24 additions and 7 deletions.
diff --git a/crates/pool/src/mempool/uo_pool.rs b/crates/pool/src/mempool/uo_pool.rs
@@ -591,8 +591,12 @@ where
                         .div_mod(self.config.num_shards.into())
                         .1
                         == shard_index.into())) &&
-                // filter out ops from senders we've already seen
-                senders.insert(op.uo.sender())
+                // filter out ops from unstaked senders we've already seen
+                if !op.account_is_staked {
+                    senders.insert(op.uo.sender())
+                } else {
+                    true
+                }
             })
             .take(max)
             .map(Into::into)
@@ -944,13 +948,13 @@ mod tests {
     async fn test_account_reputation() {
         let address = Address::random();
         let (pool, uos) = create_pool_insert_ops(vec![
-            create_op_with_errors(address, 0, 2, None, None, true),
-            create_op_with_errors(address, 1, 2, None, None, true),
-            create_op_with_errors(address, 1, 2, None, None, true),
+            create_op_with_errors(address, 0, 2, None, None, true), // accept
+            create_op_with_errors(address, 1, 2, None, None, true), // accept
+            create_op_with_errors(address, 1, 2, None, None, true), // reject
         ])
         .await;
-        // Only return 1 op per sender
-        check_ops(pool.best_operations(3, 0).unwrap(), vec![uos[0].clone()]);
+        // staked, so include all ops
+        check_ops(pool.best_operations(3, 0).unwrap(), uos[0..2].to_vec());
 
         let rep = pool.dump_reputation();
         assert_eq!(rep.len(), 1);
@@ -1383,6 +1387,19 @@ mod tests {
             .is_err());
     }
 
+    #[tokio::test]
+    async fn test_best_staked() {
+        let address = Address::random();
+        let (pool, uos) = create_pool_insert_ops(vec![
+            create_op_with_errors(address, 0, 2, None, None, true),
+            create_op_with_errors(address, 1, 2, None, None, true),
+            create_op_with_errors(address, 2, 2, None, None, true),
+        ])
+        .await;
+        // staked, so include all ops
+        check_ops(pool.best_operations(3, 0).unwrap(), uos);
+    }
+
     #[derive(Clone, Debug)]
     struct OpWithErrors {
         op: UserOperationVariant,