From 7f73dccebc0b58057920500c84833a43e839a591 Mon Sep 17 00:00:00 2001
From: Michael Sproul <michael@sigmaprime.io>
Date: Thu, 22 Oct 2020 04:47:29 +0000
Subject: [PATCH] Refine op pool pruning (#1805)

## Issue Addressed

Closes #1769
Closes #1708

## Proposed Changes

Tweaks the op pool pruning so that the attestation pool is pruned against the wall-clock epoch instead of the finalized state's epoch. This should reduce the unbounded growth that we've seen during periods without finality.

Also fixes up the voluntary exit pruning as raised in #1708.
---
 beacon_node/beacon_chain/src/beacon_chain.rs |  9 +++--
 beacon_node/operation_pool/src/lib.rs        | 37 +++++++++++---------
 2 files changed, 27 insertions(+), 19 deletions(-)

diff --git a/beacon_node/beacon_chain/src/beacon_chain.rs b/beacon_node/beacon_chain/src/beacon_chain.rs
index 37bb801aba..7f70c8f97b 100644
--- a/beacon_node/beacon_chain/src/beacon_chain.rs
+++ b/beacon_node/beacon_chain/src/beacon_chain.rs
@@ -1959,6 +1959,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             || is_reorg
         {
             self.persist_head_and_fork_choice()?;
+            self.op_pool.prune_attestations(self.epoch()?);
         }
 
         let update_head_timer = metrics::start_timer(&metrics::UPDATE_HEAD_TIMES);
@@ -2097,8 +2098,12 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             .get_state(&new_finalized_state_root, None)?
             .ok_or_else(|| Error::MissingBeaconState(new_finalized_state_root))?;
 
-        self.op_pool
-            .prune_all(&finalized_state, self.head_info()?.fork);
+        self.op_pool.prune_all(
+            &finalized_state,
+            self.epoch()?,
+            self.head_info()?.fork,
+            &self.spec,
+        );
 
         self.store_migrator.process_finalization(
             new_finalized_state_root.into(),
diff --git a/beacon_node/operation_pool/src/lib.rs b/beacon_node/operation_pool/src/lib.rs
index df1d6f834d..a22c7b085f 100644
--- a/beacon_node/operation_pool/src/lib.rs
+++ b/beacon_node/operation_pool/src/lib.rs
@@ -21,8 +21,8 @@ use std::marker::PhantomData;
 use std::ptr;
 use types::{
     typenum::Unsigned, Attestation, AttesterSlashing, BeaconState, BeaconStateError, ChainSpec,
-    EthSpec, Fork, ForkVersion, Hash256, ProposerSlashing, RelativeEpoch, SignedVoluntaryExit,
-    Validator,
+    Epoch, EthSpec, Fork, ForkVersion, Hash256, ProposerSlashing, RelativeEpoch,
+    SignedVoluntaryExit, Validator,
 };
 #[derive(Default, Debug)]
 pub struct OperationPool<T: EthSpec + Default> {
@@ -156,17 +156,14 @@ impl<T: EthSpec> OperationPool<T> {
     }
 
     /// Remove attestations which are too old to be included in a block.
-    pub fn prune_attestations(&self, finalized_state: &BeaconState<T>) {
-        // We know we can include an attestation if:
-        // state.slot <= attestation_slot + SLOTS_PER_EPOCH
-        // We approximate this check using the attestation's epoch, to avoid computing
-        // the slot or relying on the committee cache of the finalized state.
+    pub fn prune_attestations(&self, current_epoch: Epoch) {
+        // Prune attestations that are from before the previous epoch.
         self.attestations.write().retain(|_, attestations| {
             // All the attestations in this bucket have the same data, so we only need to
             // check the first one.
-            attestations.first().map_or(false, |att| {
-                finalized_state.current_epoch() <= att.data.target.epoch + 1
-            })
+            attestations
+                .first()
+                .map_or(false, |att| current_epoch <= att.data.target.epoch + 1)
         });
     }
 
@@ -299,20 +296,26 @@ impl<T: EthSpec> OperationPool<T> {
     }
 
     /// Prune if validator has already exited at the last finalized state.
-    pub fn prune_voluntary_exits(&self, finalized_state: &BeaconState<T>) {
+    pub fn prune_voluntary_exits(&self, finalized_state: &BeaconState<T>, spec: &ChainSpec) {
         prune_validator_hash_map(
             &mut self.voluntary_exits.write(),
-            |validator| validator.is_exited_at(finalized_state.current_epoch()),
+            |validator| validator.exit_epoch != spec.far_future_epoch,
             finalized_state,
         );
     }
 
     /// Prune all types of transactions given the latest finalized state and head fork.
-    pub fn prune_all(&self, finalized_state: &BeaconState<T>, head_fork: Fork) {
-        self.prune_attestations(finalized_state);
+    pub fn prune_all(
+        &self,
+        finalized_state: &BeaconState<T>,
+        current_epoch: Epoch,
+        head_fork: Fork,
+        spec: &ChainSpec,
+    ) {
+        self.prune_attestations(current_epoch);
         self.prune_proposer_slashings(finalized_state);
         self.prune_attester_slashings(finalized_state, head_fork);
-        self.prune_voluntary_exits(finalized_state);
+        self.prune_voluntary_exits(finalized_state, spec);
     }
 
     /// Total number of voluntary exits in the pool.
@@ -612,13 +615,13 @@ mod release_tests {
         );
 
         // Prune attestations shouldn't do anything at this point.
-        op_pool.prune_attestations(state);
+        op_pool.prune_attestations(state.current_epoch());
         assert_eq!(op_pool.num_attestations(), committees.len());
 
         // But once we advance to more than an epoch after the attestation, it should prune it
         // out of existence.
         state.slot += 2 * MainnetEthSpec::slots_per_epoch();
-        op_pool.prune_attestations(state);
+        op_pool.prune_attestations(state.current_epoch());
         assert_eq!(op_pool.num_attestations(), 0);
     }