Wallet store interface (#217)

[bytom/vapor.git] / account / utxo_keeper.go

diff --git a/account/utxo_keeper.go b/account/utxo_keeper.go
index 2705526..4af1a5d 100644 (file)
--- a/account/utxo_keeper.go
+++ b/account/utxo_keeper.go
@@ -1,16 +1,14 @@
 package account
 
 import (
+       "bytes"
        "container/list"
-       "encoding/json"
        "sort"
        "sync"
        "sync/atomic"
        "time"
 
        log "github.com/sirupsen/logrus"
-
-       dbm "github.com/vapor/database/leveldb"
        "github.com/vapor/errors"
        "github.com/vapor/protocol/bc"
 )
@@ -21,6 +19,7 @@ const desireUtxoCount = 5
 var (
        ErrInsufficient = errors.New("reservation found insufficient funds")
        ErrImmature     = errors.New("reservation found immature funds")
+       ErrVoteLock     = errors.New("Locked by the vote")
        ErrReserved     = errors.New("reservation found outputs already reserved")
        ErrMatchUTXO    = errors.New("can't find utxo with given hash")
        ErrReservation  = errors.New("couldn't find reservation")
@@ -54,7 +53,7 @@ type utxoKeeper struct {
        // `sync/atomic` expects the first word in an allocated struct to be 64-bit
        // aligned on both ARM and x86-32. See https://goo.gl/zW7dgq for more details.
        nextIndex     uint64
-       db            dbm.DB
+       store         AccountStore
        mtx           sync.RWMutex
        currentHeight func() uint64
 
@@ -63,9 +62,9 @@ type utxoKeeper struct {
        reservations map[uint64]*reservation
 }
 
-func newUtxoKeeper(f func() uint64, walletdb dbm.DB) *utxoKeeper {
+func newUtxoKeeper(f func() uint64, store AccountStore) *utxoKeeper {
        uk := &utxoKeeper{
-               db:            walletdb,
+               store:         store,
                currentHeight: f,
                unconfirmed:   make(map[bc.Hash]*UTXO),
                reserved:      make(map[bc.Hash]uint64),
@@ -112,17 +111,20 @@ func (uk *utxoKeeper) RemoveUnconfirmedUtxo(hashes []*bc.Hash) {
        }
 }
 
-func (uk *utxoKeeper) Reserve(accountID string, assetID *bc.AssetID, amount uint64, useUnconfirmed bool, exp time.Time) (*reservation, error) {
+func (uk *utxoKeeper) Reserve(accountID string, assetID *bc.AssetID, amount uint64, useUnconfirmed bool, vote []byte, exp time.Time) (*reservation, error) {
        uk.mtx.Lock()
        defer uk.mtx.Unlock()
 
-       utxos, immatureAmount := uk.findUtxos(accountID, assetID, useUnconfirmed)
+       utxos, immatureAmount := uk.findUtxos(accountID, assetID, useUnconfirmed, vote)
        optUtxos, optAmount, reservedAmount := uk.optUTXOs(utxos, amount)
        if optAmount+reservedAmount+immatureAmount < amount {
                return nil, ErrInsufficient
        }
 
        if optAmount+reservedAmount < amount {
+               if vote != nil {
+                       return nil, ErrVoteLock
+               }
                return nil, ErrImmature
        }
 
@@ -203,14 +205,15 @@ func (uk *utxoKeeper) expireReservation(t time.Time) {
        }
 }
 
-func (uk *utxoKeeper) findUtxos(accountID string, assetID *bc.AssetID, useUnconfirmed bool) ([]*UTXO, uint64) {
+func (uk *utxoKeeper) findUtxos(accountID string, assetID *bc.AssetID, useUnconfirmed bool, vote []byte) ([]*UTXO, uint64) {
        immatureAmount := uint64(0)
        currentHeight := uk.currentHeight()
        utxos := []*UTXO{}
        appendUtxo := func(u *UTXO) {
-               if u.AccountID != accountID || u.AssetID != *assetID {
+               if u.AccountID != accountID || u.AssetID != *assetID || !bytes.Equal(u.Vote, vote) {
                        return
                }
+
                if u.ValidHeight > currentHeight {
                        immatureAmount += u.Amount
                } else {
@@ -218,16 +221,15 @@ func (uk *utxoKeeper) findUtxos(accountID string, assetID *bc.AssetID, useUnconf
                }
        }
 
-       utxoIter := uk.db.IteratorPrefix([]byte(UTXOPreFix))
-       defer utxoIter.Release()
-       for utxoIter.Next() {
-               u := &UTXO{}
-               if err := json.Unmarshal(utxoIter.Value(), u); err != nil {
-                       log.WithFields(log.Fields{"module": logModule, "err": err}).Error("utxoKeeper findUtxos fail on unmarshal utxo")
-                       continue
-               }
-               appendUtxo(u)
+       UTXOs, err := uk.store.ListUTXOs()
+       if err != nil {
+               log.WithFields(log.Fields{"module": logModule, "err": err}).Error("utxoKeeper findUtxos fail on unmarshal utxo")
        }
+
+       for _, UTXO := range UTXOs {
+               appendUtxo(UTXO)
+       }
+
        if !useUnconfirmed {
                return utxos, immatureAmount
        }
@@ -242,15 +244,7 @@ func (uk *utxoKeeper) findUtxo(outHash bc.Hash, useUnconfirmed bool) (*UTXO, err
        if u, ok := uk.unconfirmed[outHash]; useUnconfirmed && ok {
                return u, nil
        }
-
-       u := &UTXO{}
-       if data := uk.db.Get(StandardUTXOKey(outHash)); data != nil {
-               return u, json.Unmarshal(data, u)
-       }
-       if data := uk.db.Get(ContractUTXOKey(outHash)); data != nil {
-               return u, json.Unmarshal(data, u)
-       }
-       return nil, ErrMatchUTXO
+       return uk.store.GetUTXO(outHash)
 }
 
 func (uk *utxoKeeper) optUTXOs(utxos []*UTXO, amount uint64) ([]*UTXO, uint64, uint64) {