From 7f8a5845431b80be62eb367aea26cedbdbd130cc Mon Sep 17 00:00:00 2001
From: Giulio rebuffo <giulio.rebuffo@gmail.com>
Date: Fri, 22 Nov 2024 08:37:58 +0100
Subject: [PATCH] Caplin: keep reusing the same beacon state in
 `ForkchoiceStore` across reorgs (#12833)

---
 .../forkchoice/fork_graph/fork_graph_disk.go  | 21 ++++++++++++-------
 .../fork_graph/fork_graph_disk_fs.go          |  8 +++++--
 2 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/cl/phase1/forkchoice/fork_graph/fork_graph_disk.go b/cl/phase1/forkchoice/fork_graph/fork_graph_disk.go
index 91bb6ed7671..c8847f7097e 100644
--- a/cl/phase1/forkchoice/fork_graph/fork_graph_disk.go
+++ b/cl/phase1/forkchoice/fork_graph/fork_graph_disk.go
@@ -187,7 +187,7 @@ func (f *forkGraphDisk) AddChainSegment(signedBlock *cltypes.SignedBeaconBlock,
 		return nil, BelowAnchor, nil
 	}
 
-	newState, err := f.GetState(block.ParentRoot, false)
+	newState, err := f.getState(block.ParentRoot, false, true)
 	if err != nil {
 		return nil, LogisticError, fmt.Errorf("AddChainSegment: %w, parentRoot; %x", err, block.ParentRoot)
 	}
@@ -249,7 +249,7 @@ func (f *forkGraphDisk) AddChainSegment(signedBlock *cltypes.SignedBeaconBlock,
 		// Add block to list of invalid blocks
 		log.Warn("Invalid beacon block", "slot", block.Slot, "blockRoot", blockRoot, "reason", invalidBlockErr)
 		f.badBlocks.Store(libcommon.Hash(blockRoot), struct{}{})
-		f.currentState = nil
+		//f.currentState = nil
 
 		return nil, InvalidBlock, invalidBlockErr
 	}
@@ -296,8 +296,8 @@ func (f *forkGraphDisk) AddChainSegment(signedBlock *cltypes.SignedBeaconBlock,
 	f.finalizedCheckpoints.Store(libcommon.Hash(blockRoot), newState.FinalizedCheckpoint())
 	if newState.Slot() > f.highestSeen {
 		f.highestSeen = newState.Slot()
-		f.currentState = newState
 	}
+	f.currentState = newState
 	return newState, Success, nil
 }
 
@@ -317,8 +317,11 @@ func (f *forkGraphDisk) getBlock(blockRoot libcommon.Hash) (*cltypes.SignedBeaco
 
 	return obj.(*cltypes.SignedBeaconBlock), true
 }
-
 func (f *forkGraphDisk) GetState(blockRoot libcommon.Hash, alwaysCopy bool) (*state.CachingBeaconState, error) {
+	return f.getState(blockRoot, alwaysCopy, false)
+}
+
+func (f *forkGraphDisk) getState(blockRoot libcommon.Hash, alwaysCopy bool, addChainSegment bool) (*state.CachingBeaconState, error) {
 	if f.currentState != nil && !alwaysCopy {
 		currentStateBlockRoot, err := f.currentState.BlockRoot()
 		if err != nil {
@@ -333,8 +336,12 @@ func (f *forkGraphDisk) GetState(blockRoot libcommon.Hash, alwaysCopy bool) (*st
 	blocksInTheWay := []*cltypes.SignedBeaconBlock{}
 	// Use the parent root as a reverse iterator.
 	currentIteratorRoot := blockRoot
-	var copyReferencedState *state.CachingBeaconState
+	var copyReferencedState, outState *state.CachingBeaconState
 	var err error
+	if addChainSegment {
+		outState = f.currentState
+	}
+
 	// try and find the point of recconection
 	for copyReferencedState == nil {
 		block, isSegmentPresent := f.getBlock(currentIteratorRoot)
@@ -342,7 +349,7 @@ func (f *forkGraphDisk) GetState(blockRoot libcommon.Hash, alwaysCopy bool) (*st
 			// check if it is in the header
 			bHeader, ok := f.GetHeader(currentIteratorRoot)
 			if ok && bHeader.Slot%dumpSlotFrequency == 0 {
-				copyReferencedState, err = f.readBeaconStateFromDisk(currentIteratorRoot)
+				copyReferencedState, err = f.readBeaconStateFromDisk(currentIteratorRoot, outState)
 				if err != nil {
 					log.Trace("Could not retrieve state: Missing header", "missing", currentIteratorRoot, "err", err)
 					copyReferencedState = nil
@@ -353,7 +360,7 @@ func (f *forkGraphDisk) GetState(blockRoot libcommon.Hash, alwaysCopy bool) (*st
 			return nil, nil
 		}
 		if block.Block.Slot%dumpSlotFrequency == 0 {
-			copyReferencedState, err = f.readBeaconStateFromDisk(currentIteratorRoot)
+			copyReferencedState, err = f.readBeaconStateFromDisk(currentIteratorRoot, outState)
 			if err != nil {
 				log.Trace("Could not retrieve state: Missing header", "missing", currentIteratorRoot, "err", err)
 			}
diff --git a/cl/phase1/forkchoice/fork_graph/fork_graph_disk_fs.go b/cl/phase1/forkchoice/fork_graph/fork_graph_disk_fs.go
index 811e6fd48e2..51b7ca140ff 100644
--- a/cl/phase1/forkchoice/fork_graph/fork_graph_disk_fs.go
+++ b/cl/phase1/forkchoice/fork_graph/fork_graph_disk_fs.go
@@ -39,7 +39,7 @@ func getBeaconStateCacheFilename(blockRoot libcommon.Hash) string {
 	return fmt.Sprintf("%x.cache", blockRoot)
 }
 
-func (f *forkGraphDisk) readBeaconStateFromDisk(blockRoot libcommon.Hash) (bs *state.CachingBeaconState, err error) {
+func (f *forkGraphDisk) readBeaconStateFromDisk(blockRoot libcommon.Hash, out *state.CachingBeaconState) (bs *state.CachingBeaconState, err error) {
 	var file afero.File
 	f.stateDumpLock.Lock()
 	defer f.stateDumpLock.Unlock()
@@ -77,7 +77,11 @@ func (f *forkGraphDisk) readBeaconStateFromDisk(blockRoot libcommon.Hash) (bs *s
 		return nil, fmt.Errorf("failed to read snappy buffer: %w, root: %x", err, blockRoot)
 	}
 	f.sszBuffer = f.sszBuffer[:n]
-	bs = state.New(f.beaconCfg)
+	if out == nil {
+		bs = state.New(f.beaconCfg)
+	} else {
+		bs = out
+	}
 
 	if err = bs.DecodeSSZ(f.sszBuffer, int(v[0])); err != nil {
 		return nil, fmt.Errorf("failed to decode beacon state: %w, root: %x, len: %d, decLen: %d, bs: %+v", err, blockRoot, n, len(f.sszBuffer), bs)