From a0657dbe7b60d7c00d50bdff3f7ecec0dd6ecc52 Mon Sep 17 00:00:00 2001 From: Cody Littley Date: Wed, 15 Jul 2026 16:36:08 -0500 Subject: [PATCH 01/10] implement littDB compression --- go.mod | 2 +- .../litt/disktable/compression_loop.go | 80 +++++++ .../litt/disktable/compression_test.go | 215 ++++++++++++++++++ .../db_engine/litt/disktable/control_loop.go | 37 ++- .../litt/disktable/control_loop_messages.go | 6 + sei-db/db_engine/litt/disktable/disk_table.go | 61 ++++- .../disk_table_secondary_keys_test.go | 4 +- .../litt/disktable/forward_iterator.go | 6 +- .../litt/disktable/segment/metadata_file.go | 89 ++++++-- .../disktable/segment/metadata_file_test.go | 65 +++++- .../litt/disktable/segment/segment.go | 123 +++++++++- .../litt/disktable/segment/segment_reader.go | 6 +- .../litt/disktable/segment/segment_test.go | 10 +- .../litt/disktable/segment/segment_version.go | 9 +- .../segment/shard_id_validation_test.go | 1 + sei-db/db_engine/litt/table_config.go | 17 ++ sei-db/db_engine/litt/types/compression.go | 80 +++++++ .../db_engine/litt/types/compression_test.go | 72 ++++++ 18 files changed, 831 insertions(+), 52 deletions(-) create mode 100644 sei-db/db_engine/litt/disktable/compression_loop.go create mode 100644 sei-db/db_engine/litt/disktable/compression_test.go create mode 100644 sei-db/db_engine/litt/types/compression.go create mode 100644 sei-db/db_engine/litt/types/compression_test.go diff --git a/go.mod b/go.mod index 66bfbda919..576ca68f89 100644 --- a/go.mod +++ b/go.mod @@ -208,7 +208,7 @@ require ( github.com/jmespath/go-jmespath v0.4.0 // indirect github.com/jmhodges/levigo v1.0.0 // indirect github.com/keybase/go-keychain v0.0.0-20190712205309-48d3d31d256d // indirect - github.com/klauspost/compress v1.18.3 // indirect + github.com/klauspost/compress v1.18.3 github.com/klauspost/cpuid/v2 v2.3.0 // indirect github.com/kr/pretty v0.3.1 // indirect github.com/kr/text v0.2.0 // indirect diff --git a/sei-db/db_engine/litt/disktable/compression_loop.go b/sei-db/db_engine/litt/disktable/compression_loop.go new file mode 100644 index 0000000000..0119368884 --- /dev/null +++ b/sei-db/db_engine/litt/disktable/compression_loop.go @@ -0,0 +1,80 @@ +package disktable + +import ( + "fmt" + "log/slog" + + "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/types" + "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/util" +) + +// compressionLoop compresses value bytes off the control-loop goroutine. It sits in front of the control +// loop: when compression is enabled, controlLoop.enqueue sends every control message to inputChannel, +// this loop compresses write requests, and forwards all messages (compressed writes and everything else, +// verbatim) to outputChannel (the control loop's controllerChannel) in arrival order. +// +// Forwarding all message types in order is what makes flush correct: a flush request travels the same +// channel behind the writes it must follow, so the control loop applies those writes first. Because this +// loop is single-threaded and finishes compressing a write before it reads the next message, any in-flight +// compression is complete before a following flush is forwarded; the ordering barrier is automatic. +type compressionLoop struct { + // logger for the compression loop. + logger *slog.Logger + + // errorMonitor is used to react to fatal errors anywhere in the disk table. + errorMonitor *util.ErrorMonitor + + // algorithm is the compression algorithm applied to write-request values. + algorithm types.CompressionAlgorithm + + // inputChannel receives messages from controlLoop.enqueue. + inputChannel chan any + + // outputChannel forwards messages to the control loop (its controllerChannel). + outputChannel chan any +} + +// run processes messages until shutdown. It compresses write requests and forwards every message to the +// control loop in arrival order. +func (cl *compressionLoop) run() { + for { + select { + case <-cl.errorMonitor.ImmediateShutdownRequired(): + return + case message := <-cl.inputChannel: + if req, ok := message.(*controlLoopWriteRequest); ok { + if !cl.compress(req) { + // compress panicked the DB via the error monitor; stop forwarding. + return + } + } + + // Forward every message (compressed writes and all others) in arrival order. + if err := util.Send(cl.errorMonitor, cl.outputChannel, message); err != nil { + return + } + + // The shutdown request is the last message the control loop will process; stop after + // forwarding it so this goroutine does not outlive the table. + if _, ok := message.(*controlLoopShutdownRequest); ok { + return + } + } + } +} + +// compress fills req.compressedValues with the compressed form of each value. It returns false if +// compression failed (in which case it has already panicked the DB via the error monitor). +func (cl *compressionLoop) compress(req *controlLoopWriteRequest) bool { + compressed := make([][]byte, len(req.values)) + for i, kv := range req.values { + blob, err := types.Compress(cl.algorithm, kv.Value) + if err != nil { + cl.errorMonitor.Panic(fmt.Errorf("failed to compress value: %w", err)) + return false + } + compressed[i] = blob + } + req.compressedValues = compressed + return true +} diff --git a/sei-db/db_engine/litt/disktable/compression_test.go b/sei-db/db_engine/litt/disktable/compression_test.go new file mode 100644 index 0000000000..633988c5e5 --- /dev/null +++ b/sei-db/db_engine/litt/disktable/compression_test.go @@ -0,0 +1,215 @@ +package disktable + +import ( + "bytes" + "fmt" + "log/slog" + "path/filepath" + "testing" + "time" + + "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt" + "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/disktable/keymap" + "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/types" + "github.com/stretchr/testify/require" +) + +// buildCompressedMemKeyDiskTable builds a single-shard, mem-keymap disk table with the given compression +// algorithm. Single shard keeps value placement deterministic for size assertions. +func buildCompressedMemKeyDiskTable( + t *testing.T, + clock func() time.Time, + name string, + paths []string, + algorithm types.CompressionAlgorithm, +) litt.ManagedTable { + t.Helper() + logger := slog.Default() + keymapPath := filepath.Join(paths[0], keymap.KeymapDirectoryName) + keymapTypeFile, err := setupKeymapTypeFile(keymapPath, keymap.MemKeymapType) + require.NoError(t, err) + keys, _, err := keymap.NewMemKeymap(logger, "", true) + require.NoError(t, err) + + config, err := litt.DefaultConfig(paths...) + require.NoError(t, err) + config.GCPeriod = time.Hour + config.Fsync = false + config.TargetSegmentFileSize = 1 << 20 + + tableConfig := litt.DefaultTableConfig(name) + tableConfig.ShardingFactor = 1 + tableConfig.Compression = algorithm + + runtimeConfig := litt.DefaultRuntimeConfig() + runtimeConfig.Clock = clock + runtimeConfig.Logger = logger + + table, err := NewDiskTable( + config, + runtimeConfig, + name, + tableConfig, + keys, + keymapPath, + keymapTypeFile, + paths, + true, + nil, + ) + require.NoError(t, err) + return table +} + +// compressiblePayload returns a repetitive (and therefore highly compressible) value. +func compressiblePayload() []byte { + return bytes.Repeat([]byte("littDB compression makes value files smaller. "), 200) +} + +func TestCompressionEndToEnd(t *testing.T) { + t.Parallel() + dir := t.TempDir() + table := buildCompressedMemKeyDiskTable(t, time.Now, "compressed", []string{dir}, types.CompressionS2) + defer func() { require.NoError(t, table.Close()) }() + + key := []byte("key") + value := compressiblePayload() + require.NoError(t, table.Put(key, value)) + + // Before flush: served uncompressed from the write cache. + got, ok, err := table.Get(key) + require.NoError(t, err) + require.True(t, ok) + require.Equal(t, value, got) + + // After flush: served from disk and decompressed. + require.NoError(t, table.Flush()) + got, ok, err = table.Get(key) + require.NoError(t, err) + require.True(t, ok) + require.Equal(t, value, got) + + // The on-disk value file should be meaningfully smaller than the raw value. + require.Less(t, int(table.Size()), len(value), "compressed segment should be smaller than the raw value") +} + +func TestCompressionFullValueAliasSecondary(t *testing.T) { + t.Parallel() + dir := t.TempDir() + table := buildCompressedMemKeyDiskTable(t, time.Now, "aliases", []string{dir}, types.CompressionS2) + defer func() { require.NoError(t, table.Close()) }() + + primary := []byte("block-number-42") + byHash := []byte("block-hash-deadbeef") + value := compressiblePayload() + + // A full-value alias: Offset 0, Length == len(value). This is the block-by-number / block-by-hash case. + alias := &types.SecondaryKey{Key: byHash, Offset: 0, Length: uint32(len(value))} + require.NoError(t, table.Put(primary, value, alias)) + + verify := func(stage string) { + for _, k := range [][]byte{primary, byHash} { + got, ok, err := table.Get(k) + require.NoError(t, err, stage) + require.True(t, ok, stage) + require.Equal(t, value, got, stage) + } + } + + verify("before flush") + require.NoError(t, table.Flush()) + verify("after flush") +} + +func TestCompressionRejectsSubRangeSecondary(t *testing.T) { + t.Parallel() + dir := t.TempDir() + table := buildCompressedMemKeyDiskTable(t, time.Now, "subrange", []string{dir}, types.CompressionS2) + defer func() { require.NoError(t, table.Close()) }() + + value := []byte("the quick brown fox") + // A strict sub-range secondary is not supported on a compressed table. + sub := &types.SecondaryKey{Key: []byte("quick"), Offset: 4, Length: 5} + err := table.Put([]byte("primary"), value, sub) + require.Error(t, err) + require.Contains(t, err.Error(), "compressed table") +} + +func TestCompressionFlushOrdering(t *testing.T) { + t.Parallel() + dir := t.TempDir() + table := buildCompressedMemKeyDiskTable(t, time.Now, "ordering", []string{dir}, types.CompressionS2) + defer func() { require.NoError(t, table.Close()) }() + + // Interleave writes and flushes. Every key written before a Flush() must be durable and readable + // afterwards, which only holds if the flush passes through the compression stage in order. + const rounds = 20 + for i := 0; i < rounds; i++ { + key := []byte(fmt.Sprintf("key-%03d", i)) + value := append(bytes.Repeat([]byte("payload "), 50), []byte(fmt.Sprintf("-%d", i))...) + require.NoError(t, table.Put(key, value)) + require.NoError(t, table.Flush()) + + got, ok, err := table.Get(key) + require.NoError(t, err) + require.True(t, ok) + require.Equal(t, value, got) + } + + // Re-read every key from disk after all flushes. + for i := 0; i < rounds; i++ { + key := []byte(fmt.Sprintf("key-%03d", i)) + want := append(bytes.Repeat([]byte("payload "), 50), []byte(fmt.Sprintf("-%d", i))...) + got, ok, err := table.Get(key) + require.NoError(t, err) + require.True(t, ok) + require.Equal(t, want, got) + } +} + +func TestCompressionIteration(t *testing.T) { + t.Parallel() + dir := t.TempDir() + table := buildCompressedMemKeyDiskTable(t, time.Now, "iterate", []string{dir}, types.CompressionS2) + defer func() { require.NoError(t, table.Close()) }() + + type record struct { + key string + value []byte + } + records := make([]record, 0, 10) + for i := 0; i < 10; i++ { + key := fmt.Sprintf("k-%02d", i) + value := append(compressiblePayload(), byte(i)) + records = append(records, record{key: key, value: value}) + } + + // Write each record with a full-value-alias secondary so iteration exercises grouped keys too. + for _, r := range records { + alias := &types.SecondaryKey{Key: []byte(r.key + "-alias"), Offset: 0, Length: uint32(len(r.value))} + require.NoError(t, table.Put([]byte(r.key), r.value, alias)) + } + require.NoError(t, table.Flush()) + + // Forward iteration returns primary and alias, both decompressed to the full value. + it, err := table.Iterator(false) + require.NoError(t, err) + seen := make(map[string][]byte) + for { + ok, err := it.Next() + require.NoError(t, err) + if !ok { + break + } + key, _ := it.GetKey() + value, err := it.GetValue() + require.NoError(t, err) + seen[string(key)] = append([]byte(nil), value...) + } + require.NoError(t, it.Close()) + + for _, r := range records { + require.Equal(t, r.value, seen[r.key], "primary %s", r.key) + require.Equal(t, r.value, seen[r.key+"-alias"], "alias for %s", r.key) + } +} diff --git a/sei-db/db_engine/litt/disktable/control_loop.go b/sei-db/db_engine/litt/disktable/control_loop.go index 506b6cfe91..e57f6132a4 100644 --- a/sei-db/db_engine/litt/disktable/control_loop.go +++ b/sei-db/db_engine/litt/disktable/control_loop.go @@ -11,6 +11,7 @@ import ( "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/disktable/keymap" "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/disktable/segment" "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/metrics" + "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/types" "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/util" ) @@ -24,9 +25,20 @@ type controlLoop struct { // errorMonitor is used to react to fatal errors anywhere in the disk table. errorMonitor *util.ErrorMonitor - // controllerChannel is the channel for messages sent to the control loop. + // controllerChannel is the channel the control loop reads messages from. When compression is + // disabled, enqueue sends directly here (inputChannel == controllerChannel). When compression is + // enabled, the compression loop reads inputChannel and forwards messages here. controllerChannel chan any + // inputChannel is the channel enqueue sends to. It is controllerChannel when compression is disabled, + // or the compression loop's input channel when compression is enabled (so every control message, + // including flush, passes through the compression stage in order). + inputChannel chan any + + // compressionAlgorithm is the algorithm new segments are created with. types.CompressionNone means + // segments store values verbatim. + compressionAlgorithm types.CompressionAlgorithm + // The index of the lowest numbered segment. It is advanced only by the control loop, in // deleteEligibleSegments, as collected segments' files are removed. Only the control loop goroutine touches it. lowestSegmentIndex uint32 @@ -150,7 +162,7 @@ type controlLoop struct { // database being in a panicked state. Only types defined in control_loop_messages.go are permitted to be sent // to the control loop. func (c *controlLoop) enqueue(request controlLoopMessage) error { - return util.Send(c.errorMonitor, c.controllerChannel, request) + return util.Send(c.errorMonitor, c.inputChannel, request) } // run runs the control loop for the disk table. It has sole responsibility for scheduling all operations that @@ -468,14 +480,21 @@ func (c *controlLoop) updateCurrentSize() { // handleWriteRequest handles a controlLoopWriteRequest control message. func (c *controlLoop) handleWriteRequest(req *controlLoopWriteRequest) { - for _, kv := range req.values { + for i, kv := range req.values { // Do the write. seg := c.segments[c.highestSegmentIndex] + // The number of bytes actually written to the value file: the compressed blob when compression is + // enabled, otherwise the raw value. + onDiskLen := uint64(len(kv.Value)) + if req.compressedValues != nil { + onDiskLen = uint64(len(req.compressedValues[i])) + } + // Roll to a fresh segment before writing if this value's bytes would cross the 2^32 addressable // limit of the segment's value files (offsets are stored as uint32, so a value's first byte must // sit below 2^32). - if seg.GetMaxShardSize()+uint64(len(kv.Value)) > math.MaxUint32 { + if seg.GetMaxShardSize()+onDiskLen > math.MaxUint32 { if err := c.expandSegments(); err != nil { c.errorMonitor.Panic(fmt.Errorf("failed to expand segments: %w", err)) return @@ -491,7 +510,14 @@ func (c *controlLoop) handleWriteRequest(req *controlLoopWriteRequest) { } c.newestPrimaryKey = kv.Key - keyCount, keyFileSize, err := seg.Write(kv) + var keyCount uint32 + var keyFileSize uint64 + var err error + if req.compressedValues != nil { + keyCount, keyFileSize, err = seg.WriteCompressed(kv, req.compressedValues[i]) + } else { + keyCount, keyFileSize, err = seg.Write(kv) + } shardSize := seg.GetMaxShardSize() if err != nil { c.errorMonitor.Panic( @@ -564,6 +590,7 @@ func (c *controlLoop) expandSegments() error { c.segmentPaths, c.snapshottingEnabled, c.diskTable.getShardingFactor(), + c.compressionAlgorithm, c.fsync, c.shardControlChannelSize) if err != nil { diff --git a/sei-db/db_engine/litt/disktable/control_loop_messages.go b/sei-db/db_engine/litt/disktable/control_loop_messages.go index 5d5f1b50eb..ae2481983e 100644 --- a/sei-db/db_engine/litt/disktable/control_loop_messages.go +++ b/sei-db/db_engine/litt/disktable/control_loop_messages.go @@ -28,6 +28,12 @@ type controlLoopWriteRequest struct { // values is a slice of key-value pairs to write. values []*types.PutRequest + + // compressedValues holds the compressed on-disk representation of each value, populated by the + // compression stage when the table has compression enabled. When non-nil it is parallel to values + // (compressedValues[i] is the compressed form of values[i].Value) and the control loop writes those + // bytes to the value file. When nil, values are written verbatim (the uncompressed path). + compressedValues [][]byte } // controlLoopSetShardingFactorRequest is a request to set the sharding factor that is sent to the control loop. diff --git a/sei-db/db_engine/litt/disktable/disk_table.go b/sei-db/db_engine/litt/disktable/disk_table.go index 0f7432625f..35410cc5cc 100644 --- a/sei-db/db_engine/litt/disktable/disk_table.go +++ b/sei-db/db_engine/litt/disktable/disk_table.go @@ -97,6 +97,10 @@ type DiskTable struct { // If true then ensure file operations are synced to disk. fsync bool + // The algorithm used to compress values written to new segments. types.CompressionNone means values + // are stored verbatim. Held only in memory; each segment records its own algorithm for reads. + compressionAlgorithm types.CompressionAlgorithm + // Manages flush requests and flush request batching. This is a performance optimization. flushCoordinator *flushCoordinator } @@ -146,17 +150,18 @@ func NewDiskTable( errorMonitor := util.NewErrorMonitor(runtimeConfig.CTX, runtimeConfig.Logger, runtimeConfig.FatalErrorCallback) table := &DiskTable{ - logger: runtimeConfig.Logger, - errorMonitor: errorMonitor, - clock: runtimeConfig.Clock, - roots: qualifiedRoots, - segmentPaths: segmentPaths, - name: name, - keymap: keymap, - keymapPath: keymapPath, - keymapTypeFile: keymapTypeFile, - metrics: metrics, - fsync: config.Fsync, + logger: runtimeConfig.Logger, + errorMonitor: errorMonitor, + clock: runtimeConfig.Clock, + roots: qualifiedRoots, + segmentPaths: segmentPaths, + name: name, + keymap: keymap, + keymapPath: keymapPath, + keymapTypeFile: keymapTypeFile, + metrics: metrics, + fsync: config.Fsync, + compressionAlgorithm: tableConfig.Compression, } // Sharding factor is supplied at creation time and held only in memory; it is not persisted across restarts. // (TTL is likewise in-memory, but it lives on the GC manager — its only consumer — and is seeded there.) @@ -206,6 +211,7 @@ func NewDiskTable( segmentPaths, snapshottingEnabled, table.getShardingFactor(), + table.compressionAlgorithm, config.Fsync, config.ShardControlChannelSize) if err != nil { @@ -377,9 +383,29 @@ func NewDiskTable( immutableSegmentSize: immutableSegmentSize, deletionWatermarkChan: make(chan int64, config.KeymapManagerWatermarkChannelSize), keymapDeletionWatermark: initialDeletionWatermark, + compressionAlgorithm: tableConfig.Compression, } cLoop.lowestSegmentIndex = lowestSegmentIndex cLoop.threadsafeHighestSegmentIndex.Store(highestSegmentIndex) + + // Wire the compression stage in front of the control loop when compression is enabled. enqueue sends + // to inputChannel; the compression loop compresses write requests and forwards every message (in + // order, so flush stays ordered behind its writes) to controllerChannel. When compression is + // disabled, enqueue targets controllerChannel directly and no compression goroutine runs. + var cmpLoop *compressionLoop + if tableConfig.Compression == types.CompressionNone { + cLoop.inputChannel = cLoop.controllerChannel + } else { + cmpLoop = &compressionLoop{ + logger: runtimeConfig.Logger, + errorMonitor: errorMonitor, + algorithm: tableConfig.Compression, + inputChannel: make(chan any, config.ControlChannelSize), + outputChannel: cLoop.controllerChannel, + } + cLoop.inputChannel = cmpLoop.inputChannel + } + table.controlLoop = cLoop cLoop.updateCurrentSize() @@ -413,6 +439,9 @@ func NewDiskTable( go fLoop.run() go cLoop.run() go gcMgr.run() + if cmpLoop != nil { + go cmpLoop.run() + } return table, nil } @@ -998,6 +1027,16 @@ func (d *DiskTable) PutBatch(batch []*types.PutRequest) error { return fmt.Errorf( "secondary key range [%d, %d) exceeds value length %d", sk.Offset, end, len(kv.Value)) } + // On a compressed table, a secondary key may only alias the entire value. A compressed blob + // cannot be sliced, so a strict sub-range would require storing a second (duplicated) + // compressed copy, which is a documented future optimization rather than current behavior. + if d.compressionAlgorithm != types.CompressionNone && + (sk.Offset != 0 || uint64(sk.Length) != uint64(len(kv.Value))) { + return fmt.Errorf( + "secondary key range [%d, %d) is a strict sub-range of the value (length %d), which is "+ + "not supported on a compressed table; secondary keys must alias the entire value", + sk.Offset, end, len(kv.Value)) + } skKey := util.UnsafeBytesToString(sk.Key) if _, dup := seen[skKey]; dup { return fmt.Errorf("duplicate key %x within PutRequest", sk.Key) diff --git a/sei-db/db_engine/litt/disktable/disk_table_secondary_keys_test.go b/sei-db/db_engine/litt/disktable/disk_table_secondary_keys_test.go index f56d9bad88..777e749356 100644 --- a/sei-db/db_engine/litt/disktable/disk_table_secondary_keys_test.go +++ b/sei-db/db_engine/litt/disktable/disk_table_secondary_keys_test.go @@ -448,8 +448,8 @@ func TestGroupAtomicRecoveryEndToEnd(t *testing.T) { metaPath := path.Join(segmentDir, fmt.Sprintf("%d%s", segIdx, segment.MetadataFileExtension)) mdBytes, err := os.ReadFile(metaPath) require.NoError(t, err) - require.Equal(t, segment.V3MetadataSize, len(mdBytes)) - mdBytes[segment.V3MetadataSize-1] = 0 + require.Equal(t, segment.V4MetadataSize, len(mdBytes)) + mdBytes[segment.MetadataSealedByteOffset] = 0 require.NoError(t, os.WriteFile(metaPath, mdBytes, 0600)) // Reopen. diff --git a/sei-db/db_engine/litt/disktable/forward_iterator.go b/sei-db/db_engine/litt/disktable/forward_iterator.go index cb7295acbe..715f10ff8b 100644 --- a/sei-db/db_engine/litt/disktable/forward_iterator.go +++ b/sei-db/db_engine/litt/disktable/forward_iterator.go @@ -143,7 +143,11 @@ func (it *forwardIterator) GetValue() (value []byte, err error) { // The current key is a secondary key. Its primary was visited immediately before it, so we can serve // the value as a sub-slice of the (lazily loaded) primary value. A secondary always references a // sub-range of its primary's value on the same shard. - if it.groupValid && it.secondaryWithinGroup(addr) { + // + // This optimization is invalid for compressed segments: there addr.ValueSize() is the compressed blob + // length while groupValue holds the decompressed value, so the sub-slice arithmetic would be wrong. + // Fall through to a direct read (which decompresses) instead. + if !it.currentSeg.IsCompressed() && it.groupValid && it.secondaryWithinGroup(addr) { if it.groupValue == nil { v, err := reader.Read(it.groupAddr) if err != nil { diff --git a/sei-db/db_engine/litt/disktable/segment/metadata_file.go b/sei-db/db_engine/litt/disktable/segment/metadata_file.go index 1f53836ba3..cfe431be43 100644 --- a/sei-db/db_engine/litt/disktable/segment/metadata_file.go +++ b/sei-db/db_engine/litt/disktable/segment/metadata_file.go @@ -8,6 +8,7 @@ import ( "strconv" "time" + "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/types" "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/util" ) @@ -22,15 +23,29 @@ const ( // deleted. MetadataSwapExtension = MetadataFileExtension + util.SwapFileExtension - // V3MetadataSize is the size of the metadata file at the current LatestSegmentVersion (the name - // is kept for backwards compatibility; the metadata layout has not actually changed since v3). + // V3MetadataSize is the size of a version-3 metadata file. Version 4 appends one byte (see + // V4MetadataSize); the first V3MetadataSize bytes of a v4 file are identical to a v3 file, which is + // what lets deserialize read both formats. // Layout: - // - 4 bytes for version - // - 1 byte for the sharding factor - // - 8 bytes for lastValueTimestamp - // - 4 bytes for keyCount - // - 1 byte for sealed + // - 4 bytes for version (offset 0) + // - 1 byte for the sharding factor (offset 4) + // - 8 bytes for lastValueTimestamp (offset 5) + // - 4 bytes for keyCount (offset 13) + // - 1 byte for sealed (offset 17) V3MetadataSize = 18 + + // V4MetadataSize is the size of a version-4 metadata file: the v3 layout followed by one + // compression-algorithm byte (offset 18). This is the size of metadata files written by the current + // build. + V4MetadataSize = 19 + + // MetadataSealedByteOffset is the byte offset of the sealed flag within a metadata file. It is the + // same in every version. Exposed so tests can simulate a crash-before-seal by flipping this byte. + MetadataSealedByteOffset = 17 + + // metadataCompressionByteOffset is the byte offset of the compression-algorithm byte in a v4 metadata + // file. + metadataCompressionByteOffset = 18 ) // metadataFile contains metadata about a segment. This file contains metadata about the data segment, such as @@ -60,6 +75,11 @@ type metadataFile struct { // to this segment. This value is encoded in the file. sealed bool + // The algorithm used to compress values written to this segment. Values in the segment's value files are + // decompressed with this algorithm on read. CompressionNone means values are stored verbatim. This value is + // encoded in the file (v4+); a v3 file has no compression byte and is read as CompressionNone. + compressionAlgorithm types.CompressionAlgorithm + // Path data for the segment file. This information is not serialized in the metadata file. segmentPath *SegmentPath @@ -73,14 +93,16 @@ type metadataFile struct { func createMetadataFile( index uint32, shardingFactor uint8, + compressionAlgorithm types.CompressionAlgorithm, path *SegmentPath, fsync bool, ) (*metadataFile, error) { file := &metadataFile{ - index: index, - segmentPath: path, - fsync: fsync, + index: index, + segmentPath: path, + fsync: fsync, + compressionAlgorithm: compressionAlgorithm, } file.segmentVersion = LatestSegmentVersion @@ -139,7 +161,7 @@ func getMetadataFileIndex(fileName string) (uint32, error) { // Size returns the size of the metadata file in bytes. func (m *metadataFile) Size() uint64 { - return V3MetadataSize + return V4MetadataSize } // Name returns the file name for this metadata file. @@ -165,44 +187,63 @@ func (m *metadataFile) seal(now time.Time, keyCount uint32) error { return nil } -// serialize serializes the metadata file to a byte array. +// serialize serializes the metadata file to a byte array. Metadata is always written at +// LatestSegmentVersion (v4), including the trailing compression-algorithm byte. func (m *metadataFile) serialize() []byte { - data := make([]byte, V3MetadataSize) + data := make([]byte, V4MetadataSize) - binary.BigEndian.PutUint32(data[0:4], uint32(m.segmentVersion)) + binary.BigEndian.PutUint32(data[0:4], uint32(LatestSegmentVersion)) data[4] = m.shardingFactor binary.BigEndian.PutUint64(data[5:13], m.lastValueTimestamp) binary.BigEndian.PutUint32(data[13:17], m.keyCount) if m.sealed { - data[17] = 1 + data[MetadataSealedByteOffset] = 1 } else { - data[17] = 0 + data[MetadataSealedByteOffset] = 0 } + data[metadataCompressionByteOffset] = byte(m.compressionAlgorithm) return data } -// deserialize deserializes the metadata file from a byte array. +// deserialize deserializes the metadata file from a byte array. Both version 3 (no compression byte) and +// version 4 (with a compression byte) are accepted; a v3 file is read as CompressionNone. func (m *metadataFile) deserialize(data []byte) error { if len(data) < 4 { return fmt.Errorf("metadata file is not the correct size, expected at least 4 bytes, got %d", len(data)) } m.segmentVersion = SegmentVersion(binary.BigEndian.Uint32(data[0:4])) - if m.segmentVersion != LatestSegmentVersion { - return fmt.Errorf("unsupported segment version: %d (only version %d is supported)", - m.segmentVersion, LatestSegmentVersion) + + var expectedSize int + switch m.segmentVersion { + case ShardedAddressSegmentVersion: + expectedSize = V3MetadataSize + case CompressedSegmentVersion: + expectedSize = V4MetadataSize + default: + return fmt.Errorf("unsupported segment version: %d (only versions %d and %d are supported)", + m.segmentVersion, ShardedAddressSegmentVersion, CompressedSegmentVersion) } - if len(data) != V3MetadataSize { - return fmt.Errorf("metadata file is not the correct size, expected %d, got %d", - V3MetadataSize, len(data)) + if len(data) != expectedSize { + return fmt.Errorf("metadata file is not the correct size, expected %d for version %d, got %d", + expectedSize, m.segmentVersion, len(data)) } m.shardingFactor = data[4] m.lastValueTimestamp = binary.BigEndian.Uint64(data[5:13]) m.keyCount = binary.BigEndian.Uint32(data[13:17]) - m.sealed = data[17] == 1 + m.sealed = data[MetadataSealedByteOffset] == 1 + + if m.segmentVersion == CompressedSegmentVersion { + m.compressionAlgorithm = types.CompressionAlgorithm(data[metadataCompressionByteOffset]) + if err := m.compressionAlgorithm.Validate(); err != nil { + return fmt.Errorf("invalid compression algorithm in metadata file: %w", err) + } + } else { + m.compressionAlgorithm = types.CompressionNone + } return nil } diff --git a/sei-db/db_engine/litt/disktable/segment/metadata_file_test.go b/sei-db/db_engine/litt/disktable/segment/metadata_file_test.go index aaa3690768..6560071c06 100644 --- a/sei-db/db_engine/litt/disktable/segment/metadata_file_test.go +++ b/sei-db/db_engine/litt/disktable/segment/metadata_file_test.go @@ -1,9 +1,11 @@ package segment import ( + "encoding/binary" "os" "testing" + "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/types" "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/util" "github.com/stretchr/testify/require" ) @@ -53,6 +55,65 @@ func TestUnsealedSerialization(t *testing.T) { require.True(t, os.IsNotExist(err)) } +func TestCompressionAlgorithmSerialization(t *testing.T) { + t.Parallel() + rand := util.NewTestRandom() + directory := t.TempDir() + + index := rand.Uint32() + segmentPath, err := NewSegmentPath(directory, "", "table") + require.NoError(t, err) + err = segmentPath.MakeDirectories(false) + require.NoError(t, err) + + m, err := createMetadataFile(index, 4, types.CompressionS2, segmentPath, false) + require.NoError(t, err) + require.Equal(t, types.CompressionS2, m.compressionAlgorithm) + require.Equal(t, LatestSegmentVersion, m.segmentVersion) + + // The on-disk file is the v4 size, and the algorithm survives a round trip. + stat, err := os.Stat(m.path()) + require.NoError(t, err) + require.Equal(t, int64(V4MetadataSize), stat.Size()) + + deserialized, err := loadMetadataFile(index, []*SegmentPath{segmentPath}, false) + require.NoError(t, err) + require.Equal(t, types.CompressionS2, deserialized.compressionAlgorithm) + require.Equal(t, *m, *deserialized) +} + +// TestV3MetadataReadsAsUncompressed verifies that a legacy version-3 metadata file (which has no +// compression byte) still loads, and is interpreted as CompressionNone. +func TestV3MetadataReadsAsUncompressed(t *testing.T) { + t.Parallel() + rand := util.NewTestRandom() + directory := t.TempDir() + + index := rand.Uint32() + segmentPath, err := NewSegmentPath(directory, "", "table") + require.NoError(t, err) + err = segmentPath.MakeDirectories(false) + require.NoError(t, err) + + // Write a v4 file, then rewrite it in the legacy v3 layout: version 3 and no trailing compression + // byte (18 bytes instead of 19). + m, err := createMetadataFile(index, 7, types.CompressionNone, segmentPath, false) + require.NoError(t, err) + + data, err := os.ReadFile(m.path()) + require.NoError(t, err) + require.Equal(t, V4MetadataSize, len(data)) + v3 := data[:V3MetadataSize] + binary.BigEndian.PutUint32(v3[0:4], uint32(ShardedAddressSegmentVersion)) + require.NoError(t, os.WriteFile(m.path(), v3, 0600)) + + deserialized, err := loadMetadataFile(index, []*SegmentPath{segmentPath}, false) + require.NoError(t, err) + require.Equal(t, ShardedAddressSegmentVersion, deserialized.segmentVersion) + require.Equal(t, types.CompressionNone, deserialized.compressionAlgorithm) + require.EqualValues(t, 7, deserialized.shardingFactor) +} + func TestSealedSerialization(t *testing.T) { t.Parallel() rand := util.NewTestRandom() @@ -108,7 +169,7 @@ func TestFreshFileSerialization(t *testing.T) { require.NoError(t, err) err = segmentPath.MakeDirectories(false) require.NoError(t, err) - m, err := createMetadataFile(index, 123, segmentPath, false) + m, err := createMetadataFile(index, 123, types.CompressionNone, segmentPath, false) require.NoError(t, err) require.Equal(t, index, m.index) @@ -148,7 +209,7 @@ func TestSealing(t *testing.T) { require.NoError(t, err) err = segmentPath.MakeDirectories(false) require.NoError(t, err) - m, err := createMetadataFile(index, 123, segmentPath, false) + m, err := createMetadataFile(index, 123, types.CompressionNone, segmentPath, false) require.NoError(t, err) // seal the file diff --git a/sei-db/db_engine/litt/disktable/segment/segment.go b/sei-db/db_engine/litt/disktable/segment/segment.go index 699fce41d0..09f6a8b0d6 100644 --- a/sei-db/db_engine/litt/disktable/segment/segment.go +++ b/sei-db/db_engine/litt/disktable/segment/segment.go @@ -109,6 +109,7 @@ func CreateSegment( segmentPaths []*SegmentPath, snapshottingEnabled bool, shardingFactor uint8, + compressionAlgorithm types.CompressionAlgorithm, fsync bool, shardChannelCapacity int, ) (*Segment, error) { @@ -117,7 +118,7 @@ func CreateSegment( return nil, errors.New("no segment paths provided") } - metadata, err := createMetadataFile(index, shardingFactor, segmentPaths[0], fsync) + metadata, err := createMetadataFile(index, shardingFactor, compressionAlgorithm, segmentPaths[0], fsync) if err != nil { return nil, fmt.Errorf("failed to open metadata file: %v", err) } @@ -564,6 +565,103 @@ func (s *Segment) Write(data *types.PutRequest) (keyCount uint32, keyFileSize ui return s.keyCount, s.keyFileSize, nil } +// WriteCompressed writes a value whose on-disk representation is the pre-compressed blob compressedValue +// (produced by the compression stage from data.Value). It behaves like Write except that the bytes +// written to the value file are compressedValue, and every key in the group is addressed to that single +// blob: reads fetch the blob and decompress it back to the whole value. +// +// Only full-value-alias secondary keys (Offset == 0 && Length == len(data.Value)) are valid on a +// compressed segment; PutBatch rejects any other secondary before the request reaches here. A +// compressed blob cannot be sliced, so a full-value alias shares the primary's blob (and therefore its +// Address) rather than storing a second copy. +func (s *Segment) WriteCompressed( + data *types.PutRequest, + compressedValue []byte, +) (keyCount uint32, keyFileSize uint64, err error) { + if s.metadata.sealed { + return 0, 0, fmt.Errorf("segment is sealed, cannot write data") + } + + // Shard assignment is round-robin, exactly as in Write; see that method for why this needs no lock. + shard := s.nextShard + s.nextShard++ + if s.nextShard == s.metadata.shardingFactor { + s.nextShard = 0 + } + currentSize := s.shardSizes[shard] + + if currentSize > math.MaxUint32 { + return 0, 0, + fmt.Errorf("value file already contains %d bytes, cannot add a new value", currentSize) + } + firstByteIndex := uint32(currentSize) + blobLen := uint64(len(compressedValue)) + + n := len(data.SecondaryKeys) + totalKeys := uint32(1 + n) //nolint:gosec // n bounded by caller validation + + primaryKind := types.KeyKindStandalone + if n > 0 { + primaryKind = types.KeyKindPrimary + } + + // Update accounting before sending so that callers observe consistent state. The shard grows by the + // compressed length, since that is what is written to disk. + s.unflushedKeyCount.Add(int64(totalKeys)) + s.shardSizes[shard] += blobLen + if s.shardSizes[shard] > s.maxShardSize { + s.maxShardSize = s.shardSizes[shard] + } + s.keyCount += totalKeys + s.keyFileSize += keyRecordSize(data.Key) + for _, sk := range data.SecondaryKeys { + s.keyFileSize += keyRecordSize(sk.Key) + } + + shardRequest := &valueToWrite{ + value: compressedValue, + expectedFirstByteIndex: firstByteIndex, + } + err = util.Send(s.errorMonitor, s.shardChannels[shard], shardRequest) + if err != nil { + return 0, 0, + fmt.Errorf("failed to send value to shard control loop: %v", err) + } + + // The whole compressed blob is the on-disk representation of the value; the primary and every + // (full-value-alias) secondary are addressed to it. Reads decompress the blob to recover the value. + blobAddress := types.NewAddress(s.index, firstByteIndex, shard, uint32(blobLen)) //nolint:gosec // bounded above + + primaryRequest := &types.ScopedKey{ + Key: data.Key, + Address: blobAddress, + Kind: primaryKind, + } + err = util.Send(s.errorMonitor, s.keyFileChannel, primaryRequest) + if err != nil { + return 0, 0, + fmt.Errorf("failed to send key to key file control loop: %v", err) + } + + for i, sk := range data.SecondaryKeys { + kind := types.KeyKindSecondary + if i == n-1 { + kind = types.KeyKindFinalSecondary + } + secondaryRequest := &types.ScopedKey{ + Key: sk.Key, + Address: blobAddress, // full-value alias: shares the primary's compressed blob + Kind: kind, + } + err = util.Send(s.errorMonitor, s.keyFileChannel, secondaryRequest) + if err != nil { + return 0, 0, fmt.Errorf("failed to send secondary key to key file control loop: %v", err) + } + } + + return s.keyCount, s.keyFileSize, nil +} + // keyRecordSize returns the number of bytes a key file record consumes given a key of the supplied // length. Includes the kind byte (1), the uint16 key-length prefix (2), the key bytes, and the // fixed-width serialized address. @@ -600,7 +698,28 @@ func (s *Segment) Read(key []byte, dataAddress types.Address) ([]byte, error) { if err != nil { return nil, fmt.Errorf("failed to read value: %w", err) } - return value, nil + return s.maybeDecompress(value) +} + +// maybeDecompress decompresses value using the segment's compression algorithm, or returns it unchanged +// if the segment is not compressed. All value reads (Segment.Read and SegmentReader.Read) pass through +// here so the on-disk compressed representation is never surfaced to callers. +func (s *Segment) maybeDecompress(value []byte) ([]byte, error) { + if s.metadata.compressionAlgorithm == types.CompressionNone { + return value, nil + } + decompressed, err := types.Decompress(s.metadata.compressionAlgorithm, value) + if err != nil { + return nil, fmt.Errorf("failed to decompress value: %w", err) + } + return decompressed, nil +} + +// IsCompressed reports whether values in this segment are stored compressed. Callers that read values +// through a path other than Segment.Read/SegmentReader.Read use this to avoid treating compressed bytes +// as raw value bytes. +func (s *Segment) IsCompressed() bool { + return s.metadata.compressionAlgorithm != types.CompressionNone } // GetKeys returns all keys in the data segment. Only permitted to be called after the segment has been sealed. diff --git a/sei-db/db_engine/litt/disktable/segment/segment_reader.go b/sei-db/db_engine/litt/disktable/segment/segment_reader.go index 8883fba35f..9f2ff0ac41 100644 --- a/sei-db/db_engine/litt/disktable/segment/segment_reader.go +++ b/sei-db/db_engine/litt/disktable/segment/segment_reader.go @@ -55,7 +55,11 @@ func (r *SegmentReader) Read(address types.Address) ([]byte, error) { r.readers[shardID] = reader } - return r.readers[shardID].read(address.Offset(), address.ValueSize()) + value, err := r.readers[shardID].read(address.Offset(), address.ValueSize()) + if err != nil { + return nil, err + } + return r.segment.maybeDecompress(value) } // Close releases all file handles held by the reader. It is safe to call more than once. diff --git a/sei-db/db_engine/litt/disktable/segment/segment_test.go b/sei-db/db_engine/litt/disktable/segment/segment_test.go index c25847e536..cec619a5ee 100644 --- a/sei-db/db_engine/litt/disktable/segment/segment_test.go +++ b/sei-db/db_engine/litt/disktable/segment/segment_test.go @@ -59,6 +59,7 @@ func TestWriteAndReadSegmentSingleShard(t *testing.T) { []*SegmentPath{segmentPath}, false, 1, + types.CompressionNone, false, 32) @@ -209,6 +210,7 @@ func TestWriteAndReadSegmentMultiShard(t *testing.T) { []*SegmentPath{segmentPath}, false, shardCount, + types.CompressionNone, false, 32) @@ -368,6 +370,7 @@ func TestWriteAndReadColdShard(t *testing.T) { []*SegmentPath{segmentPath}, false, shardCount, + types.CompressionNone, false, 32) @@ -478,6 +481,7 @@ func TestGetFilePaths(t *testing.T) { []*SegmentPath{segmentPath}, false, shardingFactor, + types.CompressionNone, false, 32) require.NoError(t, err) @@ -546,6 +550,7 @@ func TestRoundRobinShardAssignment(t *testing.T) { []*SegmentPath{segmentPath}, false, shardingFactor, + types.CompressionNone, false, 32) require.NoError(t, err) @@ -617,6 +622,7 @@ func newSingleShardSegment(t *testing.T) (*Segment, *SegmentPath, uint32) { []*SegmentPath{segmentPath}, false, 1, + types.CompressionNone, false, 32, ) @@ -776,8 +782,8 @@ func markSegmentUnsealed(t *testing.T, segmentPath *SegmentPath, index uint32) { metaPath := path.Join(segmentPath.SegmentDirectory(), fmt.Sprintf("%d%s", index, MetadataFileExtension)) data, err := os.ReadFile(metaPath) require.NoError(t, err) - require.Equal(t, V3MetadataSize, len(data)) - data[V3MetadataSize-1] = 0 + require.Equal(t, V4MetadataSize, len(data)) + data[MetadataSealedByteOffset] = 0 require.NoError(t, os.WriteFile(metaPath, data, 0600)) } diff --git a/sei-db/db_engine/litt/disktable/segment/segment_version.go b/sei-db/db_engine/litt/disktable/segment/segment_version.go index a2c406e9bf..95bfa49545 100644 --- a/sei-db/db_engine/litt/disktable/segment/segment_version.go +++ b/sei-db/db_engine/litt/disktable/segment/segment_version.go @@ -25,7 +25,14 @@ const ( // instance of this codebase has been deployed to production, so there is no compatibility cost to // folding the new format into the same version number rather than bumping it. ShardedAddressSegmentVersion SegmentVersion = 3 + + // CompressedSegmentVersion adds a single compression-algorithm byte to the end of the metadata file + // (see V4MetadataSize). The key-file and value-file layouts are unchanged from + // ShardedAddressSegmentVersion, so a v4 segment written with CompressionNone is byte-compatible with + // a v3 segment; the version bump exists only to make the wider metadata format explicit and to keep + // reads of pre-existing v3 metadata files working. + CompressedSegmentVersion SegmentVersion = 4 ) // LatestSegmentVersion always refers to the latest version of the segment serialization format. -const LatestSegmentVersion = ShardedAddressSegmentVersion +const LatestSegmentVersion = CompressedSegmentVersion diff --git a/sei-db/db_engine/litt/disktable/segment/shard_id_validation_test.go b/sei-db/db_engine/litt/disktable/segment/shard_id_validation_test.go index b0a29740f5..6eccbd087d 100644 --- a/sei-db/db_engine/litt/disktable/segment/shard_id_validation_test.go +++ b/sei-db/db_engine/litt/disktable/segment/shard_id_validation_test.go @@ -34,6 +34,7 @@ func TestSegmentReadRejectsOutOfRangeShardID(t *testing.T) { []*SegmentPath{segmentPath}, false, shardingFactor, + types.CompressionNone, false, 32) require.NoError(t, err) diff --git a/sei-db/db_engine/litt/table_config.go b/sei-db/db_engine/litt/table_config.go index 9805710f44..7b01f5f84e 100644 --- a/sei-db/db_engine/litt/table_config.go +++ b/sei-db/db_engine/litt/table_config.go @@ -3,6 +3,8 @@ package litt import ( "fmt" "time" + + "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/types" ) // TableConfig is the per-table configuration supplied at table creation time via DB.BuildTable. With the @@ -46,6 +48,17 @@ type TableConfig struct { // A function that is called to determine if a key is eligible for garbage collection. Keys are GC eligible once // their TTL has expired, and once this function returns true. If nil, only TTL determines GC eligibility. GCFilter GCFilter + + // The algorithm used to compress values before they are written to disk. The default is + // types.CompressionNone (no compression). This setting governs only segments created while it is in + // effect; each segment records the algorithm it was written with and is always read back with that + // algorithm, so the setting may be changed across restarts without invalidating existing data. + // + // When compression is enabled, secondary keys may only alias the entire value + // (Offset == 0 && Length == len(value)); a secondary key that indexes a strict sub-range of a value + // is rejected at write time, because a compressed blob cannot be sliced. Like the other fields here + // (except Name), this setting is not persisted and must be supplied again after each restart. + Compression types.CompressionAlgorithm } // GCFilter is a function that is called to determine if a key is eligible for garbage collection. Keys are GC @@ -69,6 +82,7 @@ func DefaultTableConfig(name string) TableConfig { ShardingFactor: 8, WriteCacheSize: 0, ReadCacheSize: 0, + Compression: types.CompressionNone, } } @@ -83,5 +97,8 @@ func (c *TableConfig) Validate() error { if c.ShardingFactor < 1 { return fmt.Errorf("sharding factor must be at least 1") } + if err := c.Compression.Validate(); err != nil { + return fmt.Errorf("invalid compression algorithm: %w", err) + } return nil } diff --git a/sei-db/db_engine/litt/types/compression.go b/sei-db/db_engine/litt/types/compression.go new file mode 100644 index 0000000000..5f6d24d7fc --- /dev/null +++ b/sei-db/db_engine/litt/types/compression.go @@ -0,0 +1,80 @@ +package types + +import ( + "fmt" + + "github.com/klauspost/compress/s2" +) + +// CompressionAlgorithm identifies the algorithm used to compress values before they are written to a +// segment's value file. It is stored (as a single byte) in each segment's metadata file so that a +// segment is always decompressed with the algorithm it was written with, independent of the table's +// current configuration. +// +// The zero value is CompressionNone, so a table or segment that never opts in behaves exactly as an +// uncompressed one. +type CompressionAlgorithm uint8 + +const ( + // CompressionNone means values are stored verbatim, with no compression. This is the default. + CompressionNone CompressionAlgorithm = 0 + + // CompressionS2 compresses each value with the S2 block codec (github.com/klauspost/compress/s2), + // a high-throughput, Snappy-compatible algorithm. The encoded block is self-describing: the + // decompressed length is recoverable from the block itself, so it is not stored separately. + CompressionS2 CompressionAlgorithm = 1 +) + +// Validate returns an error if the algorithm is not a value this build understands. It is called both +// when validating table configuration and when deserializing segment metadata, so an unknown byte in a +// metadata file (e.g. one written by a newer build) is reported rather than silently mishandled. +func (a CompressionAlgorithm) Validate() error { + switch a { + case CompressionNone, CompressionS2: + return nil + default: + return fmt.Errorf("unknown compression algorithm: %d", uint8(a)) + } +} + +// String returns a human-readable name for the algorithm. +func (a CompressionAlgorithm) String() string { + switch a { + case CompressionNone: + return "none" + case CompressionS2: + return "s2" + default: + return fmt.Sprintf("unknown(%d)", uint8(a)) + } +} + +// Compress returns the compressed form of src using the given algorithm. For CompressionNone it returns +// src unchanged. The returned slice is a freshly allocated buffer that does not alias src. +func Compress(algorithm CompressionAlgorithm, src []byte) ([]byte, error) { + switch algorithm { + case CompressionNone: + return src, nil + case CompressionS2: + return s2.Encode(nil, src), nil + default: + return nil, fmt.Errorf("cannot compress with unknown algorithm: %d", uint8(algorithm)) + } +} + +// Decompress returns the decompressed form of src, which must have been produced by Compress with the +// same algorithm. For CompressionNone it returns src unchanged. +func Decompress(algorithm CompressionAlgorithm, src []byte) ([]byte, error) { + switch algorithm { + case CompressionNone: + return src, nil + case CompressionS2: + decompressed, err := s2.Decode(nil, src) + if err != nil { + return nil, fmt.Errorf("failed to s2-decode value: %w", err) + } + return decompressed, nil + default: + return nil, fmt.Errorf("cannot decompress with unknown algorithm: %d", uint8(algorithm)) + } +} diff --git a/sei-db/db_engine/litt/types/compression_test.go b/sei-db/db_engine/litt/types/compression_test.go new file mode 100644 index 0000000000..5177350c44 --- /dev/null +++ b/sei-db/db_engine/litt/types/compression_test.go @@ -0,0 +1,72 @@ +package types + +import ( + "bytes" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestCompressionRoundTrip(t *testing.T) { + t.Parallel() + + // A compressible payload so the S2 case actually shrinks. + payload := bytes.Repeat([]byte("the quick brown fox jumps over the lazy dog. "), 100) + + cases := []struct { + name string + algo CompressionAlgorithm + }{ + {"none", CompressionNone}, + {"s2", CompressionS2}, + } + + for _, tc := range cases { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + compressed, err := Compress(tc.algo, payload) + require.NoError(t, err) + + decompressed, err := Decompress(tc.algo, compressed) + require.NoError(t, err) + require.Equal(t, payload, decompressed) + + if tc.algo == CompressionS2 { + require.Less(t, len(compressed), len(payload), "s2 should shrink a repetitive payload") + } + }) + } +} + +func TestCompressionRoundTripEmptyValue(t *testing.T) { + t.Parallel() + + for _, algo := range []CompressionAlgorithm{CompressionNone, CompressionS2} { + compressed, err := Compress(algo, []byte{}) + require.NoError(t, err) + + decompressed, err := Decompress(algo, compressed) + require.NoError(t, err) + require.Empty(t, decompressed) + } +} + +func TestCompressionAlgorithmValidate(t *testing.T) { + t.Parallel() + + require.NoError(t, CompressionNone.Validate()) + require.NoError(t, CompressionS2.Validate()) + require.Error(t, CompressionAlgorithm(99).Validate()) +} + +func TestCompressUnknownAlgorithm(t *testing.T) { + t.Parallel() + + _, err := Compress(CompressionAlgorithm(99), []byte("data")) + require.Error(t, err) + + _, err = Decompress(CompressionAlgorithm(99), []byte("data")) + require.Error(t, err) +} From d037b8895371313e130d034869ced1ef29f2dd19 Mon Sep 17 00:00:00 2001 From: Cody Littley Date: Wed, 15 Jul 2026 16:42:32 -0500 Subject: [PATCH 02/10] cleanup --- sei-db/db_engine/litt/disktable/control_loop.go | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/sei-db/db_engine/litt/disktable/control_loop.go b/sei-db/db_engine/litt/disktable/control_loop.go index e57f6132a4..481bf23201 100644 --- a/sei-db/db_engine/litt/disktable/control_loop.go +++ b/sei-db/db_engine/litt/disktable/control_loop.go @@ -27,12 +27,13 @@ type controlLoop struct { // controllerChannel is the channel the control loop reads messages from. When compression is // disabled, enqueue sends directly here (inputChannel == controllerChannel). When compression is - // enabled, the compression loop reads inputChannel and forwards messages here. + // enabled, the compression loop reads its input channel and forwards messages here. controllerChannel chan any - // inputChannel is the channel enqueue sends to. It is controllerChannel when compression is disabled, - // or the compression loop's input channel when compression is enabled (so every control message, - // including flush, passes through the compression stage in order). + // inputChannel is the entry point of the control loop's message pipeline: the channel enqueue sends + // to. It is controllerChannel when compression is disabled, or the compression loop's input channel + // when compression is enabled (so every control message, including flush, passes through the + // compression stage in order). inputChannel chan any // compressionAlgorithm is the algorithm new segments are created with. types.CompressionNone means From 9c02ca081b91c3c15bcc4a2dfa68c8d5a6f2e934 Mon Sep 17 00:00:00 2001 From: Cody Littley Date: Wed, 15 Jul 2026 16:49:00 -0500 Subject: [PATCH 03/10] unit tests --- .../litt/disktable/compression_test.go | 90 +++++++++++++++++++ 1 file changed, 90 insertions(+) diff --git a/sei-db/db_engine/litt/disktable/compression_test.go b/sei-db/db_engine/litt/disktable/compression_test.go index 633988c5e5..1eca9f63be 100644 --- a/sei-db/db_engine/litt/disktable/compression_test.go +++ b/sei-db/db_engine/litt/disktable/compression_test.go @@ -167,6 +167,61 @@ func TestCompressionFlushOrdering(t *testing.T) { } } +// TestCompressionToggleAcrossRestarts writes into the same table under alternating compression settings +// (off -> on -> off -> on), restarting between each phase, then reads everything back. It exercises the +// core guarantee that each segment is decoded with the algorithm it was written with, independent of the +// table's current configuration. +func TestCompressionToggleAcrossRestarts(t *testing.T) { + t.Parallel() + dir := t.TempDir() + name := "toggle" + + phases := []struct { + label string + algo types.CompressionAlgorithm + }{ + {"off1", types.CompressionNone}, + {"on1", types.CompressionS2}, + {"off2", types.CompressionNone}, + {"on2", types.CompressionS2}, + } + + written := make(map[string][]byte) + + for _, phase := range phases { + table := buildCompressedMemKeyDiskTable(t, time.Now, name, []string{dir}, phase.algo) + + // Every key written in a prior phase (under a possibly different algorithm) must still read back + // correctly after this restart. + for k, v := range written { + got, ok, err := table.Get([]byte(k)) + require.NoError(t, err, "reading %q in phase %s", k, phase.label) + require.True(t, ok, "missing %q in phase %s", k, phase.label) + require.Equal(t, v, got, "value mismatch for %q in phase %s", k, phase.label) + } + + // Write this phase's keys into a fresh segment created under this phase's algorithm. + for i := 0; i < 5; i++ { + key := fmt.Sprintf("%s-%d", phase.label, i) + value := append(compressiblePayload(), []byte(key)...) + require.NoError(t, table.Put([]byte(key), value)) + written[key] = value + } + require.NoError(t, table.Flush()) + require.NoError(t, table.Close()) + } + + // Final restart: read everything back, spanning segments written under both algorithms. + final := buildCompressedMemKeyDiskTable(t, time.Now, name, []string{dir}, types.CompressionNone) + defer func() { require.NoError(t, final.Close()) }() + for k, v := range written { + got, ok, err := final.Get([]byte(k)) + require.NoError(t, err) + require.True(t, ok, "missing %q after final restart", k) + require.Equal(t, v, got, "value mismatch for %q after final restart", k) + } +} + func TestCompressionIteration(t *testing.T) { t.Parallel() dir := t.TempDir() @@ -213,3 +268,38 @@ func TestCompressionIteration(t *testing.T) { require.Equal(t, r.value, seen[r.key+"-alias"], "alias for %s", r.key) } } + +// TestCompressionReverseIteration verifies that reverse iteration over a compressed segment returns the +// correct decompressed value for each key (reverse iteration reads through Segment.Read). +func TestCompressionReverseIteration(t *testing.T) { + t.Parallel() + dir := t.TempDir() + table := buildCompressedMemKeyDiskTable(t, time.Now, "reverse", []string{dir}, types.CompressionS2) + defer func() { require.NoError(t, table.Close()) }() + + want := make(map[string][]byte) + for i := 0; i < 10; i++ { + key := fmt.Sprintf("k-%02d", i) + value := append(compressiblePayload(), byte(i)) + require.NoError(t, table.Put([]byte(key), value)) + want[key] = value + } + require.NoError(t, table.Flush()) + + it, err := table.Iterator(true) + require.NoError(t, err) + seen := make(map[string][]byte) + for { + ok, err := it.Next() + require.NoError(t, err) + if !ok { + break + } + key, _ := it.GetKey() + value, err := it.GetValue() + require.NoError(t, err) + seen[string(key)] = append([]byte(nil), value...) + } + require.NoError(t, it.Close()) + require.Equal(t, want, seen) +} From b386096d3552911afab11a33e86ae214fb2f1390 Mon Sep 17 00:00:00 2001 From: Cody Littley Date: Thu, 16 Jul 2026 09:23:05 -0500 Subject: [PATCH 04/10] godoc update --- .../litt/disktable/compression_loop.go | 24 ++++++ sei-db/db_engine/litt/disktable/disk_table.go | 3 + .../db_engine/litt/metrics/littdb_metrics.go | 73 ++++++++++++++++++- sei-db/db_engine/litt/table_config.go | 6 +- 4 files changed, 102 insertions(+), 4 deletions(-) diff --git a/sei-db/db_engine/litt/disktable/compression_loop.go b/sei-db/db_engine/litt/disktable/compression_loop.go index 0119368884..64cc6e10f1 100644 --- a/sei-db/db_engine/litt/disktable/compression_loop.go +++ b/sei-db/db_engine/litt/disktable/compression_loop.go @@ -3,7 +3,9 @@ package disktable import ( "fmt" "log/slog" + "time" + "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/metrics" "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/types" "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/util" ) @@ -32,6 +34,15 @@ type compressionLoop struct { // outputChannel forwards messages to the control loop (its controllerChannel). outputChannel chan any + + // metrics encapsulates metrics for the DB. May be nil, in which case no metrics are reported. + metrics *metrics.LittDBMetrics + + // name is the table name, used to tag metrics. + name string + + // clock provides the current time, used to measure compression latency. + clock func() time.Time } // run processes messages until shutdown. It compresses write requests and forwards every message to the @@ -66,7 +77,14 @@ func (cl *compressionLoop) run() { // compress fills req.compressedValues with the compressed form of each value. It returns false if // compression failed (in which case it has already panicked the DB via the error monitor). func (cl *compressionLoop) compress(req *controlLoopWriteRequest) bool { + var start time.Time + if cl.metrics != nil { + start = cl.clock() + } + compressed := make([][]byte, len(req.values)) + var uncompressedBytes uint64 + var compressedBytes uint64 for i, kv := range req.values { blob, err := types.Compress(cl.algorithm, kv.Value) if err != nil { @@ -74,7 +92,13 @@ func (cl *compressionLoop) compress(req *controlLoopWriteRequest) bool { return false } compressed[i] = blob + uncompressedBytes += uint64(len(kv.Value)) + compressedBytes += uint64(len(blob)) } req.compressedValues = compressed + + if cl.metrics != nil { + cl.metrics.ReportCompression(cl.name, cl.clock().Sub(start), uncompressedBytes, compressedBytes) + } return true } diff --git a/sei-db/db_engine/litt/disktable/disk_table.go b/sei-db/db_engine/litt/disktable/disk_table.go index 35410cc5cc..07ec367229 100644 --- a/sei-db/db_engine/litt/disktable/disk_table.go +++ b/sei-db/db_engine/litt/disktable/disk_table.go @@ -402,6 +402,9 @@ func NewDiskTable( algorithm: tableConfig.Compression, inputChannel: make(chan any, config.ControlChannelSize), outputChannel: cLoop.controllerChannel, + metrics: metrics, + name: name, + clock: runtimeConfig.Clock, } cLoop.inputChannel = cmpLoop.inputChannel } diff --git a/sei-db/db_engine/litt/metrics/littdb_metrics.go b/sei-db/db_engine/litt/metrics/littdb_metrics.go index 72b81dc422..537d1bf881 100644 --- a/sei-db/db_engine/litt/metrics/littdb_metrics.go +++ b/sei-db/db_engine/litt/metrics/littdb_metrics.go @@ -88,6 +88,19 @@ type LittDBMetrics struct { // The latency of garbage collection operations. garbageCollectionLatency metric.Float64Histogram + // Reports on the latency of compressing a batch of values before they are written. + compressionLatency metric.Float64Histogram + + // The number of uncompressed value bytes submitted to compression since startup. + compressionUncompressedBytes metric.Int64Counter + + // The number of compressed value bytes produced by compression since startup. Compared against + // compressionUncompressedBytes, this gives the aggregate compression ratio and total bytes saved. + compressionCompressedBytes metric.Int64Counter + + // The per-batch compression ratio (compressed bytes / uncompressed bytes); lower is better. + compressionRatio metric.Float64Histogram + // Metrics for the write cache. writeCacheMetrics *util.CacheMetrics @@ -219,6 +232,32 @@ func NewLittDBMetrics() *LittDBMetrics { metric.WithExplicitBucketBoundaries(commonmetrics.LatencyBuckets...), ) + compressionLatency, _ := meter.Float64Histogram( + "litt_compression_latency_seconds", + metric.WithDescription("Reports on the latency of compressing a batch of values before they are written."), + metric.WithUnit("s"), + metric.WithExplicitBucketBoundaries(commonmetrics.LatencyBuckets...), + ) + + compressionUncompressedBytes, _ := meter.Int64Counter( + "litt_compression_uncompressed_bytes", + metric.WithDescription("The number of uncompressed value bytes submitted to compression since startup."), + metric.WithUnit("By"), + ) + + compressionCompressedBytes, _ := meter.Int64Counter( + "litt_compression_compressed_bytes", + metric.WithDescription("The number of compressed value bytes produced by compression since startup."), + metric.WithUnit("By"), + ) + + compressionRatio, _ := meter.Float64Histogram( + "litt_compression_ratio", + metric.WithDescription("The per-batch compression ratio (compressed bytes / uncompressed bytes); lower is better."), + metric.WithUnit("1"), + metric.WithExplicitBucketBoundaries(0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.25, 1.5), + ) + writeCacheMetrics := util.NewCacheMetrics("chunk_write") readCacheMetrics := util.NewCacheMetrics("chunk_read") @@ -240,8 +279,14 @@ func NewLittDBMetrics() *LittDBMetrics { garbageCollectionLatency: garbageCollectionLatency, segmentFlushLatency: segmentFlushLatency, keymapFlushLatency: keymapFlushLatency, - writeCacheMetrics: writeCacheMetrics, - readCacheMetrics: readCacheMetrics, + + compressionLatency: compressionLatency, + compressionUncompressedBytes: compressionUncompressedBytes, + compressionCompressedBytes: compressionCompressedBytes, + compressionRatio: compressionRatio, + + writeCacheMetrics: writeCacheMetrics, + readCacheMetrics: readCacheMetrics, } } @@ -367,6 +412,30 @@ func (m *LittDBMetrics) ReportGarbageCollectionLatency(tableName string, latency m.garbageCollectionLatency.Record(context.Background(), latency.Seconds(), tableAttr(tableName)) } +// ReportCompression reports the results of compressing a batch of values: the time taken, the number of +// uncompressed bytes submitted, and the number of compressed bytes produced. The per-batch ratio is +// derived and recorded when at least one byte was submitted. +func (m *LittDBMetrics) ReportCompression( + tableName string, + latency time.Duration, + uncompressedBytes uint64, + compressedBytes uint64) { + + if m == nil { + return + } + + ctx := context.Background() + attrs := tableAttr(tableName) + + m.compressionLatency.Record(ctx, latency.Seconds(), attrs) + m.compressionUncompressedBytes.Add(ctx, int64(uncompressedBytes), attrs) //nolint:gosec // byte count fits int64 + m.compressionCompressedBytes.Add(ctx, int64(compressedBytes), attrs) //nolint:gosec // byte count fits int64 + if uncompressedBytes > 0 { + m.compressionRatio.Record(ctx, float64(compressedBytes)/float64(uncompressedBytes), attrs) + } +} + func (m *LittDBMetrics) GetWriteCacheMetrics() *util.CacheMetrics { if m == nil { return nil diff --git a/sei-db/db_engine/litt/table_config.go b/sei-db/db_engine/litt/table_config.go index 7b01f5f84e..447e03f752 100644 --- a/sei-db/db_engine/litt/table_config.go +++ b/sei-db/db_engine/litt/table_config.go @@ -56,8 +56,10 @@ type TableConfig struct { // // When compression is enabled, secondary keys may only alias the entire value // (Offset == 0 && Length == len(value)); a secondary key that indexes a strict sub-range of a value - // is rejected at write time, because a compressed blob cannot be sliced. Like the other fields here - // (except Name), this setting is not persisted and must be supplied again after each restart. + // is rejected at write time, because a compressed blob cannot be sliced. This restriction could be + // lifted with additional engineering, but we should have a real use case before making the attempt. + // Like the other fields here (except Name), this setting is not persisted and must be supplied again + // after each restart. Compression types.CompressionAlgorithm } From 23cfd53cb7019d75ac8750e9c4546ecc34144913 Mon Sep 17 00:00:00 2001 From: Cody Littley Date: Thu, 16 Jul 2026 10:17:31 -0500 Subject: [PATCH 05/10] unit tests --- .../litt/disktable/compression_test.go | 256 ++++++++++++++++++ .../litt/disktable/forward_iterator.go | 5 +- sei-db/db_engine/litt/types/compression.go | 5 +- 3 files changed, 263 insertions(+), 3 deletions(-) diff --git a/sei-db/db_engine/litt/disktable/compression_test.go b/sei-db/db_engine/litt/disktable/compression_test.go index 1eca9f63be..da47cd259d 100644 --- a/sei-db/db_engine/litt/disktable/compression_test.go +++ b/sei-db/db_engine/litt/disktable/compression_test.go @@ -4,12 +4,17 @@ import ( "bytes" "fmt" "log/slog" + "math" + "os" + "path" "path/filepath" + "sync/atomic" "testing" "time" "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt" "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/disktable/keymap" + "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/disktable/segment" "github.com/sei-protocol/sei-chain/sei-db/db_engine/litt/types" "github.com/stretchr/testify/require" ) @@ -303,3 +308,254 @@ func TestCompressionReverseIteration(t *testing.T) { require.NoError(t, it.Close()) require.Equal(t, want, seen) } + +// TestCompressionRecoveryUnsealedSegment exercises the crash-recovery path (Segment.sealLoadedSegment) +// over on-disk *compressed* blobs. Every other compression test does a clean Flush+Close (which seals) +// before reopening, so this is the only coverage of recovery re-sealing a segment and validating +// value-file completeness against compressed-blob addresses. +func TestCompressionRecoveryUnsealedSegment(t *testing.T) { + t.Parallel() + + // unsealDataSegment flips the sealed byte of the data-bearing segment back to 0, simulating a crash + // that happened before the segment was sealed. This is what makes LoadSegment run recovery on reopen. + // It returns the segments directory and the segment index so callers can also corrupt the value file. + unsealDataSegment := func(t *testing.T, dir string, name string) (string, uint32) { + t.Helper() + segmentsDir := findLatestSegmentDir(t, dir, name) + idx := segmentIndexWithLargestValueFile(t, segmentsDir) + metaPath := path.Join(segmentsDir, fmt.Sprintf("%d%s", idx, segment.MetadataFileExtension)) + mdBytes, err := os.ReadFile(metaPath) + require.NoError(t, err) + require.Equal(t, segment.V4MetadataSize, len(mdBytes)) // a compressed table writes v4 metadata + mdBytes[segment.MetadataSealedByteOffset] = 0 + require.NoError(t, os.WriteFile(metaPath, mdBytes, 0600)) + return segmentsDir, idx + } + + t.Run("clean crash before seal, all groups survive", func(t *testing.T) { + t.Parallel() + dir := t.TempDir() + name := "recover-clean" + + table := buildCompressedMemKeyDiskTable(t, time.Now, name, []string{dir}, types.CompressionS2) + + written := make(map[string][]byte) + for i := 0; i < 6; i++ { + key := fmt.Sprintf("key-%02d", i) + value := append(compressiblePayload(), byte(i)) + require.NoError(t, table.Put([]byte(key), value)) + written[key] = value + } + // One group carries a full-value-alias secondary so recovery re-seals a multi-record group too. + aliasValue := append(compressiblePayload(), []byte("-alias")...) + alias := &types.SecondaryKey{Key: []byte("alias-secondary"), Offset: 0, Length: uint32(len(aliasValue))} + require.NoError(t, table.Put([]byte("aliased-primary"), aliasValue, alias)) + written["aliased-primary"] = aliasValue + + require.NoError(t, table.Flush()) + require.NoError(t, table.Close()) + + unsealDataSegment(t, dir, name) + + table = buildCompressedMemKeyDiskTable(t, time.Now, name, []string{dir}, types.CompressionS2) + defer func() { require.NoError(t, table.Close()) }() + + for key, want := range written { + got, ok, err := table.Get([]byte(key)) + require.NoError(t, err) + require.True(t, ok, "expected %q to survive recovery", key) + require.Equal(t, want, got, "value for %q must decompress correctly after recovery", key) + } + // The full-value-alias secondary resolves to the whole decompressed value. + got, ok, err := table.Get([]byte("alias-secondary")) + require.NoError(t, err) + require.True(t, ok, "alias secondary must survive recovery") + require.Equal(t, aliasValue, got) + }) + + t.Run("torn compressed value, last group dropped", func(t *testing.T) { + t.Parallel() + dir := t.TempDir() + name := "recover-torn" + + table := buildCompressedMemKeyDiskTable(t, time.Now, name, []string{dir}, types.CompressionS2) + + survivors := make(map[string][]byte) + for i := 0; i < 4; i++ { + key := fmt.Sprintf("survivor-%02d", i) + value := append(compressiblePayload(), byte(i)) + require.NoError(t, table.Put([]byte(key), value)) + survivors[key] = value + } + // The last-written group lands at the tail of the single-shard value file. It carries a + // full-value alias so we can confirm the whole group is dropped — primary and secondary alike. + tornValue := append(compressiblePayload(), []byte("-torn")...) + tornAlias := &types.SecondaryKey{ + Key: []byte("torn-secondary"), + Offset: 0, + Length: uint32(len(tornValue)), + } + require.NoError(t, table.Put([]byte("torn-primary"), tornValue, tornAlias)) + + require.NoError(t, table.Flush()) + require.NoError(t, table.Close()) + + segmentsDir := findLatestSegmentDir(t, dir, name) + idx := segmentIndexWithLargestValueFile(t, segmentsDir) + + // Lop a few bytes off the tail of the value file so the last group's blob is incomplete: its + // address end now exceeds the flushed size, so group-atomic recovery discards the whole group. + valPath := path.Join(segmentsDir, fmt.Sprintf("%d-0%s", idx, segment.ValuesFileExtension)) + data, err := os.ReadFile(valPath) + require.NoError(t, err) + require.GreaterOrEqual(t, len(data), 4) + require.NoError(t, os.WriteFile(valPath, data[:len(data)-3], 0600)) + + metaPath := path.Join(segmentsDir, fmt.Sprintf("%d%s", idx, segment.MetadataFileExtension)) + mdBytes, err := os.ReadFile(metaPath) + require.NoError(t, err) + mdBytes[segment.MetadataSealedByteOffset] = 0 + require.NoError(t, os.WriteFile(metaPath, mdBytes, 0600)) + + table = buildCompressedMemKeyDiskTable(t, time.Now, name, []string{dir}, types.CompressionS2) + defer func() { require.NoError(t, table.Close()) }() + + for key, want := range survivors { + got, ok, err := table.Get([]byte(key)) + require.NoError(t, err) + require.True(t, ok, "expected %q to survive recovery", key) + require.Equal(t, want, got) + } + // The torn group is gone in its entirety. + for _, key := range []string{"torn-primary", "torn-secondary"} { + _, ok, err := table.Get([]byte(key)) + require.NoError(t, err) + require.False(t, ok, "expected torn group member %q to be dropped by recovery", key) + } + }) +} + +// buildCompressedGCTable builds a single-shard, mem-keymap, S2-compressed table wired for deterministic +// GC testing: size-based sealing is disabled so segments seal exactly every maxSegmentKeyCount keys, +// background GC is effectively disabled (driven explicitly via RunGC), and the clock is injectable so a +// test can advance time past a segment's TTL. +func buildCompressedGCTable( + t *testing.T, + clock func() time.Time, + name string, + dir string, + maxSegmentKeyCount uint32, +) litt.ManagedTable { + t.Helper() + logger := slog.Default() + keymapPath := filepath.Join(dir, keymap.KeymapDirectoryName) + keymapTypeFile, err := setupKeymapTypeFile(keymapPath, keymap.MemKeymapType) + require.NoError(t, err) + keys, _, err := keymap.NewMemKeymap(logger, "", true) + require.NoError(t, err) + + config, err := litt.DefaultConfig(dir) + require.NoError(t, err) + config.TargetSegmentFileSize = math.MaxUint32 + config.MaxSegmentKeyCount = maxSegmentKeyCount + config.GCPeriod = time.Hour + config.Fsync = false + + tableConfig := litt.DefaultTableConfig(name) + tableConfig.ShardingFactor = 1 + tableConfig.Compression = types.CompressionS2 + + runtimeConfig := litt.DefaultRuntimeConfig() + runtimeConfig.Clock = clock + runtimeConfig.Logger = logger + + table, err := NewDiskTable( + config, + runtimeConfig, + name, + tableConfig, + keys, + keymapPath, + keymapTypeFile, + []string{dir}, + true, + nil, + ) + require.NoError(t, err) + return table +} + +// TestCompressionGarbageCollection collects TTL-expired *compressed* segments and confirms that the +// surviving compressed data still reads back (and decompresses) correctly, both via Get and via a +// forward iterator. +func TestCompressionGarbageCollection(t *testing.T) { + t.Parallel() + dir := t.TempDir() + + start := time.Unix(1_700_000_000, 0) + var fakeTime atomic.Pointer[time.Time] + fakeTime.Store(&start) + clock := func() time.Time { return *fakeTime.Load() } + + // Seal after every 2 keys so both batches below land in sealed (hence collectable) segments rather + // than the always-live mutable segment. + table := buildCompressedGCTable(t, clock, "gc", dir, 2) + defer func() { require.NoError(t, table.Close()) }() + + ttl := time.Hour + require.NoError(t, table.SetTTL(ttl)) + + // Batch A, written at t0. + oldKeys := make(map[string][]byte) + for i := 0; i < 4; i++ { + key := fmt.Sprintf("old-%02d", i) + value := append(compressiblePayload(), byte(i)) + require.NoError(t, table.Put([]byte(key), value)) + oldKeys[key] = value + } + require.NoError(t, table.Flush()) + + // Batch B, written far enough after batch A that its segments are much younger. + tB := start.Add(10 * ttl) + fakeTime.Store(&tB) + newKeys := make(map[string][]byte) + for i := 0; i < 4; i++ { + key := fmt.Sprintf("new-%02d", i) + value := append(compressiblePayload(), []byte(fmt.Sprintf("-new-%d", i))...) + require.NoError(t, table.Put([]byte(key), value)) + newKeys[key] = value + } + require.NoError(t, table.Flush()) + + // Advance just past batch A's TTL: batch A's segments are now expired, batch B's are still young. + now := tB.Add(ttl / 2) + fakeTime.Store(&now) + require.NoError(t, table.RunGC()) + + // Batch A is gone; batch B survives and decompresses to the original bytes. + for key := range oldKeys { + requireKeyAbsent(t, table, key) + } + for key, want := range newKeys { + got, ok, err := table.Get([]byte(key)) + require.NoError(t, err) + require.True(t, ok, "expected %q to survive GC", key) + require.Equal(t, want, got, "survivor %q must decompress correctly", key) + } + + // A forward iterator sees exactly the survivors, each decompressed correctly. + it, err := table.Iterator(false) + require.NoError(t, err) + entries := drainIterator(t, it) + require.NoError(t, it.Close()) + + expectedKeys := make([]string, 0, len(newKeys)) + for key := range newKeys { + expectedKeys = append(expectedKeys, key) + } + require.ElementsMatch(t, expectedKeys, entryKeys(entries)) + for _, e := range entries { + require.Equal(t, newKeys[e.key], []byte(e.value), + "iterated survivor %q must decompress correctly", e.key) + } +} diff --git a/sei-db/db_engine/litt/disktable/forward_iterator.go b/sei-db/db_engine/litt/disktable/forward_iterator.go index 715f10ff8b..faae8c7cc1 100644 --- a/sei-db/db_engine/litt/disktable/forward_iterator.go +++ b/sei-db/db_engine/litt/disktable/forward_iterator.go @@ -146,7 +146,10 @@ func (it *forwardIterator) GetValue() (value []byte, err error) { // // This optimization is invalid for compressed segments: there addr.ValueSize() is the compressed blob // length while groupValue holds the decompressed value, so the sub-slice arithmetic would be wrong. - // Fall through to a direct read (which decompresses) instead. + // Fall through to a direct read (which decompresses) instead. That re-reads and re-decodes the blob + // rather than reusing the primary's decoded value, but this path is intentionally not optimized: the + // only legal secondary on a compressed segment is a full-value alias, so reading its value returns the + // same bytes as the primary — a redundant workload not worth extra machinery. if !it.currentSeg.IsCompressed() && it.groupValid && it.secondaryWithinGroup(addr) { if it.groupValue == nil { v, err := reader.Read(it.groupAddr) diff --git a/sei-db/db_engine/litt/types/compression.go b/sei-db/db_engine/litt/types/compression.go index 5f6d24d7fc..25d06fe556 100644 --- a/sei-db/db_engine/litt/types/compression.go +++ b/sei-db/db_engine/litt/types/compression.go @@ -50,7 +50,7 @@ func (a CompressionAlgorithm) String() string { } // Compress returns the compressed form of src using the given algorithm. For CompressionNone it returns -// src unchanged. The returned slice is a freshly allocated buffer that does not alias src. +// src unchanged (the result may alias src). The caller must not mutate src or the returned slice. func Compress(algorithm CompressionAlgorithm, src []byte) ([]byte, error) { switch algorithm { case CompressionNone: @@ -63,7 +63,8 @@ func Compress(algorithm CompressionAlgorithm, src []byte) ([]byte, error) { } // Decompress returns the decompressed form of src, which must have been produced by Compress with the -// same algorithm. For CompressionNone it returns src unchanged. +// same algorithm. For CompressionNone it returns src unchanged (the result may alias src). The caller +// must not mutate src or the returned slice. func Decompress(algorithm CompressionAlgorithm, src []byte) ([]byte, error) { switch algorithm { case CompressionNone: From 448ea4956a4fd99c8ac08edafa09636a623655dd Mon Sep 17 00:00:00 2001 From: Cody Littley Date: Fri, 17 Jul 2026 14:11:15 -0500 Subject: [PATCH 06/10] bugfix --- .../litt/disktable/segment/metadata_file.go | 12 ++++++++++-- sei-db/ledger_db/block/littblock_crash_test.go | 4 ++-- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/sei-db/db_engine/litt/disktable/segment/metadata_file.go b/sei-db/db_engine/litt/disktable/segment/metadata_file.go index cfe431be43..f68c5890ac 100644 --- a/sei-db/db_engine/litt/disktable/segment/metadata_file.go +++ b/sei-db/db_engine/litt/disktable/segment/metadata_file.go @@ -159,8 +159,14 @@ func getMetadataFileIndex(fileName string) (uint32, error) { return uint32(index), nil //nolint:gosec // segment index fits uint32 } -// Size returns the size of the metadata file in bytes. +// Size returns the size of the metadata file in bytes. A metadata file loaded from disk as a legacy v3 +// file (no compression byte) is 18 bytes until it is rewritten; every file written by the current build +// is v4 (19 bytes). serialize() bumps segmentVersion to LatestSegmentVersion on write, so this reflects +// the actual on-disk size in both cases. func (m *metadataFile) Size() uint64 { + if m.segmentVersion == ShardedAddressSegmentVersion { + return V3MetadataSize + } return V4MetadataSize } @@ -188,8 +194,10 @@ func (m *metadataFile) seal(now time.Time, keyCount uint32) error { } // serialize serializes the metadata file to a byte array. Metadata is always written at -// LatestSegmentVersion (v4), including the trailing compression-algorithm byte. +// LatestSegmentVersion (v4), including the trailing compression-algorithm byte. A file loaded as legacy +// v3 is therefore promoted to v4 the moment it is rewritten, so segmentVersion is updated to match. func (m *metadataFile) serialize() []byte { + m.segmentVersion = LatestSegmentVersion data := make([]byte, V4MetadataSize) binary.BigEndian.PutUint32(data[0:4], uint32(LatestSegmentVersion)) diff --git a/sei-db/ledger_db/block/littblock_crash_test.go b/sei-db/ledger_db/block/littblock_crash_test.go index bb402aa324..e50307ec0e 100644 --- a/sei-db/ledger_db/block/littblock_crash_test.go +++ b/sei-db/ledger_db/block/littblock_crash_test.go @@ -149,8 +149,8 @@ func markSegmentUnsealedOnDisk(t *testing.T, metaPath string) { t.Helper() data, err := os.ReadFile(metaPath) require.NoError(t, err) - require.Equal(t, segment.V3MetadataSize, len(data), "unexpected metadata size for %s", metaPath) - data[segment.V3MetadataSize-1] = 0 + require.Equal(t, segment.V4MetadataSize, len(data), "unexpected metadata size for %s", metaPath) + data[segment.MetadataSealedByteOffset] = 0 require.NoError(t, os.WriteFile(metaPath, data, 0600)) } From 817b3f53f029acc6f508c8ba5df519cc2da58329 Mon Sep 17 00:00:00 2001 From: Cody Littley Date: Mon, 20 Jul 2026 08:24:44 -0500 Subject: [PATCH 07/10] handle edge case where compressed data exceeds size of original data --- sei-db/db_engine/litt/README.md | 11 +-- .../litt/disktable/compression_loop.go | 7 +- .../litt/disktable/compression_test.go | 36 ++++++++ sei-db/db_engine/litt/disktable/disk_table.go | 6 +- .../litt/disktable/segment/segment.go | 17 ++-- sei-db/db_engine/litt/table.go | 8 +- sei-db/db_engine/litt/types/compression.go | 60 +++++++++++++ .../db_engine/litt/types/compression_test.go | 90 +++++++++++++++++++ 8 files changed, 214 insertions(+), 21 deletions(-) diff --git a/sei-db/db_engine/litt/README.md b/sei-db/db_engine/litt/README.md index 6337f06ac7..ecb3109b71 100644 --- a/sei-db/db_engine/litt/README.md +++ b/sei-db/db_engine/litt/README.md @@ -44,7 +44,7 @@ The following features are currently supported by LittDB: - [tables](#table) with non-overlapping namespaces - multi-drive support (data can be spread across multiple physical volumes) - incremental backups (both local and remote) -- keys up to 64 KiB (2^16 - 1 bytes) and values up to 2^32 - 1 bytes (~4 GiB) in size +- keys up to 64 KiB (2^16 - 1 bytes) and values up to 2^32 - 2 bytes (~4 GiB) in size - incremental snapshots - incremental remote backups @@ -133,8 +133,8 @@ type Table interface { } ``` -Both primary keys and secondary keys must not exceed 64 KiB (2^16 - 1 bytes). Values may be up to 2^32 - 1 bytes -(`math.MaxUint32`, ~4 GiB); larger values are rejected. +Both primary keys and secondary keys must not exceed 64 KiB (2^16 - 1 bytes). Values may be up to 2^32 - 2 bytes +(`math.MaxUint32 - 1`, ~4 GiB); larger values are rejected. Source: [put_request.go](types/put_request.go) @@ -386,8 +386,9 @@ always written whole within a single value file: before writing a value whose by rolled and a fresh one is started, so the value lands entirely within the new file. Consequently a value file never exceeds 2^32 bytes and no value is ever split across the boundary. -This bounds the maximum size of a single [value](#value) to 2^32 - 1 bytes (`math.MaxUint32`, ~4 GiB); larger values are -rejected. (Before secondary keys were introduced, values could span past this boundary; that is no longer supported.) +This bounds the maximum size of a single [value](#value) to 2^32 - 2 bytes (`math.MaxUint32 - 1`, ~4 GiB); larger values +are rejected. One byte below the addressing boundary is reserved for a per-value header. (Before secondary keys were +introduced, values could span past this boundary; that is no longer supported.) Each segment may split its data into multiple [shards](#shard). The number of shards in a segment is called the [sharding factor](#sharding-factor). The [sharding factor](#sharding-factor) is configurable, and different segments diff --git a/sei-db/db_engine/litt/disktable/compression_loop.go b/sei-db/db_engine/litt/disktable/compression_loop.go index 64cc6e10f1..60dc124a4e 100644 --- a/sei-db/db_engine/litt/disktable/compression_loop.go +++ b/sei-db/db_engine/litt/disktable/compression_loop.go @@ -74,8 +74,9 @@ func (cl *compressionLoop) run() { } } -// compress fills req.compressedValues with the compressed form of each value. It returns false if -// compression failed (in which case it has already panicked the DB via the error monitor). +// compress fills req.compressedValues with the on-disk encoded form of each value (a one-byte algorithm +// tag plus the smaller of the compressed and raw bodies; see types.EncodeValue). It returns false if +// encoding failed (in which case it has already panicked the DB via the error monitor). func (cl *compressionLoop) compress(req *controlLoopWriteRequest) bool { var start time.Time if cl.metrics != nil { @@ -86,7 +87,7 @@ func (cl *compressionLoop) compress(req *controlLoopWriteRequest) bool { var uncompressedBytes uint64 var compressedBytes uint64 for i, kv := range req.values { - blob, err := types.Compress(cl.algorithm, kv.Value) + blob, err := types.EncodeValue(cl.algorithm, kv.Value) if err != nil { cl.errorMonitor.Panic(fmt.Errorf("failed to compress value: %w", err)) return false diff --git a/sei-db/db_engine/litt/disktable/compression_test.go b/sei-db/db_engine/litt/disktable/compression_test.go index da47cd259d..c4b63dac19 100644 --- a/sei-db/db_engine/litt/disktable/compression_test.go +++ b/sei-db/db_engine/litt/disktable/compression_test.go @@ -5,6 +5,7 @@ import ( "fmt" "log/slog" "math" + "math/rand" "os" "path" "path/filepath" @@ -71,6 +72,15 @@ func compressiblePayload() []byte { return bytes.Repeat([]byte("littDB compression makes value files smaller. "), 200) } +// incompressiblePayload returns deterministic pseudo-random bytes that S2 cannot shrink, so the +// store-smaller path keeps them raw (tagged CompressionNone) even on a compressed segment. +func incompressiblePayload() []byte { + payload := make([]byte, 4096) + rng := rand.New(rand.NewSource(1)) //nolint:gosec // test fixture, not security-sensitive + _, _ = rng.Read(payload) + return payload +} + func TestCompressionEndToEnd(t *testing.T) { t.Parallel() dir := t.TempDir() @@ -98,6 +108,32 @@ func TestCompressionEndToEnd(t *testing.T) { require.Less(t, int(table.Size()), len(value), "compressed segment should be smaller than the raw value") } +// TestCompressionMixedValues writes a compressible and an incompressible value into the same compressed +// segment and confirms both survive a flush. The incompressible value exercises the per-value store-raw +// tag (CompressionNone) alongside the compressible value's S2 tag on one segment. +func TestCompressionMixedValues(t *testing.T) { + t.Parallel() + dir := t.TempDir() + table := buildCompressedMemKeyDiskTable(t, time.Now, "mixed", []string{dir}, types.CompressionS2) + defer func() { require.NoError(t, table.Close()) }() + + compressible := compressiblePayload() + incompressible := incompressiblePayload() + require.NoError(t, table.Put([]byte("compressible"), compressible)) + require.NoError(t, table.Put([]byte("incompressible"), incompressible)) + require.NoError(t, table.Flush()) + + got, ok, err := table.Get([]byte("compressible")) + require.NoError(t, err) + require.True(t, ok) + require.Equal(t, compressible, got) + + got, ok, err = table.Get([]byte("incompressible")) + require.NoError(t, err) + require.True(t, ok) + require.Equal(t, incompressible, got) +} + func TestCompressionFullValueAliasSecondary(t *testing.T) { t.Parallel() dir := t.TempDir() diff --git a/sei-db/db_engine/litt/disktable/disk_table.go b/sei-db/db_engine/litt/disktable/disk_table.go index 07ec367229..119de9fb53 100644 --- a/sei-db/db_engine/litt/disktable/disk_table.go +++ b/sei-db/db_engine/litt/disktable/disk_table.go @@ -1005,8 +1005,10 @@ func (d *DiskTable) PutBatch(batch []*types.PutRequest) error { if len(kv.Key) > math.MaxUint16 { return fmt.Errorf("key is too large, length must not exceed 2^16 bytes: %d bytes", len(kv.Key)) } - if len(kv.Value) > math.MaxUint32 { - return fmt.Errorf("value is too large, length must not exceed 2^32 - 1 bytes: %d bytes", len(kv.Value)) + // One byte below 2^32-1: a compressed segment prefixes each value with a one-byte algorithm tag, + // and the resulting on-disk blob length must still fit the uint32 value-size field of an Address. + if len(kv.Value) >= math.MaxUint32 { + return fmt.Errorf("value is too large, length must not exceed 2^32 - 2 bytes: %d bytes", len(kv.Value)) } // Validate every secondary key in this request, and detect duplicate keys (primary vs diff --git a/sei-db/db_engine/litt/disktable/segment/segment.go b/sei-db/db_engine/litt/disktable/segment/segment.go index 09f6a8b0d6..b76c935ad1 100644 --- a/sei-db/db_engine/litt/disktable/segment/segment.go +++ b/sei-db/db_engine/litt/disktable/segment/segment.go @@ -628,9 +628,11 @@ func (s *Segment) WriteCompressed( fmt.Errorf("failed to send value to shard control loop: %v", err) } - // The whole compressed blob is the on-disk representation of the value; the primary and every - // (full-value-alias) secondary are addressed to it. Reads decompress the blob to recover the value. - blobAddress := types.NewAddress(s.index, firstByteIndex, shard, uint32(blobLen)) //nolint:gosec // bounded above + // The whole encoded blob is the on-disk representation of the value; the primary and every + // (full-value-alias) secondary are addressed to it. Reads decode the blob to recover the value. + // blobLen fits a uint32: the blob is a one-byte tag plus a body no larger than the raw value, and + // PutBatch caps the raw value at MaxUint32-1, so blobLen <= MaxUint32. + blobAddress := types.NewAddress(s.index, firstByteIndex, shard, uint32(blobLen)) //nolint:gosec // see above primaryRequest := &types.ScopedKey{ Key: data.Key, @@ -701,14 +703,15 @@ func (s *Segment) Read(key []byte, dataAddress types.Address) ([]byte, error) { return s.maybeDecompress(value) } -// maybeDecompress decompresses value using the segment's compression algorithm, or returns it unchanged -// if the segment is not compressed. All value reads (Segment.Read and SegmentReader.Read) pass through -// here so the on-disk compressed representation is never surfaced to callers. +// maybeDecompress decodes an on-disk value from a compressed segment (stripping the per-value algorithm +// tag and decompressing the body; see types.EncodeValue), or returns it unchanged if the segment is not +// compressed. All value reads (Segment.Read and SegmentReader.Read) pass through here so the on-disk +// representation is never surfaced to callers. func (s *Segment) maybeDecompress(value []byte) ([]byte, error) { if s.metadata.compressionAlgorithm == types.CompressionNone { return value, nil } - decompressed, err := types.Decompress(s.metadata.compressionAlgorithm, value) + decompressed, err := types.DecodeValue(value) if err != nil { return nil, fmt.Errorf("failed to decompress value: %w", err) } diff --git a/sei-db/db_engine/litt/table.go b/sei-db/db_engine/litt/table.go index d2d5d8d832..be219dda6d 100644 --- a/sei-db/db_engine/litt/table.go +++ b/sei-db/db_engine/litt/table.go @@ -30,7 +30,7 @@ type Table interface { // primary keys, and must not collide with the primary key or other secondaries. // // The maximum size of a key (primary or secondary) is 64 KiB (2^16 - 1 bytes). The maximum size of a - // value is 2^32 - 1 bytes (math.MaxUint32, ~4 GiB); a larger value is rejected with an error. (This + // value is 2^32 - 2 bytes (math.MaxUint32 - 1, ~4 GiB); a larger value is rejected with an error. (This // limit is stricter than it was before secondary keys existed: a value's byte offsets are stored as // uint32, and each value is written whole within a single segment file below the 2^32 boundary.) This // database has been optimized under the assumption that values are generally much larger than keys. @@ -51,7 +51,7 @@ type Table interface { // Each PutRequest may include zero or more secondary keys (see Put for semantics). // // The maximum size of a key (primary or secondary) is 64 KiB (2^16 - 1 bytes). The maximum size of a - // value is 2^32 - 1 bytes (math.MaxUint32, ~4 GiB); a larger value is rejected with an error. This + // value is 2^32 - 2 bytes (math.MaxUint32 - 1, ~4 GiB); a larger value is rejected with an error. This // database has been optimized under the assumption that values are generally much larger than keys. // This affects performance, but not correctness. // @@ -69,8 +69,8 @@ type Table interface { // (returns false if the key does not exist). If an error is returned, the value of the other returned values are // undefined. // - // The maximum size of a key is 64 KiB (2^16 - 1 bytes). The maximum size of a value is 2^32 - 1 bytes - // (math.MaxUint32, ~4 GiB). This database has been optimized under the assumption that values are + // The maximum size of a key is 64 KiB (2^16 - 1 bytes). The maximum size of a value is 2^32 - 2 bytes + // (math.MaxUint32 - 1, ~4 GiB). This database has been optimized under the assumption that values are // generally much larger than keys. This affects performance, but not correctness. // // For the sake of performance, the returned data is NOT safe to mutate. If you need to modify the data, diff --git a/sei-db/db_engine/litt/types/compression.go b/sei-db/db_engine/litt/types/compression.go index 25d06fe556..0917e1e379 100644 --- a/sei-db/db_engine/litt/types/compression.go +++ b/sei-db/db_engine/litt/types/compression.go @@ -2,10 +2,19 @@ package types import ( "fmt" + "math" "github.com/klauspost/compress/s2" ) +// s2WorstCaseExpansion bounds the number of bytes by which s2.Encode's output can exceed its input. +// S2 prefixes a block with a varint of the decompressed length (<=5 bytes for any input <= 4 GiB) and, +// for incompressible input, emits it as one literal run with a header (<=5 bytes); the true worst case +// is ~10 bytes. s2.Encode panics (ErrTooLarge) rather than return output larger than MaxUint32, so a +// value within this many bytes of the addressable ceiling must be stored uncompressed. Rounded up from +// the true bound for margin; TestS2MaxCompressibleSizeIsSafe pins it to the library's real behavior. +const s2WorstCaseExpansion = 16 + // CompressionAlgorithm identifies the algorithm used to compress values before they are written to a // segment's value file. It is stored (as a single byte) in each segment's metadata file so that a // segment is always decompressed with the algorithm it was written with, independent of the table's @@ -49,6 +58,57 @@ func (a CompressionAlgorithm) String() string { } } +// MaxCompressibleSize returns the largest input length this algorithm will compress. Inputs longer than +// this are stored uncompressed because the algorithm's worst-case output could reach the uint32 +// addressable value ceiling (for S2, Encode panics rather than emit such output). CompressionNone +// imposes no limit of its own. +func (a CompressionAlgorithm) MaxCompressibleSize() uint64 { + switch a { + case CompressionS2: + return math.MaxUint32 - s2WorstCaseExpansion + default: + return math.MaxUint64 + } +} + +// EncodeValue returns the on-disk representation of a value for a compressed segment: a one-byte +// algorithm tag followed by the value body. The tag records how the body was encoded, so each value is +// self-describing and decodes independently of the segment's configured algorithm. +// +// The body is the smaller of the compressed and raw forms (never store an expanded blob), and the value +// is stored raw whenever compression is disabled, the input exceeds the algorithm's MaxCompressibleSize +// (so Compress is never called on an unencodable input), or compression does not shrink it. The caller +// must not mutate src. +func EncodeValue(algorithm CompressionAlgorithm, src []byte) ([]byte, error) { + if algorithm != CompressionNone && uint64(len(src)) <= algorithm.MaxCompressibleSize() { + compressed, err := Compress(algorithm, src) + if err != nil { + return nil, err + } + if len(compressed) < len(src) { + return frameValue(algorithm, compressed), nil + } + } + return frameValue(CompressionNone, src), nil +} + +// DecodeValue reverses EncodeValue: it reads the leading algorithm tag and returns the decompressed +// body. The caller must not mutate src or the returned slice. +func DecodeValue(src []byte) ([]byte, error) { + if len(src) == 0 { + return nil, fmt.Errorf("cannot decode value: missing algorithm tag byte") + } + return Decompress(CompressionAlgorithm(src[0]), src[1:]) +} + +// frameValue builds the [algorithm-tag][body] on-disk blob. +func frameValue(algorithm CompressionAlgorithm, body []byte) []byte { + blob := make([]byte, 1+len(body)) + blob[0] = byte(algorithm) + copy(blob[1:], body) + return blob +} + // Compress returns the compressed form of src using the given algorithm. For CompressionNone it returns // src unchanged (the result may alias src). The caller must not mutate src or the returned slice. func Compress(algorithm CompressionAlgorithm, src []byte) ([]byte, error) { diff --git a/sei-db/db_engine/litt/types/compression_test.go b/sei-db/db_engine/litt/types/compression_test.go index 5177350c44..1e2b9bb7bf 100644 --- a/sei-db/db_engine/litt/types/compression_test.go +++ b/sei-db/db_engine/litt/types/compression_test.go @@ -2,8 +2,11 @@ package types import ( "bytes" + "math" + "math/rand" "testing" + "github.com/klauspost/compress/s2" "github.com/stretchr/testify/require" ) @@ -70,3 +73,90 @@ func TestCompressUnknownAlgorithm(t *testing.T) { _, err = Decompress(CompressionAlgorithm(99), []byte("data")) require.Error(t, err) } + +// TestEncodeValueCompressibleStoresCompressed verifies a shrinkable payload is stored under the S2 tag +// and round-trips. +func TestEncodeValueCompressibleStoresCompressed(t *testing.T) { + t.Parallel() + + payload := bytes.Repeat([]byte("the quick brown fox jumps over the lazy dog. "), 100) + + blob, err := EncodeValue(CompressionS2, payload) + require.NoError(t, err) + require.Equal(t, byte(CompressionS2), blob[0], "compressible payload should be tagged S2") + require.Less(t, len(blob), len(payload), "S2 blob (incl. tag) should be smaller than the raw payload") + + decoded, err := DecodeValue(blob) + require.NoError(t, err) + require.Equal(t, payload, decoded) +} + +// TestEncodeValueIncompressibleStoresRaw verifies that when compression does not shrink the value it is +// stored raw (tagged None) rather than in its larger compressed form, and still round-trips. +func TestEncodeValueIncompressibleStoresRaw(t *testing.T) { + t.Parallel() + + // Deterministic pseudo-random bytes: S2 cannot shrink these, so store-smaller must keep them raw. + payload := make([]byte, 4096) + rng := rand.New(rand.NewSource(1)) //nolint:gosec // test fixture, not security-sensitive + _, _ = rng.Read(payload) + + blob, err := EncodeValue(CompressionS2, payload) + require.NoError(t, err) + require.Equal(t, byte(CompressionNone), blob[0], "incompressible payload should be tagged None") + require.Equal(t, len(payload)+1, len(blob), "raw blob should be the payload plus a one-byte tag") + require.Equal(t, payload, blob[1:], "raw body should be the payload verbatim") + + decoded, err := DecodeValue(blob) + require.NoError(t, err) + require.Equal(t, payload, decoded) +} + +// TestEncodeValueEmpty verifies an empty value encodes to a lone tag byte and decodes back to empty. +func TestEncodeValueEmpty(t *testing.T) { + t.Parallel() + + blob, err := EncodeValue(CompressionS2, []byte{}) + require.NoError(t, err) + require.Len(t, blob, 1, "an empty value encodes to just the algorithm tag") + require.Equal(t, byte(CompressionNone), blob[0]) + + decoded, err := DecodeValue(blob) + require.NoError(t, err) + require.Empty(t, decoded) +} + +// TestDecodeValueEmptyInput verifies DecodeValue rejects a blob missing the tag byte. +func TestDecodeValueEmptyInput(t *testing.T) { + t.Parallel() + + _, err := DecodeValue(nil) + require.Error(t, err) + + _, err = DecodeValue([]byte{}) + require.Error(t, err) +} + +// TestDecodeValueUnknownTag verifies DecodeValue surfaces an unknown per-value algorithm tag. +func TestDecodeValueUnknownTag(t *testing.T) { + t.Parallel() + + _, err := DecodeValue([]byte{99, 'x'}) + require.Error(t, err) +} + +// TestS2MaxCompressibleSizeIsSafe pins s2WorstCaseExpansion to the library's real behavior: the largest +// input EncodeValue will hand to s2.Encode must have a defined (non-negative) MaxEncodedLen, so Encode +// can never panic with ErrTooLarge. Guards against the constant drifting too small. +func TestS2MaxCompressibleSizeIsSafe(t *testing.T) { + t.Parallel() + + maxSize := CompressionS2.MaxCompressibleSize() + require.Equal(t, uint64(math.MaxUint32-s2WorstCaseExpansion), maxSize) + require.LessOrEqual(t, maxSize, uint64(math.MaxInt), "cap must be representable as an int length") + require.GreaterOrEqual(t, s2.MaxEncodedLen(int(maxSize)), 0, + "s2.Encode would panic for the largest value EncodeValue compresses") + + require.Equal(t, uint64(math.MaxUint64), CompressionNone.MaxCompressibleSize(), + "CompressionNone imposes no size limit of its own") +} From 82e66792005a86ed5a3c33ea50660dfe5c1d4d7f Mon Sep 17 00:00:00 2001 From: Cody Littley Date: Mon, 20 Jul 2026 09:35:04 -0500 Subject: [PATCH 08/10] added note about performance hotspot --- sei-db/db_engine/litt/disktable/compression_loop.go | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/sei-db/db_engine/litt/disktable/compression_loop.go b/sei-db/db_engine/litt/disktable/compression_loop.go index 60dc124a4e..35656f6eb3 100644 --- a/sei-db/db_engine/litt/disktable/compression_loop.go +++ b/sei-db/db_engine/litt/disktable/compression_loop.go @@ -83,6 +83,12 @@ func (cl *compressionLoop) compress(req *controlLoopWriteRequest) bool { start = cl.clock() } + // Future work: if the compression thread ever becomes a bottleneck, consider the following tactics: + // - Fan out batches, handle normal work serially. + // Simple, but only helps with batch workloads. + // - Fan out all requests, have this thread only responsible for maintaining ordering. + // Complex but improves all worloads. + compressed := make([][]byte, len(req.values)) var uncompressedBytes uint64 var compressedBytes uint64 From ba73f62ba723c7ea7abc3f6525a74f35b9898b6e Mon Sep 17 00:00:00 2001 From: Cody Littley Date: Mon, 20 Jul 2026 12:07:18 -0500 Subject: [PATCH 09/10] properly bound unflushed data cache memory footprint with compression enabled --- .../db_engine/litt/disktable/control_loop.go | 3 +- sei-db/db_engine/litt/disktable/flush_loop.go | 8 +- .../litt/disktable/keymap_manager.go | 36 +++--- .../litt/disktable/keymap_manager_test.go | 28 ++++- .../litt/disktable/segment/segment.go | 104 ++++++++++++------ .../litt/disktable/segment/segment_test.go | 38 +++---- .../segment/shard_id_validation_test.go | 2 +- 7 files changed, 137 insertions(+), 82 deletions(-) diff --git a/sei-db/db_engine/litt/disktable/control_loop.go b/sei-db/db_engine/litt/disktable/control_loop.go index 481bf23201..a334aaf1fd 100644 --- a/sei-db/db_engine/litt/disktable/control_loop.go +++ b/sei-db/db_engine/litt/disktable/control_loop.go @@ -721,8 +721,7 @@ func (c *controlLoop) handleShutdownRequest(req *controlLoopShutdownRequest) { return } - // Seal the mutable segment - durableKeys, err := c.segments[c.highestSegmentIndex].Seal(c.clock()) + durableKeys, _, err := c.segments[c.highestSegmentIndex].Seal(c.clock()) if err != nil { c.errorMonitor.Panic(fmt.Errorf("failed to seal mutable segment: %w", err)) return diff --git a/sei-db/db_engine/litt/disktable/flush_loop.go b/sei-db/db_engine/litt/disktable/flush_loop.go index 0dd8a8099d..d0cb4433e7 100644 --- a/sei-db/db_engine/litt/disktable/flush_loop.go +++ b/sei-db/db_engine/litt/disktable/flush_loop.go @@ -75,7 +75,7 @@ func (f *flushLoop) run() { // the seal is finished and a new mutable segment has been created. This means that no future flush requests will be // sent to the segment that is being sealed, since only the control loop can schedule work for the flush loop. func (f *flushLoop) handleSealRequest(req *flushLoopSealRequest) { - durableKeys, err := req.segmentToSeal.Seal(req.now) + durableKeys, rawValueBytes, err := req.segmentToSeal.Seal(req.now) if err != nil { f.errorMonitor.Panic(fmt.Errorf("failed to seal segment %s: %w", req.segmentToSeal.String(), err)) return @@ -84,7 +84,7 @@ func (f *flushLoop) handleSealRequest(req *flushLoopSealRequest) { // Schedule the now-durable keys to be written into the keymap asynchronously. The segment data is already // crash durable, so we don't block the seal on the keymap write; a crash that loses the scheduled write is // repaired from the segment key files on the next startup. A full manager channel backpressures here. - err = f.keymapManager.scheduleWrite(durableKeys) + err = f.keymapManager.scheduleWrite(durableKeys, rawValueBytes) if err != nil { // The only error path is a panic-induced shutdown; the awaiting caller is released via the error // monitor, so there is nothing more to do here. @@ -117,7 +117,7 @@ func (f *flushLoop) handleFlushRequest(req *flushLoopFlushRequest) { segmentFlushStart = f.clock() } - durableKeys, err := req.flushWaitFunction() + durableKeys, rawValueBytes, err := req.flushWaitFunction() if err != nil { f.errorMonitor.Panic(fmt.Errorf("failed to flush mutable segment: %w", err)) return @@ -133,7 +133,7 @@ func (f *flushLoop) handleFlushRequest(req *flushLoopFlushRequest) { // caller without waiting for the keymap write. The segment data is already crash durable; a crash that // loses the scheduled write is repaired from the segment key files on the next startup. A full manager // channel backpressures here, which propagates back to Flush(). - err = f.keymapManager.scheduleWrite(durableKeys) + err = f.keymapManager.scheduleWrite(durableKeys, rawValueBytes) if err != nil { // The only error path is a panic-induced shutdown; the awaiting caller is released via the error // monitor, so there is nothing more to do here. diff --git a/sei-db/db_engine/litt/disktable/keymap_manager.go b/sei-db/db_engine/litt/disktable/keymap_manager.go index 9f80a99db9..9d6dadb5f6 100644 --- a/sei-db/db_engine/litt/disktable/keymap_manager.go +++ b/sei-db/db_engine/litt/disktable/keymap_manager.go @@ -14,6 +14,11 @@ import ( // keymapWriteRequest carries a batch of newly-durable keys to be written (put) into the keymap. type keymapWriteRequest struct { keys []*types.ScopedKey + + // uncompressedPutBytes is the total raw (pre-compression) value bytes represented by the primary keys in + // this batch. It feeds pendingPutBytes to bound the raw memory footprint of the unflushed-data cache, + // independent of how the values are represented on disk. + uncompressedPutBytes uint64 } // keymapDeleteRequest carries the keys of a garbage-collected segment to be deleted from the keymap. segment @@ -93,8 +98,10 @@ type keymapManager struct { // the maximum number of keys coalesced into a single keymap Put or Delete maxBatchSize int - // the maximum number of value bytes a coalescing put batch may represent before it is applied, - // independent of key count. Bounds the unflushed-data cache when values are large (few keys, many bytes). + // the maximum number of raw (pre-compression) value bytes a coalescing put batch may represent before it is + // applied, independent of key count. Bounds the raw memory footprint of the unflushed-data cache when values + // are large (few keys, many bytes). Measured in raw bytes because that is what the cache holds, independent + // of the compressed on-disk size. maxBatchBytes uint64 // the maximum number of keys deleted from the keymap in a single keymap.Delete call @@ -113,9 +120,10 @@ type keymapManager struct { // sub-batch, so a put and a same-key delete in flight resolve to deleted (the delete wins). puts []*types.ScopedKey - // pendingPutBytes is the total value bytes represented by the primary keys in puts. Secondary keys are - // zero-copy sub-ranges of their primary's value in the unflushed-data cache, so counting primaries alone - // tracks the cache memory this batch will free when applied. Reset to zero whenever puts is applied. + // pendingPutBytes is the total raw (pre-compression) value bytes represented by the primary keys in puts. + // Secondary keys are zero-copy sub-ranges of their primary's value in the unflushed-data cache, so counting + // primaries alone tracks the raw cache memory this batch will free when applied. Raw (not on-disk) bytes, + // since that is the memory the cache actually holds. Reset to zero whenever puts is applied. pendingPutBytes uint64 // deleteBacklog holds garbage-collected segments' keys awaiting deletion, in arrival (FIFO) order — @@ -175,13 +183,15 @@ func newKeymapManager( } } -// scheduleWrite enqueues a batch of newly-durable keys to be put into the keymap. Blocks if the manager's -// channel is full (backpressure); returns an error only if the DB is panicking. -func (m *keymapManager) scheduleWrite(keys []*types.ScopedKey) error { +// scheduleWrite enqueues a batch of newly-durable keys to be put into the keymap. uncompressedPutBytes is the +// total raw (pre-compression) value bytes the batch's primary keys represent, used to bound the unflushed-data +// cache. Blocks if the manager's channel is full (backpressure); returns an error only if the DB is panicking. +func (m *keymapManager) scheduleWrite(keys []*types.ScopedKey, uncompressedPutBytes uint64) error { if len(keys) == 0 { return nil } - return util.Send(m.errorMonitor, m.requestChan, &keymapWriteRequest{keys: keys}) + return util.Send(m.errorMonitor, m.requestChan, + &keymapWriteRequest{keys: keys, uncompressedPutBytes: uncompressedPutBytes}) } // scheduleDelete enqueues a garbage-collected segment's keys to be deleted from the keymap. segment is the @@ -308,13 +318,7 @@ func (m *keymapManager) routeRequest(msg any) bool { switch req := msg.(type) { case *keymapWriteRequest: m.puts = append(m.puts, req.keys...) - for _, k := range req.keys { - // Only primaries own value bytes in the unflushed-data cache; secondaries alias a sub-range of - // their primary's value, so counting them too would double-count the cache memory. - if k.Kind.IsPrimary() { - m.pendingPutBytes += uint64(k.Address.ValueSize()) - } - } + m.pendingPutBytes += req.uncompressedPutBytes case *keymapDeleteRequest: m.enqueueDelete(req) case *keymapManagerSyncRequest: diff --git a/sei-db/db_engine/litt/disktable/keymap_manager_test.go b/sei-db/db_engine/litt/disktable/keymap_manager_test.go index 0e6d6919fe..3276ca2fcb 100644 --- a/sei-db/db_engine/litt/disktable/keymap_manager_test.go +++ b/sei-db/db_engine/litt/disktable/keymap_manager_test.go @@ -113,7 +113,7 @@ func drainWatermark(ch chan int64) int64 { func TestKeymapManagerPutThenDeleteResolvesToDeleted(t *testing.T) { m, wmChan := newTestKeymapManager(t, 1024, 1 /* split every delete key */, 1_000_000, 1024) - require.NoError(t, m.scheduleWrite(scopedKeys("K", "A", "B", "X", "Y"))) + require.NoError(t, m.scheduleWrite(scopedKeys("K", "A", "B", "X", "Y"), 0)) // Delete group for segment 5 contains K, X, Y (enqueued after the put of the same keys). require.NoError(t, m.scheduleDelete(scopedKeys("X", "K", "Y"), 5)) require.NoError(t, m.drain()) @@ -128,7 +128,7 @@ func TestKeymapManagerPutThenDeleteResolvesToDeleted(t *testing.T) { func TestKeymapManagerDrainsBacklogAndAdvancesWatermark(t *testing.T) { m, wmChan := newTestKeymapManager(t, 1024, 2 /* force splitting */, 1_000_000, 1024) - require.NoError(t, m.scheduleWrite(scopedKeys("a1", "a2", "a3", "b1", "b2", "c1"))) + require.NoError(t, m.scheduleWrite(scopedKeys("a1", "a2", "a3", "b1", "b2", "c1"), 0)) require.NoError(t, m.scheduleDelete(scopedKeys("a1", "a2", "a3"), 1)) require.NoError(t, m.scheduleDelete(scopedKeys("b1", "b2"), 2)) require.NoError(t, m.scheduleDelete(scopedKeys("c1"), 3)) @@ -173,7 +173,7 @@ func TestKeymapManagerBackpressureDoesNotDeadlock(t *testing.T) { for i := 0; i < 10; i++ { ks = append(ks, fmt.Sprintf("g%d_k%d", g, i)) } - require.NoError(t, m.scheduleWrite(scopedKeys(ks...))) + require.NoError(t, m.scheduleWrite(scopedKeys(ks...), 0)) require.NoError(t, m.scheduleDelete(scopedKeys(ks...), int64(g))) allKeys = append(allKeys, ks...) } @@ -188,14 +188,32 @@ func TestKeymapManagerBackpressureDoesNotDeadlock(t *testing.T) { func TestKeymapManagerSyncAppliesThenContinues(t *testing.T) { m, _ := newTestKeymapManager(t, 1024, 1024, 1_000_000, 1024) - require.NoError(t, m.scheduleWrite(scopedKeys("k1"))) + require.NoError(t, m.scheduleWrite(scopedKeys("k1"), 0)) require.NoError(t, m.scheduleDelete(scopedKeys("k1"), 1)) require.NoError(t, m.sync()) assertAbsent(t, m, "k1") - require.NoError(t, m.scheduleWrite(scopedKeys("k2"))) + require.NoError(t, m.scheduleWrite(scopedKeys("k2"), 0)) require.NoError(t, m.sync()) assertPresent(t, m, "k2") require.NoError(t, m.drain()) } + +// TestKeymapManagerPendingPutBytesTracksRawSize pins the cache-bound accounting: pendingPutBytes must reflect the +// raw (pre-compression) value bytes carried on the write request, NOT Address.ValueSize() (which is the compressed +// on-disk size for a compressed table). This is what keeps maxBatchBytes bounding the raw unflushed-data cache. +func TestKeymapManagerPendingPutBytesTracksRawSize(t *testing.T) { + m, _ := buildTestKeymapManager(t, 1024, 1, 1_000_000, 1024) + + // Address encodes a small (compressed) on-disk size; the request carries a much larger raw size. The + // accounting must follow the raw size, proving it is decoupled from Address.ValueSize(). + keys := []*types.ScopedKey{ + {Key: []byte("k"), Address: types.NewAddress(0, 0, 0, 10), Kind: types.KeyKindStandalone}, + } + const rawBytes uint64 = 100_000 + + require.False(t, m.routeRequest(&keymapWriteRequest{keys: keys, uncompressedPutBytes: rawBytes})) + require.Equal(t, rawBytes, m.pendingPutBytes, + "pendingPutBytes must track the request's raw value bytes, not Address.ValueSize()") +} diff --git a/sei-db/db_engine/litt/disktable/segment/segment.go b/sei-db/db_engine/litt/disktable/segment/segment.go index b76c935ad1..ac55f5353d 100644 --- a/sei-db/db_engine/litt/disktable/segment/segment.go +++ b/sei-db/db_engine/litt/disktable/segment/segment.go @@ -534,12 +534,14 @@ func (s *Segment) Write(data *types.PutRequest) (keyCount uint32, keyFileSize ui } // Forward the primary key to the key file control loop, which asynchronously writes it to the - // key file. Primary always goes first; recovery relies on this ordering. - primaryRequest := &types.ScopedKey{ + // key file. Primary always goes first; recovery relies on this ordering. The primary carries the raw + // value bytes it contributes to the unflushed-data cache; secondaries alias it and contribute nothing. + primaryKey := &types.ScopedKey{ Key: data.Key, Address: types.NewAddress(s.index, firstByteIndex, shard, uint32(valueLen)), //nolint:gosec // bounded above Kind: primaryKind, } + primaryRequest := &keyFileWriteRequest{key: primaryKey, rawValueBytes: valueLen} err = util.Send(s.errorMonitor, s.keyFileChannel, primaryRequest) if err != nil { return 0, 0, @@ -551,10 +553,12 @@ func (s *Segment) Write(data *types.PutRequest) (keyCount uint32, keyFileSize ui if i == n-1 { kind = types.KeyKindFinalSecondary } - secondaryRequest := &types.ScopedKey{ - Key: sk.Key, - Address: types.NewAddress(s.index, firstByteIndex+sk.Offset, shard, sk.Length), - Kind: kind, + secondaryRequest := &keyFileWriteRequest{ + key: &types.ScopedKey{ + Key: sk.Key, + Address: types.NewAddress(s.index, firstByteIndex+sk.Offset, shard, sk.Length), + Kind: kind, + }, } err = util.Send(s.errorMonitor, s.keyFileChannel, secondaryRequest) if err != nil { @@ -634,10 +638,15 @@ func (s *Segment) WriteCompressed( // PutBatch caps the raw value at MaxUint32-1, so blobLen <= MaxUint32. blobAddress := types.NewAddress(s.index, firstByteIndex, shard, uint32(blobLen)) //nolint:gosec // see above - primaryRequest := &types.ScopedKey{ - Key: data.Key, - Address: blobAddress, - Kind: primaryKind, + // The primary carries the raw (pre-compression) value bytes it contributes to the unflushed-data + // cache; full-value-alias secondaries share the blob and contribute no additional cache memory. + primaryRequest := &keyFileWriteRequest{ + key: &types.ScopedKey{ + Key: data.Key, + Address: blobAddress, + Kind: primaryKind, + }, + rawValueBytes: uint64(len(data.Value)), } err = util.Send(s.errorMonitor, s.keyFileChannel, primaryRequest) if err != nil { @@ -650,10 +659,12 @@ func (s *Segment) WriteCompressed( if i == n-1 { kind = types.KeyKindFinalSecondary } - secondaryRequest := &types.ScopedKey{ - Key: sk.Key, - Address: blobAddress, // full-value alias: shares the primary's compressed blob - Kind: kind, + secondaryRequest := &keyFileWriteRequest{ + key: &types.ScopedKey{ + Key: sk.Key, + Address: blobAddress, // full-value alias: shares the primary's compressed blob + Kind: kind, + }, } err = util.Send(s.errorMonitor, s.keyFileChannel, secondaryRequest) if err != nil { @@ -739,8 +750,9 @@ func (s *Segment) GetKeys() ([]*types.ScopedKey, error) { } // FlushWaitFunction is a function that waits for a flush operation to complete. It returns the addresses of the data -// that was flushed, or an error if the flush operation failed. -type FlushWaitFunction func() ([]*types.ScopedKey, error) +// that was flushed and the total raw (pre-compression) value bytes those keys represent, or an error if the flush +// operation failed. +type FlushWaitFunction func() ([]*types.ScopedKey, uint64, error) // Flush schedules a flush operation. Flush operations are performed serially in the order they are scheduled. // This method returns a function that, when called, will block until the flush operation is complete. The function @@ -776,22 +788,22 @@ func (s *Segment) flush(seal bool) (FlushWaitFunction, error) { return nil, fmt.Errorf("failed to send flush request to key file: %w", err) } - return func() ([]*types.ScopedKey, error) { + return func() ([]*types.ScopedKey, uint64, error) { // Wait for each shard to finish flushing. for i := range s.shardChannels { _, err := util.Await(s.errorMonitor, shardResponseChannels[i]) if err != nil { - return nil, fmt.Errorf("failed to flush shard %d: %w", i, err) + return nil, 0, fmt.Errorf("failed to flush shard %d: %w", i, err) } } keyFlushResponse, err := util.Await(s.errorMonitor, keyResponseChannel) if err != nil { - return nil, fmt.Errorf("failed to flush key file: %w", err) + return nil, 0, fmt.Errorf("failed to flush key file: %w", err) } s.unflushedKeyCount.Add(-int64(len(keyFlushResponse.addresses))) - return keyFlushResponse.addresses, nil + return keyFlushResponse.addresses, keyFlushResponse.rawValueBytes, nil }, nil } @@ -835,30 +847,31 @@ func (s *Segment) IsSnapshot() (bool, error) { return fileInfo.Mode()&os.ModeSymlink != 0, nil } -// Seal flushes all data to disk and finalizes the metadata. Returns addresses that became durable as a result of -// this method call. After this method is called, no more data can be written to this segment. -func (s *Segment) Seal(now time.Time) ([]*types.ScopedKey, error) { +// Seal flushes all data to disk and finalizes the metadata. Returns the addresses that became durable as a result +// of this method call and the total raw (pre-compression) value bytes they represent. After this method is called, +// no more data can be written to this segment. +func (s *Segment) Seal(now time.Time) ([]*types.ScopedKey, uint64, error) { flushWaitFunction, err := s.flush(true) if err != nil { - return nil, fmt.Errorf("failed to flush segment: %w", err) + return nil, 0, fmt.Errorf("failed to flush segment: %w", err) } - addresses, err := flushWaitFunction() + addresses, rawValueBytes, err := flushWaitFunction() if err != nil { - return nil, fmt.Errorf("failed to flush segment: %w", err) + return nil, 0, fmt.Errorf("failed to flush segment: %w", err) } // Seal the metadata file. err = s.metadata.seal(now, s.keyCount) if err != nil { - return nil, fmt.Errorf("failed to seal metadata file: %w", err) + return nil, 0, fmt.Errorf("failed to seal metadata file: %w", err) } unflushedKeyCount := s.unflushedKeyCount.Load() if s.unflushedKeyCount.Load() != 0 { - return nil, fmt.Errorf("segment %d has %d unflushedKeyCount keys", s.index, unflushedKeyCount) + return nil, 0, fmt.Errorf("segment %d has %d unflushedKeyCount keys", s.index, unflushedKeyCount) } - return addresses, nil + return addresses, rawValueBytes, nil } // IsSealed returns true if the segment is sealed, and false otherwise. @@ -1003,7 +1016,11 @@ func (s *Segment) handleKeyFileWrite(data *types.ScopedKey) { } // handleKeyFileFlushRequest handles a request to flush the key file to disk. -func (s *Segment) handleKeyFileFlushRequest(request *keyFileFlushRequest, unflushedKeys []*types.ScopedKey) { +func (s *Segment) handleKeyFileFlushRequest( + request *keyFileFlushRequest, + unflushedKeys []*types.ScopedKey, + unflushedRawBytes uint64, +) { if request.seal { err := s.keys.seal() if err != nil { @@ -1017,7 +1034,8 @@ func (s *Segment) handleKeyFileFlushRequest(request *keyFileFlushRequest, unflus } request.completionChannel <- &keyFileFlushResponse{ - addresses: unflushedKeys, + addresses: unflushedKeys, + rawValueBytes: unflushedRawBytes, } } @@ -1078,13 +1096,27 @@ type keyFileFlushRequest struct { // keyFileFlushResponse is a message sent from the key file control loop to the caller of Flush to indicate that the // key file has been flushed. type keyFileFlushResponse struct { + // The keys that became durable as a result of this flush. addresses []*types.ScopedKey + + // The total raw (pre-compression) value bytes represented by the primary keys in addresses. Used by the keymap + // manager to bound the raw memory footprint of the unflushed-data cache. + rawValueBytes uint64 +} + +// keyFileWriteRequest is a message sent to the key file control loop to append one durable key to the key file. It +// also carries the raw (pre-compression) value bytes the key contributes to the unflushed-data cache: nonzero only +// for a write group's primary key, since secondaries alias the primary's value and add no cache memory. +type keyFileWriteRequest struct { + key *types.ScopedKey + rawValueBytes uint64 } // keyFileControlLoop is the main loop for performing modifications to the key file. This goroutine is responsible // for writing key-address pairs to the key file. func (s *Segment) keyFileControlLoop() { unflushedKeys := make([]*types.ScopedKey, 0, unflushedKeysInitialCapacity) + var unflushedRawBytes uint64 for { select { @@ -1094,17 +1126,19 @@ func (s *Segment) keyFileControlLoop() { case operation := <-s.keyFileChannel: if flushRequest, ok := operation.(*keyFileFlushRequest); ok { - s.handleKeyFileFlushRequest(flushRequest, unflushedKeys) + s.handleKeyFileFlushRequest(flushRequest, unflushedKeys, unflushedRawBytes) unflushedKeys = make([]*types.ScopedKey, 0, unflushedKeysInitialCapacity) + unflushedRawBytes = 0 if flushRequest.seal { // After sealing, we can exit the control loop. return } - } else if data, ok := operation.(*types.ScopedKey); ok { - s.handleKeyFileWrite(data) - unflushedKeys = append(unflushedKeys, data) + } else if data, ok := operation.(*keyFileWriteRequest); ok { + s.handleKeyFileWrite(data.key) + unflushedKeys = append(unflushedKeys, data.key) + unflushedRawBytes += data.rawValueBytes } else { s.errorMonitor.Panic( diff --git a/sei-db/db_engine/litt/disktable/segment/segment_test.go b/sei-db/db_engine/litt/disktable/segment/segment_test.go index cec619a5ee..c02684a9c3 100644 --- a/sei-db/db_engine/litt/disktable/segment/segment_test.go +++ b/sei-db/db_engine/litt/disktable/segment/segment_test.go @@ -82,7 +82,7 @@ func TestWriteAndReadSegmentSingleShard(t *testing.T) { if rand.BoolWithProbability(0.25) { flushFunction, err := seg.Flush() require.NoError(t, err) - flushedKeys, err := flushFunction() + flushedKeys, _, err := flushFunction() require.NoError(t, err) for _, flushedKey := range flushedKeys { addressMap[string(flushedKey.Key)] = flushedKey.Address @@ -96,7 +96,7 @@ func TestWriteAndReadSegmentSingleShard(t *testing.T) { if rand.BoolWithProbability(0.1) { flushFunction, err := seg.Flush() require.NoError(t, err) - flushedKeys, err := flushFunction() + flushedKeys, _, err := flushFunction() require.NoError(t, err) for _, flushedKey := range flushedKeys { addressMap[string(flushedKey.Key)] = flushedKey.Address @@ -116,7 +116,7 @@ func TestWriteAndReadSegmentSingleShard(t *testing.T) { // Seal the segment and read all keys and values. require.False(t, seg.IsSealed()) sealTime := rand.Time() - flushedKeys, err := seg.Seal(sealTime) + flushedKeys, _, err := seg.Seal(sealTime) require.NoError(t, err) require.True(t, seg.IsSealed()) @@ -231,7 +231,7 @@ func TestWriteAndReadSegmentMultiShard(t *testing.T) { if rand.BoolWithProbability(0.25) { flushFunction, err := seg.Flush() require.NoError(t, err) - flushedKeys, err := flushFunction() + flushedKeys, _, err := flushFunction() require.NoError(t, err) for _, flushedKey := range flushedKeys { addressMap[string(flushedKey.Key)] = flushedKey.Address @@ -245,7 +245,7 @@ func TestWriteAndReadSegmentMultiShard(t *testing.T) { if rand.BoolWithProbability(0.1) { flushFunction, err := seg.Flush() require.NoError(t, err) - flushedKeys, err := flushFunction() + flushedKeys, _, err := flushFunction() require.NoError(t, err) for _, flushedKey := range flushedKeys { addressMap[string(flushedKey.Key)] = flushedKey.Address @@ -265,7 +265,7 @@ func TestWriteAndReadSegmentMultiShard(t *testing.T) { // Seal the segment and read all keys and values. require.False(t, seg.IsSealed()) sealTime := rand.Time() - flushedKeys, err := seg.Seal(sealTime) + flushedKeys, _, err := seg.Seal(sealTime) require.NoError(t, err) require.True(t, seg.IsSealed()) @@ -391,7 +391,7 @@ func TestWriteAndReadColdShard(t *testing.T) { // Seal the segment and read all keys and values. require.False(t, seg.IsSealed()) sealTime := rand.Time() - flushedKeys, err := seg.Seal(sealTime) + flushedKeys, _, err := seg.Seal(sealTime) require.NoError(t, err) require.True(t, seg.IsSealed()) @@ -566,7 +566,7 @@ func TestRoundRobinShardAssignment(t *testing.T) { flushFn, err := seg.Flush() require.NoError(t, err) - flushed, err := flushFn() + flushed, _, err := flushFn() require.NoError(t, err) // Each iteration above should produce exactly one new flushed key (the one we just wrote). require.Len(t, flushed, 1) @@ -667,7 +667,7 @@ func TestSegmentSecondaryKeyAddresses(t *testing.T) { _, _, err = seg.Write(&types.PutRequest{Key: standaloneKey, Value: standaloneValue}) require.NoError(t, err) - flushedKeys, err := seg.Seal(time.Now()) + flushedKeys, _, err := seg.Seal(time.Now()) require.NoError(t, err) require.Len(t, flushedKeys, 5) @@ -736,7 +736,7 @@ func TestKeyFileKindRoundTrip(t *testing.T) { }, }) - _, err := seg.Seal(time.Now()) + _, _, err := seg.Seal(time.Now()) require.NoError(t, err) // Reload from disk and verify the on-disk record kinds. @@ -850,7 +850,7 @@ func TestSealLoadedSegmentSingleShardPrefix(t *testing.T) { for i := 0; i < n; i++ { writeNoErr(t, seg, &types.PutRequest{Key: keyFor(i), Value: valueFor(i)}) } - _, err := seg.Seal(time.Now()) + _, _, err := seg.Seal(time.Now()) require.NoError(t, err) markSegmentUnsealed(t, segmentPath, index) return segmentPath, index @@ -906,7 +906,7 @@ func TestSealLoadedSegmentGroupAtomicity(t *testing.T) { t.Parallel() seg, segmentPath, index := newSingleShardSegment(t) writeNoErr(t, seg, &types.PutRequest{Key: []byte("k1"), Value: []byte("v1")}) - _, err := seg.Seal(time.Now()) + _, _, err := seg.Seal(time.Now()) require.NoError(t, err) markSegmentUnsealed(t, segmentPath, index) @@ -927,7 +927,7 @@ func TestSealLoadedSegmentGroupAtomicity(t *testing.T) { {Key: []byte("llo"), Offset: 2, Length: 3}, }, }) - _, err := seg.Seal(time.Now()) + _, _, err := seg.Seal(time.Now()) require.NoError(t, err) markSegmentUnsealed(t, segmentPath, index) @@ -952,7 +952,7 @@ func TestSealLoadedSegmentGroupAtomicity(t *testing.T) { {Key: []byte("llo"), Offset: 2, Length: 3}, }, }) - _, err := seg.Seal(time.Now()) + _, _, err := seg.Seal(time.Now()) require.NoError(t, err) secondaryRecBytes := int(keyRecordSize([]byte("he")) + keyRecordSize([]byte("llo"))) @@ -976,7 +976,7 @@ func TestSealLoadedSegmentGroupAtomicity(t *testing.T) { {Key: []byte("llo"), Offset: 2, Length: 3}, }, }) - _, err := seg.Seal(time.Now()) + _, _, err := seg.Seal(time.Now()) require.NoError(t, err) truncateKeyFileBy(t, segmentPath, index, int(keyRecordSize([]byte("llo")))) @@ -1002,7 +1002,7 @@ func TestSealLoadedSegmentGroupAtomicity(t *testing.T) { {Key: []byte("torn-secondary"), Offset: 0, Length: 5}, }, }) - _, err := seg.Seal(time.Now()) + _, _, err := seg.Seal(time.Now()) require.NoError(t, err) truncateKeyFileBy(t, segmentPath, index, 5) @@ -1031,7 +1031,7 @@ func TestSealLoadedSegmentGroupAtomicity(t *testing.T) { {Key: []byte("oo"), Offset: 7, Length: 2}, }, }) - _, err := seg.Seal(time.Now()) + _, _, err := seg.Seal(time.Now()) require.NoError(t, err) truncateValueFileBy(t, segmentPath, index, 0, 3) @@ -1053,7 +1053,7 @@ func TestSealLoadedSegmentGroupAtomicity(t *testing.T) { {Key: []byte("oo"), Offset: 7, Length: 2}, }, }) - _, err := seg.Seal(time.Now()) + _, _, err := seg.Seal(time.Now()) require.NoError(t, err) markSegmentUnsealed(t, segmentPath, index) @@ -1079,7 +1079,7 @@ func TestSealLoadedSegmentGroupAtomicity(t *testing.T) { {Key: []byte("third-secondary"), Offset: 0, Length: 2}, }, }) - _, err := seg.Seal(time.Now()) + _, _, err := seg.Seal(time.Now()) require.NoError(t, err) truncateKeyFileBy(t, segmentPath, index, int(keyRecordSize([]byte("third-secondary")))) diff --git a/sei-db/db_engine/litt/disktable/segment/shard_id_validation_test.go b/sei-db/db_engine/litt/disktable/segment/shard_id_validation_test.go index 6eccbd087d..8cd5d8eb24 100644 --- a/sei-db/db_engine/litt/disktable/segment/shard_id_validation_test.go +++ b/sei-db/db_engine/litt/disktable/segment/shard_id_validation_test.go @@ -48,6 +48,6 @@ func TestSegmentReadRejectsOutOfRangeShardID(t *testing.T) { }) // Wind down the segment's control-loop goroutines before the test exits. - _, err = seg.Seal(time.Now()) + _, _, err = seg.Seal(time.Now()) require.NoError(t, err) } From dab46c870a3598b65f98df6f77bab17386e582ea Mon Sep 17 00:00:00 2001 From: Cody Littley Date: Mon, 20 Jul 2026 13:20:00 -0500 Subject: [PATCH 10/10] go fmt --- sei-db/db_engine/litt/disktable/compression_loop.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sei-db/db_engine/litt/disktable/compression_loop.go b/sei-db/db_engine/litt/disktable/compression_loop.go index 35656f6eb3..37db8fc480 100644 --- a/sei-db/db_engine/litt/disktable/compression_loop.go +++ b/sei-db/db_engine/litt/disktable/compression_loop.go @@ -84,9 +84,9 @@ func (cl *compressionLoop) compress(req *controlLoopWriteRequest) bool { } // Future work: if the compression thread ever becomes a bottleneck, consider the following tactics: - // - Fan out batches, handle normal work serially. + // - Fan out batches, handle normal work serially. // Simple, but only helps with batch workloads. - // - Fan out all requests, have this thread only responsible for maintaining ordering. + // - Fan out all requests, have this thread only responsible for maintaining ordering. // Complex but improves all worloads. compressed := make([][]byte, len(req.values))