Skip to content

Commit

Permalink
wallet: Move definintely unusable TXOs to a separate container
Browse files Browse the repository at this point in the history 
Definitely unusable TXOs are those that are spent by a confirmed
transaction or were produced by a now conflicted transaction. However,
we still need them for GetDebit, so we store them in a separate
m_unusable_txos container. MarkConflicted, AbandonTransaction, and
loading (via PruneSpentTXOs) will ensure that these unusable TXOs are
properly moved.
  • Loading branch information
@achow101
achow101 committed Sep 15, 2023
1 parent b27b7fa commit 7c4edf2
Show file tree
Hide file tree
Showing 5 changed files with 151 additions and 14 deletions.
2 changes: 2 additions & 0 deletions src/wallet/receive.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -280,6 +280,8 @@ Balance GetBalance(const CWallet& wallet, const int min_depth, bool avoid_reuse)

const bool is_trusted{CachedTxIsTrusted(wallet, txo.GetState(), outpoint.hash)};
const int tx_depth{wallet.GetTxStateDepthInMainChain(txo.GetState())};
Assert(tx_depth >= 0);
Assert(!wallet.IsSpent(outpoint, /*min_depth=*/1));

if (!wallet.IsSpent(outpoint) && (allow_used_addresses || !wallet.IsSpentKey(txo.GetTxOut().scriptPubKey))) {
// Get the amounts for mine and watchonly
Expand Down
8 changes: 4 additions & 4 deletions src/wallet/spend.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -331,6 +331,10 @@ CoinsResult AvailableCoins(const CWallet& wallet,
if (wallet.IsLockedCoin(outpoint) && params.skip_locked)
continue;

int nDepth = wallet.GetTxStateDepthInMainChain(txo.GetState());
Assert(nDepth >= 0);
Assert(!wallet.IsSpent(outpoint, /*min_depth=*/1));

if (wallet.IsSpent(outpoint))
continue;

Expand All @@ -348,10 +352,6 @@ CoinsResult AvailableCoins(const CWallet& wallet,

assert(mine != ISMINE_NO);

int nDepth = wallet.GetTxStateDepthInMainChain(txo.GetState());
if (nDepth < 0)
continue;

if (tx_safe_cache.count(outpoint.hash) == 0 ) {
tx_safe_cache[outpoint.hash] = {false, false};
const CWalletTx& wtx = *wallet.GetWalletTx(outpoint.hash);
Expand Down
113 changes: 105 additions & 8 deletions src/wallet/wallet.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1151,6 +1151,20 @@ CWalletTx* CWallet::AddToWallet(CTransactionRef tx, const TxState& state, const
// Break debit/credit balance caches:
wtx.MarkDirty();

// Remove or add back the inputs from m_txos to match the state of this tx.
if (wtx.isConfirmed())
{
// When a transaction becomes confirmed, we can remove all of the txos that were spent
// in its inputs as they are no longer relevant.
for (const CTxIn& txin : wtx.tx->vin) {
MarkTXOUnusable(txin.prevout);
}
} else if (wtx.isInactive()) {
// When a transaction becomes inactive, we need to mark its inputs as usable again
for (const CTxIn& txin : wtx.tx->vin) {
MarkTXOUsable(txin.prevout);
}
}
// Get the outputs that belong to the wallet
RefreshWalletTxTXOs(wtx);

Expand Down Expand Up @@ -1407,6 +1421,13 @@ void CWallet::RecursiveUpdateTxState(const uint256& tx_hash, const TryUpdatingSt
todo.insert(iter->second);
}
}
if (wtx.state<TxStateConflicted>() || wtx.state<TxStateConfirmed>()) {
// If the state applied is conflicted or confirmed, the outputs are unusable
MarkTXOUnusable(COutPoint(now, i));
} else {
// Otherwise make the outputs usable
MarkTXOUsable(COutPoint(now, i));
}
}

if (update_state == TxUpdate::NOTIFY_CHANGED) {
Expand All @@ -1416,6 +1437,21 @@ void CWallet::RecursiveUpdateTxState(const uint256& tx_hash, const TryUpdatingSt
// If a transaction changes its tx state, that usually changes the balance
// available of the outputs it spends. So force those to be recomputed
MarkInputsDirty(wtx.tx);
// Make the non-conflicted inputs usable again
for (unsigned int i = 0; i < wtx.tx->vin.size(); ++i) {
const CTxIn& txin = wtx.tx->vin.at(i);
auto unusable_txo_it = m_unusable_txos.find(txin.prevout);
if (unusable_txo_it == m_unusable_txos.end()) {
continue;
}

if (std::get_if<TxStateConflicted>(&unusable_txo_it->second.GetState()) ||
std::get_if<TxStateConfirmed>(&unusable_txo_it->second.GetState())) {
continue;
}

MarkTXOUsable(txin.prevout);
}
}
}
}
Expand Down Expand Up @@ -3248,6 +3284,10 @@ bool CWallet::AttachChain(const std::shared_ptr<CWallet>& walletInstance, interf
}
walletInstance->m_attaching_chain = false;

// Remove TXOs that have already been spent
// We do this here as we need to have an attached chain to figure out what has actually been spent.
walletInstance->PruneSpentTXOs();

return true;
}

Expand Down Expand Up @@ -4306,26 +4346,34 @@ util::Result<MigrationResult> MigrateLegacyToDescriptor(const std::string& walle
return res;
}

using TXOMap = std::unordered_map<COutPoint, WalletTXO, SaltedOutpointHasher>;
void CWallet::RefreshWalletTxTXOs(const CWalletTx& wtx)
{
AssertLockHeld(cs_wallet);
for (uint32_t i = 0; i < wtx.tx->vout.size(); ++i) {
const CTxOut& txout = wtx.tx->vout.at(i);
COutPoint outpoint(wtx.GetHash(), i);

auto it = m_txos.find(outpoint);

isminetype ismine = IsMine(txout);
if (ismine == ISMINE_NO) {
continue;
}

if (it != m_txos.end()) {
auto it = wtx.m_txos.find(i);
if (it != wtx.m_txos.end()) {
it->second.SetIsMine(ismine);
it->second.SetState(wtx.GetState());
} else {
auto [txo_it, _] = m_txos.emplace(outpoint, WalletTXO{txout, ismine, wtx.GetState(), wtx.IsCoinBase(), wtx.m_from_me, wtx.GetTxTime()});
wtx.m_txos.emplace(i, txo_it->second);
TXOMap::iterator txo_it;
bool txos_inserted = false;
if (m_last_block_processed_height >= 0 && IsSpent(outpoint, /*min_depth=*/1)) {
std::tie(txo_it, txos_inserted) = m_unusable_txos.emplace(outpoint, WalletTXO{txout, ismine, wtx.GetState(), wtx.IsCoinBase(), wtx.m_from_me, wtx.GetTxTime()});
assert(txos_inserted);
} else {
std::tie(txo_it, txos_inserted) = m_txos.emplace(outpoint, WalletTXO{txout, ismine, wtx.GetState(), wtx.IsCoinBase(), wtx.m_from_me, wtx.GetTxTime()});
}
auto [_, wtx_inserted] = wtx.m_txos.emplace(i, txo_it->second);
assert(wtx_inserted);
}
}
}
Expand All @@ -4342,9 +4390,58 @@ std::optional<WalletTXO> CWallet::GetTXO(const COutPoint& outpoint) const
{
AssertLockHeld(cs_wallet);
const auto& it = m_txos.find(outpoint);
if (it == m_txos.end()) {
return std::nullopt;
if (it != m_txos.end()) {
return it->second;
}
return it->second;
const auto& u_it = m_unusable_txos.find(outpoint);
if (u_it != m_unusable_txos.end()) {
return u_it->second;
}
return std::nullopt;
}

void CWallet::PruneSpentTXOs()
{
AssertLockHeld(cs_wallet);
auto it = m_txos.begin();
while (it != m_txos.end()) {
if (std::get_if<TxStateConflicted>(&it->second.GetState()) || IsSpent(it->first, /*min_depth=*/1)) {
it = MarkTXOUnusable(it->first).first;
} else {
it++;
}
}
}

std::pair<TXOMap::iterator, TXOMap::iterator> CWallet::MarkTXOUnusable(const COutPoint& outpoint)
{
AssertLockHeld(cs_wallet);
auto txos_it = m_txos.find(outpoint);
auto unusable_txos_it = m_unusable_txos.end();
if (txos_it != m_txos.end()) {
auto next_txo_it = std::next(txos_it);
auto nh = m_txos.extract(txos_it);
txos_it = next_txo_it;
auto [position, inserted, _] = m_unusable_txos.insert(std::move(nh));
unusable_txos_it = position;
assert(inserted);
}
return {txos_it, unusable_txos_it};
}

std::pair<TXOMap::iterator, TXOMap::iterator> CWallet::MarkTXOUsable(const COutPoint& outpoint)
{
AssertLockHeld(cs_wallet);
auto txos_it = m_txos.end();
auto unusable_txos_it = m_unusable_txos.find(outpoint);
if (unusable_txos_it != m_unusable_txos.end()) {
auto next_unusable_txo_it = std::next(unusable_txos_it);
auto nh = m_unusable_txos.extract(unusable_txos_it);
unusable_txos_it = next_unusable_txo_it;
auto [position, inserted, _] = m_txos.insert(std::move(nh));
assert(inserted);
txos_it = position;
}
return {unusable_txos_it, txos_it};
}
} // namespace wallet
11 changes: 9 additions & 2 deletions src/wallet/wallet.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -420,7 +420,11 @@ class CWallet final : public WalletStorage, public interfaces::Chain::Notificati
void AddScriptPubKeyMan(const uint256& id, std::unique_ptr<ScriptPubKeyMan> spkm_man);

//! Set of transaction outputs owned by this wallet
std::unordered_map<COutPoint, WalletTXO, SaltedOutpointHasher> m_txos GUARDED_BY(cs_wallet);
using TXOMap = std::unordered_map<COutPoint, WalletTXO, SaltedOutpointHasher>;
TXOMap m_txos GUARDED_BY(cs_wallet);
//! Set of transaction outputs that are definitely no longer usuable
//! These outputs may already be spent in a confirmed tx, or are the outputs of a conflicted tx
TXOMap m_unusable_txos GUARDED_BY(cs_wallet);

/**
* Catch wallet up to current chain, scanning new blocks, updating the best
Expand Down Expand Up @@ -501,11 +505,14 @@ class CWallet final : public WalletStorage, public interfaces::Chain::Notificati

std::set<uint256> GetTxConflicts(const CWalletTx& wtx) const EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);

const std::unordered_map<COutPoint, WalletTXO, SaltedOutpointHasher>& GetTXOs() const EXCLUSIVE_LOCKS_REQUIRED(cs_wallet) { AssertLockHeld(cs_wallet); return m_txos; };
const TXOMap& GetTXOs() const EXCLUSIVE_LOCKS_REQUIRED(cs_wallet) { AssertLockHeld(cs_wallet); return m_txos; };
std::optional<WalletTXO> GetTXO(const COutPoint& outpoint) const EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);

void RefreshWalletTxTXOs(const CWalletTx& wtx) EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
void RefreshWalletTXOs() EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
void PruneSpentTXOs() EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
std::pair<TXOMap::iterator, TXOMap::iterator> MarkTXOUnusable(const COutPoint& outpoint) EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
std::pair<TXOMap::iterator, TXOMap::iterator> MarkTXOUsable(const COutPoint& outpoint) EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);

/**
* Return depth of transaction in blockchain:
Expand Down
31 changes: 31 additions & 0 deletions test/functional/wallet_balance.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
assert_equal,
assert_is_hash_string,
assert_raises_rpc_error,
find_vout_for_address,
)
from test_framework.wallet_util import get_generate_key

Expand Down Expand Up @@ -400,5 +401,35 @@ def test_balances(*, fee_node_1=0):
balances = wallet.getbalances()
assert_equal(balances["mine"]["trusted"], amount * 2)

wallet.unloadwallet()

self.log.info("Test that the balance is unchanged by an import that makes an already spent output in an existing tx \"mine\"")
self.nodes[0].createwallet("importalreadyspent")
wallet = self.nodes[0].get_wallet_rpc("importalreadyspent")

import_change_key = get_generate_key()
addr1 = wallet.getnewaddress()
addr2 = wallet.getnewaddress()

default.importprivkey(privkey=import_change_key.privkey, rescan=False)

res = default.send(outputs=[{addr1: amount}], options={"change_address": import_change_key.p2wpkh_addr})
inputs = [{"txid":res["txid"], "vout": find_vout_for_address(default, res["txid"], import_change_key.p2wpkh_addr)}]
default.send(outputs=[{addr2: amount}], options={"inputs": inputs, "add_inputs": True})

# Mock the time forward by another day so that "now" will exclude the block we just mined
self.nodes[0].setmocktime(int(time.time()) + 86400 * 2)
# Mine 11 blocks to move the MTP past the block we just mined
self.generate(self.nodes[0], 11, sync_fun=self.no_op)

balances = wallet.getbalances()
assert_equal(balances["mine"]["trusted"], amount * 2)

# Don't rescan to make sure that the import updates the wallet txos
# The balance should not change because the output for this key is already spent
wallet.importprivkey(privkey=import_change_key.privkey, rescan=False)
balances = wallet.getbalances()
assert_equal(balances["mine"]["trusted"], amount * 2)

if __name__ == '__main__':
WalletTest().main()

0 comments on commit 7c4edf2

Please sign in to comment.