Optimise stateless prototype

cmd/state/stateless.go

@@ -174,7 +174,7 @@ func stateless(genLag, consLag int) {
 	var proofGen *state.Stateless  // Generator of proofs
 	var proofCons *state.Stateless // Consumer of proofs
 	for !interrupt {
-		trace := false //blockNum == 318335
+		trace := false // blockNum == 1807
 		if trace {
 			filename := fmt.Sprintf("right_%d.txt", blockNum-1)
 			f, err1 := os.Create(filename)
@@ -279,7 +279,6 @@ func stateless(genLag, consLag int) {
 						return
 					}
 					writeStats(wf, blockNum, pBlockProof)
-					proofCons.Prune(blockNum-uint64(consLag), false)
 				} else {
 					if err := proofCons.ApplyProof(preRoot, blockProof, block.NumberU64()-1, false); err != nil {
 						fmt.Printf("Error applying proof to consumer: %v\n", err)
@@ -289,6 +288,9 @@ func stateless(genLag, consLag int) {
 				if err := runBlock(tds, proofCons, chainConfig, bcb, header, block, trace, false); err != nil {
 					fmt.Printf("Error running block %d through proof consumer: %v\n", blockNum, err)
 				}
+				if blockNum > uint64(consLag) {
+					proofCons.Prune(blockNum-uint64(consLag), false)
+				}
 			}
 			if proofGen != nil {
 				if err := proofGen.ApplyProof(preRoot, blockProof, block.NumberU64()-1, false); err != nil {

core/blockchain.go

@@ -225,12 +225,14 @@ func (bc *BlockChain) GetTrieDbState() *state.TrieDbState {
 		log.Info("Creating StateDB from latest state", "block", currentBlockNr)
 		var err error
 		bc.trieDbState, err = state.NewTrieDbState(bc.CurrentBlock().Header().Root, bc.db, currentBlockNr)
+		if err != nil {
+			panic(err)
+		}
 		bc.trieDbState.SetNoHistory(bc.noHistory)
 		bc.trieDbState.SetResolveReads(bc.resolveReads)
-		if err != nil {
+		if err := bc.trieDbState.Rebuild(); err != nil {
 			panic(err)
 		}
-		bc.trieDbState.Rebuild()
 	}
 	return bc.trieDbState
 }
@@ -1186,12 +1188,14 @@ func (bc *BlockChain) insertChain(chain types.Blocks, verifySeals bool) (int, []
 			currentBlockNr := bc.CurrentBlock().NumberU64()
 			log.Info("Creating StateDB from latest state", "block", currentBlockNr)
 			bc.trieDbState, err = state.NewTrieDbState(bc.CurrentBlock().Header().Root, bc.db, currentBlockNr)
+			if err != nil {
+				return k, events, coalescedLogs, err
+			}
 			bc.trieDbState.SetNoHistory(bc.noHistory)
 			bc.trieDbState.SetResolveReads(bc.resolveReads)
-			if err != nil {
+			if err := bc.trieDbState.Rebuild(); err != nil {
 				return k, events, coalescedLogs, err
 			}
-			bc.trieDbState.Rebuild()
 		}
 		root = bc.trieDbState.LastRoot()
 		var parentRoot common.Hash

core/blockchain_test.go

@@ -72,7 +72,7 @@ func newCanonical(engine consensus.Engine, n int, full bool) (ethdb.Database, *B
 // Test fork of length N starting from block i
 func testFork(t *testing.T, blockchain *BlockChain, i, n int, full bool, comparator func(td1, td2 *big.Int)) {
 	// Copy old chain up to #i into a new db
-	db, blockchain2, err := newCanonical(ethash.NewFaker(), i, full)
+	db, blockchain2, err := newCanonical(ethash.NewFaker(), i, true)
 	if err != nil {
 		t.Fatal("could not make new canonical in testFork", err)
 	}
@@ -350,7 +350,7 @@ func TestBrokenBlockChain(t *testing.T)  { testBrokenChain(t, true) }
 
 func testBrokenChain(t *testing.T, full bool) {
 	// Make chain starting from genesis
-	db, blockchain, err := newCanonical(ethash.NewFaker(), 10, full)
+	db, blockchain, err := newCanonical(ethash.NewFaker(), 10, true)
 	if err != nil {
 		t.Fatalf("failed to make new canonical chain: %v", err)
 	}

core/chain_makers.go

@@ -232,7 +232,9 @@ func GenerateChain(config *params.ChainConfig, parent *types.Block, engine conse
 	if err != nil {
 		panic(err)
 	}
-	tds.Rebuild()
+	if err := tds.Rebuild(); err != nil {
+		panic(err)
+	}
 	for i := 0; i < n; i++ {
 		statedb := state.New(tds)
 		block, receipt := genblock(i, parent, statedb, tds)

core/state/database.go

@@ -35,7 +35,7 @@ import (
 )
 
 // Trie cache generation limit after which to evict trie nodes from memory.
-var MaxTrieCacheGen = uint32(4 * 1024 * 1024)
+var MaxTrieCacheGen = uint32(1024 * 1024)
 
 var AccountsBucket = []byte("AT")
 var AccountsHistoryBucket = []byte("hAT")
@@ -189,22 +189,20 @@ func (b *Buffer) merge(other *Buffer) {
 
 // Implements StateReader by wrapping a trie and a database, where trie acts as a cache for the database
 type TrieDbState struct {
-	t                *trie.Trie
-	db               ethdb.Database
-	blockNr          uint64
-	storageTries     map[common.Address]*trie.Trie
-	buffers          []*Buffer
-	aggregateBuffer  *Buffer // Merge of all buffers
-	currentBuffer    *Buffer
-	codeCache        *lru.Cache
-	codeSizeCache    *lru.Cache
-	historical       bool
-	generationCounts map[uint64]int
-	nodeCount        int
-	oldestGeneration uint64
-	noHistory        bool
-	resolveReads     bool
-	pg               *trie.ProofGenerator
+	t               *trie.Trie
+	db              ethdb.Database
+	blockNr         uint64
+	storageTries    map[common.Address]*trie.Trie
+	buffers         []*Buffer
+	aggregateBuffer *Buffer // Merge of all buffers
+	currentBuffer   *Buffer
+	codeCache       *lru.Cache
+	codeSizeCache   *lru.Cache
+	historical      bool
+	noHistory       bool
+	resolveReads    bool
+	pg              *trie.ProofGenerator
+	tp              *trie.TriePruning
 }
 
 func NewTrieDbState(root common.Hash, db ethdb.Database, blockNr uint64) (*TrieDbState, error) {
@@ -217,6 +215,10 @@ func NewTrieDbState(root common.Hash, db ethdb.Database, blockNr uint64) (*TrieD
 		return nil, err
 	}
 	t := trie.New(root, false)
+	tp, err := trie.NewTriePruning(blockNr)
+	if err != nil {
+		return nil, err
+	}
 	tds := TrieDbState{
 		t:             t,
 		db:            db,
@@ -225,10 +227,11 @@ func NewTrieDbState(root common.Hash, db ethdb.Database, blockNr uint64) (*TrieD
 		codeCache:     cc,
 		codeSizeCache: csc,
 		pg:            trie.NewProofGenerator(),
+		tp:            tp,
 	}
-	t.MakeListed(tds.joinGeneration, tds.leftGeneration)
-	tds.generationCounts = make(map[uint64]int, 4096)
-	tds.oldestGeneration = blockNr
+	t.SetTouchFunc(func(hex []byte, del bool) {
+		tp.Touch(hex, del)
+	})
 	return &tds, nil
 }
 
@@ -550,13 +553,13 @@ func (tds *TrieDbState) clearUpdates() {
 	tds.aggregateBuffer = nil
 }
 
-func (tds *TrieDbState) Rebuild() {
-	tr := tds.AccountTrie()
-	tr.Rebuild(tds.db, tds.blockNr)
+func (tds *TrieDbState) Rebuild() error {
+	return tds.AccountTrie().Rebuild(tds.db, tds.blockNr)
 }
 
 func (tds *TrieDbState) SetBlockNr(blockNr uint64) {
 	tds.blockNr = blockNr
+	tds.tp.SetBlockNr(blockNr)
 }
 
 func (tds *TrieDbState) UnwindTo(blockNr uint64) error {
@@ -713,17 +716,6 @@ func encodingToAccount(enc []byte) (*Account, error) {
 	return &data, nil
 }
 
-func (tds *TrieDbState) joinGeneration(gen uint64) {
-	tds.nodeCount++
-	tds.generationCounts[gen]++
-
-}
-
-func (tds *TrieDbState) leftGeneration(gen uint64) {
-	tds.nodeCount--
-	tds.generationCounts[gen]--
-}
-
 func (tds *TrieDbState) ReadAccountData(address common.Address) (*Account, error) {
 	h := newHasher()
 	defer returnHasherToPool(h)
@@ -799,7 +791,9 @@ func (tds *TrieDbState) getStorageTrie(address common.Address, create bool) (*tr
 		} else {
 			t = trie.New(account.Root, true)
 		}
-		t.MakeListed(tds.joinGeneration, tds.leftGeneration)
+		t.SetTouchFunc(func(hex []byte, del bool) {
+			tds.tp.TouchContract(address, hex, del)
+		})
 		tds.storageTries[address] = t
 	}
 	return t, nil
@@ -902,32 +896,40 @@ func (tds *TrieDbState) ReadAccountCodeSize(codeHash common.Hash) (codeSize int,
 var prevMemStats runtime.MemStats
 
 func (tds *TrieDbState) PruneTries(print bool) {
-	if tds.nodeCount > int(MaxTrieCacheGen) {
-		toRemove := 0
-		excess := tds.nodeCount - int(MaxTrieCacheGen)
-		gen := tds.oldestGeneration
-		for excess > 0 {
-			excess -= tds.generationCounts[gen]
-			toRemove += tds.generationCounts[gen]
-			delete(tds.generationCounts, gen)
-			gen++
+	if print {
+		mainPrunable := tds.t.CountPrunableNodes()
+		prunableNodes := mainPrunable
+		for _, storageTrie := range tds.storageTries {
+			prunableNodes += storageTrie.CountPrunableNodes()
 		}
-		// Unload all nodes with touch timestamp < gen
-		for address, storageTrie := range tds.storageTries {
-			empty := storageTrie.UnloadOlderThan(gen, false)
-			if empty {
-				delete(tds.storageTries, address)
-			}
-		}
-		tds.t.UnloadOlderThan(gen, false)
-		tds.oldestGeneration = gen
-		tds.nodeCount -= toRemove
-		var m runtime.MemStats
-		runtime.ReadMemStats(&m)
-		log.Info("Memory", "nodes", tds.nodeCount, "alloc", int(m.Alloc/1024), "sys", int(m.Sys/1024), "numGC", int(m.NumGC))
-		if print {
-			fmt.Printf("Pruning done. Nodes: %d, alloc: %d, sys: %d, numGC: %d\n", tds.nodeCount, int(m.Alloc/1024), int(m.Sys/1024), int(m.NumGC))
+		fmt.Printf("[Before] Actual prunable nodes: %d (main %d), accounted: %d\n", prunableNodes, mainPrunable, tds.tp.NodeCount())
+	}
+	pruned, emptyAddresses, err := tds.tp.PruneTo(tds.t, int(MaxTrieCacheGen), func(contract common.Address) (*trie.Trie, error) {
+		return tds.getStorageTrie(contract, false)
+	})
+	if err != nil {
+		fmt.Printf("Error while pruning: %v\n", err)
+	}
+	if !pruned {
+		//return
+	}
+	if print {
+		mainPrunable := tds.t.CountPrunableNodes()
+		prunableNodes := mainPrunable
+		for _, storageTrie := range tds.storageTries {
+			prunableNodes += storageTrie.CountPrunableNodes()
 		}
+		fmt.Printf("[After] Actual prunable nodes: %d (main %d), accounted: %d\n", prunableNodes, mainPrunable, tds.tp.NodeCount())
+	}
+	// Storage tries that were completely pruned
+	for _, address := range emptyAddresses {
+		delete(tds.storageTries, address)
+	}
+	var m runtime.MemStats
+	runtime.ReadMemStats(&m)
+	log.Info("Memory", "nodes", tds.tp.NodeCount(), "alloc", int(m.Alloc/1024), "sys", int(m.Sys/1024), "numGC", int(m.NumGC))
+	if print {
+		fmt.Printf("Pruning done. Nodes: %d, alloc: %d, sys: %d, numGC: %d\n", tds.tp.NodeCount(), int(m.Alloc/1024), int(m.Sys/1024), int(m.NumGC))
 	}
 }
 

core/state/repair.go

@@ -235,7 +235,6 @@ func (rds *RepairDbState) getStorageTrie(address common.Address, create bool) (*
 		} else {
 			t = trie.New(account.Root, true)
 		}
-		t.MakeListed(rds.joinGeneration, rds.leftGeneration)
 		rds.storageTries[address] = t
 	}
 	return t, nil
@@ -427,38 +426,8 @@ func (rds *RepairDbState) WriteAccountStorage(address common.Address, key, origi
 	return nil
 }
 
-func (rds *RepairDbState) joinGeneration(gen uint64) {
-	rds.nodeCount++
-	rds.generationCounts[gen]++
-
-}
-
-func (rds *RepairDbState) leftGeneration(gen uint64) {
-	rds.nodeCount--
-	rds.generationCounts[gen]--
-}
-
 func (rds *RepairDbState) PruneTries() {
-	if rds.nodeCount > int(MaxTrieCacheGen) {
-		toRemove := 0
-		excess := rds.nodeCount - int(MaxTrieCacheGen)
-		gen := rds.oldestGeneration
-		for excess > 0 {
-			excess -= rds.generationCounts[gen]
-			toRemove += rds.generationCounts[gen]
-			delete(rds.generationCounts, gen)
-			gen++
-		}
-		// Unload all nodes with touch timestamp < gen
-		for address, storageTrie := range rds.storageTries {
-			empty := storageTrie.UnloadOlderThan(gen, false)
-			if empty {
-				delete(rds.storageTries, address)
-			}
-		}
-		rds.oldestGeneration = gen
-		rds.nodeCount -= toRemove
-	}
+	// TODO Reintroduce pruning if necessary
 	var m runtime.MemStats
 	runtime.ReadMemStats(&m)
 	fmt.Printf("Memory: nodes=%d, alloc=%d, sys=%d\n", rds.nodeCount, int(m.Alloc/1024), int(m.Sys/1024))