db: add tests for EstimateDiskUsage

disk_usage.go

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

disk_usage_test.go

@@ -0,0 +1,146 @@
+// Copyright 2025 The LevelDB-Go and Pebble Authors. All rights reserved. Use
+// of this source code is governed by a BSD-style license that can be found in
+// the LICENSE file.
+
+package pebble
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/cockroachdb/datadriven"
+	"github.com/cockroachdb/pebble/objstorage/remote"
+	"github.com/cockroachdb/pebble/vfs"
+	"github.com/stretchr/testify/require"
+)
+
+// Test that the EstimateDiskUsage and EstimateDiskUsageByBackingType should panic when the DB is closed
+func TestEstimateDiskUsageClosedDB(t *testing.T) {
+	mem := vfs.NewMem()
+	d, err := Open("", &Options{FS: mem})
+	require.NoError(t, err)
+	require.NoError(t, d.Set([]byte("key"), []byte("value"), nil))
+	require.NoError(t, d.Close())
+	// Attempting to estimate on closed DB should panic
+	require.Panics(t, func() {
+		d.EstimateDiskUsage([]byte("a"), []byte("z"))
+	})
+	require.Panics(t, func() {
+		d.EstimateDiskUsageByBackingType([]byte("a"), []byte("z"))
+	})
+}
+
+// Test the EstimateDiskUsage and EstimateDiskUsageByBackingType data driven tests
+func TestEstimateDiskUsageDataDriven(t *testing.T) {
+	fs := vfs.NewMem()
+	remoteStorage := remote.NewInMem()
+	var d *DB
+	defer func() {
+		if d != nil {
+			_ = d.Close()
+		}
+	}()
+
+	datadriven.RunTest(t, "testdata/disk_usage", func(t *testing.T, td *datadriven.TestData) string {
+		switch td.Cmd {
+		case "open":
+			if d != nil {
+				require.NoError(t, d.Close())
+				d = nil
+			}
+			opts := &Options{FS: fs, FormatMajorVersion: FormatExciseBoundsRecord, DisableAutomaticCompactions: true}
+			opts.Experimental.RemoteStorage = remote.MakeSimpleFactory(map[remote.Locator]remote.Storage{
+				"external-locator": remoteStorage,
+			})
+			require.NoError(t, parseDBOptionsArgs(opts, td.CmdArgs))
+			var err error
+			d, err = Open("", opts)
+			require.NoError(t, err)
+			return ""
+
+		case "close":
+			if d != nil {
+				require.NoError(t, d.Close())
+				d = nil
+			}
+			return ""
+
+		case "batch":
+			b := d.NewBatch()
+			if err := runBatchDefineCmd(td, b); err != nil {
+				return err.Error()
+			}
+			if err := b.Commit(nil); err != nil {
+				return err.Error()
+			}
+			return ""
+
+		case "flush":
+			if err := d.Flush(); err != nil {
+				return err.Error()
+			}
+			return ""
+
+		case "build":
+			if err := runBuildCmd(td, d, fs); err != nil {
+				return err.Error()
+			}
+			return ""
+
+		case "ingest":
+			if err := runIngestCmd(td, d, fs); err != nil {
+				return err.Error()
+			}
+			return ""
+		case "build-remote":
+			if err := runBuildRemoteCmd(td, d, remoteStorage); err != nil {
+				return err.Error()
+			}
+			return ""
+		case "ingest-external":
+			if err := runIngestExternalCmd(t, td, d, remoteStorage, "external-locator"); err != nil {
+				return err.Error()
+			}
+			return ""
+		case "compact":
+			if err := runCompactCmd(td, d); err != nil {
+				return err.Error()
+			}
+			return runLSMCmd(td, d)
+
+		case "estimate-disk-usage":
+			// Parse range arguments, default to "a" and "z" if not specified
+			start := []byte("a")
+			end := []byte("z")
+			if len(td.CmdArgs) >= 2 {
+				start = []byte(td.CmdArgs[0].Key)
+				end = []byte(td.CmdArgs[1].Key)
+			}
+			size, err := d.EstimateDiskUsage(start, end)
+			if err != nil {
+				return err.Error()
+			}
+			return fmt.Sprintf("size: %d", size)
+
+		case "estimate-disk-usage-by-backing-type":
+			// Parse range arguments, default to "a" and "z" if not specified
+			start := []byte("a")
+			end := []byte("z")
+			if len(td.CmdArgs) >= 2 {
+				start = []byte(td.CmdArgs[0].Key)
+				end = []byte(td.CmdArgs[1].Key)
+			}
+			total, remote, external, err := d.EstimateDiskUsageByBackingType(start, end)
+			if err != nil {
+				return err.Error()
+			}
+			// remote and external should be less than or equal to total
+			require.LessOrEqual(t, remote, total)
+			require.LessOrEqual(t, external, remote)
+			return fmt.Sprintf("total: %d, remote: %d, external: %d", total, remote, external)
+
+		default:
+			return fmt.Sprintf("unknown command: %s", td.Cmd)
+		}
+	})
+}

metamorphic/config.go

@@ -28,6 +28,7 @@ const (
 	OpDBFlush
 	OpDBRatchetFormatMajorVersion
 	OpDBRestart
+	OpDBEstimateDiskUsage
 	OpIterClose
 	OpIterFirst
 	OpIterLast
@@ -158,6 +159,7 @@ func DefaultOpConfig() OpConfig {
 			OpDBFlush:                     2,
 			OpDBRatchetFormatMajorVersion: 1,
 			OpDBRestart:                   2,
+			OpDBEstimateDiskUsage:         1,
 			OpIterClose:                   5,
 			OpIterFirst:                   100,
 			OpIterLast:                    100,
@@ -219,6 +221,7 @@ func ReadOpConfig() OpConfig {
 			OpDBFlush:                     0,
 			OpDBRatchetFormatMajorVersion: 0,
 			OpDBRestart:                   0,
+			OpDBEstimateDiskUsage:         0,
 			OpIterClose:                   5,
 			OpIterFirst:                   100,
 			OpIterLast:                    100,
@@ -277,6 +280,7 @@ func WriteOpConfig() OpConfig {
 			OpDBFlush:                     2,
 			OpDBRatchetFormatMajorVersion: 1,
 			OpDBRestart:                   2,
+			OpDBEstimateDiskUsage:         1,
 			OpIterClose:                   0,
 			OpIterFirst:                   0,
 			OpIterLast:                    0,

metamorphic/generator.go

@@ -168,6 +168,7 @@ func (g *generator) generate(count uint64) []op {
 		OpDBFlush:                     g.dbFlush,
 		OpDBRatchetFormatMajorVersion: g.dbRatchetFormatMajorVersion,
 		OpDBRestart:                   g.dbRestart,
+		OpDBEstimateDiskUsage:         g.dbEstimateDiskUsage,
 		OpIterClose:                   g.randIter(g.iterClose),
 		OpIterFirst:                   g.randIter(g.iterFirst),
 		OpIterLast:                    g.randIter(g.iterLast),
@@ -414,6 +415,21 @@ func (g *generator) dbCompact() {
 	})
 }
 
+func (g *generator) dbEstimateDiskUsage() {
+	// Generate new key(s) with a 1% probability.
+	start := g.keyGenerator.RandKey(0.01)
+	end := g.keyGenerator.RandKey(0.01)
+	if g.cmp(start, end) > 0 {
+		start, end = end, start
+	}
+	dbID := g.dbs.rand(g.rng)
+	g.add(&estimateDiskUsageOp{
+		dbID:  dbID,
+		start: start,
+		end:   end,
+	})
+}
+
 func (g *generator) dbDownload() {
 	numSpans := 1 + g.expRandInt(1)
 	spans := make([]pebble.DownloadSpan, numSpans)

metamorphic/ops.go

@@ -2246,3 +2246,38 @@ func (r *replicateOp) rewriteKeys(fn func(UserKey) UserKey) {
 func (r *replicateOp) diagramKeyRanges() []pebble.KeyRange {
 	return []pebble.KeyRange{{Start: r.start, End: r.end}}
 }
+
+// estimateDiskUsageOp models DB.EstimateDiskUsage and DB.EstimateDiskUsageByBackingType operations.
+type estimateDiskUsageOp struct {
+	dbID  objID
+	start UserKey
+	end   UserKey
+}
+
+func (o *estimateDiskUsageOp) run(t *Test, h historyRecorder) {
+	db := t.getDB(o.dbID)
+	_, err := db.EstimateDiskUsage(o.start, o.end)
+	if err != nil {
+		h.Recordf("%s // %v", o.formattedString(t.testOpts.KeyFormat), err)
+	} else {
+		h.Recordf("%s // <OK>", o.formattedString(t.testOpts.KeyFormat))
+	}
+
+}
+
+func (o *estimateDiskUsageOp) formattedString(kf KeyFormat) string {
+	return fmt.Sprintf("%s.EstimateDiskUsage(%q, %q)",
+		o.dbID, kf.FormatKey(o.start), kf.FormatKey(o.end))
+}
+
+func (o *estimateDiskUsageOp) receiver() objID      { return o.dbID }
+func (o *estimateDiskUsageOp) syncObjs() objIDSlice { return nil }
+
+func (o *estimateDiskUsageOp) rewriteKeys(fn func(UserKey) UserKey) {
+	o.start = fn(o.start)
+	o.end = fn(o.end)
+}
+
+func (o *estimateDiskUsageOp) diagramKeyRanges() []pebble.KeyRange {
+	return []pebble.KeyRange{{Start: o.start, End: o.end}}
+}

metamorphic/parser.go

@@ -130,6 +130,8 @@ func opArgs(op op) (receiverID *objID, targetID *objID, args []interface{}) {
 		return &t.writerID, nil, []interface{}{&t.start, &t.end, &t.suffix}
 	case *replicateOp:
 		return &t.source, nil, []interface{}{&t.dest, &t.start, &t.end}
+	case *estimateDiskUsageOp:
+		return &t.dbID, nil, []interface{}{&t.start, &t.end}
 	}
 	panic(fmt.Sprintf("unsupported op type: %T", op))
 }
@@ -146,6 +148,7 @@ var methods = map[string]*methodInfo{
 	"Close":                     makeMethod(closeOp{}, dbTag, batchTag, iterTag, snapTag),
 	"Commit":                    makeMethod(batchCommitOp{}, batchTag),
 	"Compact":                   makeMethod(compactOp{}, dbTag),
+	"EstimateDiskUsage":         makeMethod(estimateDiskUsageOp{}, dbTag),
 	"Delete":                    makeMethod(deleteOp{}, dbTag, batchTag),
 	"DeleteRange":               makeMethod(deleteRangeOp{}, dbTag, batchTag),
 	"Download":                  makeMethod(downloadOp{}, dbTag),