1 files changed, 76 insertions, 16 deletions

diff --git a/vendor/github.com/dexon-foundation/dexon-consensus/core/db/memory.go b/vendor/github.com/dexon-foundation/dexon-consensus/core/db/memory.go
index 4246e4fe1..7393de9db 100644
--- a/vendor/github.com/dexon-foundation/dexon-consensus/core/db/memory.go
+++ b/vendor/github.com/dexon-foundation/dexon-consensus/core/db/memory.go
@@ -24,6 +24,7 @@ import (
 	"sync"
 
 	"github.com/dexon-foundation/dexon-consensus/common"
+	"github.com/dexon-foundation/dexon-consensus/core/crypto/dkg"
 	"github.com/dexon-foundation/dexon-consensus/core/types"
 )
 
@@ -41,10 +42,15 @@ func (seq *blockSeqIterator) NextBlock() (types.Block, error) {
 
 // MemBackedDB is a memory backed DB implementation.
 type MemBackedDB struct {
-	blocksMutex        sync.RWMutex
-	blockHashSequence  common.Hashes
-	blocksByHash       map[common.Hash]*types.Block
-	persistantFilePath string
+	blocksLock               sync.RWMutex
+	blockHashSequence        common.Hashes
+	blocksByHash             map[common.Hash]*types.Block
+	compactionChainTipLock   sync.RWMutex
+	compactionChainTipHash   common.Hash
+	compactionChainTipHeight uint64
+	dkgPrivateKeysLock       sync.RWMutex
+	dkgPrivateKeys           map[uint64]*dkg.PrivateKey
+	persistantFilePath       string
 }
 
 // NewMemBackedDB initialize a memory-backed database.
@@ -53,6 +59,7 @@ func NewMemBackedDB(persistantFilePath ...string) (
 	dbInst = &MemBackedDB{
 		blockHashSequence: common.Hashes{},
 		blocksByHash:      make(map[common.Hash]*types.Block),
+		dkgPrivateKeys:    make(map[uint64]*dkg.PrivateKey),
 	}
 	if len(persistantFilePath) == 0 || len(persistantFilePath[0]) == 0 {
 		return
@@ -87,8 +94,8 @@ func NewMemBackedDB(persistantFilePath ...string) (
 
 // HasBlock returns wheter or not the DB has a block identified with the hash.
 func (m *MemBackedDB) HasBlock(hash common.Hash) bool {
-	m.blocksMutex.RLock()
-	defer m.blocksMutex.RUnlock()
+	m.blocksLock.RLock()
+	defer m.blocksLock.RUnlock()
 
 	_, ok := m.blocksByHash[hash]
 	return ok
@@ -96,8 +103,8 @@ func (m *MemBackedDB) HasBlock(hash common.Hash) bool {
 
 // GetBlock returns a block given a hash.
 func (m *MemBackedDB) GetBlock(hash common.Hash) (types.Block, error) {
-	m.blocksMutex.RLock()
-	defer m.blocksMutex.RUnlock()
+	m.blocksLock.RLock()
+	defer m.blocksLock.RUnlock()
 
 	return m.internalGetBlock(hash)
 }
@@ -116,8 +123,8 @@ func (m *MemBackedDB) PutBlock(block types.Block) error {
 		return ErrBlockExists
 	}
 
-	m.blocksMutex.Lock()
-	defer m.blocksMutex.Unlock()
+	m.blocksLock.Lock()
+	defer m.blocksLock.Unlock()
 
 	m.blockHashSequence = append(m.blockHashSequence, block.Hash)
 	m.blocksByHash[block.Hash] = &block
@@ -130,13 +137,66 @@ func (m *MemBackedDB) UpdateBlock(block types.Block) error {
 		return ErrBlockDoesNotExist
 	}
 
-	m.blocksMutex.Lock()
-	defer m.blocksMutex.Unlock()
+	m.blocksLock.Lock()
+	defer m.blocksLock.Unlock()
 
 	m.blocksByHash[block.Hash] = &block
 	return nil
 }
 
+// PutCompactionChainTipInfo saves tip of compaction chain into the database.
+func (m *MemBackedDB) PutCompactionChainTipInfo(
+	blockHash common.Hash, height uint64) error {
+	m.compactionChainTipLock.Lock()
+	defer m.compactionChainTipLock.Unlock()
+	if m.compactionChainTipHeight >= height {
+		return ErrInvalidCompactionChainTipHeight
+	}
+	m.compactionChainTipHeight = height
+	m.compactionChainTipHash = blockHash
+	return nil
+}
+
+// GetCompactionChainTipInfo get the tip info of compaction chain into the
+// database.
+func (m *MemBackedDB) GetCompactionChainTipInfo() (
+	hash common.Hash, height uint64) {
+	m.compactionChainTipLock.RLock()
+	defer m.compactionChainTipLock.RUnlock()
+	return m.compactionChainTipHash, m.compactionChainTipHeight
+}
+
+// HasDKGPrivateKey check existence of DKG private key of one round.
+func (m *MemBackedDB) HasDKGPrivateKey(round uint64) (bool, error) {
+	m.dkgPrivateKeysLock.RLock()
+	defer m.dkgPrivateKeysLock.RUnlock()
+	_, exists := m.dkgPrivateKeys[round]
+	return exists, nil
+}
+
+// GetDKGPrivateKey get DKG private key of one round.
+func (m *MemBackedDB) GetDKGPrivateKey(round uint64) (
+	dkg.PrivateKey, error) {
+	m.dkgPrivateKeysLock.RLock()
+	defer m.dkgPrivateKeysLock.RUnlock()
+	if prv, exists := m.dkgPrivateKeys[round]; exists {
+		return *prv, nil
+	}
+	return dkg.PrivateKey{}, ErrDKGPrivateKeyDoesNotExist
+}
+
+// PutDKGPrivateKey save DKG private key of one round.
+func (m *MemBackedDB) PutDKGPrivateKey(
+	round uint64, prv dkg.PrivateKey) error {
+	m.dkgPrivateKeysLock.Lock()
+	defer m.dkgPrivateKeysLock.Unlock()
+	if _, exists := m.dkgPrivateKeys[round]; exists {
+		return ErrDKGPrivateKeyExists
+	}
+	m.dkgPrivateKeys[round] = &prv
+	return nil
+}
+
 // Close implement Closer interface, which would release allocated resource.
 func (m *MemBackedDB) Close() (err error) {
 	// Save internal state to a pretty-print json file. It's a temporary way
@@ -145,8 +205,8 @@ func (m *MemBackedDB) Close() (err error) {
 		return
 	}
 
-	m.blocksMutex.RLock()
-	defer m.blocksMutex.RUnlock()
+	m.blocksLock.RLock()
+	defer m.blocksLock.RUnlock()
 
 	toDump := struct {
 		Sequence common.Hashes
@@ -167,8 +227,8 @@ func (m *MemBackedDB) Close() (err error) {
 }
 
 func (m *MemBackedDB) getBlockByIndex(idx int) (types.Block, error) {
-	m.blocksMutex.RLock()
-	defer m.blocksMutex.RUnlock()
+	m.blocksLock.RLock()
+	defer m.blocksLock.RUnlock()
 
 	if idx >= len(m.blockHashSequence) {
 		return types.Block{}, ErrIterationFinished