Use lru cache (#2433)

* use lru cache

* use new func

* cache block after write is done

* dont iterate as orederd dict

chia/consensus/blockchain.py

@@ -269,6 +269,7 @@ async def receive_block(
                         ] = fetched_block_record.sub_epoch_summary_included
                 if peak_height is not None:
                     self._peak_height = peak_height
+                self.block_store.cache_block(block)
             except BaseException:
                 await self.block_store.db_wrapper.rollback_transaction()
                 raise

chia/full_node/block_store.py

@@ -69,7 +69,7 @@ async def add_full_block(self, block: FullBlock, block_record: BlockRecord) -> N
         if cached is not None:
             # Since write to db can fail, we remove from cache here to avoid potential inconsistency
             # Adding to cache only from reading
-            self.block_cache.put(block.header_hash, None)
+            self.block_cache.remove(block.header_hash)
         cursor_1 = await self.db.execute(
             "INSERT OR REPLACE INTO full_blocks VALUES(?, ?, ?, ?, ?)",
             (
@@ -123,6 +123,9 @@ async def get_sub_epoch_challenge_segments(
             return SubEpochSegments.from_bytes(row[0]).challenge_segments
         return None
 
+    def cache_block(self, block: FullBlock):
+        self.block_cache.put(block.header_hash, block)
+
     async def get_full_block(self, header_hash: bytes32) -> Optional[FullBlock]:
         cached = self.block_cache.get(header_hash)
         if cached is not None:

chia/full_node/coin_store.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Optional
+from typing import List, Optional
 
 import aiosqlite
 
@@ -8,6 +8,7 @@
 from chia.types.full_block import FullBlock
 from chia.util.db_wrapper import DBWrapper
 from chia.util.ints import uint32, uint64
+from chia.util.lru_cache import LRUCache
 
 
 class CoinStore:
@@ -17,7 +18,7 @@ class CoinStore:
     """
 
     coin_record_db: aiosqlite.Connection
-    coin_record_cache: Dict[str, CoinRecord]
+    coin_record_cache: LRUCache
     cache_size: uint32
     db_wrapper: DBWrapper
 
@@ -57,7 +58,7 @@ async def create(cls, db_wrapper: DBWrapper, cache_size: uint32 = uint32(60000))
         await self.coin_record_db.execute("CREATE INDEX IF NOT EXISTS coin_spent on coin_record(puzzle_hash)")
 
         await self.coin_record_db.commit()
-        self.coin_record_cache = dict()
+        self.coin_record_cache = LRUCache(cache_size)
         return self
 
     async def new_block(self, block: FullBlock, additions: List[Coin], removals: List[bytes32]):
@@ -101,8 +102,9 @@ async def new_block(self, block: FullBlock, additions: List[Coin], removals: Lis
 
     # Checks DB and DiffStores for CoinRecord with coin_name and returns it
     async def get_coin_record(self, coin_name: bytes32) -> Optional[CoinRecord]:
-        if coin_name.hex() in self.coin_record_cache:
-            return self.coin_record_cache[coin_name.hex()]
+        cached = self.coin_record_cache.get(coin_name.hex())
+        if cached is not None:
+            return cached
         cursor = await self.coin_record_db.execute("SELECT * from coin_record WHERE coin_name=?", (coin_name.hex(),))
         row = await cursor.fetchone()
         await cursor.close()
@@ -189,7 +191,7 @@ async def rollback_to_block(self, block_index: int):
         """
         # Update memory cache
         delete_queue: bytes32 = []
-        for coin_name, coin_record in self.coin_record_cache.items():
+        for coin_name, coin_record in list(self.coin_record_cache.cache.items()):
             if int(coin_record.spent_block_index) > block_index:
                 new_record = CoinRecord(
                     coin_record.coin,
@@ -199,12 +201,12 @@ async def rollback_to_block(self, block_index: int):
                     coin_record.coinbase,
                     coin_record.timestamp,
                 )
-                self.coin_record_cache[coin_record.coin.name().hex()] = new_record
+                self.coin_record_cache.put(coin_record.coin.name().hex(), new_record)
             if int(coin_record.confirmed_block_index) > block_index:
                 delete_queue.append(coin_name)
 
         for coin_name in delete_queue:
-            del self.coin_record_cache[coin_name]
+            self.coin_record_cache.remove(coin_name)
 
         # Delete from storage
         c1 = await self.coin_record_db.execute("DELETE FROM coin_record WHERE confirmed_index>?", (block_index,))
@@ -232,11 +234,7 @@ async def _add_coin_record(self, record: CoinRecord) -> None:
             ),
         )
         await cursor.close()
-        self.coin_record_cache[record.coin.name().hex()] = record
-        if len(self.coin_record_cache) > self.cache_size:
-            while len(self.coin_record_cache) > self.cache_size:
-                first_in = list(self.coin_record_cache.keys())[0]
-                del self.coin_record_cache[first_in]
+        self.coin_record_cache.put(record.coin.name().hex(), record)
 
     # Update coin_record to be spent in DB
     async def _set_spent(self, coin_name: bytes32, index: uint32):

chia/util/lru_cache.py

@@ -19,3 +19,6 @@ def put(self, key: Any, value: Any) -> None:
         self.cache.move_to_end(key)
         if len(self.cache) > self.capacity:
             self.cache.popitem(last=False)
+
+    def remove(self, key: Any) -> None:
+        self.cache.pop(key)