Add lock, keep cache consistent (#3051)

* execute task decorator

* use blockchain lock

* indentation

* lint

* execute_task

chia/full_node/coin_store.py

@@ -55,7 +55,7 @@ class CoinStore:
 
         await self.coin_record_db.execute("CREATE INDEX IF NOT EXISTS coin_spent on coin_record(spent)")
 
-        await self.coin_record_db.execute("CREATE INDEX IF NOT EXISTS coin_spent on coin_record(puzzle_hash)")
+        await self.coin_record_db.execute("CREATE INDEX IF NOT EXISTS coin_puzzle_hash on coin_record(puzzle_hash)")
 
         await self.coin_record_db.commit()
         self.coin_record_cache = LRUCache(cache_size)
@@ -106,7 +106,7 @@ class CoinStore:
 
     # Checks DB and DiffStores for CoinRecord with coin_name and returns it
     async def get_coin_record(self, coin_name: bytes32) -> Optional[CoinRecord]:
-        cached = self.coin_record_cache.get(coin_name.hex())
+        cached = self.coin_record_cache.get(coin_name)
         if cached is not None:
             return cached
         cursor = await self.coin_record_db.execute("SELECT * from coin_record WHERE coin_name=?", (coin_name.hex(),))
@@ -114,7 +114,9 @@ class CoinStore:
         await cursor.close()
         if row is not None:
             coin = Coin(bytes32(bytes.fromhex(row[6])), bytes32(bytes.fromhex(row[5])), uint64.from_bytes(row[7]))
-            return CoinRecord(coin, row[1], row[2], row[3], row[4], row[8])
+            record = CoinRecord(coin, row[1], row[2], row[3], row[4], row[8])
+            self.coin_record_cache.put(record.coin.name(), record)
+            return record
         return None
 
     async def get_coins_added_at_height(self, height: uint32) -> List[CoinRecord]:
@@ -205,7 +207,7 @@ class CoinStore:
                     coin_record.coinbase,
                     coin_record.timestamp,
                 )
-                self.coin_record_cache.put(coin_record.coin.name().hex(), new_record)
+                self.coin_record_cache.put(coin_record.coin.name(), new_record)
             if int(coin_record.confirmed_block_index) > block_index:
                 delete_queue.append(coin_name)
 
@@ -223,6 +225,9 @@ class CoinStore:
 
     # Store CoinRecord in DB and ram cache
     async def _add_coin_record(self, record: CoinRecord, allow_replace: bool) -> None:
+        if self.coin_record_cache.get(record.coin.name()) is not None:
+            self.coin_record_cache.remove(record.coin.name())
+
         cursor = await self.coin_record_db.execute(
             f"INSERT {'OR REPLACE ' if allow_replace else ''}INTO coin_record VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?)",
             (
@@ -238,7 +243,6 @@ class CoinStore:
             ),
         )
         await cursor.close()
-        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) -> uint64:

chia/full_node/full_node_api.py

@@ -32,7 +32,7 @@ from chia.types.mempool_inclusion_status import MempoolInclusionStatus
 from chia.types.mempool_item import MempoolItem
 from chia.types.peer_info import PeerInfo
 from chia.types.unfinished_block import UnfinishedBlock
-from chia.util.api_decorators import api_request, peer_required, bytes_required
+from chia.util.api_decorators import api_request, peer_required, bytes_required, execute_task
 from chia.util.generator_tools import get_block_header
 from chia.util.hash import std_hash
 from chia.util.ints import uint8, uint32, uint64, uint128
@@ -92,6 +92,7 @@ class FullNodeAPI:
         await peer.close()
         return None
 
+    @execute_task
     @peer_required
     @api_request
     async def new_peak(self, request: full_node_protocol.NewPeak, peer: ws.WSChiaConnection) -> Optional[Message]:

chia/server/server.py

@@ -139,6 +139,8 @@ class ChiaServer:
         self.app_shut_down_task: Optional[asyncio.Task] = None
         self.received_message_callback: Optional[Callable] = None
         self.api_tasks: Dict[bytes32, asyncio.Task] = {}
+        self.execute_tasks: Set[bytes32] = set()
+
         self.tasks_from_peer: Dict[bytes32, Set[bytes32]] = {}
         self.banned_peers: Dict[str, float] = {}
         self.invalid_protocol_ban_seconds = 10
@@ -465,6 +467,8 @@ class ChiaServer:
 
         task_ids = self.tasks_from_peer[peer_id]
         for task_id in task_ids:
+            if task_id in self.execute_tasks:
+                continue
             task = self.api_tasks[task_id]
             task.cancel()
 
@@ -501,6 +505,9 @@ class ChiaServer:
                         if self.api.api_ready is False:
                             return None
 
+                    if hasattr(f, "execute_task"):
+                        self.execute_tasks.add(task_id)
+
                     if hasattr(f, "peer_required"):
                         coroutine = f(full_message.data, connection)
                     else:
@@ -542,6 +549,8 @@ class ChiaServer:
                         self.api_tasks.pop(task_id)
                     if task_id in self.tasks_from_peer[connection.peer_node_id]:
                         self.tasks_from_peer[connection.peer_node_id].remove(task_id)
+                    if task_id in self.execute_tasks:
+                        self.execute_tasks.remove(task_id)
 
             task_id = token_bytes()
             api_task = asyncio.create_task(api_call(payload_inc, connection_inc, task_id))

chia/simulator/full_node_simulator.py

@@ -1,5 +1,3 @@
-import asyncio
-import traceback
 from typing import List, Optional
 
 from chia.consensus.block_record import BlockRecord
@@ -16,7 +14,6 @@ class FullNodeSimulator(FullNodeAPI):
         super().__init__(full_node)
         self.bt = block_tools
         self.full_node = full_node
-        self.lock = asyncio.Lock()
         self.config = full_node.config
         self.time_per_block = None
         if "simulation" in self.config and self.config["simulation"] is True:
@@ -47,8 +44,7 @@ class FullNodeSimulator(FullNodeAPI):
 
     @api_request
     async def farm_new_transaction_block(self, request: FarmNewBlockProtocol):
-        await self.lock.acquire()
-        try:
+        async with self.full_node.blockchain.lock:
             self.log.info("Farming new block!")
             current_blocks = await self.get_all_full_blocks()
             if len(current_blocks) == 0:
@@ -77,16 +73,11 @@ class FullNodeSimulator(FullNodeAPI):
                 previous_generator=self.full_node.full_node_store.previous_generator,
             )
             rr = RespondBlock(more[-1])
-            await self.full_node.respond_block(rr)
-        except Exception as e:
-            error_stack = traceback.format_exc()
-            self.log.error(f"Error while farming block: {error_stack} {e}")
-        finally:
-            self.lock.release()
+        await self.full_node.respond_block(rr)
 
     @api_request
     async def farm_new_block(self, request: FarmNewBlockProtocol):
-        async with self.lock:
+        async with self.full_node.blockchain.lock:
             self.log.info("Farming new block!")
             current_blocks = await self.get_all_full_blocks()
             if len(current_blocks) == 0:
@@ -111,7 +102,7 @@ class FullNodeSimulator(FullNodeAPI):
                 current_time=self.use_current_time,
             )
             rr: RespondBlock = RespondBlock(more[-1])
-            await self.full_node.respond_block(rr)
+        await self.full_node.respond_block(rr)
 
     @api_request
     async def reorg_from_index_to_new_index(self, request: ReorgProtocol):

chia/util/api_decorators.py

@@ -51,3 +51,11 @@ def bytes_required(func):
         return func
 
     return inner()
+
+
+def execute_task(func):
+    def inner():
+        setattr(func, "execute_task", True)
+        return func
+
+    return inner()

chia/wallet/wallet_node_api.py

@@ -2,7 +2,7 @@ from chia.protocols import full_node_protocol, introducer_protocol, wallet_proto
 from chia.server.outbound_message import NodeType
 from chia.server.ws_connection import WSChiaConnection
 from chia.types.mempool_inclusion_status import MempoolInclusionStatus
-from chia.util.api_decorators import api_request, peer_required
+from chia.util.api_decorators import api_request, peer_required, execute_task
 from chia.util.errors import Err
 from chia.wallet.wallet_node import WalletNode
 
@@ -39,6 +39,7 @@ class WalletNodeAPI:
         """
         pass
 
+    @execute_task
     @peer_required
     @api_request
     async def new_peak_wallet(self, peak: wallet_protocol.NewPeakWallet, peer: WSChiaConnection):