From 1bd5819dad47ce478d7b4fcd6a7e24485edb159e Mon Sep 17 00:00:00 2001
From: Etan Kissling <etan@status.im>
Date: Fri, 15 Mar 2024 22:48:18 +0100
Subject: [PATCH] cache head block eligibility for fork choice (#6076)

When there are long periods of non-finality, `nodeIsViableForHead` has
been observed to consume significant time as it repeatedly walks the
non-finalized check graph as part of determining what heads are eligible
for fork choice. Caching the result resolves that.

Overall, it may still be better to prune fork choice more aggressively
when finality advances, to fully avoid the case specced out using the
linear scan. The current implementation is very close to spec, though,
so such a change should not be introduced without thorough testing.

The simple cache should allow significantly better performance on Goerli
while the network is still supported (Mid April).
---
 .../fork_choice/fork_choice_types.nim         |  1 +
 beacon_chain/fork_choice/proto_array.nim      | 43 +++++++++++++++----
 2 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/beacon_chain/fork_choice/fork_choice_types.nim b/beacon_chain/fork_choice/fork_choice_types.nim
index 26b1be20d..ab776cb40 100644
--- a/beacon_chain/fork_choice/fork_choice_types.nim
+++ b/beacon_chain/fork_choice/fork_choice_types.nim
@@ -109,6 +109,7 @@ type
     bid*: BlockId
     parent*: Option[Index]
     checkpoints*: FinalityCheckpoints
+    sharedFinalizedEpoch*: Epoch
     weight*: int64
     invalid*: bool
     bestChild*: Option[Index]
diff --git a/beacon_chain/fork_choice/proto_array.nim b/beacon_chain/fork_choice/proto_array.nim
index ed3179efc..381279055 100644
--- a/beacon_chain/fork_choice/proto_array.nim
+++ b/beacon_chain/fork_choice/proto_array.nim
@@ -82,9 +82,9 @@ func maybeUpdateBestChildAndDescendant(self: var ProtoArray,
                                        childIdx: Index): FcResult[void]
 
 func nodeIsViableForHead(
-    self: ProtoArray, node: ProtoNode, nodeIdx: Index): bool
+    self: var ProtoArray, node: ProtoNode, nodeIdx: Index): bool
 func nodeLeadsToViableHead(
-    self: ProtoArray, node: ProtoNode, nodeIdx: Index): FcResult[bool]
+    self: var ProtoArray, node: ProtoNode, nodeIdx: Index): FcResult[bool]
 
 # ProtoArray routines
 # ----------------------------------------------------------------------
@@ -158,6 +158,11 @@ func applyScoreChanges*(self: var ProtoArray,
       indicesLen: self.indices.len)
 
   self.currentEpoch = currentEpoch
+
+  doAssert checkpoints.finalized.epoch >= self.checkpoints.finalized.epoch
+  doAssert checkpoints.finalized.epoch > self.checkpoints.finalized.epoch or
+    checkpoints.finalized.root == self.checkpoints.finalized.root or
+    self.checkpoints.finalized.epoch == GENESIS_EPOCH
   self.checkpoints = checkpoints
 
   ## Alias
@@ -499,7 +504,7 @@ func maybeUpdateBestChildAndDescendant(self: var ProtoArray,
   ok()
 
 func nodeLeadsToViableHead(
-    self: ProtoArray, node: ProtoNode, nodeIdx: Index): FcResult[bool] =
+    self: var ProtoArray, node: ProtoNode, nodeIdx: Index): FcResult[bool] =
   ## Indicates if the node itself or its best-descendant are viable
   ## for blockchain head
   let bestDescendantIsViableForHead = block:
@@ -517,7 +522,7 @@ func nodeLeadsToViableHead(
   ok(bestDescendantIsViableForHead or self.nodeIsViableForHead(node, nodeIdx))
 
 func nodeIsViableForHead(
-    self: ProtoArray, node: ProtoNode, nodeIdx: Index): bool =
+    self: var ProtoArray, node: ProtoNode, nodeIdx: Index): bool =
   ## This is the equivalent of `filter_block_tree` function in consensus specs
   ## https://github.com/ethereum/consensus-specs/blob/v1.3.0/specs/phase0/fork-choice.md#filter_block_tree
 
@@ -554,18 +559,38 @@ func nodeIsViableForHead(
       false
     elif self.checkpoints.finalized.epoch == GENESIS_EPOCH:
       true
+    elif node.sharedFinalizedEpoch == self.checkpoints.finalized.epoch:
+      # Already checked to share history with `self.checkpoints.finalized`.
+      # `self.checkpoints.finalized` cannot reorg, see `applyScoreChanges`
+      true
     else:
-      let finalizedSlot = self.checkpoints.finalized.epoch.start_slot
+      # Check that this node is not going to be pruned
+      let 
+        finalizedEpoch = self.checkpoints.finalized.epoch
+        finalizedSlot = finalizedEpoch.start_slot
       var ancestor = some node
-      while ancestor.isSome and ancestor.unsafeGet.bid.slot > finalizedSlot:
+      while ancestor.isSome and ancestor.unsafeGet.bid.slot > finalizedSlot and
+          ancestor.unsafeGet.sharedFinalizedEpoch != finalizedEpoch:
         if ancestor.unsafeGet.parent.isSome:
           ancestor = self.nodes[ancestor.unsafeGet.parent.unsafeGet]
         else:
           ancestor.reset()
-      if ancestor.isSome:
-        ancestor.unsafeGet.bid.root == self.checkpoints.finalized.root
-      else:
+      if ancestor.isNone:
         false
+      elif ancestor.unsafeGet.sharedFinalizedEpoch != finalizedEpoch and
+          ancestor.unsafeGet.bid.root != self.checkpoints.finalized.root:
+        false
+      else:
+        # An ancestor was already checked to share canonical history,
+        # or we reached the `finalizedSlot` and the root matches expectations
+        let ancestorSlot = ancestor.unsafeGet.bid.slot
+        var idx = nodeIdx
+        template mutableNode: untyped = self.nodes.buf[idx - self.nodes.offset]
+        while mutableNode.bid.slot > ancestorSlot:
+          mutableNode.sharedFinalizedEpoch = finalizedEpoch
+          idx = mutableNode.parent.unsafeGet
+        mutableNode.sharedFinalizedEpoch = finalizedEpoch
+        true
 
 func propagateInvalidity*(
     self: var ProtoArray, startPhysicalIdx: Index) =