Refactor MVCCStore: Add Compact, Snapshot consistency, and Logger options

- Add Compact method for garbage collection of old versions.
- Include LastCommitTS in snapshots to ensure consistency after restore.
- Refactor NewMVCCStore to accept options, enabling custom logger injection.
- Add tests for Compact functionality.

Author: bootjp

store/mvcc_store.go

@@ -27,6 +27,18 @@ const (
 	checksumSize = 4
 )
 
+// mvccSnapshot is used solely for gob snapshot serialization.
+type mvccSnapshot struct {
+	LastCommitTS uint64
+	Entries      []mvccSnapshotEntry
+}
+
+// mvccSnapshotEntry is used solely for gob snapshot serialization.
+type mvccSnapshotEntry struct {
+	Key      []byte
+	Versions []VersionedValue
+}
+
 func byteSliceComparator(a, b interface{}) int {
 	ab, okA := a.([]byte)
 	bb, okB := b.([]byte)
@@ -69,14 +81,28 @@ func (s *mvccStore) LastCommitTS() uint64 {
 	return s.lastCommitTS
 }
 
+// MVCCStoreOption configures the MVCCStore.
+type MVCCStoreOption func(*mvccStore)
+
+// WithLogger sets a custom logger for the store.
+func WithLogger(l *slog.Logger) MVCCStoreOption {
+	return func(s *mvccStore) {
+		s.log = l
+	}
+}
+
 // NewMVCCStore creates a new MVCC-enabled in-memory store.
-func NewMVCCStore() MVCCStore {
-	return &mvccStore{
+func NewMVCCStore(opts ...MVCCStoreOption) MVCCStore {
+	s := &mvccStore{
 		tree: treemap.NewWith(byteSliceComparator),
 		log: slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{
 			Level: slog.LevelWarn,
 		})),
 	}
+	for _, opt := range opts {
+		opt(s)
+	}
+	return s
 }
 
 var _ MVCCStore = (*mvccStore)(nil)
@@ -345,7 +371,7 @@ func (s *mvccStore) Snapshot() (io.ReadWriter, error) {
 	s.mtx.RLock()
 	defer s.mtx.RUnlock()
 
-	state := make([]mvccSnapshotEntry, 0, s.tree.Size())
+	entries := make([]mvccSnapshotEntry, 0, s.tree.Size())
 	s.tree.Each(func(key interface{}, value interface{}) {
 		k, ok := key.([]byte)
 		if !ok {
@@ -355,14 +381,19 @@ func (s *mvccStore) Snapshot() (io.ReadWriter, error) {
 		if !ok {
 			return
 		}
-		state = append(state, mvccSnapshotEntry{
+		entries = append(entries, mvccSnapshotEntry{
 			Key:      bytes.Clone(k),
 			Versions: append([]VersionedValue(nil), versions...),
 		})
 	})
 
+	snapshot := mvccSnapshot{
+		LastCommitTS: s.lastCommitTS,
+		Entries:      entries,
+	}
+
 	buf := &bytes.Buffer{}
-	if err := gob.NewEncoder(buf).Encode(state); err != nil {
+	if err := gob.NewEncoder(buf).Encode(snapshot); err != nil {
 		return nil, errors.WithStack(err)
 	}
 
@@ -388,35 +419,90 @@ func (s *mvccStore) Restore(r io.Reader) error {
 		return errors.WithStack(ErrInvalidChecksum)
 	}
 
-	var state []mvccSnapshotEntry
-	if err := gob.NewDecoder(bytes.NewReader(payload)).Decode(&state); err != nil {
+	var snapshot mvccSnapshot
+	if err := gob.NewDecoder(bytes.NewReader(payload)).Decode(&snapshot); err != nil {
 		return errors.WithStack(err)
 	}
 
 	s.mtx.Lock()
 	defer s.mtx.Unlock()
 
 	s.tree.Clear()
-	for _, entry := range state {
+	s.lastCommitTS = snapshot.LastCommitTS
+	for _, entry := range snapshot.Entries {
 		versions := append([]VersionedValue(nil), entry.Versions...)
 		s.tree.Put(bytes.Clone(entry.Key), versions)
-		if len(versions) > 0 {
-			last := versions[len(versions)-1].TS
-			if last > s.lastCommitTS {
-				s.lastCommitTS = last
+	}
+
+	return nil
+}
+
+func compactVersions(versions []VersionedValue, minTS uint64) ([]VersionedValue, bool) {
+	if len(versions) == 0 {
+		return versions, false
+	}
+
+	// Find the latest version that is <= minTS
+	keepIdx := -1
+	for i := len(versions) - 1; i >= 0; i-- {
+		if versions[i].TS <= minTS {
+			keepIdx = i
+			break
+		}
+	}
+
+	// If all versions are newer than minTS, keep everything
+	if keepIdx == -1 {
+		return versions, false
+	}
+
+	// If the oldest version is the one to keep, we can't remove anything before it
+	if keepIdx == 0 {
+		return versions, false
+	}
+
+	// We keep versions starting from keepIdx
+	// The version at keepIdx represents the state at minTS.
+	newVersions := make([]VersionedValue, len(versions)-keepIdx)
+	copy(newVersions, versions[keepIdx:])
+	return newVersions, true
+}
+
+func (s *mvccStore) Compact(ctx context.Context, minTS uint64) error {
+	s.mtx.Lock()
+	defer s.mtx.Unlock()
+
+	var updates map[string][]VersionedValue = make(map[string][]VersionedValue)
+
+	it := s.tree.Iterator()
+	for it.Next() {
+		versions, ok := it.Value().([]VersionedValue)
+		if !ok {
+			continue
+		}
+
+		newVersions, changed := compactVersions(versions, minTS)
+		if changed {
+			// tree keys are []byte, need string for map key
+			keyBytes, ok := it.Key().([]byte)
+			if !ok {
+				continue
 			}
+			updates[string(keyBytes)] = newVersions
 		}
 	}
 
+	for k, v := range updates {
+		s.tree.Put([]byte(k), v)
+	}
+
+	s.log.InfoContext(ctx, "compact",
+		slog.Uint64("min_ts", minTS),
+		slog.Int("updated_keys", len(updates)),
+	)
 	return nil
 }
 
 func (s *mvccStore) Close() error {
 	return nil
 }
-
-// mvccSnapshotEntry is used solely for gob snapshot serialization.
-type mvccSnapshotEntry struct {
-	Key      []byte
-	Versions []VersionedValue
-}

store/mvcc_store_compact_test.go

@@ -0,0 +1,102 @@
+package store
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestMVCCStore_Compact(t *testing.T) {
+	s := NewMVCCStore()
+	ctx := context.Background()
+
+	key := []byte("key1")
+
+	// Helper to add versions
+	require.NoError(t, s.PutAt(ctx, key, []byte("v10"), 10, 0))
+	require.NoError(t, s.PutAt(ctx, key, []byte("v20"), 20, 0))
+	require.NoError(t, s.PutAt(ctx, key, []byte("v30"), 30, 0))
+	require.NoError(t, s.PutAt(ctx, key, []byte("v40"), 40, 0))
+
+	// Verify initial state
+	val, err := s.GetAt(ctx, key, 15)
+	require.NoError(t, err)
+	assert.Equal(t, []byte("v10"), val)
+
+	// Compact at 25
+	// Should keep: v20 (latest <= 25), v30, v40
+	// Should remove: v10
+	err = s.Compact(ctx, 25)
+	require.NoError(t, err)
+
+	// v10 should be gone physically, but logically checking at TS 15
+	// depends on how GetAt is implemented.
+	// Current GetAt implementation:
+	// It iterates versions backwards. If we removed v10, and only have v20, v30, v40...
+	// querying at TS 15:
+	// v40 > 15 -> skip
+	// v30 > 15 -> skip
+	// v20 > 15 -> skip
+	// No version <= 15 found -> ErrKeyNotFound
+	// This is the expected behavior of Compaction: you cannot query older than minTS.
+
+	_, err = s.GetAt(ctx, key, 15)
+	assert.Equal(t, ErrKeyNotFound, err, "Should not find version older than compacted minTS")
+
+	// Query at 25 should return v20
+	val, err = s.GetAt(ctx, key, 25)
+	require.NoError(t, err)
+	assert.Equal(t, []byte("v20"), val)
+
+	// Query at 35 should return v30
+	val, err = s.GetAt(ctx, key, 35)
+	require.NoError(t, err)
+	assert.Equal(t, []byte("v30"), val)
+}
+
+func TestMVCCStore_Compact_Delete(t *testing.T) {
+	s := NewMVCCStore()
+	ctx := context.Background()
+
+	key := []byte("key_del")
+
+	require.NoError(t, s.PutAt(ctx, key, []byte("v10"), 10, 0))
+	require.NoError(t, s.DeleteAt(ctx, key, 20)) // Tombstone at 20
+
+	// Compact at 25
+	// Should keep Tombstone at 20 (as it is latest <= 25)
+	// Should remove v10
+	err := s.Compact(ctx, 25)
+	require.NoError(t, err)
+
+	// Query at 15 -> Not found (compacted)
+	_, err = s.GetAt(ctx, key, 15)
+	assert.Equal(t, ErrKeyNotFound, err)
+
+	// Query at 25 -> Not found (Tombstone at 20)
+	exists, err := s.ExistsAt(ctx, key, 25)
+	require.NoError(t, err)
+	assert.False(t, exists)
+}
+
+func TestMVCCStore_Compact_KeepAll(t *testing.T) {
+	s := NewMVCCStore()
+	ctx := context.Background()
+	key := []byte("key_keep")
+
+	require.NoError(t, s.PutAt(ctx, key, []byte("v50"), 50, 0))
+	require.NoError(t, s.PutAt(ctx, key, []byte("v60"), 60, 0))
+
+	// Compact at 40
+	// All versions are > 40. None <= 40.
+	// Logic says: keepIdx = -1 (not found).
+	// Should return all versions.
+	err := s.Compact(ctx, 40)
+	require.NoError(t, err)
+
+	val, err := s.GetAt(ctx, key, 55)
+	require.NoError(t, err)
+	assert.Equal(t, []byte("v50"), val)
+}

store/store.go

@@ -61,6 +61,8 @@ type MVCCStore interface {
 	ApplyMutations(ctx context.Context, mutations []*KVPairMutation, startTS, commitTS uint64) error
 	// LastCommitTS returns the highest commit timestamp applied on this node.
 	LastCommitTS() uint64
+	// Compact removes versions older than minTS that are no longer needed.
+	Compact(ctx context.Context, minTS uint64) error
 	Snapshot() (io.ReadWriter, error)
 	Restore(buf io.Reader) error
 	Close() error