Skip to content

Commit

Permalink
balance update refactoring and update on reconfig
Browse files Browse the repository at this point in the history
  • Loading branch information
@chappjc
chappjc committed Jan 11, 2021
1 parent f8c47a1 commit 8ff4d94
Show file tree
Hide file tree
Showing 2 changed files with 81 additions and 56 deletions.
124 changes: 74 additions & 50 deletions client/core/core.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1025,6 +1025,7 @@ func (c *Core) encryptionKey(pw []byte) (encrypt.Crypter, error) {
}
return crypter, nil
}

func (c *Core) storeDepositAddress(wdbID []byte, addr string) error {
// Store the new address in the DB.
dbWallet, err := c.db.Wallet(wdbID)
Expand All @@ -1051,7 +1052,7 @@ func (c *Core) connectAndUpdateWallet(w *xcWallet) error {
}
// First update balances since it is included in WalletState. Ignore errors
// because some wallets may not reveal balance until unlocked.
_, err = c.walletBalances(w)
_, err = c.updateWalletBalance(w)
if err != nil {
// Warn because the balances will be stale.
c.log.Warnf("Could not retrieve balances from %s wallet: %v", unbip(assetID), err)
Expand Down Expand Up @@ -1160,28 +1161,40 @@ func (c *Core) connectAndUnlock(crypter encrypt.Crypter, wallet *xcWallet) error
return nil
}

// walletBalances retrieves balances for the wallet.
func (c *Core) walletBalances(wallet *xcWallet) (*WalletBalance, error) {
c.connMtx.RLock()
defer c.connMtx.RUnlock()
// walletBalance gets the xcWallet's current WalletBalance, which includes the
// db.Balance plus order/contract locked amounts. The data is not stored. Use
// updateWalletBalance instead to also update xcWallet.balance and the DB.
func (c *Core) walletBalance(wallet *xcWallet) (*WalletBalance, error) {
bal, err := wallet.Balance()
if err != nil {
return nil, err
}
contractLockedAmt, orderLockedAmt := c.lockedAmounts(wallet.AssetID)
walletBal := &WalletBalance{
return &WalletBalance{
Balance: &db.Balance{
Balance: *bal,
Stamp: time.Now(),
},
ContractLocked: contractLockedAmt,
OrderLocked: orderLockedAmt,
ContractLocked: contractLockedAmt,
}, nil
}

// updateWalletBalance retrieves balances for the wallet, updates
// xcWallet.balance and the balance in the DB, and emits a BalanceNote.
func (c *Core) updateWalletBalance(wallet *xcWallet) (*WalletBalance, error) {
walletBal, err := c.walletBalance(wallet)
if err != nil {
return nil, err
}
wallet.setBalance(walletBal)

// Store the db.Balance.
err = c.db.UpdateBalance(wallet.dbID, walletBal.Balance)
if err != nil {
return nil, fmt.Errorf("error updating %s balance in database: %v", unbip(wallet.AssetID), err)
}

c.notify(newBalanceNote(wallet.AssetID, walletBal))
return walletBal, nil
}
Expand All @@ -1190,9 +1203,9 @@ func (c *Core) walletBalances(wallet *xcWallet) (*WalletBalance, error) {
// swaps (contractLocked) and the total amount locked by orders for future
// swaps (orderLocked). Only applies to trades where the specified assetID is
// the fromAssetID.
//
// The connMtx lock MUST be held for reads.
func (c *Core) lockedAmounts(assetID uint32) (contractLocked, orderLocked uint64) {
c.connMtx.RLock()
defer c.connMtx.RUnlock()
for _, dc := range c.conns {
dc.tradeMtx.RLock()
for _, tracker := range dc.trades {
Expand Down Expand Up @@ -1222,7 +1235,7 @@ func (c *Core) updateBalances(assets assetMap) {
c.log.Errorf("non-existent %d wallet should exist", assetID)
continue
}
_, err := c.walletBalances(w)
_, err := c.updateWalletBalance(w)
if err != nil {
c.log.Errorf("error updating %q balance: %v", unbip(assetID), err)
continue
Expand Down Expand Up @@ -1353,9 +1366,9 @@ func (c *Core) CreateWallet(appPW, walletPW []byte, form *WalletForm) error {

dbWallet := &db.Wallet{
AssetID: assetID,
Balance: &db.Balance{},
Settings: form.Config,
EncryptedPW: encPW,
// Balance and Address are set after connect.
}

wallet, err := c.loadWallet(dbWallet)
Expand All @@ -1375,21 +1388,20 @@ func (c *Core) CreateWallet(appPW, walletPW []byte, form *WalletForm) error {

err = wallet.Unlock(crypter)
if err != nil {
return initErr("%s wallet authentication error: %v", symbol, err)
return initErr("%s wallet authentication error: %w", symbol, err)
}

// Store the wallet in the database.
err = c.db.UpdateWallet(dbWallet)
balances, err := c.walletBalance(wallet)
if err != nil {
return initErr("error storing wallet credentials: %v", err)
return initErr("error getting wallet balance for %s: %w", symbol, err)
}
wallet.balance = balances // update xcWallet's WalletBalance
dbWallet.Balance = balances.Balance // store the db.Balance

// walletBalances will update the database record with the current balance.
// UpdateWallet must be called to create the database record before
// walletBalances is used.
balances, err := c.walletBalances(wallet)
// Store the wallet in the database.
err = c.db.UpdateWallet(dbWallet)
if err != nil {
return initErr("error getting wallet balance for %s: %v", symbol, err)
return initErr("error storing wallet credentials: %w", err)
}

c.log.Infof("Created %s wallet. Balance available = %d / "+
Expand All @@ -1404,6 +1416,7 @@ func (c *Core) CreateWallet(appPW, walletPW []byte, form *WalletForm) error {

c.refreshUser()

c.notify(newBalanceNote(assetID, balances)) // redundant with wallet state note?
c.notify(newWalletStateNote(wallet.state()))

return nil
Expand All @@ -1412,34 +1425,35 @@ func (c *Core) CreateWallet(appPW, walletPW []byte, form *WalletForm) error {
// loadWallet uses the data from the database to construct a new exchange
// wallet. The returned wallet is running but not connected.
func (c *Core) loadWallet(dbWallet *db.Wallet) (*xcWallet, error) {
c.connMtx.RLock() // required to calculate contractlocked amount
defer c.connMtx.RUnlock()
contractLockedAmt, orderLockedAmt := c.lockedAmounts(dbWallet.AssetID)
wallet := &xcWallet{
AssetID: dbWallet.AssetID,
balance: &WalletBalance{
Balance: dbWallet.Balance,
ContractLocked: contractLockedAmt,
OrderLocked: orderLockedAmt,
},
encPW: dbWallet.EncryptedPW,
address: dbWallet.Address,
dbID: dbWallet.ID(),
}
// Create the client/asset.Wallet.
assetID := dbWallet.AssetID
walletCfg := &asset.WalletConfig{
Settings: dbWallet.Settings,
TipChange: func(err error) {
c.tipChange(dbWallet.AssetID, err)
c.tipChange(assetID, err)
},
}
logger := c.log.SubLogger(unbip(dbWallet.AssetID))
w, err := asset.Setup(dbWallet.AssetID, walletCfg, logger, c.net)
logger := c.log.SubLogger(unbip(assetID))
w, err := asset.Setup(assetID, walletCfg, logger, c.net)
if err != nil {
return nil, fmt.Errorf("error creating wallet: %v", err)
}
wallet.Wallet = w
wallet.connector = dex.NewConnectionMaster(w)
return wallet, nil

// Construct the unconnected xcWallet.
contractLockedAmt, orderLockedAmt := c.lockedAmounts(assetID)
return &xcWallet{
Wallet: w,
connector: dex.NewConnectionMaster(w),
AssetID: assetID,
balance: &WalletBalance{
Balance: dbWallet.Balance,
OrderLocked: orderLockedAmt,
ContractLocked: contractLockedAmt,
},
encPW: dbWallet.EncryptedPW,
address: dbWallet.Address,
dbID: dbWallet.ID(),
}, nil
}

// WalletState returns the *WalletState for the asset ID.
Expand Down Expand Up @@ -1470,7 +1484,7 @@ func (c *Core) OpenWallet(assetID uint32, appPW []byte) error {
}

state := wallet.state()
balances, err := c.walletBalances(wallet)
balances, err := c.updateWalletBalance(wallet)
if err != nil {
return err
}
Expand Down Expand Up @@ -1551,12 +1565,10 @@ func (c *Core) ReconfigureWallet(appPW []byte, assetID uint32, cfg map[string]st
dbWallet := &db.Wallet{
AssetID: oldWallet.AssetID,
Settings: cfg,
Balance: &db.Balance{}, // in case retrieving new balance after connect fails
EncryptedPW: oldWallet.encPW,
Address: oldWallet.address,
}
if oldWallet.balance != nil {
dbWallet.Balance = oldWallet.balance.Balance
}
// Reload the wallet with the new settings.
wallet, err := c.loadWallet(dbWallet)
if err != nil {
Expand All @@ -1576,16 +1588,27 @@ func (c *Core) ReconfigureWallet(appPW []byte, assetID uint32, cfg map[string]st
err := wallet.Unlock(crypter)
if err != nil {
wallet.Disconnect()
return newError(walletAuthErr, "wallet successfully connected, but errored unlocking. reconfiguration not saved: %v", err)
return newError(walletAuthErr, "wallet successfully connected, but failed to unlock. "+
"reconfiguration not saved: %v", err)
}
}

balances, err := c.walletBalance(wallet)
if err != nil {
c.log.Warnf("Error getting balance for wallet %s: %v", unbip(assetID), err)
// Do not fail in case this requires an unlocked wallet.
} else {
wallet.balance = balances // update xcWallet's WalletBalance
dbWallet.Balance = balances.Balance // store the db.Balance
}

err = c.db.UpdateWallet(dbWallet)
if err != nil {
wallet.Disconnect()
return newError(dbErr, "error saving wallet configuration: %v", err)
}

// Update all relevant trackedTrades' toWallet and fromWallet.
c.connMtx.RLock()
for _, dc := range c.conns {
dc.tradeMtx.RLock()
Expand All @@ -1607,6 +1630,7 @@ func (c *Core) ReconfigureWallet(appPW []byte, assetID uint32, cfg map[string]st
}
c.wallets[assetID] = wallet

c.notify(newBalanceNote(assetID, balances)) // redundant with wallet config note?
details := fmt.Sprintf("Configuration for %s wallet has been updated. Deposit address = %s", unbip(assetID), wallet.address)
c.notify(newWalletConfigNote("Wallet Configuration Updated", details, db.Success, wallet.state()))

Expand Down Expand Up @@ -2051,9 +2075,9 @@ func (c *Core) Login(pw []byte) (*LoginResult, error) {
// Attempt to connect to and retrieve balance from all known wallets. It is
// not an error if we can't connect, unless we need the wallet for active
// trades, but that condition is checked later in resolveActiveTrades.
// Ignoring walletBalances errors here too, to accommodate wallets that must
// be unlocked to get the balance. We won't try to unlock here, but if the
// wallet is needed for active trades, it will be unlocked in
// Ignore updateWalletBalance errors here too, to accommodate wallets that
// must be unlocked to get the balance. We won't try to unlock here, but if
// the wallet is needed for active trades, it will be unlocked in
// resolveActiveTrades and the balance updated there.
var wg sync.WaitGroup
var connectCount uint32
Expand Down Expand Up @@ -2987,13 +3011,13 @@ func (c *Core) authDEX(dc *dexConnection) error {
return nil
}

// AssetBalance retrieves the current wallet balance.
// AssetBalance retrieves and updates the current wallet balance.
func (c *Core) AssetBalance(assetID uint32) (*WalletBalance, error) {
wallet, err := c.connectedWallet(assetID)
if err != nil {
return nil, fmt.Errorf("%d -> %s wallet error: %v", assetID, unbip(assetID), err)
}
return c.walletBalances(wallet)
return c.updateWalletBalance(wallet)
}

// initialize pulls the known DEXes from the database and attempts to connect
Expand Down
13 changes: 7 additions & 6 deletions client/core/wallet.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,18 +13,19 @@ import (
"decred.org/dcrdex/dex/encrypt"
)

// xcWallet is a wallet.
// xcWallet is a wallet. Use (*Core).loadWallet to construct a xcWallet.
type xcWallet struct {
asset.Wallet
connector *dex.ConnectionMaster
AssetID uint32
connector *dex.ConnectionMaster
AssetID uint32
dbID []byte
encPW []byte // empty means wallet not password protected

mtx sync.RWMutex
hookedUp bool
balance *WalletBalance
encPW []byte
pw string
address string
dbID []byte
hookedUp bool
synced bool
syncProgress float32
}
Expand Down

0 comments on commit 8ff4d94

Please sign in to comment.