Skip to content

Commit 42da76f

Browse files
Raneet10Raneet10
authored
Add batching for KVstore (cosmos#149)
* add batching to KVStore * add batch for PrefixKV
1 parent 7d2f4ae commit 42da76f

11 files changed

+138
-22
lines changed

‎block/manager_test.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -29,9 +29,9 @@ func TestInitialState(t *testing.T) {
2929
LastBlockHeight: 128,
3030
}
3131

32-
emptyStore := store.New(store.NewInMemoryKVStore())
32+
emptyStore := store.New(store.NewDefaultInMemoryKVStore())
3333

34-
fullStore := store.New(store.NewInMemoryKVStore())
34+
fullStore := store.New(store.NewDefaultInMemoryKVStore())
3535
err := fullStore.UpdateState(sampleState)
3636
require.NoError(t, err)
3737

‎da/mock/mock_test.go

+3-3
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ import (
1515

1616
func TestLifecycle(t *testing.T) {
1717
var da da.DataAvailabilityLayerClient = &MockDataAvailabilityLayerClient{}
18-
dalcKV := store.NewInMemoryKVStore()
18+
dalcKV := store.NewDefaultInMemoryKVStore()
1919

2020
require := require.New(t)
2121

@@ -31,7 +31,7 @@ func TestLifecycle(t *testing.T) {
3131

3232
func TestMockDALC(t *testing.T) {
3333
var dalc da.DataAvailabilityLayerClient = &MockDataAvailabilityLayerClient{}
34-
dalcKV := store.NewInMemoryKVStore()
34+
dalcKV := store.NewDefaultInMemoryKVStore()
3535

3636
require := require.New(t)
3737
assert := assert.New(t)
@@ -72,7 +72,7 @@ func TestRetrieve(t *testing.T) {
7272
var dalc da.DataAvailabilityLayerClient = mock
7373
var retriever da.BlockRetriever = mock
7474

75-
dalcKV := store.NewInMemoryKVStore()
75+
dalcKV := store.NewDefaultInMemoryKVStore()
7676

7777
require := require.New(t)
7878
assert := assert.New(t)

‎node/integration_test.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -156,7 +156,7 @@ func createNodes(num int, t *testing.T) ([]*Node, []*mocks.Application) {
156156
nodes := make([]*Node, num)
157157
apps := make([]*mocks.Application, num)
158158
dalc := &mockda.MockDataAvailabilityLayerClient{}
159-
_ = dalc.Init(nil, store.NewInMemoryKVStore(), log.TestingLogger())
159+
_ = dalc.Init(nil, store.NewDefaultInMemoryKVStore(), log.TestingLogger())
160160
_ = dalc.Start()
161161
nodes[0], apps[0] = createNode(0, true, dalc, keys, t)
162162
for i := 1; i < num; i++ {

‎node/node.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -82,9 +82,9 @@ func NewNode(ctx context.Context, conf config.NodeConfig, nodeKey crypto.PrivKey
8282
var baseKV store.KVStore
8383
if conf.RootDir == "" && conf.DBPath == "" { // this is used for testing
8484
logger.Info("WARNING: working in in-memory mode")
85-
baseKV = store.NewInMemoryKVStore()
85+
baseKV = store.NewDefaultInMemoryKVStore()
8686
} else {
87-
baseKV = store.NewKVStore(conf.RootDir, conf.DBPath, "optimint")
87+
baseKV = store.NewDefaultKVStore(conf.RootDir, conf.DBPath, "optimint")
8888
}
8989
mainKV := store.NewPrefixKV(baseKV, mainPrefix)
9090
dalcKV := store.NewPrefixKV(baseKV, dalcPrefix)

‎store/badger.go

+38
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
)
88

99
var _ KVStore = &BadgerKV{}
10+
var _ Batch = &BadgerBatch{}
1011

1112
var (
1213
// ErrKeyNotFound is returned if key is not found in KVStore.
@@ -53,3 +54,40 @@ func (b *BadgerKV) Delete(key []byte) error {
5354
}
5455
return txn.Commit()
5556
}
57+
58+
// NewBatch creates new batch.
59+
// Note: badger batches should be short lived as they use extra resources.
60+
func (b *BadgerKV) NewBatch() Batch {
61+
return &BadgerBatch{
62+
txn: b.db.NewTransaction(true),
63+
}
64+
}
65+
66+
// BadgerBatch encapsulates badger transaction
67+
type BadgerBatch struct {
68+
txn *badger.Txn
69+
}
70+
71+
// Set accumulates key-value entries in a transaction
72+
func (bb *BadgerBatch) Set(key, value []byte) error {
73+
if err := bb.txn.Set(key, value); err != nil {
74+
return err
75+
}
76+
77+
return nil
78+
}
79+
80+
// Delete removes the key and associated value from store
81+
func (bb *BadgerBatch) Delete(key []byte) error {
82+
return bb.txn.Delete(key)
83+
}
84+
85+
// Commit commits a transaction
86+
func (bb *BadgerBatch) Commit() error {
87+
return bb.txn.Commit()
88+
}
89+
90+
// Discard cancels a transaction
91+
func (bb *BadgerBatch) Discard() {
92+
bb.txn.Discard()
93+
}

‎store/badger_test.go

+3-3
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import (
88
)
99

1010
func TestGetErrors(t *testing.T) {
11-
dalcKV := NewInMemoryKVStore()
11+
dalcKV := NewDefaultInMemoryKVStore()
1212

1313
tc := []struct {
1414
name string
@@ -30,7 +30,7 @@ func TestGetErrors(t *testing.T) {
3030
}
3131

3232
func TestSetErrors(t *testing.T) {
33-
dalcKV := NewInMemoryKVStore()
33+
dalcKV := NewDefaultInMemoryKVStore()
3434

3535
tc := []struct {
3636
name string
@@ -53,7 +53,7 @@ func TestSetErrors(t *testing.T) {
5353
}
5454

5555
func TestDeleteErrors(t *testing.T) {
56-
dalcKV := NewInMemoryKVStore()
56+
dalcKV := NewDefaultInMemoryKVStore()
5757

5858
tc := []struct {
5959
name string

‎store/kv.go

+11-2
Original file line numberDiff line numberDiff line change
@@ -13,10 +13,19 @@ type KVStore interface {
1313
Get(key []byte) ([]byte, error) // Get gets the value for a key.
1414
Set(key []byte, value []byte) error // Set updates the value for a key.
1515
Delete(key []byte) error // Delete deletes a key.
16+
NewBatch() Batch
17+
}
18+
19+
// Batch enables batching of transactions
20+
type Batch interface {
21+
Set(key, value []byte) error // Accumulates KV entries in a transaction
22+
Delete(key []byte) error // Deletes the given key
23+
Commit() error // Commits the transaction
24+
Discard() // Discards the transaction
1625
}
1726

1827
// NewInMemoryKVStore builds KVStore that works in-memory (without accessing disk).
19-
func NewInMemoryKVStore() KVStore {
28+
func NewDefaultInMemoryKVStore() KVStore {
2029
db, err := badger.Open(badger.DefaultOptions("").WithInMemory(true))
2130
if err != nil {
2231
panic(err)
@@ -26,7 +35,7 @@ func NewInMemoryKVStore() KVStore {
2635
}
2736
}
2837

29-
func NewKVStore(rootDir, dbPath, dbName string) KVStore {
38+
func NewDefaultKVStore(rootDir, dbPath, dbName string) KVStore {
3039
path := filepath.Join(rootify(rootDir, dbPath), dbName)
3140
db, err := badger.Open(badger.DefaultOptions(path))
3241
if err != nil {

‎store/prefix.go

+29-1
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
11
package store
22

3+
var _ KVStore = &PrefixKV{}
4+
var _ Batch = &PrefixKVBatch{}
5+
36
type PrefixKV struct {
47
kv KVStore
58
prefix []byte
@@ -23,5 +26,30 @@ func (p *PrefixKV) Set(key []byte, value []byte) error {
2326
func (p *PrefixKV) Delete(key []byte) error {
2427
return p.kv.Delete(append(p.prefix, key...))
2528
}
29+
func (p *PrefixKV) NewBatch() Batch {
30+
return &PrefixKVBatch{
31+
b: p.kv.NewBatch(),
32+
prefix: p.prefix,
33+
}
34+
}
2635

27-
var _ KVStore = &PrefixKV{}
36+
type PrefixKVBatch struct {
37+
b Batch
38+
prefix []byte
39+
}
40+
41+
func (pb *PrefixKVBatch) Set(key, value []byte) error {
42+
return pb.b.Set(append(pb.prefix, key...), value)
43+
}
44+
45+
func (pb *PrefixKVBatch) Delete(key []byte) error {
46+
return pb.b.Delete(append(pb.prefix, key...))
47+
}
48+
49+
func (pb *PrefixKVBatch) Commit() error {
50+
return pb.b.Commit()
51+
}
52+
53+
func (pb *PrefixKVBatch) Discard() {
54+
pb.b.Discard()
55+
}

‎store/prefix_test.go

+34-1
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ func TestPrefixKV(t *testing.T) {
1313
assert := assert.New(t)
1414
require := require.New(t)
1515

16-
base := NewInMemoryKVStore()
16+
base := NewDefaultInMemoryKVStore()
1717

1818
p1 := NewPrefixKV(base, []byte{1})
1919
p2 := NewPrefixKV(base, []byte{2})
@@ -71,3 +71,36 @@ func TestPrefixKV(t *testing.T) {
7171
require.NoError(err)
7272
assert.Equal(val22, v)
7373
}
74+
75+
func TestPrefixKVBatch(t *testing.T) {
76+
t.Parallel()
77+
78+
assert := assert.New(t)
79+
require := require.New(t)
80+
81+
basekv := NewDefaultInMemoryKVStore()
82+
prefixkv := NewPrefixKV(basekv, []byte("prefix1"))
83+
prefixbatchkv1 := prefixkv.NewBatch()
84+
85+
keys := [][]byte{[]byte("key1"), []byte("key2"), []byte("key3"), []byte("key4")}
86+
values := [][]byte{[]byte("value1"), []byte("value2"), []byte("value3"), []byte("value4")}
87+
88+
for i := 0; i < len(keys); i++ {
89+
err := prefixbatchkv1.Set(keys[i], values[i])
90+
require.NoError(err)
91+
}
92+
93+
err := prefixbatchkv1.Commit()
94+
require.NoError(err)
95+
96+
for i := 0; i < len(keys); i++ {
97+
vals, err := prefixkv.Get(keys[i])
98+
assert.Equal(vals, values[i])
99+
require.NoError(err)
100+
}
101+
102+
prefixbatchkv2 := prefixkv.NewBatch()
103+
err = prefixbatchkv2.Delete([]byte("key1"))
104+
require.NoError(err)
105+
106+
}

‎store/store.go

+13-5
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,9 @@ var _ Store = &DefaultStore{}
3333

3434
// New returns new, default store.
3535
func New(kv KVStore) Store {
36-
return &DefaultStore{db: kv}
36+
return &DefaultStore{
37+
db: kv,
38+
}
3739
}
3840

3941
// Height returns height of the highest block saved in the Store.
@@ -59,12 +61,18 @@ func (s *DefaultStore) SaveBlock(block *types.Block, commit *types.Commit) error
5961

6062
s.mtx.Lock()
6163
defer s.mtx.Unlock()
62-
// TODO(tzdybal): use transaction for consistency of DB (https://github.com/celestiaorg/optimint/issues/80)
63-
err = multierr.Append(err, s.db.Set(getBlockKey(hash), blockBlob))
64-
err = multierr.Append(err, s.db.Set(getCommitKey(hash), commitBlob))
65-
err = multierr.Append(err, s.db.Set(getIndexKey(block.Header.Height), hash[:]))
64+
65+
bb := s.db.NewBatch()
66+
err = multierr.Append(err, bb.Set(getBlockKey(hash), blockBlob))
67+
err = multierr.Append(err, bb.Set(getCommitKey(hash), commitBlob))
68+
err = multierr.Append(err, bb.Set(getIndexKey(block.Header.Height), hash[:]))
6669

6770
if err != nil {
71+
bb.Discard()
72+
return err
73+
}
74+
75+
if err = bb.Commit(); err != nil {
6876
return err
6977
}
7078

‎store/store_test.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ func TestBlockstoreHeight(t *testing.T) {
3838
for _, c := range cases {
3939
t.Run(c.name, func(t *testing.T) {
4040
assert := assert.New(t)
41-
bstore := New(NewInMemoryKVStore())
41+
bstore := New(NewDefaultInMemoryKVStore())
4242
assert.Equal(uint64(0), bstore.Height())
4343

4444
for _, block := range c.blocks {
@@ -80,7 +80,7 @@ func TestBlockstoreLoad(t *testing.T) {
8080
}
8181
}()
8282

83-
for _, kv := range []KVStore{NewInMemoryKVStore(), NewKVStore(tmpDir, "db", "test")} {
83+
for _, kv := range []KVStore{NewDefaultInMemoryKVStore(), NewDefaultKVStore(tmpDir, "db", "test")} {
8484
for _, c := range cases {
8585
t.Run(c.name, func(t *testing.T) {
8686
assert := assert.New(t)

0 commit comments

Comments
 (0)