From 2689b35da9edfcce3c3f851ac31285d1c6598e58 Mon Sep 17 00:00:00 2001
From: Alex Sharov <AskAlexSharov@gmail.com>
Date: Wed, 8 Jul 2020 11:46:04 +0700
Subject: [PATCH] TxLookup unwind: Remove entries for blocks between
 unwindPoint+1 and stage.BlockNumber (#723)

---
 cmd/integration/stages.go        |  3 ++-
 cmd/integration/state_stages.go  |  3 ++-
 eth/stagedsync/stage_txlookup.go | 12 +++++++++---
 eth/stagedsync/stagedsync.go     |  2 +-
 4 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/cmd/integration/stages.go b/cmd/integration/stages.go
index 73c639e48..e4f2c8ea9 100644
--- a/cmd/integration/stages.go
+++ b/cmd/integration/stages.go
@@ -307,7 +307,8 @@ func stage9(ctx context.Context) error {
 
 	if unwind > 0 {
 		u := &stagedsync.UnwindState{Stage: stages.TxLookup, UnwindPoint: stage9.BlockNumber - unwind}
-		return stagedsync.UnwindTxLookup(u, db, "", ch)
+		s := progress(db, stages.TxLookup)
+		return stagedsync.UnwindTxLookup(u, s, db, "", ch)
 	}
 
 	return stagedsync.SpawnTxLookup(stage9, db, "", ch)
diff --git a/cmd/integration/state_stages.go b/cmd/integration/state_stages.go
index 70be07c9d..abb9900b5 100644
--- a/cmd/integration/state_stages.go
+++ b/cmd/integration/state_stages.go
@@ -120,7 +120,8 @@ func syncBySmallSteps(ctx context.Context, chaindata string) error {
 		to := execStage.BlockNumber - unwind
 		{
 			u := &stagedsync.UnwindState{Stage: stages.TxLookup, UnwindPoint: to}
-			if err = stagedsync.UnwindTxLookup(u, db, "", ch); err != nil {
+			stage9 := progress(db, stages.TxLookup)
+			if err = stagedsync.UnwindTxLookup(u, stage9, db, "", ch); err != nil {
 				return fmt.Errorf("unwindTxLookup: %w", err)
 			}
 		}
diff --git a/eth/stagedsync/stage_txlookup.go b/eth/stagedsync/stage_txlookup.go
index 9363f2317..d0106958e 100644
--- a/eth/stagedsync/stage_txlookup.go
+++ b/eth/stagedsync/stage_txlookup.go
@@ -64,10 +64,16 @@ func TxLookupTransform(db ethdb.Database, startKey, endKey []byte, quitCh <-chan
 	})
 }
 
-func UnwindTxLookup(u *UnwindState, db ethdb.Database, datadir string, quitCh <-chan struct{}) error {
+func UnwindTxLookup(u *UnwindState, s *StageState, db ethdb.Database, datadir string, quitCh <-chan struct{}) error {
 	collector := etl.NewCollector(datadir, etl.NewSortableBuffer(etl.BufferOptimalSize))
-	// Remove lookup entries for all blocks above unwindPoint
-	if err := db.Walk(dbutils.BlockBodyPrefix, dbutils.EncodeBlockNumber(u.UnwindPoint+1), 8*8, func(k, v []byte) (b bool, e error) {
+
+	// Remove lookup entries for blocks between unwindPoint+1 and stage.BlockNumber
+	if err := db.Walk(dbutils.BlockBodyPrefix, dbutils.EncodeBlockNumber(u.UnwindPoint+1), 0, func(k, v []byte) (b bool, e error) {
+		blockNumber := binary.BigEndian.Uint64(k[:8])
+		if blockNumber > s.BlockNumber {
+			return false, nil
+		}
+
 		if err := common.Stopped(quitCh); err != nil {
 			return false, err
 		}
diff --git a/eth/stagedsync/stagedsync.go b/eth/stagedsync/stagedsync.go
index 9cefcd5c9..1436b3d55 100644
--- a/eth/stagedsync/stagedsync.go
+++ b/eth/stagedsync/stagedsync.go
@@ -133,7 +133,7 @@ func PrepareStagedSync(
 				return SpawnTxLookup(s, stateDB, datadir, quitCh)
 			},
 			UnwindFunc: func(u *UnwindState, s *StageState) error {
-				return UnwindTxLookup(u, stateDB, datadir, quitCh)
+				return UnwindTxLookup(u, s, stateDB, datadir, quitCh)
 			},
 		},
 	}