db: datadriven testing for span policy enforcer

Author: xinhaoz

compaction_test.go

@@ -1507,6 +1507,27 @@ func runCompactionTest(
 			s := blobRewriteLog.String()
 			return s
 
+		case "scan-policy-violations":
+			// Run the span policy enforcer's scan to detect violations and mark files.
+			if d.opts.Experimental.SpanPolicyFunc == nil {
+				return "no span policy configured"
+			}
+			enforcer := newSpanPolicyEnforcer(d, SpanPolicyEnforcerOptions{})
+			enforcer.scanOnce()
+			return ""
+
+		case "pending-policy-enforcement":
+			// Show files pending policy enforcement compaction.
+			d.mu.Lock()
+			count := d.mu.compact.policyEnforcementFiles.Count()
+			var buf strings.Builder
+			fmt.Fprintf(&buf, "pending: %d\n", count)
+			for f, level := range d.mu.compact.policyEnforcementFiles.Ascending() {
+				fmt.Fprintf(&buf, "  L%d: %s\n", level, f.TableNum)
+			}
+			d.mu.Unlock()
+			return buf.String()
+
 		case "set-span-policies":
 			var spanPolicies []SpanAndPolicy
 			for line := range crstrings.LinesSeq(td.Input) {
@@ -1552,6 +1573,11 @@ func runCompactionTest(
 							td.Fatalf(t, "parsing minimum-mvcc-garbage-size: %s", err)
 						}
 						policy.ValueStoragePolicy.MinimumMVCCGarbageSize = int(size)
+					case "prefer-fast-compression":
+						if len(parts) != 1 {
+							td.Fatalf(t, "expected prefer-fast-compression with no value, got: %s", arg)
+						}
+						policy.PreferFastCompression = true
 					default:
 						td.Fatalf(t, "unknown span policy arg: %s", arg)
 					}
@@ -1660,6 +1686,11 @@ func TestCompaction(t *testing.T) {
 			maxVersion: FormatNewest,
 			verbose:    true,
 		},
+		"policy_enforcement": {
+			minVersion: FormatNewest,
+			maxVersion: FormatNewest,
+			cmp:        DefaultComparer,
+		},
 	}
 	datadriven.Walk(t, "testdata/compaction", func(t *testing.T, path string) {
 		filename := filepath.Base(path)

data_test.go

@@ -25,6 +25,7 @@ import (
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/pebble/internal/base"
 	"github.com/cockroachdb/pebble/internal/blobtest"
+	"github.com/cockroachdb/pebble/internal/compression"
 	"github.com/cockroachdb/pebble/internal/humanize"
 	"github.com/cockroachdb/pebble/internal/keyspan"
 	"github.com/cockroachdb/pebble/internal/manifest"
@@ -36,6 +37,7 @@ import (
 	"github.com/cockroachdb/pebble/objstorage/objstorageprovider"
 	"github.com/cockroachdb/pebble/objstorage/remote"
 	"github.com/cockroachdb/pebble/sstable"
+	"github.com/cockroachdb/pebble/sstable/block"
 	"github.com/cockroachdb/pebble/sstable/block/blockkind"
 	"github.com/cockroachdb/pebble/sstable/tablefilters/bloom"
 	"github.com/cockroachdb/pebble/valsep"
@@ -1899,6 +1901,32 @@ func parseDBOptionsArgs(opts *Options, args []datadriven.CmdArg) error {
 				Secondary: wal.Dir{FS: opts.FS, Dirname: cmdArg.Vals[0]},
 			}
 			opts.WALFailover.EnsureDefaults()
+		case "compression":
+			var profile block.CompressionProfile
+			switch cmdArg.Vals[0] {
+			case "zstd":
+				profile = *block.ZstdCompression
+			case "snappy":
+				profile = *block.SnappyCompression
+			case "none":
+				profile = *block.NoCompression
+			case "zstd-force":
+				// For testing: Zstd with MinReductionPercent=0 so even small
+				// values are stored compressed.
+				profile = block.CompressionProfile{
+					Name:                "test-zstd-force",
+					DataBlocks:          block.SimpleCompressionSetting(compression.ZstdLevel3),
+					ValueBlocks:         block.SimpleCompressionSetting(compression.ZstdLevel3),
+					OtherBlocks:         compression.ZstdLevel3,
+					MinReductionPercent: 0,
+				}
+			default:
+				return errors.Newf("unrecognized compression %q", cmdArg.Vals[0])
+			}
+			for i := range opts.Levels {
+				p := profile
+				opts.Levels[i].Compression = func() *block.CompressionProfile { return &p }
+			}
 		}
 	}
 	if len(spanPolicies) > 0 {

span_policy_enforcer.go

@@ -373,3 +373,29 @@ func (p *spanPolicyEnforcer) markForEnforcement(f *manifest.TableMetadata, level
 	// Trigger compaction scheduling.
 	p.db.maybeScheduleCompaction()
 }
+
+// scanOnce performs a complete scan through all files in the LSM, checking for
+// policy violations and marking files for enforcement. This is useful for
+// testing where we want to run a full scan without timing concerns.
+func (p *spanPolicyEnforcer) scanOnce() {
+	p.db.mu.Lock()
+	vers := p.db.mu.versions.currentVersion()
+	p.db.mu.Unlock()
+
+	// Reset cursor to start from the beginning.
+	p.cursor = policyEnforcerCursor{level: 0}
+
+	for {
+		nextFile, level := p.nextFile(vers, p.cursor)
+		if nextFile == nil {
+			// Reached end of scan.
+			break
+		}
+
+		if p.checkPolicyViolation(nextFile) {
+			p.markForEnforcement(nextFile, level)
+		}
+
+		p.cursor = makePolicyCursorAfterFile(nextFile, level)
+	}
+}

testdata/compaction/policy_enforcement

@@ -0,0 +1,139 @@
+# Test span policy enforcer detecting compression violations and triggering
+# policy enforcement compactions.
+
+# Create files with Zstd compression using zstd-force (MinReductionPercent=0).
+# The span policy requires fast compression for keys >= "m".
+
+define compression=zstd-force
+L1
+  a#10,SET:a b#10,SET:b
+L1
+  m#10,SET:m n#10,SET:n
+L2
+  x#5,SET:x z#5,SET:z
+----
+L1:
+  000004:[a#10,SET-b#10,SET]
+  000005:[m#10,SET-n#10,SET]
+L2:
+  000006:[x#5,SET-z#5,SET]
+
+# Configure span policy: keys >= "m" require fast compression.
+set-span-policies
+m,zzz prefer-fast-compression
+----
+
+# Initially no pending enforcement files.
+pending-policy-enforcement
+----
+pending: 0
+
+# Run the enforcer scan to detect violations and mark files.
+# Files 000005 (m-n) and 000006 (x-z) are in the policy span and use Zstd.
+scan-policy-violations
+----
+
+# Verify files are now pending enforcement. The enforcer's run loop would pause
+# here until these are processed.
+pending-policy-enforcement
+----
+pending: 2
+  L2: 000006
+  L1: 000005
+
+# Run compaction. The scheduler should pick up policy enforcement compactions.
+auto-compact
+----
+L1:
+  000004:[a#10,SET-b#10,SET]
+  000005:[m#10,SET-n#10,SET]
+L2:
+  000007:[x#0,SET-z#0,SET]
+
+# After auto-compact processes enforcement compactions, pending should be cleared.
+# (auto-compact may process one or both files depending on scheduling)
+pending-policy-enforcement
+----
+pending: 0
+
+# Scan again to check for any remaining violations.
+scan-policy-violations
+----
+
+# One file may still violate policy if not recompacted in first pass.
+pending-policy-enforcement
+----
+pending: 1
+  L1: 000005
+
+auto-compact
+----
+L1:
+  000004:[a#10,SET-b#10,SET]
+  000008:[m#0,SET-n#0,SET]
+L2:
+  000007:[x#0,SET-z#0,SET]
+
+# All enforcement compactions complete.
+pending-policy-enforcement
+----
+pending: 0
+
+# Test interaction with manual compaction: a file marked for enforcement
+# gets moved by a manual compaction. Since the move keeps the same file
+# reference, the enforcement mark remains valid.
+
+define compression=zstd-force
+L1
+  a#10,SET:a b#10,SET:b
+L2
+  m#10,SET:m n#10,SET:n
+L3
+  x#5,SET:x z#5,SET:z
+----
+L1:
+  000004:[a#10,SET-b#10,SET]
+L2:
+  000005:[m#10,SET-n#10,SET]
+L3:
+  000006:[x#5,SET-z#5,SET]
+
+# Set policy and scan to mark files.
+set-span-policies
+m,zzz prefer-fast-compression
+----
+
+scan-policy-violations
+----
+
+# Manually compact the L2 file. This moves 000005 to L3 but keeps the same
+# file reference, so the enforcement mark remains valid.
+compact m-o L2
+----
+L1:
+  000004:[a#10,SET-b#10,SET]
+L3:
+  000005:[m#10,SET-n#10,SET]
+  000006:[x#5,SET-z#5,SET]
+
+# Auto-compact picks up the enforcement compaction for one of the marked files.
+# File 000006 is outside the policy span (x-z < m), so only 000005 was marked.
+auto-compact
+----
+L1:
+  000004:[a#10,SET-b#10,SET]
+L3:
+  000005:[m#10,SET-n#10,SET]
+  000007:[x#0,SET-z#0,SET]
+
+# Scan and compact again - this should compact file 000005.
+scan-policy-violations
+----
+
+auto-compact
+----
+L1:
+  000004:[a#10,SET-b#10,SET]
+L3:
+  000008:[m#0,SET-n#0,SET]
+  000007:[x#0,SET-z#0,SET]