State Trie clearing journal edge case fix (#14)

* Fixed dirty object journal edge case

* Removed duplicate logic

* Fixed mutex locks for tests and fixed variable for dirty objects

* Removed redundant code

* Readded bug fix and added related test

core/state/dump.go

@@ -26,6 +26,7 @@ import (
 	"sync"
 
 	"fmt"
+
 	"github.com/eth-classic/go-ethereum/common"
 	"github.com/eth-classic/go-ethereum/rlp"
 	"github.com/eth-classic/go-ethereum/trie"
@@ -80,7 +81,7 @@ func (self *StateDB) RawDump(addresses []common.Address) Dump {
 			panic(err)
 		}
 
-		obj := newObject(nil, addrA, data, nil)
+		obj := newObject(nil, addrA, data)
 		account := DumpAccount{
 			Balance:  data.Balance.String(),
 			Nonce:    data.Nonce,
@@ -178,7 +179,7 @@ func iterator(sdb *StateDB, addresses []common.Address, c chan *AddressedRawAcco
 			panic(err)
 		}
 
-		obj := newObject(nil, addrA, data, nil)
+		obj := newObject(nil, addrA, data)
 		account := AddressedRawAccount{
 			DumpAccount: DumpAccount{
 				Balance:  data.Balance.String(),

core/state/journal.go

@@ -22,11 +22,67 @@ import (
 	"github.com/eth-classic/go-ethereum/common"
 )
 
+// journalEntry is a modification entry in the state change journal that can be
+// reverted on demand.
 type journalEntry interface {
-	undo(*StateDB)
+	// revert undoes the changes introduced by this journal entry.
+	revert(*StateDB)
+
+	// dirtied returns the Ethereum address modified by this journal entry.
+	dirtied() *common.Address
+}
+
+// journal contains the list of state modifications applied since the last state
+// commit. These are tracked to be able to be reverted in case of an execution
+// exception or revertal request.
+type journal struct {
+	entries []journalEntry         // Current changes tracked by the journal
+	dirties map[common.Address]int // Dirty accounts and the number of changes
+}
+
+// newJournal create a new initialized journal.
+func newJournal() *journal {
+	return &journal{
+		dirties: make(map[common.Address]int),
+	}
+}
+
+// append inserts a new modification entry to the end of the change journal.
+func (j *journal) append(entry journalEntry) {
+	j.entries = append(j.entries, entry)
+	if addr := entry.dirtied(); addr != nil {
+		j.dirties[*addr]++
+	}
+}
+
+// revert undoes a batch of journalled modifications along with any reverted
+// dirty handling too.
+func (j *journal) revert(statedb *StateDB, snapshot int) {
+	for i := len(j.entries) - 1; i >= snapshot; i-- {
+		// Undo the changes made by the operation
+		j.entries[i].revert(statedb)
+
+		// Drop any dirty tracking induced by the change
+		if addr := j.entries[i].dirtied(); addr != nil {
+			if j.dirties[*addr]--; j.dirties[*addr] == 0 {
+				delete(j.dirties, *addr)
+			}
+		}
+	}
+	j.entries = j.entries[:snapshot]
+}
+
+// dirty explicitly sets an address to dirty, even if the change entries would
+// otherwise suggest it as clean. This method is an ugly hack to handle the RIPEMD
+// precompile consensus exception.
+func (j *journal) dirty(addr common.Address) {
+	j.dirties[addr]++
 }
 
-type journal []journalEntry
+// length returns the current number of entries in the journal.
+func (j *journal) length() int {
+	return len(j.entries)
+}
 
 type (
 	// Changes to the account trie.
@@ -77,55 +133,85 @@ type (
 	}
 )
 
-func (ch createObjectChange) undo(s *StateDB) {
+func (ch createObjectChange) revert(s *StateDB) {
 	delete(s.stateObjects, *ch.account)
 	delete(s.stateObjectsDirty, *ch.account)
 }
 
-func (ch resetObjectChange) undo(s *StateDB) {
+func (ch createObjectChange) dirtied() *common.Address {
+	return ch.account
+}
+
+func (ch resetObjectChange) revert(s *StateDB) {
 	s.setStateObject(ch.prev)
 }
 
-func (ch suicideChange) undo(s *StateDB) {
+func (ch resetObjectChange) dirtied() *common.Address {
+	return nil
+}
+
+func (ch suicideChange) revert(s *StateDB) {
 	obj := s.getStateObject(*ch.account)
 	if obj != nil {
 		obj.suicided = ch.prev
 		obj.setBalance(ch.prevbalance)
 	}
 }
 
+func (ch suicideChange) dirtied() *common.Address {
+	return ch.account
+}
+
 var ripemd = common.HexToAddress("0000000000000000000000000000000000000003")
 
-func (ch touchChange) undo(s *StateDB) {
-	if !ch.prev && *ch.account != ripemd {
-		s.getStateObject(*ch.account).touched = ch.prev
-		if !ch.prevDirty {
-			delete(s.stateObjectsDirty, *ch.account)
-		}
-	}
+func (ch touchChange) revert(s *StateDB) {
+}
+
+func (ch touchChange) dirtied() *common.Address {
+	return ch.account
 }
 
-func (ch balanceChange) undo(s *StateDB) {
+func (ch balanceChange) revert(s *StateDB) {
 	s.getStateObject(*ch.account).setBalance(ch.prev)
 }
 
-func (ch nonceChange) undo(s *StateDB) {
+func (ch balanceChange) dirtied() *common.Address {
+	return ch.account
+}
+
+func (ch nonceChange) revert(s *StateDB) {
 	s.getStateObject(*ch.account).setNonce(ch.prev)
 }
 
-func (ch codeChange) undo(s *StateDB) {
+func (ch nonceChange) dirtied() *common.Address {
+	return ch.account
+}
+
+func (ch codeChange) revert(s *StateDB) {
 	s.getStateObject(*ch.account).setCode(common.BytesToHash(ch.prevhash), ch.prevcode)
 }
 
-func (ch storageChange) undo(s *StateDB) {
+func (ch codeChange) dirtied() *common.Address {
+	return ch.account
+}
+
+func (ch storageChange) revert(s *StateDB) {
 	s.getStateObject(*ch.account).setState(ch.key, ch.prevalue)
 }
 
-func (ch refundChange) undo(s *StateDB) {
+func (ch storageChange) dirtied() *common.Address {
+	return ch.account
+}
+
+func (ch refundChange) revert(s *StateDB) {
 	s.refund = ch.prev
 }
 
-func (ch addLogChange) undo(s *StateDB) {
+func (ch refundChange) dirtied() *common.Address {
+	return nil
+}
+
+func (ch addLogChange) revert(s *StateDB) {
 	logs := s.logs[ch.txhash]
 	if len(logs) == 1 {
 		delete(s.logs, ch.txhash)
@@ -135,6 +221,14 @@ func (ch addLogChange) undo(s *StateDB) {
 	s.logSize--
 }
 
-func (ch addPreimageChange) undo(s *StateDB) {
+func (ch addLogChange) dirtied() *common.Address {
+	return nil
+}
+
+func (ch addPreimageChange) revert(s *StateDB) {
 	delete(s.preimages, ch.hash)
 }
+
+func (ch addPreimageChange) dirtied() *common.Address {
+	return nil
+}

core/state/managed_state.go

@@ -84,8 +84,8 @@ func (ms *ManagedState) NewNonce(addr common.Address) uint64 {
 
 // GetNonce returns the canonical nonce for the managed or unmanaged account
 func (ms *ManagedState) GetNonce(addr common.Address) uint64 {
-	ms.mu.RLock()
-	defer ms.mu.RUnlock()
+	ms.mu.Lock()
+	defer ms.mu.Unlock()
 
 	if ms.hasAccount(addr) {
 		account := ms.getAccount(addr)

core/state/state_object.go

@@ -88,9 +88,7 @@ type StateObject struct {
 	// during the "update" phase of the state transition.
 	dirtyCode bool // true if the code was updated
 	suicided  bool
-	touched   bool
 	deleted   bool
-	onDirty   func(addr common.Address) // Callback method to mark a state object newly dirty
 }
 
 // empty returns whether the account is considered empty.
@@ -108,7 +106,7 @@ type Account struct {
 }
 
 // newObject creates a state object.
-func newObject(db *StateDB, address common.Address, data Account, onDirty func(addr common.Address)) *StateObject {
+func newObject(db *StateDB, address common.Address, data Account) *StateObject {
 	if data.Balance == nil {
 		data.Balance = new(big.Int)
 	}
@@ -122,7 +120,6 @@ func newObject(db *StateDB, address common.Address, data Account, onDirty func(a
 		data:          data,
 		cachedStorage: make(Storage),
 		dirtyStorage:  make(Storage),
-		onDirty:       onDirty,
 	}
 }
 
@@ -140,23 +137,17 @@ func (self *StateObject) setError(err error) {
 
 func (self *StateObject) markSuicided() {
 	self.suicided = true
-	if self.onDirty != nil {
-		self.onDirty(self.Address())
-		self.onDirty = nil
-	}
 }
 
 func (c *StateObject) touch() {
-	c.db.journal = append(c.db.journal, touchChange{
-		account:   &c.address,
-		prev:      c.touched,
-		prevDirty: c.onDirty == nil,
+	c.db.journal.append(touchChange{
+		account: &c.address,
 	})
-	if c.onDirty != nil {
-		c.onDirty(c.Address())
-		c.onDirty = nil
+	if c.address == ripemd {
+		// Explicitly put it in the dirty-cache, which is otherwise generated from
+		// flattened journals.
+		c.db.journal.dirty(c.address)
 	}
-	c.touched = true
 }
 
 func (c *StateObject) getTrie(db Database) Trie {
@@ -198,7 +189,7 @@ func (self *StateObject) GetState(db Database, key common.Hash) common.Hash {
 
 // SetState updates a value in account storage.
 func (self *StateObject) SetState(db Database, key, value common.Hash) {
-	self.db.journal = append(self.db.journal, storageChange{
+	self.db.journal.append(storageChange{
 		account:  &self.address,
 		key:      key,
 		prevalue: self.GetState(db, key),
@@ -209,11 +200,6 @@ func (self *StateObject) SetState(db Database, key, value common.Hash) {
 func (self *StateObject) setState(key, value common.Hash) {
 	self.cachedStorage[key] = value
 	self.dirtyStorage[key] = value
-
-	if self.onDirty != nil {
-		self.onDirty(self.Address())
-		self.onDirty = nil
-	}
 }
 
 // updateTrie writes cached storage modifications into the object's storage trie.
@@ -299,7 +285,7 @@ func (c *StateObject) SubBalance(amount *big.Int) {
 }
 
 func (self *StateObject) SetBalance(amount *big.Int) {
-	self.db.journal = append(self.db.journal, balanceChange{
+	self.db.journal.append(balanceChange{
 		account: &self.address,
 		prev:    new(big.Int).Set(self.data.Balance),
 	})
@@ -308,17 +294,13 @@ func (self *StateObject) SetBalance(amount *big.Int) {
 
 func (self *StateObject) setBalance(amount *big.Int) {
 	self.data.Balance = amount
-	if self.onDirty != nil {
-		self.onDirty(self.Address())
-		self.onDirty = nil
-	}
 }
 
 // Return the gas back to the origin. Used by the Virtual machine or Closures
 func (c *StateObject) ReturnGas(*big.Int, *big.Int) {}
 
-func (self *StateObject) deepCopy(db *StateDB, onDirty func(addr common.Address)) *StateObject {
-	stateObject := newObject(db, self.address, self.data, onDirty)
+func (self *StateObject) deepCopy(db *StateDB) *StateObject {
+	stateObject := newObject(db, self.address, self.data)
 	if self.trie != nil {
 		stateObject.trie = db.db.CopyTrie(self.trie)
 	}
@@ -358,7 +340,7 @@ func (self *StateObject) Code(db Database) []byte {
 
 func (self *StateObject) SetCode(codeHash common.Hash, code []byte) {
 	prevcode := self.Code(self.db.db)
-	self.db.journal = append(self.db.journal, codeChange{
+	self.db.journal.append(codeChange{
 		account:  &self.address,
 		prevhash: self.CodeHash(),
 		prevcode: prevcode,
@@ -370,14 +352,10 @@ func (self *StateObject) setCode(codeHash common.Hash, code []byte) {
 	self.code = code
 	self.data.CodeHash = codeHash[:]
 	self.dirtyCode = true
-	if self.onDirty != nil {
-		self.onDirty(self.Address())
-		self.onDirty = nil
-	}
 }
 
 func (self *StateObject) SetNonce(nonce uint64) {
-	self.db.journal = append(self.db.journal, nonceChange{
+	self.db.journal.append(nonceChange{
 		account: &self.address,
 		prev:    self.data.Nonce,
 	})
@@ -386,10 +364,6 @@ func (self *StateObject) SetNonce(nonce uint64) {
 
 func (self *StateObject) setNonce(nonce uint64) {
 	self.data.Nonce = nonce
-	if self.onDirty != nil {
-		self.onDirty(self.Address())
-		self.onDirty = nil
-	}
 }
 
 func (self *StateObject) CodeHash() []byte {