Separate locking and non-locking cases for get_confirmed_balance_for_wallet to allow calling a few wallet_state_manager methods while already under the wallet_state_manager lock (#7790)

This commit is contained in:
Adam Kelly 2021-08-03 09:55:57 -07:00 committed by GitHub
parent c4b6014e01
commit d029a2a594
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 26 additions and 9 deletions

View File

@ -73,7 +73,7 @@ class DIDWallet:
if self.wallet_info is None:
raise ValueError("Internal Error")
self.wallet_id = self.wallet_info.id
bal = await self.standard_wallet.get_confirmed_balance()
bal = await self.wallet_state_manager.get_confirmed_balance_for_wallet(self.standard_wallet.id())
if amount > bal:
raise ValueError("Not enough balance")

View File

@ -63,6 +63,13 @@ from chia.server.server import ChiaServer
from chia.wallet.did_wallet.did_wallet import DIDWallet
def get_balance_from_coin_records(coin_records: Set[WalletCoinRecord]) -> uint128:
amount: uint128 = uint128(0)
for record in coin_records:
amount = uint128(amount + record.coin.amount)
return uint128(amount)
class WalletStateManager:
constants: ConsensusConstants
config: Dict
@ -499,21 +506,29 @@ class WalletStateManager:
return False
async def get_confirmed_balance_for_wallet_already_locked(self, wallet_id: int) -> uint128:
# This is a workaround to be able to call la locking operation when already locked
# for example, in the create method of DID wallet
assert self.lock.locked() is False
unspent_coin_records = await self.coin_store.get_unspent_coins_for_wallet(wallet_id)
return get_balance_from_coin_records(unspent_coin_records)
async def get_confirmed_balance_for_wallet(
self, wallet_id: int, unspent_coin_records: Optional[Set[WalletCoinRecord]] = None
) -> uint128:
"""
Returns the confirmed balance, including coinbase rewards that are not spendable.
"""
# lock only if unspent_coin_records is None
# lock only if unspent_coin_records is None.
# This API should change so that get_balance_from_coin_records is called for Set[WalletCoinRecord]
# and this method is called only for the unspent_coin_records==None case.
if unspent_coin_records is None:
async with self.lock:
if unspent_coin_records is None:
unspent_coin_records = await self.coin_store.get_unspent_coins_for_wallet(wallet_id)
amount: uint128 = uint128(0)
for record in unspent_coin_records:
amount = uint128(amount + record.coin.amount)
return uint128(amount)
unspent_coin_records = await self.get_confirmed_balance_for_wallet_with_lock(wallet_id)
return get_balance_from_coin_records(unspent_coin_records)
async def get_confirmed_balance_for_wallet_with_lock(self, wallet_id: int) -> Set[WalletCoinRecord]:
async with self.lock:
return await self.coin_store.get_unspent_coins_for_wallet(wallet_id)
async def get_unconfirmed_balance(
self, wallet_id, unspent_coin_records: Optional[Set[WalletCoinRecord]] = None
@ -522,6 +537,8 @@ class WalletStateManager:
Returns the balance, including coinbase rewards that are not spendable, and unconfirmed
transactions.
"""
# This API should change so that get_balance_from_coin_records is called for Set[WalletCoinRecord]
# and this method is called only for the unspent_coin_records==None case.
confirmed = await self.get_confirmed_balance_for_wallet(wallet_id, unspent_coin_records)
unconfirmed_tx: List[TransactionRecord] = await self.tx_store.get_unconfirmed_for_wallet(wallet_id)
removal_amount: int = 0