added caching for account/wallet balance and removed --reserved_subtotals argument instead always returning the subtotals

lbry/lbry/extras/daemon/Daemon.py

@@ -1159,27 +1159,24 @@ class Daemon(metaclass=JSONRPCServerType):
         return wallet
 
     @requires("wallet")
-    async def jsonrpc_wallet_balance(self, wallet_id=None, confirmations=0, reserved_subtotals=False):
+    async def jsonrpc_wallet_balance(self, wallet_id=None, confirmations=0):
         """
         Return the balance of a wallet
 
         Usage:
-            wallet_balance [--wallet_id=<wallet_id>] [--confirmations=<confirmations>] [--reserved_subtotals]
+            wallet_balance [--wallet_id=<wallet_id>] [--confirmations=<confirmations>]
 
         Options:
             --wallet_id=<wallet_id>         : (str) balance for specific wallet
             --confirmations=<confirmations> : (int) Only include transactions with this many
                                               confirmed blocks.
-            --reserved_subtotals            : (bool) Include detailed reserved balances on
-                                              claims, tips and supports.
 
         Returns:
             (decimal) amount of lbry credits in wallet
         """
         wallet = self.wallet_manager.get_wallet_or_default(wallet_id)
         balance = await self.ledger.get_detailed_balance(
-            accounts=wallet.accounts, confirmations=confirmations,
+            accounts=wallet.accounts, confirmations=confirmations
-            reserved_subtotals=reserved_subtotals
         )
         return dict_values_to_lbc(balance)
 
@@ -1362,13 +1359,13 @@ class Daemon(metaclass=JSONRPCServerType):
             return paginate_list(await wallet.get_detailed_accounts(**kwargs), page, page_size)
 
     @requires("wallet")
-    async def jsonrpc_account_balance(self, account_id=None, wallet_id=None, confirmations=0, reserved_subtotals=False):
+    async def jsonrpc_account_balance(self, account_id=None, wallet_id=None, confirmations=0):
         """
         Return the balance of an account
 
         Usage:
             account_balance [<account_id>] [<address> | --address=<address>] [--wallet_id=<wallet_id>]
-                            [<confirmations> | --confirmations=<confirmations>] [--reserved_subtotals]
+                            [<confirmations> | --confirmations=<confirmations>]
 
         Options:
             --account_id=<account_id>       : (str) If provided only the balance for this
@@ -1376,8 +1373,6 @@ class Daemon(metaclass=JSONRPCServerType):
             --wallet_id=<wallet_id>         : (str) balance for specific wallet
             --confirmations=<confirmations> : (int) Only include transactions with this many
                                               confirmed blocks.
-            --reserved_subtotals            : (bool) Include detailed reserved balances on
-                                              claims, tips and supports.
 
         Returns:
             (decimal) amount of lbry credits in wallet
@@ -1385,7 +1380,7 @@ class Daemon(metaclass=JSONRPCServerType):
         wallet = self.wallet_manager.get_wallet_or_default(wallet_id)
         account = wallet.get_account_or_default(account_id)
         balance = await account.get_detailed_balance(
-            confirmations=confirmations, reserved_subtotals=reserved_subtotals
+            confirmations=confirmations, reserved_subtotals=True
         )
         return dict_values_to_lbc(balance)
 

lbry/lbry/wallet/ledger.py

@@ -5,7 +5,8 @@ from functools import partial
 from typing import Tuple, List
 from datetime import datetime
 
-from torba.client.baseledger import BaseLedger
+import pylru
+from torba.client.baseledger import BaseLedger, TransactionEvent
 from torba.client.baseaccount import SingleKey
 from lbry.schema.result import Outputs
 from lbry.schema.url import URL
@@ -52,6 +53,7 @@ class MainNetLedger(BaseLedger):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.fee_per_name_char = self.config.get('fee_per_name_char', self.default_fee_per_name_char)
+        self._balance_cache = pylru.lrucache(100000)
 
     async def _inflate_outputs(self, query, accounts):
         outputs = Outputs.from_base64(await query)
@@ -105,6 +107,7 @@ class MainNetLedger(BaseLedger):
         await asyncio.gather(*(a.maybe_migrate_certificates() for a in self.accounts))
         await asyncio.gather(*(a.save_max_gap() for a in self.accounts))
         await self._report_state()
+        self.on_transaction.listen(self._reset_balance_cache)
 
     async def _report_state(self):
         try:
@@ -128,6 +131,14 @@ class MainNetLedger(BaseLedger):
                 'Failed to display wallet state, please file issue '
                 'for this bug along with the traceback you see below:')
 
+    async def _reset_balance_cache(self, e: TransactionEvent):
+        account_ids = [
+            r['account'] for r in await self.db.get_addresses(('account',), address=e.address)
+        ]
+        for account_id in account_ids:
+            if account_id in self._balance_cache:
+                del self._balance_cache[account_id]
+
     @staticmethod
     def constraint_spending_utxos(constraints):
         constraints['txo_type'] = 0
@@ -288,17 +299,18 @@ class MainNetLedger(BaseLedger):
     def get_transaction_history_count(self, **constraints):
         return self.db.get_transaction_count(**constraints)
 
-    @staticmethod
+    async def get_detailed_balance(self, accounts, confirmations=0):
-    async def get_detailed_balance(accounts, confirmations=0, reserved_subtotals=False):
         result = {}
         for account in accounts:
-            balance = await account.get_detailed_balance(confirmations, reserved_subtotals)
+            balance = self._balance_cache.get(account.id)
+            if not balance:
+                balance = self._balance_cache[account.id] =\
+                    await account.get_detailed_balance(confirmations, reserved_subtotals=True)
             if result:
                 for key, value in balance.items():
                     if key == 'reserved_subtotals':
-                        if value is not None:
+                        for subkey, subvalue in value.items():
-                            for subkey, subvalue in value.items():
+                            result['reserved_subtotals'][subkey] += subvalue
-                                result['reserved_subtotals'][subkey] += subvalue
                     else:
                         result[key] += value
             else:

lbry/tests/integration/test_transaction_commands.py

@@ -36,112 +36,3 @@ class TransactionCommandsTestCase(CommandTestCase):
         await self.assertBalance(self.account, '0.0')
         await self.daemon.jsonrpc_utxo_release()
         await self.assertBalance(self.account, '11.0')
-
-    async def test_granular_balances(self):
-        account2 = await self.daemon.jsonrpc_account_create("Tip-er")
-
-        account_balance = self.daemon.jsonrpc_account_balance
-        wallet_balance = self.daemon.jsonrpc_wallet_balance
-
-        expected = {
-            'total': '10.0',
-            'available': '10.0',
-            'reserved': '0.0',
-            'reserved_subtotals': None
-        }
-        self.assertEqual(await account_balance(reserved_subtotals=False), expected)
-        self.assertEqual(await wallet_balance(reserved_subtotals=False), expected)
-
-        expected = {
-            'total': '10.0',
-            'available': '10.0',
-            'reserved': '0.0',
-            'reserved_subtotals': {'claims': '0.0', 'supports': '0.0', 'tips': '0.0'}
-        }
-        self.assertEqual(await account_balance(reserved_subtotals=True), expected)
-        self.assertEqual(await wallet_balance(reserved_subtotals=True), expected)
-
-        # claim with update + supporting our own claim
-        stream1 = await self.stream_create('granularity', '3.0')
-        await self.stream_update(self.get_claim_id(stream1), data=b'news', bid='1.0')
-        await self.support_create(self.get_claim_id(stream1), '2.0')
-        expected = {
-            'total': '9.977534',
-            'available': '6.977534',
-            'reserved': '3.0',
-            'reserved_subtotals': {'claims': '1.0', 'supports': '2.0', 'tips': '0.0'}
-        }
-        self.assertEqual(await account_balance(reserved_subtotals=True), expected)
-        self.assertEqual(await wallet_balance(reserved_subtotals=True), expected)
-
-        address2 = await self.daemon.jsonrpc_address_unused(account2.id)
-
-        # send lbc to someone else
-        tx = await self.daemon.jsonrpc_account_send('1.0', address2)
-        await self.confirm_tx(tx.id)
-        self.assertEqual(await account_balance(reserved_subtotals=True), {
-            'total': '8.97741',
-            'available': '5.97741',
-            'reserved': '3.0',
-            'reserved_subtotals': {'claims': '1.0', 'supports': '2.0', 'tips': '0.0'}
-        })
-        self.assertEqual(await wallet_balance(reserved_subtotals=True), {
-            'total': '9.97741',
-            'available': '6.97741',
-            'reserved': '3.0',
-            'reserved_subtotals': {'claims': '1.0', 'supports': '2.0', 'tips': '0.0'}
-        })
-
-        # tip received
-        support1 = await self.support_create(
-            self.get_claim_id(stream1), '0.3', tip=True, funding_account_ids=[account2.id]
-        )
-        self.assertEqual(await account_balance(reserved_subtotals=True), {
-            'total': '9.27741',
-            'available': '5.97741',
-            'reserved': '3.3',
-            'reserved_subtotals': {'claims': '1.0', 'supports': '2.0', 'tips': '0.3'}
-        })
-        self.assertEqual(await wallet_balance(reserved_subtotals=True), {
-            'total': '9.977268',
-            'available': '6.677268',
-            'reserved': '3.3',
-            'reserved_subtotals': {'claims': '1.0', 'supports': '2.0', 'tips': '0.3'}
-        })
-
-        # tip claimed
-        tx = await self.daemon.jsonrpc_support_abandon(txid=support1['txid'], nout=0)
-        await self.confirm_tx(tx.id)
-        self.assertEqual(await account_balance(reserved_subtotals=True), {
-            'total': '9.277303',
-            'available': '6.277303',
-            'reserved': '3.0',
-            'reserved_subtotals': {'claims': '1.0', 'supports': '2.0', 'tips': '0.0'}
-        })
-        self.assertEqual(await wallet_balance(reserved_subtotals=True), {
-            'total': '9.977161',
-            'available': '6.977161',
-            'reserved': '3.0',
-            'reserved_subtotals': {'claims': '1.0', 'supports': '2.0', 'tips': '0.0'}
-        })
-
-        stream2 = await self.stream_create(
-            'granularity-is-cool', '0.1', account_id=account2.id, funding_account_ids=[account2.id]
-        )
-
-        # tip another claim
-        await self.support_create(
-            self.get_claim_id(stream2), '0.2', tip=True, funding_account_ids=[self.account.id]
-        )
-        self.assertEqual(await account_balance(reserved_subtotals=True), {
-            'total': '9.077157',
-            'available': '6.077157',
-            'reserved': '3.0',
-            'reserved_subtotals': {'claims': '1.0', 'supports': '2.0', 'tips': '0.0'}
-        })
-        self.assertEqual(await wallet_balance(reserved_subtotals=True), {
-            'total': '9.938908',
-            'available': '6.638908',
-            'reserved': '3.3',
-            'reserved_subtotals': {'claims': '1.1', 'supports': '2.0', 'tips': '0.2'}
-        })

lbry/tests/integration/test_wallet_commands.py

@@ -16,6 +16,142 @@ class WalletCommands(CommandTestCase):
         await self.daemon.jsonrpc_wallet_add(wallet.id)
         self.assertEqual(len(session.hashX_subs), 28)
 
+    async def test_balance_caching(self):
+        self.merchant_address = await self.blockchain.get_raw_change_address()
+        wallet_balance = self.daemon.jsonrpc_wallet_balance
+        ledger = self.ledger
+        query_count = self.ledger.db.db.query_count
+
+        expected = {
+            'total': '10.0',
+            'available': '10.0',
+            'reserved': '0.0',
+            'reserved_subtotals': {'claims': '0.0', 'supports': '0.0', 'tips': '0.0'}
+        }
+        self.assertIsNone(ledger._balance_cache.get(self.account.id))
+
+        query_count += 3
+        self.assertEqual(await wallet_balance(), expected)
+        self.assertEqual(self.ledger.db.db.query_count, query_count)
+        self.assertEqual(ledger._balance_cache.get(self.account.id), expected)
+
+        # calling again uses cache
+        self.assertEqual(await wallet_balance(), expected)
+        self.assertEqual(self.ledger.db.db.query_count, query_count)
+        self.assertEqual(ledger._balance_cache.get(self.account.id), expected)
+
+        await self.stream_create()
+
+        expected = {
+            'total': '9.979893',
+            'available': '8.979893',
+            'reserved': '1.0',
+            'reserved_subtotals': {'claims': '1.0', 'supports': '0.0', 'tips': '0.0'}
+        }
+        # on_transaction event reset balance cache
+        self.assertEqual(await wallet_balance(), expected)
+        self.assertEqual(ledger._balance_cache.get(self.account.id), expected)
+
+    async def test_granular_balances(self):
+        account2 = await self.daemon.jsonrpc_account_create("Tip-er")
+
+        account_balance = self.daemon.jsonrpc_account_balance
+        wallet_balance = self.daemon.jsonrpc_wallet_balance
+
+        expected = {
+            'total': '10.0',
+            'available': '10.0',
+            'reserved': '0.0',
+            'reserved_subtotals': {'claims': '0.0', 'supports': '0.0', 'tips': '0.0'}
+        }
+        self.assertEqual(await account_balance(), expected)
+        self.assertEqual(await wallet_balance(), expected)
+
+        # claim with update + supporting our own claim
+        stream1 = await self.stream_create('granularity', '3.0')
+        await self.stream_update(self.get_claim_id(stream1), data=b'news', bid='1.0')
+        await self.support_create(self.get_claim_id(stream1), '2.0')
+        expected = {
+            'total': '9.977534',
+            'available': '6.977534',
+            'reserved': '3.0',
+            'reserved_subtotals': {'claims': '1.0', 'supports': '2.0', 'tips': '0.0'}
+        }
+        self.assertEqual(await account_balance(), expected)
+        self.assertEqual(await wallet_balance(), expected)
+
+        address2 = await self.daemon.jsonrpc_address_unused(account2.id)
+
+        # send lbc to someone else
+        tx = await self.daemon.jsonrpc_account_send('1.0', address2)
+        await self.confirm_tx(tx.id)
+        self.assertEqual(await account_balance(), {
+            'total': '8.97741',
+            'available': '5.97741',
+            'reserved': '3.0',
+            'reserved_subtotals': {'claims': '1.0', 'supports': '2.0', 'tips': '0.0'}
+        })
+        self.assertEqual(await wallet_balance(), {
+            'total': '9.97741',
+            'available': '6.97741',
+            'reserved': '3.0',
+            'reserved_subtotals': {'claims': '1.0', 'supports': '2.0', 'tips': '0.0'}
+        })
+
+        # tip received
+        support1 = await self.support_create(
+            self.get_claim_id(stream1), '0.3', tip=True, funding_account_ids=[account2.id]
+        )
+        self.assertEqual(await account_balance(), {
+            'total': '9.27741',
+            'available': '5.97741',
+            'reserved': '3.3',
+            'reserved_subtotals': {'claims': '1.0', 'supports': '2.0', 'tips': '0.3'}
+        })
+        self.assertEqual(await wallet_balance(), {
+            'total': '9.977268',
+            'available': '6.677268',
+            'reserved': '3.3',
+            'reserved_subtotals': {'claims': '1.0', 'supports': '2.0', 'tips': '0.3'}
+        })
+
+        # tip claimed
+        tx = await self.daemon.jsonrpc_support_abandon(txid=support1['txid'], nout=0)
+        await self.confirm_tx(tx.id)
+        self.assertEqual(await account_balance(), {
+            'total': '9.277303',
+            'available': '6.277303',
+            'reserved': '3.0',
+            'reserved_subtotals': {'claims': '1.0', 'supports': '2.0', 'tips': '0.0'}
+        })
+        self.assertEqual(await wallet_balance(), {
+            'total': '9.977161',
+            'available': '6.977161',
+            'reserved': '3.0',
+            'reserved_subtotals': {'claims': '1.0', 'supports': '2.0', 'tips': '0.0'}
+        })
+
+        stream2 = await self.stream_create(
+            'granularity-is-cool', '0.1', account_id=account2.id, funding_account_ids=[account2.id]
+        )
+
+        # tip another claim
+        await self.support_create(
+            self.get_claim_id(stream2), '0.2', tip=True, funding_account_ids=[self.account.id]
+        )
+        self.assertEqual(await account_balance(), {
+            'total': '9.077157',
+            'available': '6.077157',
+            'reserved': '3.0',
+            'reserved_subtotals': {'claims': '1.0', 'supports': '2.0', 'tips': '0.0'}
+        })
+        self.assertEqual(await wallet_balance(), {
+            'total': '9.938908',
+            'available': '6.638908',
+            'reserved': '3.3',
+            'reserved_subtotals': {'claims': '1.1', 'supports': '2.0', 'tips': '0.2'}
+        })
+
 
 class WalletEncryptionAndSynchronization(CommandTestCase):
 

torba/torba/client/basedatabase.py

@@ -20,6 +20,7 @@ class AIOSQLite:
         self.executor = ThreadPoolExecutor(max_workers=1)
         self.connection: sqlite3.Connection = None
         self._closing = False
+        self.query_count = 0
 
     @classmethod
     async def connect(cls, path: Union[bytes, str], *args, **kwargs):
@@ -65,6 +66,7 @@ class AIOSQLite:
     def __run_transaction(self, fun: Callable[[sqlite3.Connection, Any, Any], Any], *args, **kwargs):
         self.connection.execute('begin')
         try:
+            self.query_count += 1
             result = fun(self.connection, *args, **kwargs)  # type: ignore
             self.connection.commit()
             return result
@@ -605,11 +607,12 @@ class BaseDatabase(SQLiteMixin):
             'pubkey', 'chain_code', 'n', 'depth'
         )
         addresses = rows_to_dict(await self.select_addresses(', '.join(cols), **constraints), cols)
-        for address in addresses:
+        if 'pubkey' in cols:
-            address['pubkey'] = PubKey(
+            for address in addresses:
-                self.ledger, address.pop('pubkey'), address.pop('chain_code'),
+                address['pubkey'] = PubKey(
-                address.pop('n'), address.pop('depth')
+                    self.ledger, address.pop('pubkey'), address.pop('chain_code'),
-            )
+                    address.pop('n'), address.pop('depth')
+                )
         return addresses
 
     async def get_address_count(self, cols=None, **constraints):