fix(core): clean canonical metadata during rollback

core/blockchain_impl.go

@@ -1319,37 +1319,38 @@ func (bc *BlockChainImpl) Rollback(chain []common.Hash) error {
 	bc.chainmu.Lock()
 	defer bc.chainmu.Unlock()
 
+	oldHeader, newHeader := bc.hc.CurrentHeader(), bc.hc.CurrentHeader()
+	oldFastBlock, newFastBlock := bc.CurrentFastBlock(), bc.CurrentFastBlock()
+	oldBlock, newBlock := bc.CurrentBlock(), bc.CurrentBlock()
 	valsToRemove := map[common.Address]struct{}{}
+	var canonicalStart *uint64
+
 	for i := len(chain) - 1; i >= 0; i-- {
 		hash := chain[i]
 
-		currentHeader := bc.hc.CurrentHeader()
-		if currentHeader != nil && currentHeader.Hash() == hash {
-			parentHeader := bc.GetHeader(currentHeader.ParentHash(), currentHeader.Number().Uint64()-1)
-			if parentHeader != nil {
-				if err := bc.hc.SetCurrentHeader(parentHeader); err != nil {
-					return errors.Wrap(err, "HeaderChain SetCurrentHeader")
-				}
+		if newHeader != nil && newHeader.Hash() == hash {
+			parent := bc.GetHeader(newHeader.ParentHash(), newHeader.Number().Uint64()-1)
+			if parent != nil {
+				newHeader = parent
 			}
 		}
-		if currentFastBlock := bc.CurrentFastBlock(); currentFastBlock != nil && currentFastBlock.Hash() == hash {
-			newFastBlock := bc.GetBlock(currentFastBlock.ParentHash(), currentFastBlock.NumberU64()-1)
-			if newFastBlock != nil {
-				bc.currentFastBlock.Store(newFastBlock)
-				headFastBlockGauge.Update(int64(newFastBlock.NumberU64()))
-				rawdb.WriteHeadFastBlockHash(bc.db, newFastBlock.Hash())
+		if newFastBlock != nil && newFastBlock.Hash() == hash {
+			parent := bc.GetBlock(newFastBlock.ParentHash(), newFastBlock.NumberU64()-1)
+			if parent != nil {
+				newFastBlock = parent
 			}
 		}
-		if currentBlock := bc.CurrentBlock(); currentBlock != nil && currentBlock.Hash() == hash {
-			newBlock := bc.GetBlock(currentBlock.ParentHash(), currentBlock.NumberU64()-1)
-			if newBlock != nil {
-				bc.currentBlock.Store(newBlock)
-				headBlockGauge.Update(int64(newBlock.NumberU64()))
-				if err := rawdb.WriteHeadBlockHash(bc.db, newBlock.Hash()); err != nil {
-					return err
+		if newBlock != nil && newBlock.Hash() == hash {
+			rolledBlock := newBlock
+			parent := bc.GetBlock(rolledBlock.ParentHash(), rolledBlock.NumberU64()-1)
+			if parent != nil {
+				number := rolledBlock.NumberU64()
+				if canonicalStart == nil || number < *canonicalStart {
+					canonicalStart = &number
 				}
+				newBlock = parent
 
-				for _, stkTxn := range currentBlock.StakingTransactions() {
+				for _, stkTxn := range rolledBlock.StakingTransactions() {
 					if stkTxn.StakingType() == staking.DirectiveCreateValidator {
 						if addr, err := stkTxn.SenderAddress(); err == nil {
 							valsToRemove[addr] = struct{}{}
@@ -1359,7 +1360,96 @@ func (bc *BlockChainImpl) Rollback(chain []common.Hash) error {
 			}
 		}
 	}
-	return bc.removeInValidatorList(valsToRemove)
+
+	headerChanged := oldHeader != newHeader
+	fastBlockChanged := oldFastBlock != newFastBlock
+	blockChanged := oldBlock != newBlock
+	if !headerChanged && !fastBlockChanged && !blockChanged {
+		return nil
+	}
+
+	var validatorList []common.Address
+	if len(valsToRemove) > 0 {
+		existing, err := bc.ReadValidatorList()
+		if err != nil {
+			return err
+		}
+		validatorList = make([]common.Address, 0, len(existing))
+		for _, addr := range existing {
+			if _, remove := valsToRemove[addr]; !remove {
+				validatorList = append(validatorList, addr)
+			}
+		}
+	}
+
+	batch := bc.db.NewBatch()
+	if headerChanged {
+		if err := rawdb.WriteHeadHeaderHash(batch, newHeader.Hash()); err != nil {
+			return errors.Wrap(err, "write head header hash")
+		}
+	}
+	if fastBlockChanged {
+		if err := rawdb.WriteHeadFastBlockHash(batch, newFastBlock.Hash()); err != nil {
+			return errors.Wrap(err, "write fast head block hash")
+		}
+	}
+	if blockChanged {
+		if err := rawdb.WriteHeadBlockHash(batch, newBlock.Hash()); err != nil {
+			return errors.Wrap(err, "write head block hash")
+		}
+	}
+	if validatorList != nil {
+		if err := rawdb.WriteValidatorList(batch, validatorList); err != nil {
+			return errors.Wrap(err, "write validator list")
+		}
+	}
+
+	canonicalLocked := false
+	if canonicalStart != nil {
+		bc.hc.canonicalMu.Lock()
+		canonicalLocked = true
+		if err := rawdb.DeleteCanonicalHash(batch, *canonicalStart); err != nil {
+			bc.hc.canonicalMu.Unlock()
+			return err
+		}
+		for number := *canonicalStart + 1; number > *canonicalStart; number++ {
+			if rawdb.ReadCanonicalHash(bc.db, number) == (common.Hash{}) {
+				break
+			}
+			if err := rawdb.DeleteCanonicalHash(batch, number); err != nil {
+				bc.hc.canonicalMu.Unlock()
+				return err
+			}
+		}
+	}
+	if err := batch.Write(); err != nil {
+		if canonicalLocked {
+			bc.hc.canonicalMu.Unlock()
+		}
+		return err
+	}
+	if canonicalLocked {
+		bc.hc.canonicalCache.Purge()
+		bc.hc.canonicalMu.Unlock()
+	}
+
+	if headerChanged {
+		bc.hc.currentHeader.Store(newHeader)
+		bc.hc.currentHeaderHash = newHeader.Hash()
+		headHeaderGauge.Update(newHeader.Number().Int64())
+	}
+	if fastBlockChanged {
+		bc.currentFastBlock.Store(newFastBlock)
+		headFastBlockGauge.Update(int64(newFastBlock.NumberU64()))
+	}
+	if blockChanged {
+		bc.currentBlock.Store(newBlock)
+		headBlockGauge.Update(int64(newBlock.NumberU64()))
+	}
+	if validatorList != nil {
+		bc.validatorListCache.Add("validatorList", validatorList)
+	}
+	return nil
 }
 
 // SetReceiptsData computes all the non-consensus fields of the receipts

core/blockchain_impl_test.go

@@ -2,15 +2,182 @@ package core
 
 import (
 	"crypto/ecdsa"
+	"errors"
 	"math/big"
 	"testing"
+	"time"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/ethereum/go-ethereum/ethdb"
+	"github.com/harmony-one/harmony/core/rawdb"
 	"github.com/harmony-one/harmony/core/types"
 	staking "github.com/harmony-one/harmony/staking/types"
 )
 
+var errInjectedBatchWrite = errors.New("injected batch write failure")
+
+type failingBatchDatabase struct {
+	ethdb.Database
+}
+
+func (db failingBatchDatabase) NewBatch() ethdb.Batch {
+	return failingBatch{Batch: db.Database.NewBatch()}
+}
+
+type failingBatch struct {
+	ethdb.Batch
+}
+
+func (failingBatch) Write() error {
+	return errInjectedBatchWrite
+}
+
+func TestRollbackClearsCanonicalMappingAndCache(t *testing.T) {
+	key, _ := crypto.GenerateKey()
+	chain, _, header, database := getTestEnvironment(*key)
+	defer chain.Stop()
+
+	genesis := chain.Genesis()
+	makeBlock := func(parent common.Hash, number int64, extra string) *types.Block {
+		h := header.With().
+			ParentHash(parent).
+			Number(big.NewInt(number)).
+			Root(genesis.Root()).
+			Extra([]byte(extra)).
+			Header()
+		return types.NewBlockWithHeader(h)
+	}
+
+	oldBlock := makeBlock(genesis.Hash(), 1, "old canonical block")
+	if err := rawdb.WriteBlock(database, oldBlock); err != nil {
+		t.Fatalf("write old block: %v", err)
+	}
+	if err := chain.WriteHeadBlock(oldBlock); err != nil {
+		t.Fatalf("set old canonical head: %v", err)
+	}
+	if got := chain.GetBlockByNumber(1); got == nil || got.Hash() != oldBlock.Hash() {
+		t.Fatalf("failed to warm canonical cache with old block: got %v", got)
+	}
+
+	staleBlock2 := makeBlock(oldBlock.Hash(), 2, "stale block above current head")
+	if err := rawdb.WriteCanonicalHash(database, staleBlock2.Hash(), 2); err != nil {
+		t.Fatalf("write stale canonical mapping: %v", err)
+	}
+	if got := chain.GetCanonicalHash(2); got != staleBlock2.Hash() {
+		t.Fatalf("failed to warm stale canonical cache: got %s want %s", got, staleBlock2.Hash())
+	}
+
+	if err := chain.Rollback([]common.Hash{oldBlock.Hash()}); err != nil {
+		t.Fatalf("rollback old block: %v", err)
+	}
+	if got := chain.CurrentBlock().Hash(); got != genesis.Hash() {
+		t.Fatalf("current block after rollback: got %s want genesis %s", got, genesis.Hash())
+	}
+	if got := rawdb.ReadCanonicalHash(database, 1); got != (common.Hash{}) {
+		t.Fatalf("persistent canonical mapping survived rollback: got %s", got)
+	}
+	if got := chain.GetCanonicalHash(1); got != (common.Hash{}) {
+		t.Fatalf("cached canonical mapping survived rollback: got %s", got)
+	}
+	if got := rawdb.ReadCanonicalHash(database, 2); got != (common.Hash{}) {
+		t.Fatalf("persistent canonical mapping above rolled-back head survived: got %s", got)
+	}
+	if got := chain.GetCanonicalHash(2); got != (common.Hash{}) {
+		t.Fatalf("cached canonical mapping above rolled-back head survived: got %s", got)
+	}
+
+	replacement := makeBlock(genesis.Hash(), 1, "replacement canonical block")
+	if err := rawdb.WriteBlock(database, replacement); err != nil {
+		t.Fatalf("write replacement block: %v", err)
+	}
+	if err := chain.WriteHeadBlock(replacement); err != nil {
+		t.Fatalf("set replacement canonical head: %v", err)
+	}
+	if got := chain.GetBlockByNumber(1); got == nil || got.Hash() != replacement.Hash() {
+		t.Fatalf("canonical lookup after replacement: got %v want %s", got, replacement.Hash())
+	}
+}
+
+func TestRollbackBatchFailureLeavesHeadsAndCanonicalMappingUnchanged(t *testing.T) {
+	key, _ := crypto.GenerateKey()
+	chain, _, header, database := getTestEnvironment(*key)
+	defer chain.Stop()
+
+	genesis := chain.Genesis()
+	oldBlock := types.NewBlockWithHeader(header.With().
+		ParentHash(genesis.Hash()).
+		Number(big.NewInt(1)).
+		Root(genesis.Root()).
+		Extra([]byte("old canonical block")).
+		Header())
+	if err := rawdb.WriteBlock(database, oldBlock); err != nil {
+		t.Fatalf("write old block: %v", err)
+	}
+	if err := chain.WriteHeadBlock(oldBlock); err != nil {
+		t.Fatalf("write old head: %v", err)
+	}
+
+	failingDB := failingBatchDatabase{Database: database}
+	chain.db = failingDB
+	chain.hc.chainDb = failingDB
+
+	err := chain.Rollback([]common.Hash{oldBlock.Hash()})
+	if !errors.Is(err, errInjectedBatchWrite) {
+		t.Fatalf("rollback error = %v, want %v", err, errInjectedBatchWrite)
+	}
+	if got := chain.CurrentBlock().Hash(); got != oldBlock.Hash() {
+		t.Fatalf("in-memory block head changed after failed rollback: got %s want %s", got, oldBlock.Hash())
+	}
+	if got := chain.CurrentFastBlock().Hash(); got != oldBlock.Hash() {
+		t.Fatalf("in-memory fast head changed after failed rollback: got %s want %s", got, oldBlock.Hash())
+	}
+	if got := chain.CurrentHeader().Hash(); got != oldBlock.Hash() {
+		t.Fatalf("in-memory header head changed after failed rollback: got %s want %s", got, oldBlock.Hash())
+	}
+	if got := rawdb.ReadHeadBlockHash(database); got != oldBlock.Hash() {
+		t.Fatalf("persistent block head changed after failed rollback: got %s want %s", got, oldBlock.Hash())
+	}
+	if got := rawdb.ReadHeadFastBlockHash(database); got != oldBlock.Hash() {
+		t.Fatalf("persistent fast head changed after failed rollback: got %s want %s", got, oldBlock.Hash())
+	}
+	if got := rawdb.ReadHeadHeaderHash(database); got != oldBlock.Hash() {
+		t.Fatalf("persistent header head changed after failed rollback: got %s want %s", got, oldBlock.Hash())
+	}
+	if got := rawdb.ReadCanonicalHash(database, 1); got != oldBlock.Hash() {
+		t.Fatalf("canonical mapping changed after failed rollback: got %s want %s", got, oldBlock.Hash())
+	}
+}
+
+func TestGetCanonicalHashWaitsForCanonicalMutation(t *testing.T) {
+	key, _ := crypto.GenerateKey()
+	chain, _, _, _ := getTestEnvironment(*key)
+	defer chain.Stop()
+
+	chain.hc.canonicalMu.Lock()
+	done := make(chan common.Hash, 1)
+	go func() {
+		done <- chain.GetCanonicalHash(0)
+	}()
+
+	select {
+	case <-done:
+		chain.hc.canonicalMu.Unlock()
+		t.Fatal("canonical read bypassed mutation lock")
+	case <-time.After(25 * time.Millisecond):
+	}
+
+	chain.hc.canonicalMu.Unlock()
+	select {
+	case got := <-done:
+		if got != chain.Genesis().Hash() {
+			t.Fatalf("canonical read after mutation lock: got %s want %s", got, chain.Genesis().Hash())
+		}
+	case <-time.After(time.Second):
+		t.Fatal("canonical read did not resume after mutation lock")
+	}
+}
+
 // TestIsSpentIgnoresMutatedMerkleProofIdentity guards against replaying a
 // genuine, already-applied CXReceiptsProof by mutating the unauthenticated
 // MerkleProof.ShardID/BlockNum while keeping the same signed Header: the

core/headerchain.go

@@ -22,6 +22,7 @@ import (
 	"math"
 	"math/big"
 	mrand "math/rand"
+	"sync"
 	"sync/atomic"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -61,6 +62,7 @@ type HeaderChain struct {
 	tdCache        *lru.Cache // Cache for the most recent block total difficulties
 	numberCache    *lru.Cache // Cache for the most recent block numbers
 	canonicalCache *lru.Cache // number -> Hash
+	canonicalMu    sync.RWMutex
 
 	procInterrupt func() bool
 
@@ -402,8 +404,9 @@ func (hc *HeaderChain) getHashByNumber(number uint64) common.Hash {
 }
 
 func (hc *HeaderChain) GetCanonicalHash(number uint64) common.Hash {
-	// Since canonical chain is immutable, it's safe to read header
-	// hash by number from cache.
+	hc.canonicalMu.RLock()
+	defer hc.canonicalMu.RUnlock()
+
 	if hash, ok := hc.canonicalCache.Get(number); ok {
 		return hash.(common.Hash)
 	}