optimize snapshot create and load

Author: Jan Xie

ethereum/pruning_trie.py

@@ -884,6 +884,49 @@ def to_dict(self):
             res[key] = value
         return res
 
+    def iter_branch(self):
+        for key_str, value in self._iter_branch(self.root_node):
+            if key_str:
+                nibbles = [int(x) for x in key_str.split(b'+')]
+            else:
+                nibbles = []
+            key = nibbles_to_bin(without_terminator(nibbles))
+            yield key, value
+
+    def _iter_branch(self, node):
+        '''yield (key, value) stored in this and the descendant nodes
+        :param node: node in form of list, or BLANK_NODE
+
+        .. note::
+            Here key is in full form, rather than key of the individual node
+        '''
+        if node == BLANK_NODE:
+            raise StopIteration
+
+        node_type = self._get_node_type(node)
+
+        if is_key_value_type(node_type):
+            nibbles = without_terminator(unpack_to_nibbles(node[0]))
+            key = b'+'.join([to_string(x) for x in nibbles])
+            if node_type == NODE_TYPE_EXTENSION:
+                sub_tree = self._iter_branch(self._decode_to_node(node[1]))
+            else:
+                sub_tree = [(to_string(NIBBLE_TERMINATOR), node[1])]
+
+            # prepend key of this node to the keys of children
+            for sub_key, sub_value in sub_tree:
+                full_key = (key + b'+' + sub_key).strip(b'+')
+                yield (full_key, sub_value)
+
+        elif node_type == NODE_TYPE_BRANCH:
+            for i in range(16):
+                sub_tree = self._iter_branch(self._decode_to_node(node[i]))
+                for sub_key, sub_value in sub_tree:
+                    full_key = (str_to_bytes(str(i)) + b'+' + sub_key).strip(b'+')
+                    yield (full_key, sub_value)
+            if node[16]:
+                yield (to_string(NIBBLE_TERMINATOR), node[-1])
+
     def get(self, key):
         return self._get(self.root_node, bin_to_nibbles(to_string(key)))
 

ethereum/snapshot.py

@@ -3,7 +3,8 @@
 from ethereum.blocks import Account, BlockHeader, Block, CachedBlock
 from ethereum.utils import is_numeric, is_string, encode_hex, decode_hex, zpad, scan_bin, big_endian_to_int
 from ethereum.securetrie import SecureTrie
-from ethereum.trie import Trie, BLANK_NODE, BLANK_ROOT
+from ethereum.trie import BLANK_NODE, BLANK_ROOT
+from ethereum.pruning_trie import Trie
 
 
 class FakeHeader(object):
@@ -68,18 +69,18 @@ def create_base_snapshot(base):
 def create_state_snapshot(env, state_trie):
     alloc = dict()
     count = 0
-    for addr, account_rlp in state_trie.to_dict().items():
+    for addr, account_rlp in state_trie.iter_branch():
         alloc[encode_hex(addr)] = create_account_snapshot(env, account_rlp)
         count += 1
-        print "[%d] created account snapshot %s" % encode_hex(addr)
+        print "[%d] created account snapshot %s" % (count, encode_hex(addr))
     return alloc
 
 
 def create_account_snapshot(env, rlpdata):
     account = get_account(env, rlpdata)
     storage_trie = SecureTrie(Trie(env.db, account.storage))
     storage = dict()
-    for k, v in storage_trie.to_dict().items():
+    for k, v in storage_trie.iter_branch():
         storage[encode_hex(k.lstrip('\x00') or '\x00')] = encode_hex(v)
     return {
         'nonce': snapshot_form(account.nonce),
@@ -132,6 +133,7 @@ def get_block(env, blockhash):
     def validate_uncles():
         return True
 
+    print "Start loading recent blocks from snapshot"
     first_block = rlp.decode(first_block_rlp, Block, env=chain.env)
     chain.index.add_block(first_block)
     chain._store_block(first_block)
@@ -157,11 +159,15 @@ def validate_uncles():
 def load_state(env, alloc):
     db = env.db
     state = SecureTrie(Trie(db, BLANK_ROOT))
-    for addr, account in alloc.items():
+    count = 0
+    print "Start loading state from snapshot"
+    for addr in alloc:
+        account = alloc[addr]
         acct = Account.blank_account(db, env.config['ACCOUNT_INITIAL_NONCE'])
         if len(account['storage']) > 0:
             t = SecureTrie(Trie(db, BLANK_ROOT))
-            for k, v in account['storage'].items():
+            for k in account['storage']:
+                v = account['storage'][k]
                 enckey = zpad(decode_hex(k), 32)
                 t.update(enckey, decode_hex(v))
             acct.storage = t.root_hash
@@ -172,6 +178,10 @@ def load_state(env, alloc):
         if account['code']:
             acct.code = decode_hex(account['code'])
         state.update(decode_hex(addr), rlp.encode(acct))
+        count += 1
+        if count % 1000 == 0:
+            db.commit()
+        print "[%d] loaded account %s" % (count, addr)
     db.commit()
     return state