Rollback on errors from forkchoice insertion (#14556)

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
 

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.

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))
+}

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.

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))
+}