From 4d78c8ea658989e3b297d840adddf562ad725a30 Mon Sep 17 00:00:00 2001
From: Jannik Luhn <jannik@brainbot.com>
Date: Tue, 17 Dec 2019 16:57:35 +0700
Subject: [PATCH] Cache recent states and blocks by root

---
 eth2/beacon/db/chain.py | 41 ++++++++++++++++++++++++-----------------
 1 file changed, 24 insertions(+), 17 deletions(-)

diff --git a/eth2/beacon/db/chain.py b/eth2/beacon/db/chain.py
index c66e15607a..2f296a42b3 100644
--- a/eth2/beacon/db/chain.py
+++ b/eth2/beacon/db/chain.py
@@ -1,5 +1,4 @@
 from abc import ABC, abstractmethod
-import functools
 from typing import Iterable, Optional, Tuple, Type, cast
 
 from cytoolz import concat, first, sliding_window
@@ -9,6 +8,7 @@
 from eth_typing import Hash32
 from eth_utils import ValidationError, encode_hex, to_tuple
 import ssz
+from lru import LRU
 
 from eth2.beacon.constants import ZERO_SIGNING_ROOT
 from eth2.beacon.db.exceptions import (
@@ -28,6 +28,14 @@
 from eth2.beacon.typing import Epoch, HashTreeRoot, SigningRoot, Slot
 from eth2.configs import Eth2GenesisConfig
 
+# When performing a chain sync (either fast or regular modes), we'll very often need to look
+# up recent blocks to validate the chain, and decoding their SSZ representation is
+# relatively expensive so we cache that here, but use a small cache because we *should* only
+# be looking up recent blocks. We cache by root instead of ssz representation as ssz
+# representation is not unique if different length configs are considered
+state_cache = LRU(128)
+block_cache = LRU(128)
+
 
 class AttestationKey(ssz.Serializable):
     fields = [("block_root", ssz.sedes.bytes32), ("index", ssz.sedes.uint8)]
@@ -368,13 +376,20 @@ def _get_block_by_root(
         Raise BlockNotFound if it is not present in the db.
         """
         validate_word(block_root, title="block root")
+
+        if block_root in block_cache and block_root in db:
+            return block_cache[block_root]
+
         try:
             block_ssz = db[block_root]
         except KeyError:
             raise BlockNotFound(
                 "No block with signing root {0} found".format(encode_hex(block_root))
             )
-        return _decode_block(block_ssz, block_class)
+
+        block = ssz.decode(block_ssz, block_class)
+        block_cache[block_root] = block
+        return block
 
     def get_slot_by_root(self, block_root: SigningRoot) -> Slot:
         """
@@ -737,11 +752,17 @@ def _get_state_by_root(
         Raises StateNotFound if it is not present in the db.
         """
         # TODO: validate_state_root
+        if state_root in state_cache and state_root in db:
+            return state_cache[state_root]
+
         try:
             state_ssz = db[state_root]
         except KeyError:
             raise StateNotFound(f"No state with root {encode_hex(state_root)} found")
-        return _decode_state(state_ssz, state_class)
+
+        state = ssz.decode(state_ssz, state_class)
+        state_cache[state] = state
+        return state
 
     def persist_state(self, state: BeaconState) -> None:
         """
@@ -911,17 +932,3 @@ def get(self, key: bytes) -> bytes:
         Return the value for the given key or a KeyError if it doesn't exist in the database.
         """
         return self.db[key]
-
-
-# When performing a chain sync (either fast or regular modes), we'll very often need to look
-# up recent blocks to validate the chain, and decoding their SSZ representation is
-# relatively expensive so we cache that here, but use a small cache because we *should* only
-# be looking up recent blocks.
-@functools.lru_cache(128)
-def _decode_block(block_ssz: bytes, sedes: Type[BaseBeaconBlock]) -> BaseBeaconBlock:
-    return ssz.decode(block_ssz, sedes=sedes)
-
-
-@functools.lru_cache(128)
-def _decode_state(state_ssz: bytes, state_class: Type[BeaconState]) -> BeaconState:
-    return ssz.decode(state_ssz, sedes=state_class)