CoinStore benchmark and performance improvements (#8514)

* add coin_store benchmark

* optimize _set_spent in CoinStore

* simplify _add_coin_record, since we won't use it to replace an entry anymore

* use executemany in add_coin_record and set_spent

* disable sqlite sync

benchmarks/coin_store.py

@@ -0,0 +1,244 @@
+import asyncio
+import random
+from time import time
+from pathlib import Path
+from chia.full_node.coin_store import CoinStore
+from typing import List
+import os
+import sys
+
+import aiosqlite
+from chia.util.db_wrapper import DBWrapper
+from chia.consensus.coinbase import create_farmer_coin, create_pool_coin
+from chia.consensus.default_constants import DEFAULT_CONSTANTS
+from chia.types.blockchain_format.sized_bytes import bytes32
+from chia.types.blockchain_format.coin import Coin
+from chia.util.ints import uint64
+
+
+NUM_ITERS = 200
+
+
+async def setup_db() -> DBWrapper:
+    db_filename = Path("coin-store-benchmark.db")
+    try:
+        os.unlink(db_filename)
+    except FileNotFoundError:
+        pass
+    connection = await aiosqlite.connect(db_filename)
+    await connection.execute("pragma journal_mode=wal")
+    await connection.execute("pragma synchronous=OFF")
+    return DBWrapper(connection)
+
+
+def rand_hash() -> bytes32:
+    return random.randbytes(32)
+
+
+def make_coin() -> Coin:
+    return Coin(rand_hash(), rand_hash(), uint64(1))
+
+
+async def run_new_block_benchmark():
+
+    db_wrapper: DBWrapper = await setup_db()
+
+    try:
+        coin_store = await CoinStore.create(db_wrapper)
+        # farmer puzzle hash
+        ph = bytes32(b"a" * 32)
+
+        all_added: List[bytes32] = []
+
+        block_height = 1
+        timestamp = 1631794488
+
+        print("Building database ", end="")
+        for height in range(block_height, block_height + NUM_ITERS):
+            additions = []
+            removals = []
+
+            # add some new coins
+            for i in range(2000):
+                c = make_coin()
+                additions.append(c)
+                all_added.append(c.get_hash())
+
+            # farm rewards
+            farmer_coin = create_farmer_coin(height, ph, 250000000, DEFAULT_CONSTANTS.GENESIS_CHALLENGE)
+            pool_coin = create_pool_coin(height, ph, 1750000000, DEFAULT_CONSTANTS.GENESIS_CHALLENGE)
+            reward_coins = [pool_coin, farmer_coin]
+            all_added += [pool_coin.name(), farmer_coin.name()]
+
+            # remove some coins we've added previously
+            random.shuffle(all_added)
+            removals = all_added[:100]
+            all_added = all_added[100:]
+
+            await coin_store.new_block(
+                height,
+                timestamp,
+                set(reward_coins),
+                additions,
+                removals,
+            )
+            await db_wrapper.db.commit()
+
+            # 19 seconds per block
+            timestamp += 19
+
+            print(".", end="")
+            sys.stdout.flush()
+        block_height += NUM_ITERS
+
+        total_time = 0
+        total_add = 0
+        total_remove = 0
+        print("\nProfiling mostly additions ", end="")
+        for height in range(block_height, block_height + NUM_ITERS):
+            additions = []
+            removals = []
+
+            # add some new coins
+            for i in range(2000):
+                c = make_coin()
+                additions.append(c)
+                all_added.append(c.get_hash())
+            total_add += 2000
+
+            farmer_coin = create_farmer_coin(height, ph, 250000000, DEFAULT_CONSTANTS.GENESIS_CHALLENGE)
+            pool_coin = create_pool_coin(height, ph, 1750000000, DEFAULT_CONSTANTS.GENESIS_CHALLENGE)
+            reward_coins = [pool_coin, farmer_coin]
+            all_added += [pool_coin.name(), farmer_coin.name()]
+            total_add += 2
+
+            # remove some coins we've added previously
+            random.shuffle(all_added)
+            removals = all_added[:100]
+            all_added = all_added[100:]
+            total_remove += 100
+
+            start = time()
+            await coin_store.new_block(
+                height,
+                timestamp,
+                set(reward_coins),
+                additions,
+                removals,
+            )
+            await db_wrapper.db.commit()
+            stop = time()
+
+            # 19 seconds per block
+            timestamp += 19
+
+            total_time += stop - start
+            print(".", end="")
+            sys.stdout.flush()
+
+        block_height += NUM_ITERS
+
+        print(f"\nMOSTLY ADDITIONS, time: {total_time:0.4f}s additions: {total_add} removals: {total_remove}")
+
+        print("Profiling mostly removals ", end="")
+        total_add = 0
+        total_remove = 0
+        total_time = 0
+        for height in range(block_height, block_height + NUM_ITERS):
+            additions = []
+            removals = []
+
+            # add one new coins
+            c = make_coin()
+            additions.append(c)
+            all_added.append(c.get_hash())
+            total_add += 1
+
+            farmer_coin = create_farmer_coin(height, ph, 250000000, DEFAULT_CONSTANTS.GENESIS_CHALLENGE)
+            pool_coin = create_pool_coin(height, ph, 1750000000, DEFAULT_CONSTANTS.GENESIS_CHALLENGE)
+            reward_coins = [pool_coin, farmer_coin]
+            all_added += [pool_coin.name(), farmer_coin.name()]
+            total_add += 2
+
+            # remove some coins we've added previously
+            random.shuffle(all_added)
+            removals = all_added[:700]
+            all_added = all_added[700:]
+            total_remove += 700
+
+            start = time()
+            await coin_store.new_block(
+                height,
+                timestamp,
+                set(reward_coins),
+                additions,
+                removals,
+            )
+            await db_wrapper.db.commit()
+
+            stop = time()
+
+            # 19 seconds per block
+            timestamp += 19
+
+            total_time += stop - start
+            print(".", end="")
+            sys.stdout.flush()
+
+        block_height += NUM_ITERS
+
+        print(f"\nMOSTLY REMOVALS, time: {total_time:0.4f}s additions: {total_add} removals: {total_remove}")
+
+        print("Profiling full block transactions", end="")
+        total_add = 0
+        total_remove = 0
+        total_time = 0
+        for height in range(block_height, block_height + NUM_ITERS):
+            additions = []
+            removals = []
+
+            # add some new coins
+            for i in range(2000):
+                c = make_coin()
+                additions.append(c)
+                all_added.append(c.get_hash())
+            total_add += 2000
+
+            farmer_coin = create_farmer_coin(height, ph, 250000000, DEFAULT_CONSTANTS.GENESIS_CHALLENGE)
+            pool_coin = create_pool_coin(height, ph, 1750000000, DEFAULT_CONSTANTS.GENESIS_CHALLENGE)
+            reward_coins = [pool_coin, farmer_coin]
+            all_added += [pool_coin.name(), farmer_coin.name()]
+            total_add += 2
+
+            # remove some coins we've added previously
+            random.shuffle(all_added)
+            removals = all_added[:2000]
+            all_added = all_added[2000:]
+            total_remove += 2000
+
+            start = time()
+            await coin_store.new_block(
+                height,
+                timestamp,
+                set(reward_coins),
+                additions,
+                removals,
+            )
+            await db_wrapper.db.commit()
+            stop = time()
+
+            # 19 seconds per block
+            timestamp += 19
+
+            total_time += stop - start
+            print(".", end="")
+            sys.stdout.flush()
+
+        print(f"\nFULLBLOCKS, time: {total_time:0.4f}s additions: {total_add} removals: {total_remove}")
+
+    finally:
+        await db_wrapper.db.close()
+
+
+if __name__ == "__main__":
+    asyncio.run(run_new_block_benchmark())

chia/clvm/spend_sim.py

@@ -129,8 +129,9 @@ class SpendSim:
             uint64(calculate_base_farmer_reward(next_block_height) + fees),
             self.defaults.GENESIS_CHALLENGE,
         )
-        await self.mempool_manager.coin_store._add_coin_record(self.new_coin_record(pool_coin, True), False)
-        await self.mempool_manager.coin_store._add_coin_record(self.new_coin_record(farmer_coin, True), False)
+        await self.mempool_manager.coin_store._add_coin_records(
+            [self.new_coin_record(pool_coin, True), self.new_coin_record(farmer_coin, True)]
+        )
 
         # Coin store gets updated
         generator_bundle: Optional[SpendBundle] = None
@@ -147,10 +148,12 @@ class SpendSim:
                     return_additions = additions
                     return_removals = removals
 
-                for addition in additions:
-                    await self.mempool_manager.coin_store._add_coin_record(self.new_coin_record(addition), False)
-                for removal in removals:
-                    await self.mempool_manager.coin_store._set_spent(removal.name(), uint32(self.block_height + 1))
+                await self.mempool_manager.coin_store._add_coin_records(
+                    [self.new_coin_record(addition) for addition in additions]
+                )
+                await self.mempool_manager.coin_store._set_spent(
+                    [r.name() for r in removals], uint32(self.block_height + 1)
+                )
 
         # SimBlockRecord is created
         generator: Optional[BlockGenerator] = await self.generate_transaction_generator(generator_bundle)

chia/consensus/blockchain.py

@@ -316,7 +316,7 @@ class Blockchain(BlockchainInterface):
                     tx_removals, tx_additions = [], []
                 if block.is_transaction_block():
                     assert block.foliage_transaction_block is not None
-                    added, _ = await self.coin_store.new_block(
+                    added = await self.coin_store.new_block(
                         block.height,
                         block.foliage_transaction_block.timestamp,
                         block.get_included_reward_coins(),
@@ -381,19 +381,25 @@ class Blockchain(BlockchainInterface):
                         tx_removals, tx_additions = await self.get_tx_removals_and_additions(fetched_full_block, None)
                     if fetched_full_block.is_transaction_block():
                         assert fetched_full_block.foliage_transaction_block is not None
-                        removed_rec, added_rec = await self.coin_store.new_block(
+                        added_rec = await self.coin_store.new_block(
                             fetched_full_block.height,
                             fetched_full_block.foliage_transaction_block.timestamp,
                             fetched_full_block.get_included_reward_coins(),
                             tx_additions,
                             tx_removals,
                         )
+                        removed_rec: List[Optional[CoinRecord]] = [
+                            await self.coin_store.get_coin_record(name) for name in tx_removals
+                        ]
 
                         # Set additions first, than removals in order to handle ephemeral coin state
                         # Add in height order is also required
+                        record: Optional[CoinRecord]
                         for record in added_rec:
+                            assert record
                             lastest_coin_state[record.name] = record
                         for record in removed_rec:
+                            assert record
                             lastest_coin_state[record.name] = record
 
             # Changes the peak to be the new peak

chia/full_node/coin_store.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Set, Tuple, Dict
+from typing import List, Optional, Set, Dict
 import aiosqlite
 from chia.protocols.wallet_protocol import CoinState
 from chia.types.blockchain_format.coin import Coin
@@ -72,15 +72,15 @@ class CoinStore:
         included_reward_coins: Set[Coin],
         tx_additions: List[Coin],
         tx_removals: List[bytes32],
-    ) -> Tuple[List[CoinRecord], List[CoinRecord]]:
+    ) -> List[CoinRecord]:
         """
         Only called for blocks which are blocks (and thus have rewards and transactions)
+        Returns a list of the CoinRecords that were added by this block
         """
 
         start = time()
 
-        added_coin_records = []
-        removed_coin_records = []
+        additions = []
 
         for coin in tx_additions:
             record: CoinRecord = CoinRecord(
@@ -91,8 +91,7 @@ class CoinStore:
                 False,
                 timestamp,
             )
-            added_coin_records.append(record)
-            await self._add_coin_record(record, False)
+            additions.append(record)
 
         if height == 0:
             assert len(included_reward_coins) == 0
@@ -108,16 +107,11 @@ class CoinStore:
                 True,
                 timestamp,
             )
-            added_coin_records.append(reward_coin_r)
-            await self._add_coin_record(reward_coin_r, False)
+            additions.append(reward_coin_r)
+
+        await self._add_coin_records(additions)
+        await self._set_spent(tx_removals, height)
 
-        total_amount_spent: int = 0
-        for coin_name in tx_removals:
-            removed_coin_record = await self._set_spent(coin_name, height)
-            total_amount_spent += removed_coin_record.coin.amount
-            removed_coin_records.append(removed_coin_record)
-        # Sanity check, already checked in block_body_validation
-        assert sum([a.amount for a in tx_additions]) <= total_amount_spent
         end = time()
         if end - start > 10:
             log.warning(
@@ -126,7 +120,7 @@ class CoinStore:
                 + "blockchain database is on a fast drive"
             )
 
-        return removed_coin_records, added_coin_records
+        return additions
 
     # Checks DB and DiffStores for CoinRecord with coin_name and returns it
     async def get_coin_record(self, coin_name: bytes32) -> Optional[CoinRecord]:
@@ -394,40 +388,44 @@ class CoinStore:
         return list(coin_changes.values())
 
     # 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())
+    async def _add_coin_records(self, records: List[CoinRecord]) -> None:
 
-        cursor = await self.coin_record_db.execute(
-            f"INSERT {'OR REPLACE ' if allow_replace else ''}INTO coin_record VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?)",
-            (
-                record.coin.name().hex(),
-                record.confirmed_block_index,
-                record.spent_block_index,
-                int(record.spent),
-                int(record.coinbase),
-                str(record.coin.puzzle_hash.hex()),
-                str(record.coin.parent_coin_info.hex()),
-                bytes(record.coin.amount),
-                record.timestamp,
-            ),
+        values = []
+        for record in records:
+            self.coin_record_cache.put(record.coin.name(), record)
+            values.append(
+                (
+                    record.coin.name().hex(),
+                    record.confirmed_block_index,
+                    record.spent_block_index,
+                    int(record.spent),
+                    int(record.coinbase),
+                    str(record.coin.puzzle_hash.hex()),
+                    str(record.coin.parent_coin_info.hex()),
+                    bytes(record.coin.amount),
+                    record.timestamp,
+                )
+            )
+
+        cursor = await self.coin_record_db.executemany(
+            "INSERT INTO coin_record VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?)",
+            values,
         )
         await cursor.close()
 
     # Update coin_record to be spent in DB
-    async def _set_spent(self, coin_name: bytes32, index: uint32) -> CoinRecord:
-        current: Optional[CoinRecord] = await self.get_coin_record(coin_name)
-        if current is None:
-            raise ValueError(f"Cannot spend a coin that does not exist in db: {coin_name}")
+    async def _set_spent(self, coin_names: List[bytes32], index: uint32):
 
-        assert not current.spent  # Redundant sanity check, already checked in block_body_validation
-        spent: CoinRecord = CoinRecord(
-            current.coin,
-            current.confirmed_block_index,
-            index,
-            True,
-            current.coinbase,
-            current.timestamp,
-        )  # type: ignore # noqa
-        await self._add_coin_record(spent, True)
-        return spent
+        # if this coin is in the cache, mark it as spent in there
+        updates = []
+        for coin_name in coin_names:
+            r = self.coin_record_cache.get(coin_name)
+            if r is not None:
+                self.coin_record_cache.put(
+                    r.name, CoinRecord(r.coin, r.confirmed_block_index, index, True, r.coinbase, r.timestamp)
+                )
+            updates.append((index, coin_name.hex()))
+
+        await self.coin_record_db.executemany(
+            "UPDATE OR FAIL coin_record SET spent=1,spent_index=? WHERE coin_name=?", updates
+        )

chia/full_node/full_node.py

@@ -130,7 +130,7 @@ class FullNode:
         # create the store (db) and full node instance
         self.connection = await aiosqlite.connect(self.db_path)
         await self.connection.execute("pragma journal_mode=wal")
-        await self.connection.execute("pragma synchronous=NORMAL")
+        await self.connection.execute("pragma synchronous=OFF")
         if self.config.get("log_sqlite_cmds", False):
             sql_log_path = path_from_root(self.root_path, "log/sql.log")
             self.log.info(f"logging SQL commands to {sql_log_path}")

chia/server/node_discovery.py

@@ -93,7 +93,7 @@ class FullNodeDiscovery:
         mkdir(self.peer_db_path.parent)
         self.connection = await aiosqlite.connect(self.peer_db_path)
         await self.connection.execute("pragma journal_mode=wal")
-        await self.connection.execute("pragma synchronous=NORMAL")
+        await self.connection.execute("pragma synchronous=OFF")
         self.address_manager_store = await AddressManagerStore.create(self.connection)
         if not await self.address_manager_store.is_empty():
             self.address_manager = await self.address_manager_store.deserialize()

chia/wallet/wallet_state_manager.py

@@ -138,7 +138,7 @@ class WalletStateManager:
         self.log.debug(f"Starting in db path: {db_path}")
         self.db_connection = await aiosqlite.connect(db_path)
         await self.db_connection.execute("pragma journal_mode=wal")
-        await self.db_connection.execute("pragma synchronous=NORMAL")
+        await self.db_connection.execute("pragma synchronous=OFF")
 
         self.db_wrapper = DBWrapper(self.db_connection)
         self.coin_store = await WalletCoinStore.create(self.db_wrapper)

tests/core/full_node/test_coin_store.py

@@ -194,10 +194,7 @@ class TestCoinStoreWithBlocks:
                     coins = block.get_included_reward_coins()
                     records = [await coin_store.get_coin_record(coin.name()) for coin in coins]
 
-                    for record in records:
-                        await coin_store._set_spent(record.coin.name(), block.height)
-                        with pytest.raises(AssertionError):
-                            await coin_store._set_spent(record.coin.name(), block.height)
+                    await coin_store._set_spent([r.name for r in records], block.height)
 
                     records = [await coin_store.get_coin_record(coin.name()) for coin in coins]
                     for record in records:
@@ -212,7 +209,7 @@ class TestCoinStoreWithBlocks:
         async with DBConnection() as db_wrapper:
             coin_store = await CoinStore.create(db_wrapper, cache_size=uint32(cache_size))
 
-            records: List[Optional[CoinRecord]] = []
+            records: List[CoinRecord] = []
 
             for block in blocks:
                 if block.is_transaction_block():
@@ -232,9 +229,7 @@ class TestCoinStoreWithBlocks:
                     coins = block.get_included_reward_coins()
                     records = [await coin_store.get_coin_record(coin.name()) for coin in coins]
 
-                    for record in records:
-                        assert record is not None
-                        await coin_store._set_spent(record.coin.name(), block.height)
+                    await coin_store._set_spent([r.name for r in records], block.height)
 
                     records = [await coin_store.get_coin_record(coin.name()) for coin in coins]
                     for record in records: