From 073cf19b691a4ea6770e94e634440b0e470d15f4 Mon Sep 17 00:00:00 2001
From: Potuz <potuz@prysmaticlabs.com>
Date: Fri, 18 Oct 2024 13:49:55 -0300
Subject: [PATCH] Rollback on errors from forkchoice insertion (#14556)

---
 CHANGELOG.md                                       |  1 +
 .../forkchoice/doubly-linked-tree/forkchoice.go    |  9 ++++++++-
 .../doubly-linked-tree/forkchoice_test.go          | 14 ++++++++++++++
 .../forkchoice/doubly-linked-tree/store.go         | 10 ++++++++--
 .../forkchoice/doubly-linked-tree/store_test.go    |  9 +++++++++
 5 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index aa7c18cb6b35..22447bb5f331 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -25,6 +25,7 @@ The format is based on Keep a Changelog, and this project adheres to Semantic Ve
 - Switch to compounding when consolidating with source==target.
 - Revert block db save when saving state fails.
 - Return false from HasBlock if the block is being synced. 
+- Cleanup forkchoice on failed insertions.
 
 ### Deprecated
 
diff --git a/beacon-chain/forkchoice/doubly-linked-tree/forkchoice.go b/beacon-chain/forkchoice/doubly-linked-tree/forkchoice.go
index dea38cb49183..296dfbeeca3c 100644
--- a/beacon-chain/forkchoice/doubly-linked-tree/forkchoice.go
+++ b/beacon-chain/forkchoice/doubly-linked-tree/forkchoice.go
@@ -141,7 +141,14 @@ func (f *ForkChoice) InsertNode(ctx context.Context, state state.BeaconState, ro
 	}
 
 	jc, fc = f.store.pullTips(state, node, jc, fc)
-	return f.updateCheckpoints(ctx, jc, fc)
+	if err := f.updateCheckpoints(ctx, jc, fc); err != nil {
+		_, remErr := f.store.removeNode(ctx, node)
+		if remErr != nil {
+			log.WithError(remErr).Error("could not remove node")
+		}
+		return errors.Wrap(err, "could not update checkpoints")
+	}
+	return nil
 }
 
 // updateCheckpoints update the checkpoints when inserting a new node.
diff --git a/beacon-chain/forkchoice/doubly-linked-tree/forkchoice_test.go b/beacon-chain/forkchoice/doubly-linked-tree/forkchoice_test.go
index 54c6c53edad9..496a0dff1283 100644
--- a/beacon-chain/forkchoice/doubly-linked-tree/forkchoice_test.go
+++ b/beacon-chain/forkchoice/doubly-linked-tree/forkchoice_test.go
@@ -3,6 +3,7 @@ package doublylinkedtree
 import (
 	"context"
 	"encoding/binary"
+	"errors"
 	"testing"
 	"time"
 
@@ -887,3 +888,16 @@ func TestForkchoiceParentRoot(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, zeroHash, root)
 }
+
+func TestForkChoice_CleanupInserting(t *testing.T) {
+	f := setup(0, 0)
+	ctx := context.Background()
+	st, blkRoot, err := prepareForkchoiceState(ctx, 1, indexToHash(1), params.BeaconConfig().ZeroHash, params.BeaconConfig().ZeroHash, 2, 2)
+	f.SetBalancesByRooter(func(_ context.Context, _ [32]byte) ([]uint64, error) {
+		return f.justifiedBalances, errors.New("mock err")
+	})
+
+	require.NoError(t, err)
+	require.NotNil(t, f.InsertNode(ctx, st, blkRoot))
+	require.Equal(t, false, f.HasNode(blkRoot))
+}
diff --git a/beacon-chain/forkchoice/doubly-linked-tree/store.go b/beacon-chain/forkchoice/doubly-linked-tree/store.go
index 77c896c46caf..931d7254dc8e 100644
--- a/beacon-chain/forkchoice/doubly-linked-tree/store.go
+++ b/beacon-chain/forkchoice/doubly-linked-tree/store.go
@@ -107,7 +107,9 @@ func (s *Store) insert(ctx context.Context,
 			s.headNode = n
 			s.highestReceivedNode = n
 		} else {
-			return n, errInvalidParentRoot
+			delete(s.nodeByRoot, root)
+			delete(s.nodeByPayload, payloadHash)
+			return nil, errInvalidParentRoot
 		}
 	} else {
 		parent.children = append(parent.children, n)
@@ -128,7 +130,11 @@ func (s *Store) insert(ctx context.Context,
 		jEpoch := s.justifiedCheckpoint.Epoch
 		fEpoch := s.finalizedCheckpoint.Epoch
 		if err := s.treeRootNode.updateBestDescendant(ctx, jEpoch, fEpoch, slots.ToEpoch(currentSlot)); err != nil {
-			return n, err
+			_, remErr := s.removeNode(ctx, n)
+			if remErr != nil {
+				log.WithError(remErr).Error("could not remove node")
+			}
+			return nil, errors.Wrap(err, "could not update best descendants")
 		}
 	}
 	// Update metrics.
diff --git a/beacon-chain/forkchoice/doubly-linked-tree/store_test.go b/beacon-chain/forkchoice/doubly-linked-tree/store_test.go
index ba621b56ca6d..29b60599f6c4 100644
--- a/beacon-chain/forkchoice/doubly-linked-tree/store_test.go
+++ b/beacon-chain/forkchoice/doubly-linked-tree/store_test.go
@@ -525,3 +525,12 @@ func TestStore_TargetRootForEpoch(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, root4, target)
 }
+
+func TestStore_CleanupInserting(t *testing.T) {
+	f := setup(0, 0)
+	ctx := context.Background()
+	st, blkRoot, err := prepareForkchoiceState(ctx, 1, indexToHash(1), indexToHash(2), params.BeaconConfig().ZeroHash, 0, 0)
+	require.NoError(t, err)
+	require.NotNil(t, f.InsertNode(ctx, st, blkRoot))
+	require.Equal(t, false, f.HasNode(blkRoot))
+}