Skip to content

Commit 6f5050d

Browse files
SachumanSachuman
committed
db: add tests for EstimateDiskUsage
Added tests for EstimateDiskUsage and EstimateDiskUsageByBackingType. Added calls to metamorphic tests, not for results but for races/panics
1 parent 761a9fe commit 6f5050d

File tree

6 files changed

+418
-0
lines changed

6 files changed

+418
-0
lines changed

‎disk_usage_test.go

Lines changed: 239 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,239 @@
1+
// Copyright 2025 The LevelDB-Go and Pebble Authors. All rights reserved. Use
2+
// of this source code is governed by a BSD-style license that can be found in
3+
// the LICENSE file.
4+
5+
package pebble
6+
7+
import (
8+
"context"
9+
"fmt"
10+
"testing"
11+
12+
"github.com/cockroachdb/datadriven"
13+
"github.com/cockroachdb/pebble/vfs"
14+
"github.com/stretchr/testify/require"
15+
)
16+
17+
func TestEstimateDiskUsage(t *testing.T) {
18+
mem := vfs.NewMem()
19+
d, err := Open("", &Options{FS: mem})
20+
require.NoError(t, err)
21+
defer d.Close()
22+
23+
// Test with empty database
24+
size, err := d.EstimateDiskUsage([]byte("a"), []byte("z"))
25+
require.NoError(t, err)
26+
require.Equal(t, uint64(0), size)
27+
28+
require.NoError(t, d.Set([]byte("apple"), []byte("fruit"), nil))
29+
require.NoError(t, d.Set([]byte("banana"), []byte("fruit"), nil))
30+
require.NoError(t, d.Set([]byte("cherry"), []byte("fruit"), nil))
31+
require.NoError(t, d.Set([]byte("date"), []byte("fruit"), nil))
32+
require.NoError(t, d.Flush())
33+
34+
testCases := []struct {
35+
name string
36+
start string
37+
end string
38+
}{
39+
{"full range", "a", "z"},
40+
{"partial range", "apple", "cherry"},
41+
{"single key range", "banana", "banana"},
42+
{"empty range", "zebra", "zebra"},
43+
{"overlapping range", "b", "d"},
44+
}
45+
46+
for _, tc := range testCases {
47+
t.Run(tc.name, func(t *testing.T) {
48+
size, err := d.EstimateDiskUsage([]byte(tc.start), []byte(tc.end))
49+
require.NoError(t, err)
50+
// Size should be non-negative
51+
require.GreaterOrEqual(t, size, uint64(0))
52+
})
53+
}
54+
// Test invalid range (start > end)
55+
_, err = d.EstimateDiskUsage([]byte("z"), []byte("a"))
56+
require.Error(t, err)
57+
require.Contains(t, err.Error(), "invalid key-range")
58+
}
59+
60+
func TestEstimateDiskUsageByBackingType(t *testing.T) {
61+
mem := vfs.NewMem()
62+
d, err := Open("", &Options{FS: mem})
63+
require.NoError(t, err)
64+
defer d.Close()
65+
66+
// Add some data
67+
require.NoError(t, d.Set([]byte("key1"), []byte("value1"), nil))
68+
require.NoError(t, d.Set([]byte("key2"), []byte("value2"), nil))
69+
require.NoError(t, d.Flush())
70+
71+
total, remote, external, err := d.EstimateDiskUsageByBackingType([]byte("a"), []byte("z"))
72+
require.NoError(t, err)
73+
74+
// Total should be non-negative
75+
require.GreaterOrEqual(t, total, uint64(0))
76+
77+
// Remote and external should be subsets of total
78+
require.LessOrEqual(t, remote, total)
79+
require.LessOrEqual(t, external, total)
80+
81+
// External should be subset of remote (external implies remote)
82+
require.LessOrEqual(t, external, remote)
83+
84+
// Test invalid range
85+
_, _, _, err = d.EstimateDiskUsageByBackingType([]byte("z"), []byte("a"))
86+
require.Error(t, err)
87+
require.Contains(t, err.Error(), "invalid key-range")
88+
}
89+
90+
func TestDiskUsageDataDriven(t *testing.T) {
91+
fs := vfs.NewMem()
92+
var d *DB
93+
defer func() {
94+
if d != nil {
95+
_ = d.Close()
96+
}
97+
}()
98+
99+
datadriven.RunTest(t, "testdata/disk_usage", func(t *testing.T, td *datadriven.TestData) string {
100+
switch td.Cmd {
101+
case "open":
102+
if d != nil {
103+
require.NoError(t, d.Close())
104+
d = nil
105+
}
106+
opts := &Options{FS: fs}
107+
require.NoError(t, parseDBOptionsArgs(opts, td.CmdArgs))
108+
var err error
109+
d, err = Open("", opts)
110+
require.NoError(t, err)
111+
return ""
112+
113+
case "close":
114+
if d != nil {
115+
require.NoError(t, d.Close())
116+
d = nil
117+
}
118+
return ""
119+
120+
case "batch":
121+
b := d.NewBatch()
122+
if err := runBatchDefineCmd(td, b); err != nil {
123+
return err.Error()
124+
}
125+
if err := b.Commit(nil); err != nil {
126+
return err.Error()
127+
}
128+
return ""
129+
130+
case "flush":
131+
if err := d.Flush(); err != nil {
132+
return err.Error()
133+
}
134+
return ""
135+
136+
case "build":
137+
if err := runBuildCmd(td, d, fs); err != nil {
138+
return err.Error()
139+
}
140+
return ""
141+
142+
case "ingest":
143+
if err := runIngestCmd(td, d, fs); err != nil {
144+
return err.Error()
145+
}
146+
return ""
147+
148+
case "compact":
149+
if len(td.CmdArgs) != 2 {
150+
return "compact <start> <end>"
151+
}
152+
l := []byte(td.CmdArgs[0].Key)
153+
r := []byte(td.CmdArgs[1].Key)
154+
if err := d.Compact(context.Background(), l, r, false /* parallelize */); err != nil {
155+
return err.Error()
156+
}
157+
return ""
158+
159+
case "estimate-disk-usage":
160+
// Parse range arguments, default to "a" and "z" if not specified
161+
start := []byte("a")
162+
end := []byte("z")
163+
if len(td.CmdArgs) >= 2 {
164+
start = []byte(td.CmdArgs[0].Key)
165+
end = []byte(td.CmdArgs[1].Key)
166+
}
167+
168+
var expect string
169+
if arg, ok := td.Arg("expect"); ok && len(arg.Vals) > 0 {
170+
expect = arg.Vals[0]
171+
}
172+
173+
size, err := d.EstimateDiskUsage(start, end)
174+
if err != nil {
175+
return err.Error()
176+
}
177+
178+
switch expect {
179+
case "zero":
180+
require.Equal(t, uint64(0), size)
181+
case "non-zero":
182+
require.Greater(t, size, uint64(0))
183+
}
184+
return "ok"
185+
186+
case "estimate-disk-usage-by-backing-type":
187+
// Parse range arguments, default to "a" and "z" if not specified
188+
start := []byte("a")
189+
end := []byte("z")
190+
if len(td.CmdArgs) >= 2 {
191+
start = []byte(td.CmdArgs[0].Key)
192+
end = []byte(td.CmdArgs[1].Key)
193+
}
194+
195+
total, remote, external, err := d.EstimateDiskUsageByBackingType(start, end)
196+
if err != nil {
197+
return err.Error()
198+
}
199+
200+
require.LessOrEqual(t, remote, total)
201+
require.LessOrEqual(t, external, remote)
202+
203+
if arg, ok := td.Arg("expect-total"); ok && len(arg.Vals) > 0 && arg.Vals[0] == "non-zero" {
204+
require.Greater(t, total, uint64(0))
205+
}
206+
if arg, ok := td.Arg("expect-remote"); ok && len(arg.Vals) > 0 && arg.Vals[0] == "zero" {
207+
require.Equal(t, uint64(0), remote)
208+
}
209+
if arg, ok := td.Arg("expect-external"); ok && len(arg.Vals) > 0 && arg.Vals[0] == "zero" {
210+
require.Equal(t, uint64(0), external)
211+
}
212+
213+
return "ok"
214+
215+
default:
216+
return fmt.Sprintf("unknown command: %s", td.Cmd)
217+
}
218+
})
219+
}
220+
221+
// Test that the EstimateDiskUsage and EstimateDiskUsageByBackingType should panic when the DB is closed
222+
func TestEstimateDiskUsageClosedDB(t *testing.T) {
223+
mem := vfs.NewMem()
224+
d, err := Open("", &Options{FS: mem})
225+
require.NoError(t, err)
226+
227+
require.NoError(t, d.Set([]byte("key"), []byte("value"), nil))
228+
229+
require.NoError(t, d.Close())
230+
231+
// Attempting to estimate on closed DB should panic
232+
require.Panics(t, func() {
233+
d.EstimateDiskUsage([]byte("a"), []byte("z"))
234+
})
235+
236+
require.Panics(t, func() {
237+
d.EstimateDiskUsageByBackingType([]byte("a"), []byte("z"))
238+
})
239+
}

‎metamorphic/config.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ const (
2828
OpDBFlush
2929
OpDBRatchetFormatMajorVersion
3030
OpDBRestart
31+
OpDBEstimateDiskUsage
3132
OpIterClose
3233
OpIterFirst
3334
OpIterLast
@@ -158,6 +159,7 @@ func DefaultOpConfig() OpConfig {
158159
OpDBFlush: 2,
159160
OpDBRatchetFormatMajorVersion: 1,
160161
OpDBRestart: 2,
162+
OpDBEstimateDiskUsage: 1,
161163
OpIterClose: 5,
162164
OpIterFirst: 100,
163165
OpIterLast: 100,
@@ -219,6 +221,7 @@ func ReadOpConfig() OpConfig {
219221
OpDBFlush: 0,
220222
OpDBRatchetFormatMajorVersion: 0,
221223
OpDBRestart: 0,
224+
OpDBEstimateDiskUsage: 0,
222225
OpIterClose: 5,
223226
OpIterFirst: 100,
224227
OpIterLast: 100,
@@ -277,6 +280,7 @@ func WriteOpConfig() OpConfig {
277280
OpDBFlush: 2,
278281
OpDBRatchetFormatMajorVersion: 1,
279282
OpDBRestart: 2,
283+
OpDBEstimateDiskUsage: 1,
280284
OpIterClose: 0,
281285
OpIterFirst: 0,
282286
OpIterLast: 0,

‎metamorphic/generator.go

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -168,6 +168,7 @@ func (g *generator) generate(count uint64) []op {
168168
OpDBFlush: g.dbFlush,
169169
OpDBRatchetFormatMajorVersion: g.dbRatchetFormatMajorVersion,
170170
OpDBRestart: g.dbRestart,
171+
OpDBEstimateDiskUsage: g.dbEstimateDiskUsage,
171172
OpIterClose: g.randIter(g.iterClose),
172173
OpIterFirst: g.randIter(g.iterFirst),
173174
OpIterLast: g.randIter(g.iterLast),
@@ -414,6 +415,21 @@ func (g *generator) dbCompact() {
414415
})
415416
}
416417

418+
func (g *generator) dbEstimateDiskUsage() {
419+
// Generate new key(s) with a 1% probability.
420+
start := g.keyGenerator.RandKey(0.01)
421+
end := g.keyGenerator.RandKey(0.01)
422+
if g.cmp(start, end) > 0 {
423+
start, end = end, start
424+
}
425+
dbID := g.dbs.rand(g.rng)
426+
g.add(&estimateDiskUsageOp{
427+
dbID: dbID,
428+
start: start,
429+
end: end,
430+
})
431+
}
432+
417433
func (g *generator) dbDownload() {
418434
numSpans := 1 + g.expRandInt(1)
419435
spans := make([]pebble.DownloadSpan, numSpans)

‎metamorphic/ops.go

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2246,3 +2246,38 @@ func (r *replicateOp) rewriteKeys(fn func(UserKey) UserKey) {
22462246
func (r *replicateOp) diagramKeyRanges() []pebble.KeyRange {
22472247
return []pebble.KeyRange{{Start: r.start, End: r.end}}
22482248
}
2249+
2250+
// estimateDiskUsageOp models DB.EstimateDiskUsage and DB.EstimateDiskUsageByBackingType operations.
2251+
type estimateDiskUsageOp struct {
2252+
dbID objID
2253+
start UserKey
2254+
end UserKey
2255+
}
2256+
2257+
func (o *estimateDiskUsageOp) run(t *Test, h historyRecorder) {
2258+
db := t.getDB(o.dbID)
2259+
_, err := db.EstimateDiskUsage(o.start, o.end)
2260+
if err != nil {
2261+
h.Recordf("%s // %v", o.formattedString(t.testOpts.KeyFormat), err)
2262+
} else {
2263+
h.Recordf("%s // <OK>", o.formattedString(t.testOpts.KeyFormat))
2264+
}
2265+
2266+
}
2267+
2268+
func (o *estimateDiskUsageOp) formattedString(kf KeyFormat) string {
2269+
return fmt.Sprintf("%s.EstimateDiskUsage(%q, %q)",
2270+
o.dbID, kf.FormatKey(o.start), kf.FormatKey(o.end))
2271+
}
2272+
2273+
func (o *estimateDiskUsageOp) receiver() objID { return o.dbID }
2274+
func (o *estimateDiskUsageOp) syncObjs() objIDSlice { return nil }
2275+
2276+
func (o *estimateDiskUsageOp) rewriteKeys(fn func(UserKey) UserKey) {
2277+
o.start = fn(o.start)
2278+
o.end = fn(o.end)
2279+
}
2280+
2281+
func (o *estimateDiskUsageOp) diagramKeyRanges() []pebble.KeyRange {
2282+
return []pebble.KeyRange{{Start: o.start, End: o.end}}
2283+
}

‎metamorphic/parser.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,8 @@ func opArgs(op op) (receiverID *objID, targetID *objID, args []interface{}) {
130130
return &t.writerID, nil, []interface{}{&t.start, &t.end, &t.suffix}
131131
case *replicateOp:
132132
return &t.source, nil, []interface{}{&t.dest, &t.start, &t.end}
133+
case *estimateDiskUsageOp:
134+
return &t.dbID, nil, []interface{}{&t.start, &t.end}
133135
}
134136
panic(fmt.Sprintf("unsupported op type: %T", op))
135137
}
@@ -146,6 +148,7 @@ var methods = map[string]*methodInfo{
146148
"Close": makeMethod(closeOp{}, dbTag, batchTag, iterTag, snapTag),
147149
"Commit": makeMethod(batchCommitOp{}, batchTag),
148150
"Compact": makeMethod(compactOp{}, dbTag),
151+
"EstimateDiskUsage": makeMethod(estimateDiskUsageOp{}, dbTag),
149152
"Delete": makeMethod(deleteOp{}, dbTag, batchTag),
150153
"DeleteRange": makeMethod(deleteRangeOp{}, dbTag, batchTag),
151154
"Download": makeMethod(downloadOp{}, dbTag),

0 commit comments

Comments
 (0)