Skip to content

Commit 67288c5

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 8615b15 commit 67288c5

File tree

7 files changed

+351
-1
lines changed

7 files changed

+351
-1
lines changed

‎disk_usage.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ func (d *DB) EstimateDiskUsage(start, end []byte) (uint64, error) {
2929
}
3030

3131
// EstimateDiskUsageByBackingType is like EstimateDiskUsage but additionally
32-
// returns the subsets of that size in remote ane external files.
32+
// returns the subsets of that size in remote and external files.
3333
func (d *DB) EstimateDiskUsageByBackingType(
3434
start, end []byte,
3535
) (totalSize, remoteSize, externalSize uint64, _ error) {

‎disk_usage_test.go‎

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
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+
"fmt"
9+
"testing"
10+
11+
"github.com/cockroachdb/datadriven"
12+
"github.com/cockroachdb/pebble/objstorage/remote"
13+
"github.com/cockroachdb/pebble/vfs"
14+
"github.com/stretchr/testify/require"
15+
)
16+
17+
// Test that the EstimateDiskUsage and EstimateDiskUsageByBackingType should panic when the DB is closed
18+
func TestEstimateDiskUsageClosedDB(t *testing.T) {
19+
mem := vfs.NewMem()
20+
d, err := Open("", &Options{FS: mem})
21+
require.NoError(t, err)
22+
require.NoError(t, d.Set([]byte("key"), []byte("value"), nil))
23+
require.NoError(t, d.Close())
24+
// Attempting to estimate on closed DB should panic
25+
require.Panics(t, func() {
26+
d.EstimateDiskUsage([]byte("a"), []byte("z"))
27+
})
28+
require.Panics(t, func() {
29+
d.EstimateDiskUsageByBackingType([]byte("a"), []byte("z"))
30+
})
31+
}
32+
33+
// Test the EstimateDiskUsage and EstimateDiskUsageByBackingType data driven tests
34+
func TestEstimateDiskUsageDataDriven(t *testing.T) {
35+
fs := vfs.NewMem()
36+
remoteStorage := remote.NewInMem()
37+
var d *DB
38+
defer func() {
39+
if d != nil {
40+
_ = d.Close()
41+
}
42+
}()
43+
44+
datadriven.RunTest(t, "testdata/disk_usage", func(t *testing.T, td *datadriven.TestData) string {
45+
switch td.Cmd {
46+
case "open":
47+
if d != nil {
48+
require.NoError(t, d.Close())
49+
d = nil
50+
}
51+
opts := &Options{FS: fs, FormatMajorVersion: FormatExciseBoundsRecord, DisableAutomaticCompactions: true}
52+
opts.Experimental.RemoteStorage = remote.MakeSimpleFactory(map[remote.Locator]remote.Storage{
53+
"external-locator": remoteStorage,
54+
})
55+
require.NoError(t, parseDBOptionsArgs(opts, td.CmdArgs))
56+
var err error
57+
d, err = Open("", opts)
58+
require.NoError(t, err)
59+
return ""
60+
61+
case "close":
62+
if d != nil {
63+
require.NoError(t, d.Close())
64+
d = nil
65+
}
66+
return ""
67+
68+
case "batch":
69+
b := d.NewBatch()
70+
if err := runBatchDefineCmd(td, b); err != nil {
71+
return err.Error()
72+
}
73+
if err := b.Commit(nil); err != nil {
74+
return err.Error()
75+
}
76+
return ""
77+
78+
case "flush":
79+
if err := d.Flush(); err != nil {
80+
return err.Error()
81+
}
82+
return ""
83+
84+
case "build":
85+
if err := runBuildCmd(td, d, fs); err != nil {
86+
return err.Error()
87+
}
88+
return ""
89+
90+
case "ingest":
91+
if err := runIngestCmd(td, d, fs); err != nil {
92+
return err.Error()
93+
}
94+
return ""
95+
case "build-remote":
96+
if err := runBuildRemoteCmd(td, d, remoteStorage); err != nil {
97+
return err.Error()
98+
}
99+
return ""
100+
case "ingest-external":
101+
if err := runIngestExternalCmd(t, td, d, remoteStorage, "external-locator"); err != nil {
102+
return err.Error()
103+
}
104+
return ""
105+
case "compact":
106+
if err := runCompactCmd(td, d); err != nil {
107+
return err.Error()
108+
}
109+
return runLSMCmd(td, d)
110+
111+
case "estimate-disk-usage":
112+
// Parse range arguments, default to "a" and "z" if not specified
113+
start := []byte("a")
114+
end := []byte("z")
115+
if len(td.CmdArgs) >= 2 {
116+
start = []byte(td.CmdArgs[0].Key)
117+
end = []byte(td.CmdArgs[1].Key)
118+
}
119+
size, err := d.EstimateDiskUsage(start, end)
120+
if err != nil {
121+
return err.Error()
122+
}
123+
return fmt.Sprintf("size: %d", size)
124+
125+
case "estimate-disk-usage-by-backing-type":
126+
// Parse range arguments, default to "a" and "z" if not specified
127+
start := []byte("a")
128+
end := []byte("z")
129+
if len(td.CmdArgs) >= 2 {
130+
start = []byte(td.CmdArgs[0].Key)
131+
end = []byte(td.CmdArgs[1].Key)
132+
}
133+
total, remote, external, err := d.EstimateDiskUsageByBackingType(start, end)
134+
if err != nil {
135+
return err.Error()
136+
}
137+
// remote and external should be less than or equal to total
138+
require.LessOrEqual(t, remote, total)
139+
require.LessOrEqual(t, external, remote)
140+
return fmt.Sprintf("total: %d, remote: %d, external: %d", total, remote, external)
141+
142+
default:
143+
return fmt.Sprintf("unknown command: %s", td.Cmd)
144+
}
145+
})
146+
}

‎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)