fix: littDB consistent keyfiles (#1573)

* Save value size to keyfile.

Signed-off-by: Cody Littley <[email protected]>

* Add key count to metadata file.

Signed-off-by: Cody Littley <[email protected]>

* Fix DB size calculation.

Signed-off-by: Cody Littley <[email protected]>

* Fix key count.

Signed-off-by: Cody Littley <[email protected]>

* Consistent key files.

Signed-off-by: Cody Littley <[email protected]>

* remove old keymap reload logic

Signed-off-by: Cody Littley <[email protected]>

* fix tickers, fix broken test

Signed-off-by: Cody Littley <[email protected]>

* Fix TODOs.

Signed-off-by: Cody Littley <[email protected]>

* cleanup

Signed-off-by: Cody Littley <[email protected]>

* Made suggested change.

Signed-off-by: Cody Littley <[email protected]>

---------

Signed-off-by: Cody Littley <[email protected]>

Author: cody-littley

common/testutils/test_utils.go

@@ -8,7 +8,7 @@ import (
 	"time"
 
 	"github.com/Layr-Labs/eigensdk-go/logging"
-	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"golang.org/x/exp/rand"
 )
 
@@ -38,13 +38,21 @@ func AssertEventuallyTrue(t *testing.T, condition func() bool, duration time.Dur
 	}
 
 	ticker := time.NewTicker(1 * time.Millisecond)
-	select {
-	case <-ticker.C:
-		if condition() {
+	defer ticker.Stop()
+
+	ctx, cancel := context.WithTimeout(context.Background(), duration)
+	defer cancel()
+
+	for {
+		select {
+		case <-ticker.C:
+			if condition() {
+				return
+			}
+		case <-ctx.Done():
+			require.True(t, condition(), debugInfo...)
 			return
 		}
-	case <-time.After(duration):
-		assert.True(t, condition(), debugInfo...)
 	}
 }
 

litt/disktable/control_loop.go

@@ -60,7 +60,7 @@ type controlLoop struct {
 	targetFileSize uint32
 
 	// The maximum number of keys in a segment.
-	maxKeyCount uint64
+	maxKeyCount uint32
 
 	// The target size for key files.
 	targetKeyFileSize uint64
@@ -119,6 +119,7 @@ func (c *controlLoop) enqueue(request controlLoopMessage) error {
 // mutate the data in the disk table.
 func (c *controlLoop) run() {
 	ticker := time.NewTicker(c.garbageCollectionPeriod)
+	defer ticker.Stop()
 
 	for {
 		select {
@@ -157,13 +158,15 @@ func (c *controlLoop) doGarbageCollection() {
 		return
 	}
 
-	if c.metrics != nil {
-		defer func() {
+	defer func() {
+		if c.metrics != nil {
 			end := c.clock()
 			delta := end.Sub(start)
 			c.metrics.ReportGarbageCollectionLatency(c.name, delta)
-		}()
-	}
+
+		}
+		c.updateCurrentSize()
+	}()
 
 	for index := c.lowestSegmentIndex; index <= c.highestSegmentIndex; index++ {
 		seg := c.segments[index]
@@ -198,8 +201,13 @@ func (c *controlLoop) doGarbageCollection() {
 			}
 		}
 
+		if seg.Size() > c.immutableSegmentSize {
+			c.logger.Errorf("segment %d size %d is larger than immutable segment size %d, "+
+				"reported DB size will not be accurate", index, seg.Size(), c.immutableSegmentSize)
+		}
+
 		c.immutableSegmentSize -= seg.Size()
-		c.keyCount.Add(-1 * int64(len(keys)))
+		c.keyCount.Add(-1 * int64(seg.KeyCount()))
 
 		// Deletion of segment files will happen when the segment is released by all reservation holders.
 		seg.Release()
@@ -209,8 +217,6 @@ func (c *controlLoop) doGarbageCollection() {
 
 		c.lowestSegmentIndex++
 	}
-
-	c.updateCurrentSize()
 }
 
 // getReservedSegment returns the segment with the given index. Segment is reserved, and it is the caller's
@@ -245,7 +251,10 @@ func (c *controlLoop) getSegments() (map[uint32]*segment.Segment, error) {
 
 // updateCurrentSize updates the size of the table.
 func (c *controlLoop) updateCurrentSize() {
-	size := c.immutableSegmentSize + c.segments[c.highestSegmentIndex].Size() + segment.MetadataSize
+	size := c.immutableSegmentSize +
+		c.segments[c.highestSegmentIndex].Size() +
+		c.metadata.Size()
+
 	c.size.Store(size)
 }
 
@@ -280,8 +289,6 @@ func (c *controlLoop) handleWriteRequest(req *controlLoopWriteRequest) {
 func (c *controlLoop) expandSegments() error {
 	now := c.clock()
 
-	c.immutableSegmentSize += c.segments[c.highestSegmentIndex].Size()
-
 	// Seal the previous segment.
 	flushLoopResponseChan := make(chan struct{}, 1)
 	request := &flushLoopSealRequest{
@@ -302,6 +309,9 @@ func (c *controlLoop) expandSegments() error {
 		return fmt.Errorf("failed to seal segment: %w", err)
 	}
 
+	// Record the size of the segment.
+	c.immutableSegmentSize += c.segments[c.highestSegmentIndex].Size()
+
 	// Create a new segment.
 	salt := [16]byte{}
 	_, err = c.saltShaker.Read(salt[:])

litt/disktable/disk_table.go

@@ -198,6 +198,12 @@ func NewDiskTable(
 		return nil, fmt.Errorf("failed to gather segment files: %w", err)
 	}
 
+	keyCount := int64(0)
+	for _, seg := range segments {
+		keyCount += int64(seg.KeyCount())
+	}
+	table.keyCount.Store(keyCount)
+
 	immutableSegmentSize := uint64(0)
 	for _, seg := range segments {
 		immutableSegmentSize += seg.Size()
@@ -282,6 +288,7 @@ func NewDiskTable(
 		keymap:                  keymap,
 		flushLoop:               fLoop,
 		garbageCollectionPeriod: config.GCPeriod,
+		immutableSegmentSize:    immutableSegmentSize,
 	}
 	cLoop.threadsafeHighestSegmentIndex.Store(highestSegmentIndex)
 	table.controlLoop = cLoop
@@ -306,27 +313,14 @@ func (d *DiskTable) reloadKeymap(
 	lowestSegmentIndex uint32,
 	highestSegmentIndex uint32) error {
 
-	if len(segments) > 0 {
-		// TODO(cody.littley): there is currently a race condition that may cause key files and value files to
-		//  be inconsistent. This is not a problem as long as the keymap is not reloaded, since this inconsistency
-		//  only hurts us when we are reloading the keymap. This needs to be fixed, but is not urgent since reloading
-		//  the keymap is not a feature currently used in production.
-		d.logger.Errorf("there is currently a race condition in this method, " +
-			"changing keymap type should not be used until fixed")
-	}
-
 	start := d.clock()
 	defer func() {
 		d.logger.Infof("spent %v reloading keymap", d.clock().Sub(start))
 	}()
 
-	// It's possible that some of the data written near the end of the previous session was corrupted.
-	// Read data from the end until the first valid key/value pair is found.
-	isValid := false
+	batch := make([]*types.ScopedKey, 0, keymapReloadBatchSize)
 
-	batch := make([]*types.KAPair, 0, keymapReloadBatchSize)
-
-	for i := highestSegmentIndex; i >= lowestSegmentIndex && i+1 != 0; i-- {
+	for i := lowestSegmentIndex; i <= highestSegmentIndex; i++ {
 		if !segments[i].IsSealed() {
 			// ignore unsealed segment, this will have been created in the current session and will not
 			// yet contain any data.
@@ -340,32 +334,15 @@ func (d *DiskTable) reloadKeymap(
 		for keyIndex := len(keys) - 1; keyIndex >= 0; keyIndex-- {
 			key := keys[keyIndex]
 
-			if !isValid {
-				_, err = segments[i].Read(key.Key, key.Address)
-				if err == nil {
-					// we found a valid key/value pair. All subsequent keys are valid.
-					isValid = true
-				} else {
-					// This is not cause for alarm (probably).
-					// This can happen when the database is not cleanly shut down,
-					// and just means that some data near the end was not fully committed.
-					d.logger.Infof("truncated value for key %s with address %s for segment %d",
-						key.Key, key.Address, i)
-				}
-			}
-
-			if isValid {
-				batch = append(batch, key)
-				if len(batch) == keymapReloadBatchSize {
-					err = d.keymap.Put(batch)
-					if err != nil {
-						return fmt.Errorf("failed to put keys for segment %d: %w", i, err)
-					}
-					batch = make([]*types.KAPair, 0, keymapReloadBatchSize)
+			batch = append(batch, key)
+			if len(batch) == keymapReloadBatchSize {
+				err = d.keymap.Put(batch)
+				if err != nil {
+					return fmt.Errorf("failed to put keys for segment %d: %w", i, err)
 				}
+				batch = make([]*types.ScopedKey, 0, keymapReloadBatchSize)
 			}
 		}
-
 	}
 
 	if len(batch) > 0 {
@@ -766,7 +743,7 @@ func (d *DiskTable) RunGC() error {
 
 // writeKeysToKeymap flushes all keys to the keymap. Once they are flushed, it also removes the keys from the
 // unflushedDataCache.
-func (d *DiskTable) writeKeysToKeymap(keys []*types.KAPair) error {
+func (d *DiskTable) writeKeysToKeymap(keys []*types.ScopedKey) error {
 	if len(keys) == 0 {
 		// Nothing to flush.
 		return nil

litt/disktable/disk_table_test.go

@@ -811,7 +811,7 @@ func lastFileMissingTest(t *testing.T, tableBuilder *tableBuilder, typeToDelete
 	// to the files being artificially deleted in this test), the keymap will not hold any value that has not
 	// yet been durably flushed to disk.
 	for key := range missingKeys {
-		err = table.(*DiskTable).keymap.Delete([]*types.KAPair{{Key: []byte(key)}})
+		err = table.(*DiskTable).keymap.Delete([]*types.ScopedKey{{Key: []byte(key)}})
 		require.NoError(t, err)
 	}
 
@@ -1041,7 +1041,7 @@ func truncatedKeyFileTest(t *testing.T, tableBuilder *tableBuilder) {
 	// to the files being artificially deleted in this test), the keymap will not hold any value that has not
 	// yet been durably flushed to disk.
 	for key := range missingKeys {
-		err = table.(*DiskTable).keymap.Delete([]*types.KAPair{{Key: []byte(key)}})
+		err = table.(*DiskTable).keymap.Delete([]*types.ScopedKey{{Key: []byte(key)}})
 		require.NoError(t, err)
 	}
 
@@ -1281,7 +1281,7 @@ func truncatedValueFileTest(t *testing.T, tableBuilder *tableBuilder) {
 	// to the files being artificially deleted in this test), the keymap will not hold any value that has not
 	// yet been durably flushed to disk.
 	for key := range missingKeys {
-		err = table.(*DiskTable).keymap.Delete([]*types.KAPair{{Key: []byte(key)}})
+		err = table.(*DiskTable).keymap.Delete([]*types.ScopedKey{{Key: []byte(key)}})
 		require.NoError(t, err)
 	}
 
@@ -1485,7 +1485,7 @@ func unflushedKeysTest(t *testing.T, tableBuilder *tableBuilder) {
 	// to the files being artificially deleted in this test), the keymap will not hold any value that has not
 	// yet been durably flushed to disk.
 	for key := range missingKeys {
-		err = table.(*DiskTable).keymap.Delete([]*types.KAPair{{Key: []byte(key)}})
+		err = table.(*DiskTable).keymap.Delete([]*types.ScopedKey{{Key: []byte(key)}})
 		require.NoError(t, err)
 	}
 
@@ -2137,10 +2137,18 @@ func tableSizeTest(t *testing.T, tableBuilder *tableBuilder) {
 		// Once in a while, scan the table and verify that all expected values are present.
 		// Don't do this every time for the sake of test runtime.
 		if rand.BoolWithProbability(0.01) || i == iterations-1 /* always check on the last iteration */ {
+
+			// Force garbage collection to run in order to remove expired values from counts.
+			err = table.Flush()
+			require.NoError(t, err)
+			err = (table).(*DiskTable).RunGC()
+			require.NoError(t, err)
+
 			// Remove expired values from the expected values.
 			newlyExpiredKeys := make([]string, 0)
 			for key, creationTime := range creationTimes {
-				if newTime.Sub(creationTime) > ttl {
+				age := newTime.Sub(creationTime)
+				if age > ttl {
 					newlyExpiredKeys = append(newlyExpiredKeys, key)
 				}
 			}
@@ -2183,7 +2191,74 @@ func tableSizeTest(t *testing.T, tableBuilder *tableBuilder) {
 	err = table.RunGC()
 	require.NoError(t, err)
 
+	// disable garbage collection
+	err = table.SetTTL(0)
+	require.NoError(t, err)
+	err = table.Flush()
+	require.NoError(t, err)
+
+	// Write some data that won't expire, just to be sure that the table is not empty.
+	for i := 0; i < 10; i++ {
+		key := rand.PrintableVariableBytes(32, 64)
+		value := rand.PrintableVariableBytes(1, 128)
+		err = table.Put(key, value)
+		require.NoError(t, err)
+		expectedValues[string(key)] = value
+	}
+
+	err = table.Flush()
+	require.NoError(t, err)
+
 	reportedSize := table.Size()
+	reportedKeyCount := table.KeyCount()
+
+	// The exact key count is hard to predict for the sake of this unit test, since GC is "lazy" and may not
+	// immediately remove all values that are legal to be removed. But at the very least, all unexpired
+	// values should be present, and the key count should not exceed the number of total inserted values.
+	require.GreaterOrEqual(t, reportedKeyCount, uint64(len(expectedValues)))
+	require.LessOrEqual(t, reportedKeyCount, uint64(len(expectedValues)+len(expiredValues)))
+
+	err = table.Close()
+	require.NoError(t, err)
+
+	// Walk the "directory" file tree and calculate the actual size of the table.
+	// There is some asynchrony in file deletion, so we retry a reasonable number of times.
+	testutils.AssertEventuallyTrue(t, func() bool {
+		actualSize := uint64(0)
+
+		err = filepath.Walk(directory, func(path string, info os.FileInfo, err error) error {
+			if err != nil {
+				// files may be deleted in the middle of the walk
+				return nil
+			}
+			if info.IsDir() {
+				// directory sizes are not factored into the table size
+				return nil
+			}
+			if strings.Contains(path, "keymap") {
+				// table size does not currently include the keymap size
+				return nil
+			}
+			actualSize += uint64(info.Size())
+			return nil
+		})
+		require.NoError(t, err)
+		return actualSize == reportedSize
+	}, time.Second)
+
+	// Restart the table. The size should be accurately reported.
+	table, err = tableBuilder.builder(clock, tableName, []string{directory})
+	require.NoError(t, err)
+
+	newReportedSize := table.Size()
+	newReportedKeyCount := table.KeyCount()
+
+	// New size should be greater than the old size, since GC is disabled and
+	// we will have started a new segment upon restart.
+	require.LessOrEqual(t, reportedSize, newReportedSize)
+
+	// The number of keys should be the same as before.
+	require.Equal(t, reportedKeyCount, newReportedKeyCount)
 
 	err = table.Close()
 	require.NoError(t, err)
@@ -2211,8 +2286,7 @@ func tableSizeTest(t *testing.T, tableBuilder *tableBuilder) {
 		})
 		require.NoError(t, err)
 
-		return actualSize == reportedSize
-		//require.Equal(t, actualSize, reportedSize, "delta: %d", int(actualSize)-int(reportedSize))
+		return actualSize == newReportedSize
 	}, time.Second)
 }
 

litt/disktable/keymap/keymap.go

@@ -19,13 +19,14 @@ const KeymapInitializedFileName = "initialized"
 
 // Keymap maintains a mapping between keys and addresses. Implementations of this interface are goroutine safe.
 type Keymap interface {
-	// Put adds keys to the keymap as a batch.
+	// Put adds keys to the keymap as a batch. This method is required to store the address, but can ignore
+	// other fields in the ScopedKey struct such as the value length.
 	//
 	// A keymap provides atomicity for individual key-address pairs, but not for the batch as a whole.
 	//
 	// It is not safe to modify the contents of any slices passed to this function after the call.
 	// This includes the byte slices containing the keys.
-	Put(pairs []*types.KAPair) error
+	Put(pairs []*types.ScopedKey) error
 
 	// Get returns the address for a key. Returns true if the key exists, and false otherwise (i.e. does not
 	// return an error if the key does not exist).
@@ -39,7 +40,7 @@ type Keymap interface {
 	//
 	// It is not safe to modify the contents of any slices passed to this function after the call.
 	// This includes the byte slices containing the keys.
-	Delete(keys []*types.KAPair) error
+	Delete(keys []*types.ScopedKey) error
 
 	// Stop stops the keymap.
 	Stop() error