Skip to content

Commit 188f7e3

Browse files
DracoLiDracoLi
authored
feat: add in memory implementation of HeightIndex Database (#4212)
1 parent c85e9da commit 188f7e3

File tree

4 files changed

+329
-5
lines changed

4 files changed

+329
-5
lines changed

‎database/database.go‎

Lines changed: 17 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -95,17 +95,29 @@ type Database interface {
9595
health.Checker
9696
}
9797

98-
// HeightIndex defines the interface for storing and retrieving entries by height.
98+
// HeightIndex defines the interface for storing and retrieving values by height.
9999
type HeightIndex interface {
100-
// Put inserts the entry into the store at the given height.
101-
Put(height uint64, bytes []byte) error
100+
// Put inserts the value into the database at the given height.
101+
//
102+
// If value is nil or an empty slice, then when it's retrieved it may be nil
103+
// or an empty slice.
104+
//
105+
// value is safe to read and modify after calling Put.
106+
Put(height uint64, value []byte) error
102107

103-
// Get retrieves an entry by its height.
108+
// Get retrieves a value by its height.
109+
// Returns [ErrNotFound] if the key is not present in the database.
110+
//
111+
// Returned []byte is safe to read and modify after calling Get.
104112
Get(height uint64) ([]byte, error)
105113

106-
// Has checks if an entry exists at the given height.
114+
// Has checks if a value exists at the given height.
115+
//
116+
// Returns true even if the stored value is nil or empty.
107117
Has(height uint64) (bool, error)
108118

109119
// Close closes the database.
120+
//
121+
// Calling Close after Close returns [ErrClosed].
110122
io.Closer
111123
}

‎database/heightindexdb/dbtest/dbtest.go‎

Lines changed: 216 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,216 @@
1+
// Copyright (C) 2019-2025, Ava Labs, Inc. All rights reserved.
2+
// See the file LICENSE for licensing terms.
3+
4+
package dbtest
5+
6+
import (
7+
"bytes"
8+
"testing"
9+
10+
"github.com/stretchr/testify/require"
11+
12+
"github.com/ava-labs/avalanchego/database"
13+
)
14+
15+
// Tests is a list of all database tests
16+
var Tests = []struct {
17+
Name string
18+
Test func(t *testing.T, newDB func() database.HeightIndex)
19+
}{
20+
{"TestPutGet", TestPutGet},
21+
{"TestHas", TestHas},
22+
{"TestCloseAndPut", TestCloseAndPut},
23+
{"TestCloseAndGet", TestCloseAndGet},
24+
{"TestCloseAndHas", TestCloseAndHas},
25+
{"TestClose", TestClose},
26+
}
27+
28+
type putArgs struct {
29+
height uint64
30+
data []byte
31+
}
32+
33+
func TestPutGet(t *testing.T, newDB func() database.HeightIndex) {
34+
tests := []struct {
35+
name string
36+
puts []putArgs
37+
queryHeight uint64
38+
want []byte
39+
wantErr error
40+
}{
41+
{
42+
name: "normal operation",
43+
puts: []putArgs{
44+
{1, []byte("test data 1")},
45+
},
46+
queryHeight: 1,
47+
want: []byte("test data 1"),
48+
},
49+
{
50+
name: "not found error when getting on non-existing height",
51+
puts: []putArgs{
52+
{1, []byte("test data")},
53+
},
54+
queryHeight: 2,
55+
wantErr: database.ErrNotFound,
56+
},
57+
{
58+
name: "overwriting data on same height",
59+
puts: []putArgs{
60+
{1, []byte("original data")},
61+
{1, []byte("overwritten data")},
62+
},
63+
queryHeight: 1,
64+
want: []byte("overwritten data"),
65+
},
66+
{
67+
name: "put and get nil data",
68+
puts: []putArgs{
69+
{1, nil},
70+
},
71+
queryHeight: 1,
72+
},
73+
{
74+
name: "put and get empty bytes",
75+
puts: []putArgs{
76+
{1, []byte{}},
77+
},
78+
queryHeight: 1,
79+
},
80+
{
81+
name: "put and get large data",
82+
puts: []putArgs{
83+
{1, make([]byte, 1000)},
84+
},
85+
queryHeight: 1,
86+
want: make([]byte, 1000),
87+
},
88+
}
89+
90+
for _, tt := range tests {
91+
t.Run(tt.name, func(t *testing.T) {
92+
db := newDB()
93+
defer func() {
94+
require.NoError(t, db.Close())
95+
}()
96+
97+
// Perform all puts
98+
for _, write := range tt.puts {
99+
require.NoError(t, db.Put(write.height, write.data))
100+
}
101+
102+
// modify the original value of the put data to ensure the saved
103+
// value won't be changed after Get
104+
if len(tt.puts) > int(tt.queryHeight) && tt.puts[tt.queryHeight].data != nil {
105+
copy(tt.puts[tt.queryHeight].data, []byte("modified data"))
106+
}
107+
108+
// Query the specific height
109+
retrievedData, err := db.Get(tt.queryHeight)
110+
require.ErrorIs(t, err, tt.wantErr)
111+
require.True(t, bytes.Equal(tt.want, retrievedData))
112+
113+
// modify the data returned from Get and ensure it won't change the
114+
// data from a second Get
115+
copy(retrievedData, []byte("modified data"))
116+
newData, err := db.Get(tt.queryHeight)
117+
require.ErrorIs(t, err, tt.wantErr)
118+
require.True(t, bytes.Equal(tt.want, newData))
119+
})
120+
}
121+
}
122+
123+
func TestHas(t *testing.T, newDB func() database.HeightIndex) {
124+
tests := []struct {
125+
name string
126+
puts []putArgs
127+
queryHeight uint64
128+
want bool
129+
}{
130+
{
131+
name: "non-existent item",
132+
queryHeight: 1,
133+
},
134+
{
135+
name: "existing item with data",
136+
puts: []putArgs{{1, []byte("test data")}},
137+
queryHeight: 1,
138+
want: true,
139+
},
140+
{
141+
name: "existing item with nil data",
142+
puts: []putArgs{{1, nil}},
143+
queryHeight: 1,
144+
want: true,
145+
},
146+
{
147+
name: "existing item with empty bytes",
148+
puts: []putArgs{{1, []byte{}}},
149+
queryHeight: 1,
150+
want: true,
151+
},
152+
{
153+
name: "has returns true on overridden height",
154+
puts: []putArgs{
155+
{1, []byte("original data")},
156+
{1, []byte("overridden data")},
157+
},
158+
queryHeight: 1,
159+
want: true,
160+
},
161+
}
162+
163+
for _, tt := range tests {
164+
t.Run(tt.name, func(t *testing.T) {
165+
db := newDB()
166+
defer func() {
167+
require.NoError(t, db.Close())
168+
}()
169+
170+
// Perform all puts
171+
for _, write := range tt.puts {
172+
require.NoError(t, db.Put(write.height, write.data))
173+
}
174+
175+
ok, err := db.Has(tt.queryHeight)
176+
require.NoError(t, err)
177+
require.Equal(t, tt.want, ok)
178+
})
179+
}
180+
}
181+
182+
func TestCloseAndPut(t *testing.T, newDB func() database.HeightIndex) {
183+
db := newDB()
184+
require.NoError(t, db.Close())
185+
186+
// Try to put after close - should return error
187+
err := db.Put(1, []byte("test"))
188+
require.ErrorIs(t, err, database.ErrClosed)
189+
}
190+
191+
func TestCloseAndGet(t *testing.T, newDB func() database.HeightIndex) {
192+
db := newDB()
193+
require.NoError(t, db.Close())
194+
195+
// Try to get after close - should return error
196+
_, err := db.Get(1)
197+
require.ErrorIs(t, err, database.ErrClosed)
198+
}
199+
200+
func TestCloseAndHas(t *testing.T, newDB func() database.HeightIndex) {
201+
db := newDB()
202+
require.NoError(t, db.Close())
203+
204+
// Try to has after close - should return error
205+
_, err := db.Has(1)
206+
require.ErrorIs(t, err, database.ErrClosed)
207+
}
208+
209+
func TestClose(t *testing.T, newDB func() database.HeightIndex) {
210+
db := newDB()
211+
require.NoError(t, db.Close())
212+
213+
// Second close should return error
214+
err := db.Close()
215+
require.ErrorIs(t, err, database.ErrClosed)
216+
}

‎database/heightindexdb/memdb/database.go‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
// Copyright (C) 2019-2025, Ava Labs, Inc. All rights reserved.
2+
// See the file LICENSE for licensing terms.
3+
4+
package memdb
5+
6+
import (
7+
"slices"
8+
"sync"
9+
10+
"github.com/ava-labs/avalanchego/database"
11+
)
12+
13+
var _ database.HeightIndex = (*Database)(nil)
14+
15+
// Database is an in-memory implementation of database.HeightIndex
16+
type Database struct {
17+
mu sync.RWMutex
18+
data map[uint64][]byte
19+
closed bool
20+
}
21+
22+
func (db *Database) Put(height uint64, data []byte) error {
23+
db.mu.Lock()
24+
defer db.mu.Unlock()
25+
26+
if db.closed {
27+
return database.ErrClosed
28+
}
29+
30+
if db.data == nil {
31+
db.data = make(map[uint64][]byte)
32+
}
33+
34+
db.data[height] = slices.Clone(data)
35+
return nil
36+
}
37+
38+
func (db *Database) Get(height uint64) ([]byte, error) {
39+
db.mu.RLock()
40+
defer db.mu.RUnlock()
41+
42+
if db.closed {
43+
return nil, database.ErrClosed
44+
}
45+
46+
data, ok := db.data[height]
47+
if !ok {
48+
return nil, database.ErrNotFound
49+
}
50+
51+
return slices.Clone(data), nil
52+
}
53+
54+
func (db *Database) Has(height uint64) (bool, error) {
55+
db.mu.RLock()
56+
defer db.mu.RUnlock()
57+
58+
if db.closed {
59+
return false, database.ErrClosed
60+
}
61+
62+
_, ok := db.data[height]
63+
return ok, nil
64+
}
65+
66+
func (db *Database) Close() error {
67+
db.mu.Lock()
68+
defer db.mu.Unlock()
69+
70+
if db.closed {
71+
return database.ErrClosed
72+
}
73+
74+
db.closed = true
75+
db.data = nil
76+
return nil
77+
}

‎database/heightindexdb/memdb/database_test.go‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
// Copyright (C) 2019-2025, Ava Labs, Inc. All rights reserved.
2+
// See the file LICENSE for licensing terms.
3+
4+
package memdb
5+
6+
import (
7+
"testing"
8+
9+
"github.com/ava-labs/avalanchego/database"
10+
"github.com/ava-labs/avalanchego/database/heightindexdb/dbtest"
11+
)
12+
13+
func TestInterface(t *testing.T) {
14+
for _, test := range dbtest.Tests {
15+
t.Run("memdb_"+test.Name, func(t *testing.T) {
16+
test.Test(t, func() database.HeightIndex { return &Database{} })
17+
})
18+
}
19+
}

0 commit comments

Comments
 (0)