Skip to content

Commit

Permalink
[factory] refactore workingsetstore (#4519)
Browse files Browse the repository at this point in the history
  • Loading branch information
dustinxie authored Dec 11, 2024
1 parent 5848248 commit 6ac9264
Show file tree
Hide file tree
Showing 3 changed files with 221 additions and 197 deletions.
211 changes: 14 additions & 197 deletions state/factory/workingsetstore.go
Original file line number Diff line number Diff line change
@@ -1,25 +1,15 @@
// Copyright (c) 2022 IoTeX Foundation
// Copyright (c) 2024 IoTeX Foundation
// This source code is provided 'as is' and no warranties are given as to title or non-infringement, merchantability
// or fitness for purpose and, to the extent permitted by law, all liability for your use of the code is disclaimed.
// This source code is governed by Apache License 2.0 that can be found in the LICENSE file.

package factory

import (
"context"
"fmt"

"github.com/iotexproject/go-pkgs/hash"
"go.uber.org/zap"

"github.com/pkg/errors"

"github.com/iotexproject/iotex-core/v2/action/protocol"
"github.com/iotexproject/iotex-core/v2/db"
"github.com/iotexproject/iotex-core/v2/db/trie"
"github.com/iotexproject/iotex-core/v2/pkg/log"
"github.com/iotexproject/iotex-core/v2/pkg/util/byteutil"
"github.com/iotexproject/iotex-core/v2/state"
)

type (
Expand All @@ -35,220 +25,47 @@ type (
ReadView(string) (interface{}, error)
WriteView(string, interface{}) error
}
stateDBWorkingSetStore struct {
view protocol.View
flusher db.KVStoreFlusher
readBuffer bool
}
factoryWorkingSetStore struct {
view protocol.View
flusher db.KVStoreFlusher
tlt trie.TwoLayerTrie
trieRoots map[int][]byte
workingSetStoreCommon struct {
view protocol.View
flusher db.KVStoreFlusher
}
)

func newStateDBWorkingSetStore(view protocol.View, flusher db.KVStoreFlusher, readBuffer bool) workingSetStore {
return &stateDBWorkingSetStore{
flusher: flusher,
view: view,
readBuffer: readBuffer,
}
}

func newFactoryWorkingSetStore(view protocol.View, flusher db.KVStoreFlusher) (workingSetStore, error) {
tlt, err := newTwoLayerTrie(ArchiveTrieNamespace, flusher.KVStoreWithBuffer(), ArchiveTrieRootKey, true)
if err != nil {
return nil, err
}

return &factoryWorkingSetStore{
flusher: flusher,
view: view,
tlt: tlt,
trieRoots: make(map[int][]byte),
}, nil
}

func newFactoryWorkingSetStoreAtHeight(view protocol.View, flusher db.KVStoreFlusher, height uint64) (workingSetStore, error) {
rootKey := fmt.Sprintf("%s-%d", ArchiveTrieRootKey, height)
tlt, err := newTwoLayerTrie(ArchiveTrieNamespace, flusher.KVStoreWithBuffer(), rootKey, false)
if err != nil {
return nil, err
}

return &factoryWorkingSetStore{
flusher: flusher,
view: view,
tlt: tlt,
trieRoots: make(map[int][]byte),
}, nil
}

func (store *stateDBWorkingSetStore) Start(context.Context) error {
return nil
}

func (store *stateDBWorkingSetStore) Stop(context.Context) error {
return nil
}

func (store *stateDBWorkingSetStore) ReadView(name string) (interface{}, error) {
func (store *workingSetStoreCommon) ReadView(name string) (interface{}, error) {
return store.view.Read(name)
}

func (store *stateDBWorkingSetStore) WriteView(name string, value interface{}) error {
func (store *workingSetStoreCommon) WriteView(name string, value interface{}) error {
return store.view.Write(name, value)
}

func (store *stateDBWorkingSetStore) Get(ns string, key []byte) ([]byte, error) {
data, err := store.flusher.KVStoreWithBuffer().Get(ns, key)
if err != nil {
if errors.Cause(err) == db.ErrNotExist {
return nil, errors.Wrapf(state.ErrStateNotExist, "failed to get state of ns = %x and key = %x", ns, key)
}
return nil, err
}
return data, nil
}

func (store *stateDBWorkingSetStore) Put(ns string, key []byte, value []byte) error {
func (store *workingSetStoreCommon) Put(ns string, key []byte, value []byte) error {
store.flusher.KVStoreWithBuffer().MustPut(ns, key, value)
return nil
}

func (store *stateDBWorkingSetStore) Delete(ns string, key []byte) error {
func (store *workingSetStoreCommon) Delete(ns string, key []byte) error {
store.flusher.KVStoreWithBuffer().MustDelete(ns, key)
return nil
}

func (store *stateDBWorkingSetStore) States(ns string, keys [][]byte) ([][]byte, [][]byte, error) {
if store.readBuffer {
return readStates(store.flusher.KVStoreWithBuffer(), ns, keys)
}
return readStates(store.flusher.BaseKVStore(), ns, keys)
}

func (store *stateDBWorkingSetStore) Digest() hash.Hash256 {
func (store *workingSetStoreCommon) Digest() hash.Hash256 {
return hash.Hash256b(store.flusher.SerializeQueue())
}

func (store *stateDBWorkingSetStore) Finalize(height uint64) error {
// Persist current chain Height
store.flusher.KVStoreWithBuffer().MustPut(
AccountKVNamespace,
[]byte(CurrentHeightKey),
byteutil.Uint64ToBytes(height),
)
return nil
}

func (store *stateDBWorkingSetStore) Commit() error {
func (store *workingSetStoreCommon) Commit() error {
_dbBatchSizelMtc.WithLabelValues().Set(float64(store.flusher.KVStoreWithBuffer().Size()))
return store.flusher.Flush()
}

func (store *stateDBWorkingSetStore) Snapshot() int {
func (store *workingSetStoreCommon) Snapshot() int {
return store.flusher.KVStoreWithBuffer().Snapshot()
}

func (store *stateDBWorkingSetStore) RevertSnapshot(snapshot int) error {
func (store *workingSetStoreCommon) RevertSnapshot(snapshot int) error {
return store.flusher.KVStoreWithBuffer().RevertSnapshot(snapshot)
}

func (store *stateDBWorkingSetStore) ResetSnapshots() {
store.flusher.KVStoreWithBuffer().ResetSnapshots()
}

func (store *factoryWorkingSetStore) Start(ctx context.Context) error {
return store.tlt.Start(ctx)
}

func (store *factoryWorkingSetStore) Stop(ctx context.Context) error {
return store.tlt.Stop(ctx)
}

func (store *factoryWorkingSetStore) ReadView(name string) (interface{}, error) {
return store.view.Read(name)
}

func (store *factoryWorkingSetStore) WriteView(name string, value interface{}) error {
return store.view.Write(name, value)
}

func (store *factoryWorkingSetStore) Get(ns string, key []byte) ([]byte, error) {
return readStateFromTLT(store.tlt, ns, key)
}

func (store *factoryWorkingSetStore) Put(ns string, key []byte, value []byte) error {
store.flusher.KVStoreWithBuffer().MustPut(ns, key, value)
nsHash := hash.Hash160b([]byte(ns))

return store.tlt.Upsert(nsHash[:], toLegacyKey(key), value)
}

func (store *factoryWorkingSetStore) Delete(ns string, key []byte) error {
store.flusher.KVStoreWithBuffer().MustDelete(ns, key)
nsHash := hash.Hash160b([]byte(ns))

err := store.tlt.Delete(nsHash[:], toLegacyKey(key))
if errors.Cause(err) == trie.ErrNotExist {
return errors.Wrapf(state.ErrStateNotExist, "key %x doesn't exist in namespace %x", key, nsHash)
}
return err
}

func (store *factoryWorkingSetStore) States(ns string, keys [][]byte) ([][]byte, [][]byte, error) {
return readStatesFromTLT(store.tlt, ns, keys)
}

func (store *factoryWorkingSetStore) Digest() hash.Hash256 {
return hash.Hash256b(store.flusher.SerializeQueue())
}

func (store *factoryWorkingSetStore) Finalize(h uint64) error {
rootHash, err := store.tlt.RootHash()
if err != nil {
return err
}
store.flusher.KVStoreWithBuffer().MustPut(AccountKVNamespace, []byte(CurrentHeightKey), byteutil.Uint64ToBytes(h))
store.flusher.KVStoreWithBuffer().MustPut(ArchiveTrieNamespace, []byte(ArchiveTrieRootKey), rootHash)
// Persist the historical accountTrie's root hash
store.flusher.KVStoreWithBuffer().MustPut(
ArchiveTrieNamespace,
[]byte(fmt.Sprintf("%s-%d", ArchiveTrieRootKey, h)),
rootHash,
)
return nil
}

func (store *factoryWorkingSetStore) Commit() error {
_dbBatchSizelMtc.WithLabelValues().Set(float64(store.flusher.KVStoreWithBuffer().Size()))
return store.flusher.Flush()
}

func (store *factoryWorkingSetStore) Snapshot() int {
rh, err := store.tlt.RootHash()
if err != nil {
log.L().Panic("failed to do snapshot", zap.Error(err))
}
s := store.flusher.KVStoreWithBuffer().Snapshot()
store.trieRoots[s] = rh
return s
}

func (store *factoryWorkingSetStore) RevertSnapshot(snapshot int) error {
if err := store.flusher.KVStoreWithBuffer().RevertSnapshot(snapshot); err != nil {
return err
}
root, ok := store.trieRoots[snapshot]
if !ok {
// this should not happen, b/c we save the trie root on a successful return of Snapshot(), but check anyway
return errors.Wrapf(trie.ErrInvalidTrie, "failed to get trie root for snapshot = %d", snapshot)
}
return store.tlt.SetRootHash(root[:])
}

func (store *factoryWorkingSetStore) ResetSnapshots() {
func (store *workingSetStoreCommon) ResetSnapshots() {
store.flusher.KVStoreWithBuffer().ResetSnapshots()
store.trieRoots = make(map[int][]byte)
}
138 changes: 138 additions & 0 deletions state/factory/workingsetstore_factory.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,138 @@
// Copyright (c) 2024 IoTeX Foundation
// This source code is provided 'as is' and no warranties are given as to title or non-infringement, merchantability
// or fitness for purpose and, to the extent permitted by law, all liability for your use of the code is disclaimed.
// This source code is governed by Apache License 2.0 that can be found in the LICENSE file.

package factory

import (
"context"
"fmt"

"github.com/iotexproject/go-pkgs/hash"
"github.com/pkg/errors"
"go.uber.org/zap"

"github.com/iotexproject/iotex-core/v2/action/protocol"
"github.com/iotexproject/iotex-core/v2/db"
"github.com/iotexproject/iotex-core/v2/db/trie"
"github.com/iotexproject/iotex-core/v2/pkg/log"
"github.com/iotexproject/iotex-core/v2/pkg/util/byteutil"
"github.com/iotexproject/iotex-core/v2/state"
)

type factoryWorkingSetStore struct {
*workingSetStoreCommon
tlt trie.TwoLayerTrie
trieRoots map[int][]byte
}

func newFactoryWorkingSetStore(view protocol.View, flusher db.KVStoreFlusher) (workingSetStore, error) {
tlt, err := newTwoLayerTrie(ArchiveTrieNamespace, flusher.KVStoreWithBuffer(), ArchiveTrieRootKey, true)
if err != nil {
return nil, err
}

return &factoryWorkingSetStore{
workingSetStoreCommon: &workingSetStoreCommon{
flusher: flusher,
view: view,
},
tlt: tlt,
trieRoots: make(map[int][]byte),
}, nil
}

func newFactoryWorkingSetStoreAtHeight(view protocol.View, flusher db.KVStoreFlusher, height uint64) (workingSetStore, error) {
rootKey := fmt.Sprintf("%s-%d", ArchiveTrieRootKey, height)
tlt, err := newTwoLayerTrie(ArchiveTrieNamespace, flusher.KVStoreWithBuffer(), rootKey, false)
if err != nil {
return nil, err
}

return &factoryWorkingSetStore{
workingSetStoreCommon: &workingSetStoreCommon{
flusher: flusher,
view: view,
},
tlt: tlt,
trieRoots: make(map[int][]byte),
}, nil
}

func (store *factoryWorkingSetStore) Start(ctx context.Context) error {
return store.tlt.Start(ctx)
}

func (store *factoryWorkingSetStore) Stop(ctx context.Context) error {
return store.tlt.Stop(ctx)
}

func (store *factoryWorkingSetStore) Get(ns string, key []byte) ([]byte, error) {
return readStateFromTLT(store.tlt, ns, key)
}

func (store *factoryWorkingSetStore) Put(ns string, key []byte, value []byte) error {
store.workingSetStoreCommon.Put(ns, key, value)
nsHash := hash.Hash160b([]byte(ns))

return store.tlt.Upsert(nsHash[:], toLegacyKey(key), value)
}

func (store *factoryWorkingSetStore) Delete(ns string, key []byte) error {
store.workingSetStoreCommon.Delete(ns, key)
nsHash := hash.Hash160b([]byte(ns))

err := store.tlt.Delete(nsHash[:], toLegacyKey(key))
if errors.Cause(err) == trie.ErrNotExist {
return errors.Wrapf(state.ErrStateNotExist, "key %x doesn't exist in namespace %x", key, nsHash)
}
return err
}

func (store *factoryWorkingSetStore) States(ns string, keys [][]byte) ([][]byte, [][]byte, error) {
return readStatesFromTLT(store.tlt, ns, keys)
}

func (store *factoryWorkingSetStore) Finalize(h uint64) error {
rootHash, err := store.tlt.RootHash()
if err != nil {
return err
}
store.flusher.KVStoreWithBuffer().MustPut(AccountKVNamespace, []byte(CurrentHeightKey), byteutil.Uint64ToBytes(h))
store.flusher.KVStoreWithBuffer().MustPut(ArchiveTrieNamespace, []byte(ArchiveTrieRootKey), rootHash)
// Persist the historical accountTrie's root hash
store.flusher.KVStoreWithBuffer().MustPut(
ArchiveTrieNamespace,
[]byte(fmt.Sprintf("%s-%d", ArchiveTrieRootKey, h)),
rootHash,
)
return nil
}

func (store *factoryWorkingSetStore) Snapshot() int {
rh, err := store.tlt.RootHash()
if err != nil {
log.L().Panic("failed to do snapshot", zap.Error(err))
}
s := store.workingSetStoreCommon.Snapshot()
store.trieRoots[s] = rh
return s
}

func (store *factoryWorkingSetStore) RevertSnapshot(snapshot int) error {
if err := store.workingSetStoreCommon.RevertSnapshot(snapshot); err != nil {
return err
}
root, ok := store.trieRoots[snapshot]
if !ok {
// this should not happen, b/c we save the trie root on a successful return of Snapshot(), but check anyway
return errors.Wrapf(trie.ErrInvalidTrie, "failed to get trie root for snapshot = %d", snapshot)
}
return store.tlt.SetRootHash(root[:])
}

func (store *factoryWorkingSetStore) ResetSnapshots() {
store.workingSetStoreCommon.ResetSnapshots()
store.trieRoots = make(map[int][]byte)
}
Loading

0 comments on commit 6ac9264

Please sign in to comment.