1/**
  2 *  @file
  3 *  @copyright defined in aergo/LICENSE.txt
  4 */
  5
  6package trie
  7
  8import (
  9	"bytes"
 10	"fmt"
 11	"sync"
 12	"sync/atomic"
 13
 14	"github.com/aergoio/aergo-lib/db"
 15)
 16
 17// LoadCache loads the first layers of the merkle tree given a root
 18// This is called after a node restarts so that it doesnt become slow with db reads
 19// LoadCache also updates the Root with the given root.
 20func (s *Trie) LoadCache(root []byte) error {
 21	if s.db.Store == nil {
 22		return fmt.Errorf("DB not connected to trie")
 23	}
 24	s.db.liveCache = make(map[Hash][][]byte)
 25	ch := make(chan error, 1)
 26	s.loadCache(root, nil, 0, s.TrieHeight, ch)
 27	s.Root = root
 28	return <-ch
 29}
 30
 31// loadCache loads the first layers of the merkle tree given a root
 32func (s *Trie) loadCache(root []byte, batch [][]byte, iBatch, height int, ch chan<- (error)) {
 33	if height < s.CacheHeightLimit || len(root) == 0 {
 34		ch <- nil
 35		return
 36	}
 37	if height%4 == 0 {
 38		// Load the node from db
 39		s.db.lock.Lock()
 40		dbval := s.db.Store.Get(root[:HashLength])
 41		s.db.lock.Unlock()
 42		if len(dbval) == 0 {
 43			ch <- fmt.Errorf("the trie node %x is unavailable in the disk db, db may be corrupted", root)
 44			return
 45		}
 46		//Store node in cache.
 47		var node Hash
 48		copy(node[:], root)
 49		batch = s.parseBatch(dbval)
 50		s.db.liveMux.Lock()
 51		s.db.liveCache[node] = batch
 52		s.db.liveMux.Unlock()
 53		iBatch = 0
 54		if batch[0][0] == 1 {
 55			// if height == 0 this will also return
 56			ch <- nil
 57			return
 58		}
 59	}
 60	if iBatch != 0 && batch[iBatch][HashLength] == 1 {
 61		// Check if node is a leaf node
 62		ch <- nil
 63	} else {
 64		// Load subtree
 65		lnode, rnode := batch[2*iBatch+1], batch[2*iBatch+2]
 66
 67		lch := make(chan error, 1)
 68		rch := make(chan error, 1)
 69		go s.loadCache(lnode, batch, 2*iBatch+1, height-1, lch)
 70		go s.loadCache(rnode, batch, 2*iBatch+2, height-1, rch)
 71		if err := <-lch; err != nil {
 72			ch <- err
 73			return
 74		}
 75		if err := <-rch; err != nil {
 76			ch <- err
 77			return
 78		}
 79		ch <- nil
 80	}
 81}
 82
 83// Get fetches the value of a key by going down the current trie root.
 84func (s *Trie) Get(key []byte) ([]byte, error) {
 85	s.lock.RLock()
 86	defer s.lock.RUnlock()
 87	s.atomicUpdate = false
 88	return s.get(s.Root, key, nil, 0, s.TrieHeight)
 89}
 90
 91// GetWithRoot fetches the value of a key by going down for the specified root.
 92func (s *Trie) GetWithRoot(key []byte, root []byte) ([]byte, error) {
 93	s.lock.RLock()
 94	defer s.lock.RUnlock()
 95	s.atomicUpdate = false
 96	if root == nil {
 97		root = s.Root
 98	}
 99	return s.get(root, key, nil, 0, s.TrieHeight)
100}
101
102// get fetches the value of a key given a trie root
103func (s *Trie) get(root, key []byte, batch [][]byte, iBatch, height int) ([]byte, error) {
104	if len(root) == 0 {
105		// the trie does not contain the key
106		return nil, nil
107	}
108	// Fetch the children of the node
109	batch, iBatch, lnode, rnode, isShortcut, err := s.loadChildren(root, height, iBatch, batch)
110	if err != nil {
111		return nil, err
112	}
113	if isShortcut {
114		if bytes.Equal(lnode[:HashLength], key) {
115			return rnode[:HashLength], nil
116		}
117		// also returns nil if height 0 is not a shortcut
118		return nil, nil
119	}
120	if bitIsSet(key, s.TrieHeight-height) {
121		return s.get(rnode, key, batch, 2*iBatch+2, height-1)
122	}
123	return s.get(lnode, key, batch, 2*iBatch+1, height-1)
124}
125
126// WalkResult contains the key and value obtained with a Walk() operation
127type WalkResult struct {
128	Value []byte
129	Key   []byte
130}
131
132// Walk finds all the trie stored values from left to right and calls callback.
133// If callback returns a number diferent from 0, the walk will stop, else it will continue.
134func (s *Trie) Walk(root []byte, callback func(*WalkResult) int32) error {
135	walkc := make(chan *WalkResult)
136	s.lock.RLock()
137	defer s.lock.RUnlock()
138	if root == nil {
139		root = s.Root
140	}
141	s.atomicUpdate = false
142	finishedWalk := make(chan (bool), 1)
143	stop := int32(0)
144	wg := sync.WaitGroup{} // WaitGroup to avoid Walk() return before all callback executions are finished.
145	go func() {
146		for {
147			select {
148			case <-finishedWalk:
149				return
150			case value := <-walkc:
151				stopCallback := callback(value)
152				wg.Done()
153				// In order to avoid data races we need to check the current value of stop, while at the
154				// same time we store our callback value. If our callback value is 0 means that we have
155				// override the previous non-zero value, so we need to restore it.
156				if cv := atomic.SwapInt32(&stop, stopCallback); cv != 0 || stopCallback != 0 {
157					if stopCallback == 0 {
158						atomic.StoreInt32(&stop, cv)
159					}
160					// We need to return (instead of break) in order to stop iterating if some callback returns non zero
161					return
162				}
163			}
164		}
165	}()
166	err := s.walk(walkc, &stop, root, nil, 0, s.TrieHeight, &wg)
167	finishedWalk <- true
168	wg.Wait()
169	return err
170}
171
172// walk fetches the value of a key given a trie root
173func (s *Trie) walk(walkc chan (*WalkResult), stop *int32, root []byte, batch [][]byte, ibatch, height int, wg *sync.WaitGroup) error {
174	if len(root) == 0 || atomic.LoadInt32(stop) != 0 {
175		// The sub tree is empty or stop walking
176		return nil
177	}
178	// Fetch the children of the node
179	batch, ibatch, lnode, rnode, isShortcut, err := s.loadChildren(root, height, ibatch, batch)
180	if err != nil {
181		return err
182	}
183	if isShortcut {
184		wg.Add(1)
185		walkc <- &WalkResult{Value: rnode[:HashLength], Key: lnode[:HashLength]}
186		return nil
187	}
188	// Go left
189	if err := s.walk(walkc, stop, lnode, batch, 2*ibatch+1, height-1, wg); err != nil {
190		return err
191	}
192	// Go Right
193	if err := s.walk(walkc, stop, rnode, batch, 2*ibatch+2, height-1, wg); err != nil {
194		return err
195	}
196	return nil
197}
198
199// TrieRootExists returns true if the root exists in Database.
200func (s *Trie) TrieRootExists(root []byte) bool {
201	s.db.lock.RLock()
202	dbval := s.db.Store.Get(root)
203	s.db.lock.RUnlock()
204	if len(dbval) != 0 {
205		return true
206	}
207	return false
208}
209
210// Commit stores the updated nodes to disk.
211// Commit should be called for every block otherwise past tries
212// are not recorded and it is not possible to revert to them
213// (except if AtomicUpdate is used, which records every state).
214func (s *Trie) Commit() error {
215	if s.db.Store == nil {
216		return fmt.Errorf("DB not connected to trie")
217	}
218	// NOTE The tx interface doesnt handle ErrTxnTooBig
219	txn := s.db.Store.NewTx().(DbTx)
220	s.StageUpdates(txn)
221	txn.(db.Transaction).Commit()
222	return nil
223}
224
225// StageUpdates requires a database transaction as input
226// Unlike Commit(), it doesnt commit the transaction
227// the database transaction MUST be commited otherwise the
228// state ROOT will not exist.
229func (s *Trie) StageUpdates(txn DbTx) {
230	s.lock.Lock()
231	defer s.lock.Unlock()
232	// Commit the new nodes to database, clear updatedNodes and store the Root in pastTries for reverts.
233	if !s.atomicUpdate {
234		// if previously AtomicUpdate was called, then past tries is already updated
235		s.updatePastTries()
236	}
237	s.db.commit(&txn)
238
239	s.db.updatedNodes = make(map[Hash][][]byte)
240	s.prevRoot = s.Root
241}
242
243// Stash rolls back the changes made by previous updates
244// and loads the cache from before the rollback.
245func (s *Trie) Stash(rollbackCache bool) error {
246	s.lock.Lock()
247	defer s.lock.Unlock()
248	s.Root = s.prevRoot
249	if rollbackCache {
250		// Making a temporary liveCache requires it to be copied, so it's quicker
251		// to just load the cache from DB if a block state root was incorrect.
252		s.db.liveCache = make(map[Hash][][]byte)
253		ch := make(chan error, 1)
254		s.loadCache(s.Root, nil, 0, s.TrieHeight, ch)
255		err := <-ch
256		if err != nil {
257			return err
258		}
259	} else {
260		s.db.liveCache = make(map[Hash][][]byte)
261	}
262	s.db.updatedNodes = make(map[Hash][][]byte)
263	// also stash past tries created by Atomic update
264	for i := len(s.pastTries) - 1; i >= 0; i-- {
265		if bytes.Equal(s.pastTries[i], s.Root) {
266			break
267		} else {
268			// remove from past tries
269			s.pastTries = s.pastTries[:len(s.pastTries)-1]
270		}
271	}
272	return nil
273}