[WIP] Refactor DB persistence logic to work with iterators (#1701)

* Allow persistence of iterable of chain headers

Handle chain validation and setting of scores inside a single loop.

* Add missing typing information to the _set_hash_scores_to_db method

* Move code to persist the header outside the _set_hash_scores_to_db function

* Implement tests for HeaderDB.persist_header_chain to work with iterators

* Fix flake8 errors

* Allow persistence of iterable of chain blocks

Author: onyb

eth/beacon/db/chain.py

@@ -7,6 +7,7 @@
     Type,
 )
 from cytoolz import (
+    concat,
     first,
     sliding_window,
 )
@@ -289,61 +290,79 @@ def persist_block_chain(
         with self.db.atomic_batch() as db:
             return self._persist_block_chain(db, blocks, self.block_class)
 
+    @classmethod
+    def _set_block_scores_to_db(
+            cls,
+            db: BaseDB,
+            block: BaseBeaconBlock
+    ) -> None:
+        # TODO: It's a stub before we implement fork choice rule
+        score = block.slot
+
+        db.set(
+            SchemaV1.make_block_root_to_score_lookup_key(block.root),
+            rlp.encode(score, sedes=rlp.sedes.big_endian_int),
+        )
+
     @classmethod
     def _persist_block_chain(
             cls,
             db: BaseDB,
             blocks: Iterable[BaseBeaconBlock],
             block_class: Type[BaseBeaconBlock]
     ) -> Tuple[Tuple[BaseBeaconBlock, ...], Tuple[BaseBeaconBlock, ...]]:
+        blocks_iterator = iter(blocks)
+
         try:
-            first_block = first(blocks)
+            first_block = first(blocks_iterator)
         except StopIteration:
             return tuple(), tuple()
-        else:
-            for parent, child in sliding_window(2, blocks):
-                if parent.root != child.parent_root:
-                    raise ValidationError(
-                        "Non-contiguous chain. Expected {} to have {} as parent but was {}".format(
-                            encode_hex(child.root),
-                            encode_hex(parent.root),
-                            encode_hex(child.parent_root),
-                        )
-                    )
 
-            is_genesis = first_block.parent_root == GENESIS_PARENT_HASH
-            if not is_genesis and not cls._block_exists(db, first_block.parent_root):
-                raise ParentNotFound(
-                    "Cannot persist block ({}) with unknown parent ({})".format(
-                        encode_hex(first_block.root), encode_hex(first_block.parent_root)))
+        is_genesis = first_block.parent_root == GENESIS_PARENT_HASH
+        if not is_genesis and not cls._block_exists(db, first_block.parent_root):
+            raise ParentNotFound(
+                "Cannot persist block ({}) with unknown parent ({})".format(
+                    encode_hex(first_block.root), encode_hex(first_block.parent_root)))
 
-            if is_genesis:
-                score = 0
-            else:
-                score = cls._get_score(db, first_block.parent_root)
+        if is_genesis:
+            score = 0
+        else:
+            score = cls._get_score(db, first_block.parent_root)
 
-        for block in blocks:
-            db.set(
-                block.root,
-                rlp.encode(block),
-            )
+        curr_block_head = first_block
+        db.set(
+            curr_block_head.root,
+            rlp.encode(curr_block_head),
+        )
+        cls._set_block_scores_to_db(db, curr_block_head)
 
-            # TODO: It's a stub before we implement fork choice rule
-            score = block.slot
+        orig_blocks_seq = concat([(first_block,), blocks_iterator])
+
+        for parent, child in sliding_window(2, orig_blocks_seq):
+            if parent.root != child.parent_root:
+                raise ValidationError(
+                    "Non-contiguous chain. Expected {} to have {} as parent but was {}".format(
+                        encode_hex(child.root),
+                        encode_hex(parent.root),
+                        encode_hex(child.parent_root),
+                    )
+                )
 
+            curr_block_head = child
             db.set(
-                SchemaV1.make_block_root_to_score_lookup_key(block.root),
-                rlp.encode(score, sedes=rlp.sedes.big_endian_int),
+                curr_block_head.root,
+                rlp.encode(curr_block_head),
             )
+            cls._set_block_scores_to_db(db, curr_block_head)
 
         try:
             previous_canonical_head = cls._get_canonical_head(db, block_class).root
             head_score = cls._get_score(db, previous_canonical_head)
         except CanonicalHeadNotFound:
-            return cls._set_as_canonical_chain_head(db, block.root, block_class)
+            return cls._set_as_canonical_chain_head(db, curr_block_head.root, block_class)
 
         if score > head_score:
-            return cls._set_as_canonical_chain_head(db, block.root, block_class)
+            return cls._set_as_canonical_chain_head(db, curr_block_head.root, block_class)
         else:
             return tuple(), tuple()
 

eth/db/header.py

@@ -5,6 +5,7 @@
 import rlp
 
 from cytoolz import (
+    concat,
     first,
     sliding_window,
 )
@@ -204,62 +205,82 @@ def persist_header_chain(self,
         with self.db.atomic_batch() as db:
             return self._persist_header_chain(db, headers)
 
+    @classmethod
+    def _set_hash_scores_to_db(
+            cls,
+            db: BaseDB,
+            header: BlockHeader,
+            score: int
+    ) -> int:
+        new_score = score + header.difficulty
+
+        db.set(
+            SchemaV1.make_block_hash_to_score_lookup_key(header.hash),
+            rlp.encode(new_score, sedes=rlp.sedes.big_endian_int),
+        )
+
+        return new_score
+
     @classmethod
     def _persist_header_chain(
             cls,
             db: BaseDB,
             headers: Iterable[BlockHeader]
     ) -> Tuple[Tuple[BlockHeader, ...], Tuple[BlockHeader, ...]]:
+        headers_iterator = iter(headers)
+
         try:
-            first_header = first(headers)
+            first_header = first(headers_iterator)
         except StopIteration:
             return tuple(), tuple()
-        else:
 
-            for parent, child in sliding_window(2, headers):
-                if parent.hash != child.parent_hash:
-                    raise ValidationError(
-                        "Non-contiguous chain. Expected {} to have {} as parent but was {}".format(
-                            encode_hex(child.hash),
-                            encode_hex(parent.hash),
-                            encode_hex(child.parent_hash),
-                        )
-                    )
+        is_genesis = first_header.parent_hash == GENESIS_PARENT_HASH
+        if not is_genesis and not cls._header_exists(db, first_header.parent_hash):
+            raise ParentNotFound(
+                "Cannot persist block header ({}) with unknown parent ({})".format(
+                    encode_hex(first_header.hash), encode_hex(first_header.parent_hash)))
 
-            is_genesis = first_header.parent_hash == GENESIS_PARENT_HASH
-            if not is_genesis and not cls._header_exists(db, first_header.parent_hash):
-                raise ParentNotFound(
-                    "Cannot persist block header ({}) with unknown parent ({})".format(
-                        encode_hex(first_header.hash), encode_hex(first_header.parent_hash)))
+        if is_genesis:
+            score = 0
+        else:
+            score = cls._get_score(db, first_header.parent_hash)
 
-            if is_genesis:
-                score = 0
-            else:
-                score = cls._get_score(db, first_header.parent_hash)
+        curr_chain_head = first_header
+        db.set(
+            curr_chain_head.hash,
+            rlp.encode(curr_chain_head),
+        )
+        score = cls._set_hash_scores_to_db(db, curr_chain_head, score)
+
+        orig_headers_seq = concat([(first_header,), headers_iterator])
+        for parent, child in sliding_window(2, orig_headers_seq):
+            if parent.hash != child.parent_hash:
+                raise ValidationError(
+                    "Non-contiguous chain. Expected {} to have {} as parent but was {}".format(
+                        encode_hex(child.hash),
+                        encode_hex(parent.hash),
+                        encode_hex(child.parent_hash),
+                    )
+                )
 
-        for header in headers:
+            curr_chain_head = child
             db.set(
-                header.hash,
-                rlp.encode(header),
+                curr_chain_head.hash,
+                rlp.encode(curr_chain_head),
             )
 
-            score += header.difficulty
-
-            db.set(
-                SchemaV1.make_block_hash_to_score_lookup_key(header.hash),
-                rlp.encode(score, sedes=rlp.sedes.big_endian_int),
-            )
+            score = cls._set_hash_scores_to_db(db, curr_chain_head, score)
 
         try:
             previous_canonical_head = cls._get_canonical_head(db).hash
             head_score = cls._get_score(db, previous_canonical_head)
         except CanonicalHeadNotFound:
-            return cls._set_as_canonical_chain_head(db, header.hash)
+            return cls._set_as_canonical_chain_head(db, curr_chain_head.hash)
 
         if score > head_score:
-            return cls._set_as_canonical_chain_head(db, header.hash)
-        else:
-            return tuple(), tuple()
+            return cls._set_as_canonical_chain_head(db, curr_chain_head.hash)
+
+        return tuple(), tuple()
 
     @classmethod
     def _set_as_canonical_chain_head(cls, db: BaseDB, block_hash: Hash32

tests/database/test_header_db.py

@@ -81,6 +81,28 @@ def test_headerdb_get_canonical_head_with_header_chain(headerdb, genesis_header)
     assert_headers_eq(head, headers[-1])
 
 
+@pytest.mark.parametrize(
+    'chain_length',
+    (0, 1, 2, 3),
+)
+def test_headerdb_get_canonical_head_with_header_chain_iterator(
+        headerdb,
+        genesis_header,
+        chain_length):
+
+    headerdb.persist_header(genesis_header)
+    headers = mk_header_chain(genesis_header, length=chain_length)
+
+    headerdb.persist_header_chain(header for header in headers)
+
+    head = headerdb.get_canonical_head()
+
+    if chain_length == 0:
+        assert_headers_eq(head, genesis_header)
+    else:
+        assert_headers_eq(head, headers[-1])
+
+
 def test_headerdb_persist_header_disallows_unknown_parent(headerdb):
     header = BlockHeader(
         difficulty=GENESIS_DIFFICULTY,