From dc643c9f3213d2e12fc007f8e9a7585aeb1198b0 Mon Sep 17 00:00:00 2001
From: Nishant Das <nishdas93@gmail.com>
Date: Mon, 2 Dec 2024 21:29:36 +0800
Subject: [PATCH] Fix Deadline Again During Rollback (#14686)

* fix it again

* CHANGELOG
---
 CHANGELOG.md                                  |  1 +
 beacon-chain/blockchain/process_block.go      |  7 +--
 beacon-chain/blockchain/process_block_test.go | 56 +++++++++++++++++++
 3 files changed, 60 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 882e42f8ae..40c830d2be 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -101,6 +101,7 @@ The format is based on Keep a Changelog, and this project adheres to Semantic Ve
 - temporary solution to handling electra attesation and attester_slashing events. [pr](14655)
 - Diverse log improvements and comment additions.
 - P2P: Avoid infinite loop when looking for peers in small networks.
+- Fixed another rollback bug due to a context deadline.
 
 
 ### Security
diff --git a/beacon-chain/blockchain/process_block.go b/beacon-chain/blockchain/process_block.go
index c816c23881..39a384fff0 100644
--- a/beacon-chain/blockchain/process_block.go
+++ b/beacon-chain/blockchain/process_block.go
@@ -404,10 +404,9 @@ func (s *Service) savePostStateInfo(ctx context.Context, r [32]byte, b interface
 		return errors.Wrapf(err, "could not save block from slot %d", b.Block().Slot())
 	}
 	if err := s.cfg.StateGen.SaveState(ctx, r, st); err != nil {
-		log.Warnf("Rolling back insertion of block with root %#x", r)
-		if err := s.cfg.BeaconDB.DeleteBlock(ctx, r); err != nil {
-			log.WithError(err).Errorf("Could not delete block with block root %#x", r)
-		}
+		// Do not use parent context in the event it deadlined
+		ctx = trace.NewContext(context.Background(), span)
+		s.rollbackBlock(ctx, r)
 		return errors.Wrap(err, "could not save state")
 	}
 	return nil
diff --git a/beacon-chain/blockchain/process_block_test.go b/beacon-chain/blockchain/process_block_test.go
index 53fb11e805..9b6a99200e 100644
--- a/beacon-chain/blockchain/process_block_test.go
+++ b/beacon-chain/blockchain/process_block_test.go
@@ -2352,6 +2352,62 @@ func TestRollbackBlock(t *testing.T) {
 	require.Equal(t, false, hasState)
 }
 
+func TestRollbackBlock_SavePostStateInfo_ContextDeadline(t *testing.T) {
+	service, tr := minimalTestService(t)
+	ctx := tr.ctx
+
+	st, keys := util.DeterministicGenesisState(t, 64)
+	stateRoot, err := st.HashTreeRoot(ctx)
+	require.NoError(t, err, "Could not hash genesis state")
+
+	require.NoError(t, service.saveGenesisData(ctx, st))
+
+	genesis := blocks.NewGenesisBlock(stateRoot[:])
+	wsb, err := consensusblocks.NewSignedBeaconBlock(genesis)
+	require.NoError(t, err)
+	require.NoError(t, service.cfg.BeaconDB.SaveBlock(ctx, wsb), "Could not save genesis block")
+	parentRoot, err := genesis.Block.HashTreeRoot()
+	require.NoError(t, err, "Could not get signing root")
+	require.NoError(t, service.cfg.BeaconDB.SaveState(ctx, st, parentRoot), "Could not save genesis state")
+	require.NoError(t, service.cfg.BeaconDB.SaveHeadBlockRoot(ctx, parentRoot), "Could not save genesis state")
+	require.NoError(t, service.cfg.BeaconDB.SaveJustifiedCheckpoint(ctx, &ethpb.Checkpoint{Root: parentRoot[:]}))
+	require.NoError(t, service.cfg.BeaconDB.SaveFinalizedCheckpoint(ctx, &ethpb.Checkpoint{Root: parentRoot[:]}))
+
+	st, err = service.HeadState(ctx)
+	require.NoError(t, err)
+	b, err := util.GenerateFullBlock(st, keys, util.DefaultBlockGenConfig(), 128)
+	require.NoError(t, err)
+	wsb, err = consensusblocks.NewSignedBeaconBlock(b)
+	require.NoError(t, err)
+	root, err := b.Block.HashTreeRoot()
+	require.NoError(t, err)
+	preState, err := service.getBlockPreState(ctx, wsb.Block())
+	require.NoError(t, err)
+	postState, err := service.validateStateTransition(ctx, preState, wsb)
+	require.NoError(t, err)
+
+	// Save state summaries so that the cache is flushed and saved to disk
+	// later.
+	for i := 1; i <= 127; i++ {
+		require.NoError(t, service.cfg.BeaconDB.SaveStateSummary(ctx, &ethpb.StateSummary{
+			Slot: primitives.Slot(i),
+			Root: bytesutil.Bytes32(uint64(i)),
+		}))
+	}
+
+	// Set deadlined context when saving block and state
+	cancCtx, canc := context.WithCancel(ctx)
+	canc()
+
+	require.ErrorContains(t, context.Canceled.Error(), service.savePostStateInfo(cancCtx, root, wsb, postState))
+
+	// The block should no longer exist.
+	require.Equal(t, false, service.cfg.BeaconDB.HasBlock(ctx, root))
+	hasState, err := service.cfg.StateGen.HasState(ctx, root)
+	require.NoError(t, err)
+	require.Equal(t, false, hasState)
+}
+
 func TestRollbackBlock_ContextDeadline(t *testing.T) {
 	service, tr := minimalTestService(t)
 	ctx := tr.ctx