core, rawdb, dex: improve gov state for syncing (#393)

* core, rawdb, dex: improve gov state for syncing

5 files changed, 108 insertions, 3 deletions

diff --git a/core/blockchain.go b/core/blockchain.go
index 3f0e7203e..a1e82313b 100644
--- a/core/blockchain.go
+++ b/core/blockchain.go
@@ -117,6 +117,7 @@ type BlockChain struct {
 	mu      sync.RWMutex // global mutex for locking chain operations
 	chainmu sync.RWMutex // blockchain insertion lock
 	procmu  sync.RWMutex // block processor lock
+	govmu   sync.RWMutex // gov state lock
 
 	checkpoint       int          // checkpoint counts towards the new checkpoint
 	currentBlock     atomic.Value // Current head of the block chain
@@ -983,7 +984,7 @@ func (bc *BlockChain) WriteBlockWithoutState(block *types.Block, td *big.Int) (e
 }
 
 // WriteBlockWithState writes the block and all associated state to the database.
-func (bc *BlockChain) WriteBlockWithState(block *types.Block, receipts []*types.Receipt, state *state.StateDB) (status WriteStatus, err error) {
+func (bc *BlockChain) WriteBlockWithState(block *types.Block, receipts []*types.Receipt, statedb *state.StateDB) (status WriteStatus, err error) {
 	bc.wg.Add(1)
 	defer bc.wg.Done()
 
@@ -1006,7 +1007,7 @@ func (bc *BlockChain) WriteBlockWithState(block *types.Block, receipts []*types.
 	}
 	rawdb.WriteBlock(bc.db, block)
 
-	root, err := state.Commit(bc.chainConfig.IsEIP158(block.Number()))
+	root, err := statedb.Commit(bc.chainConfig.IsEIP158(block.Number()))
 	if err != nil {
 		return NonStatTy, err
 	}
@@ -1017,6 +1018,33 @@ func (bc *BlockChain) WriteBlockWithState(block *types.Block, receipts []*types.
 	}
 	height, _ := bc.GetRoundHeight(block.Round())
 
+	// Write gov state into disk
+	if height == block.NumberU64() {
+		// spawn a goroutine to write gov state
+		go func() {
+			retry := 3
+			n := 0
+			for n < retry {
+				// get gov state from state db
+				tt := time.Now()
+				log.Debug("Write gov state", "n", n, "t", tt)
+				govState, err := state.GetGovState(statedb, block.Header(),
+					vm.GovernanceContractAddress)
+				log.Debug("Get gov state finished", "n", n, "elapsed", time.Since(tt))
+				if err == nil {
+					bc.govmu.Lock()
+					rawdb.WriteGovState(bc.db, block.Hash(), govState)
+					bc.govmu.Unlock()
+					log.Debug("Write gov state finished", "n", n, "elapsed", time.Since(tt))
+					break
+				} else {
+					log.Warn("Get gov state fail", "n", n, "err", err)
+				}
+				n++
+			}
+		}()
+	}
+
 	// If we're running an archive node or the block is snapshot height, always flush
 	if bc.cacheConfig.Disabled || height == block.NumberU64() {
 		if err := triedb.Commit(root, false); err != nil {
@@ -1101,7 +1129,7 @@ func (bc *BlockChain) WriteBlockWithState(block *types.Block, receipts []*types.
 		}
 		// Write the positional metadata for transaction/receipt lookups and preimages
 		rawdb.WriteTxLookupEntries(batch, block)
-		rawdb.WritePreimages(batch, state.Preimages())
+		rawdb.WritePreimages(batch, statedb.Preimages())
 
 		status = CanonStatTy
 	} else {
@@ -1889,6 +1917,16 @@ func (bc *BlockChain) GetGovStateByHash(hash common.Hash) (*types.GovState, erro
 	if header == nil {
 		return nil, fmt.Errorf("header not found")
 	}
+
+	// Get gov state from disk first.
+	bc.govmu.Lock()
+	if govState := rawdb.ReadGovState(bc.db, header.Hash()); govState != nil {
+		bc.govmu.Unlock()
+		log.Debug("Read gov state from db success")
+		return govState, nil
+	}
+	bc.govmu.Unlock()
+
 	statedb, err := bc.StateAt(header.Root)
 	if err != nil {
 		return nil, err
@@ -1901,6 +1939,16 @@ func (bc *BlockChain) GetGovStateByNumber(number uint64) (*types.GovState, error
 	if header == nil {
 		return nil, fmt.Errorf("header not found")
 	}
+
+	// Get gov state from disk first.
+	bc.govmu.Lock()
+	if govState := rawdb.ReadGovState(bc.db, header.Hash()); govState != nil {
+		bc.govmu.Unlock()
+		log.Debug("Read gov state from db success")
+		return govState, nil
+	}
+	bc.govmu.Unlock()
+
 	statedb, err := bc.StateAt(header.Root)
 	if err != nil {
 		return nil, err
diff --git a/core/headerchain.go b/core/headerchain.go
index 4eff41b26..ffa005010 100644
--- a/core/headerchain.go
+++ b/core/headerchain.go
@@ -372,6 +372,7 @@ func (hc *HeaderChain) WriteDexonHeader(header *types.HeaderWithGovState) (statu
 
 	// Store the govState
 	if govState := header.GovState; govState != nil {
+		rawdb.WriteGovState(hc.chainDb, header.Hash(), header.GovState)
 		batch := hc.chainDb.NewBatch()
 		for _, node := range govState.Proof {
 			batch.Put(crypto.Keccak256(node), node)
diff --git a/core/rawdb/accessors_chain.go b/core/rawdb/accessors_chain.go
index 801ad9362..41a79b21d 100644
--- a/core/rawdb/accessors_chain.go
+++ b/core/rawdb/accessors_chain.go
@@ -20,6 +20,7 @@ import (
 	"bytes"
 	"encoding/binary"
 	"math/big"
+	"time"
 
 	"github.com/dexon-foundation/dexon/common"
 	"github.com/dexon-foundation/dexon/core/types"
@@ -382,3 +383,49 @@ func FindCommonAncestor(db DatabaseReader, a, b *types.Header) *types.Header {
 	}
 	return a
 }
+
+// ReadGovStateRLP retrieves
+func ReadGovStateRLP(db DatabaseReader, hash common.Hash) rlp.RawValue {
+	data, _ := db.Get(govStateKey(hash))
+	return data
+}
+
+// WriteGovStateRLP
+func WriteGovStateRLP(db DatabaseWriter, hash common.Hash, rlp rlp.RawValue) {
+	if err := db.Put(govStateKey(hash), rlp); err != nil {
+		log.Crit("Failed to store gov state", "err", err)
+	}
+}
+
+// ReadGovState
+func ReadGovState(db DatabaseReader, hash common.Hash) *types.GovState {
+	data := ReadGovStateRLP(db, hash)
+	if len(data) == 0 {
+		return nil
+	}
+	govState := new(types.GovState)
+	if err := rlp.Decode(bytes.NewReader(data), govState); err != nil {
+		log.Error("Invalid gov state RLP", "hash", hash, "err", err)
+		return nil
+	}
+	return govState
+}
+
+// WriteGovState
+func WriteGovState(db DatabaseWriter, hash common.Hash, govState *types.GovState) {
+	t := time.Now()
+	log.Debug("Rawdb WriteGovState", "t", t)
+	data, err := rlp.EncodeToBytes(govState)
+	if err != nil {
+		log.Crit("Failed to RLP encode gov state", "err", err)
+	}
+	log.Debug("Rawdb WriteGovState", "len", len(data), "elapsed", time.Since(t))
+	WriteGovStateRLP(db, hash, data)
+}
+
+// DeleteGovState
+func DeleteGovState(db DatabaseDeleter, hash common.Hash) {
+	if err := db.Delete(govStateKey(hash)); err != nil {
+		log.Crit("Failed to delete gov satate", "err", err)
+	}
+}
diff --git a/core/rawdb/schema.go b/core/rawdb/schema.go
index 8ca47676d..5b92891df 100644
--- a/core/rawdb/schema.go
+++ b/core/rawdb/schema.go
@@ -50,6 +50,8 @@ var (
 	blockBodyPrefix     = []byte("b") // blockBodyPrefix + num (uint64 big endian) + hash -> block body
 	blockReceiptsPrefix = []byte("r") // blockReceiptsPrefix + num (uint64 big endian) + hash -> block receipts
 
+	govStatePrefix = []byte("g")
+
 	txLookupPrefix  = []byte("l") // txLookupPrefix + hash -> transaction/receipt lookup metadata
 	bloomBitsPrefix = []byte("B") // bloomBitsPrefix + bit (uint16 big endian) + section (uint64 big endian) + hash -> bloom bits
 
@@ -118,6 +120,10 @@ func txLookupKey(hash common.Hash) []byte {
 	return append(txLookupPrefix, hash.Bytes()...)
 }
 
+func govStateKey(hash common.Hash) []byte {
+	return append(govStatePrefix, hash.Bytes()...)
+}
+
 // coreBlockKey = coreBlockPrefix + hash
 func coreBlockKey(hash common.Hash) []byte {
 	return append(coreBlockPrefix, hash.Bytes()...)
diff --git a/dex/handler.go b/dex/handler.go
index fb2ce4e4d..fed1dc8b7 100644
--- a/dex/handler.go
+++ b/dex/handler.go
@@ -588,7 +588,10 @@ func (pm *ProtocolManager) handleMsg(p *peer) error {
 
 			for _, header := range headers {
 				if _, exist := snapshotHeight[header.Number.Uint64()]; exist {
+					tt := time.Now()
+					log.Debug("Handler get gov state by hash", "t", tt)
 					s, err := pm.blockchain.GetGovStateByHash(header.Hash())
+					log.Debug("Handler get gov state by hash", "elapsed", time.Since(tt))
 					if err != nil {
 						log.Warn("Get gov state by hash fail", "number", header.Number.Uint64(), "err", err)
 						return p.SendBlockHeaders(query.Flag, []*types.HeaderWithGovState{})