updated equity frontend. (#1177)

[bytom/bytom.git] / wallet / wallet.go

diff --git a/wallet/wallet.go b/wallet/wallet.go
index b80cf1b..c00651a 100644 (file)
--- a/wallet/wallet.go
+++ b/wallet/wallet.go
@@ -2,6 +2,7 @@ package wallet
 
 import (
        "encoding/json"
+       "sync"
 
        log "github.com/sirupsen/logrus"
        "github.com/tendermint/tmlibs/db"
@@ -14,8 +15,10 @@ import (
        "github.com/bytom/protocol/bc/types"
 )
 
-//SINGLE single sign
-const SINGLE = 1
+const (
+       //SINGLE single sign
+       SINGLE = 1
+)
 
 var walletKey = []byte("walletInfo")
 
@@ -30,6 +33,7 @@ type StatusInfo struct {
 //Wallet is related to storing account unspent outputs
 type Wallet struct {
        DB         db.DB
+       rw         sync.RWMutex
        status     StatusInfo
        AccountMgr *account.Manager
        AssetReg   *asset.Registry
@@ -54,7 +58,7 @@ func NewWallet(walletDB db.DB, account *account.Manager, asset *asset.Registry,
        }
 
        go w.walletUpdater()
-
+       go w.delUnconfirmedTx()
        return w, nil
 }
 
@@ -86,6 +90,9 @@ func (w *Wallet) commitWalletInfo(batch db.Batch) error {
 
 // AttachBlock attach a new block
 func (w *Wallet) AttachBlock(block *types.Block) error {
+       w.rw.Lock()
+       defer w.rw.Unlock()
+
        if block.PreviousBlockHash != w.status.WorkHash {
                log.Warn("wallet skip attachBlock due to status hash not equal to previous hash")
                return nil
@@ -99,7 +106,7 @@ func (w *Wallet) AttachBlock(block *types.Block) error {
 
        storeBatch := w.DB.NewBatch()
        w.indexTransactions(storeBatch, block, txStatus)
-       w.buildAccountUTXOs(storeBatch, block, txStatus)
+       w.attachUtxos(storeBatch, block, txStatus)
 
        w.status.WorkHeight = block.Height
        w.status.WorkHash = block.Hash()
@@ -112,6 +119,9 @@ func (w *Wallet) AttachBlock(block *types.Block) error {
 
 // DetachBlock detach a block and rollback state
 func (w *Wallet) DetachBlock(block *types.Block) error {
+       w.rw.Lock()
+       defer w.rw.Unlock()
+
        blockHash := block.Hash()
        txStatus, err := w.chain.GetTransactionStatus(&blockHash)
        if err != nil {
@@ -119,7 +129,7 @@ func (w *Wallet) DetachBlock(block *types.Block) error {
        }
 
        storeBatch := w.DB.NewBatch()
-       w.reverseAccountUTXOs(storeBatch, block, txStatus)
+       w.detachUtxos(storeBatch, block, txStatus)
        w.deleteTransactions(storeBatch, w.status.BestHeight)
 
        w.status.BestHeight = block.Height - 1
@@ -145,24 +155,25 @@ func (w *Wallet) walletUpdater() {
                        }
 
                        if err := w.DetachBlock(block); err != nil {
-                               log.WithField("err", err).Error("walletUpdater detachBlock")
+                               log.WithField("err", err).Error("walletUpdater detachBlock stop")
                                return
                        }
                }
 
                block, _ := w.chain.GetBlockByHeight(w.status.WorkHeight + 1)
                if block == nil {
-                       <-w.chain.BlockWaiter(w.status.WorkHeight + 1)
+                       w.walletBlockWaiter()
                        continue
                }
 
                if err := w.AttachBlock(block); err != nil {
-                       log.WithField("err", err).Error("walletUpdater stop")
+                       log.WithField("err", err).Error("walletUpdater AttachBlock stop")
                        return
                }
        }
 }
 
+//RescanBlocks provide a trigger to rescan blocks
 func (w *Wallet) RescanBlocks() {
        select {
        case w.rescanCh <- struct{}{}:
@@ -174,19 +185,30 @@ func (w *Wallet) RescanBlocks() {
 func (w *Wallet) getRescanNotification() {
        select {
        case <-w.rescanCh:
-               block, _ := w.chain.GetBlockByHeight(0)
-               w.status.WorkHash = bc.Hash{}
-               w.AttachBlock(block)
+               w.setRescanStatus()
        default:
                return
        }
 }
 
-func (w *Wallet) createProgram(account *account.Account, XPub *pseudohsm.XPub, index uint64) error {
-       for i := uint64(0); i < index; i++ {
-               if _, err := w.AccountMgr.CreateAddress(nil, account.ID, false); err != nil {
-                       return err
-               }
+func (w *Wallet) setRescanStatus() {
+       block, _ := w.chain.GetBlockByHeight(0)
+       w.status.WorkHash = bc.Hash{}
+       w.AttachBlock(block)
+}
+
+func (w *Wallet) walletBlockWaiter() {
+       select {
+       case <-w.chain.BlockWaiter(w.status.WorkHeight + 1):
+       case <-w.rescanCh:
+               w.setRescanStatus()
        }
-       return nil
+}
+
+// GetWalletStatusInfo return current wallet StatusInfo
+func (w *Wallet) GetWalletStatusInfo() StatusInfo {
+       w.rw.RLock()
+       defer w.rw.RUnlock()
+
+       return w.status
 }