Merge pull request #423 from bootjp/feature/pebble-compaction

store: harden Pebble MVCC compaction

Author: bootjp

adapter/dynamodb.go

@@ -114,6 +114,7 @@ type DynamoDBServer struct {
 	coordinator     kv.Coordinator
 	httpServer      *http.Server
 	targetHandlers  map[string]func(http.ResponseWriter, *http.Request)
+	readTracker     *kv.ActiveTimestampTracker
 	requestObserver monitoring.DynamoDBRequestObserver
 	itemUpdateLocks [itemUpdateLockStripeCount]sync.Mutex
 	tableLocks      [tableLockStripeCount]sync.Mutex
@@ -126,6 +127,12 @@ func WithDynamoDBRequestObserver(observer monitoring.DynamoDBRequestObserver) Dy
 	}
 }
 
+func WithDynamoDBActiveTimestampTracker(tracker *kv.ActiveTimestampTracker) DynamoDBServerOption {
+	return func(server *DynamoDBServer) {
+		server.readTracker = tracker
+	}
+}
+
 type dynamoRequestMetricsContextKey struct{}
 
 type dynamoRequestMetricsState struct {
@@ -1888,6 +1895,8 @@ func (d *DynamoDBServer) queryItems(ctx context.Context, in queryInput) (*queryO
 	if err != nil {
 		return nil, err
 	}
+	readPin := d.pinReadTS(readTS)
+	defer readPin.Release()
 	keySchema, cond, err := resolveQueryCondition(in, schema)
 	if err != nil {
 		return nil, err
@@ -1926,6 +1935,8 @@ func (d *DynamoDBServer) scanItems(ctx context.Context, in scanInput) (*queryOut
 	if err != nil {
 		return nil, err
 	}
+	readPin := d.pinReadTS(readTS)
+	defer readPin.Release()
 	indexKeySchema := schema.PrimaryKey
 	if strings.TrimSpace(in.IndexName) != "" {
 		indexKeySchema, err = schema.keySchemaForQuery(in.IndexName)
@@ -7367,6 +7378,13 @@ func (d *DynamoDBServer) nextTxnReadTS() uint64 {
 	return clock.Next()
 }
 
+func (d *DynamoDBServer) pinReadTS(ts uint64) *kv.ActiveTimestampToken {
+	if d == nil || d.readTracker == nil {
+		return &kv.ActiveTimestampToken{}
+	}
+	return d.readTracker.Pin(ts)
+}
+
 func (d *DynamoDBServer) loadTableSchema(ctx context.Context, tableName string) (*dynamoTableSchema, bool, error) {
 	return d.loadTableSchemaAt(ctx, tableName, snapshotTS(d.coordinator.Clock(), d.store))
 }
@@ -7544,6 +7562,9 @@ func (d *DynamoDBServer) migrateLegacySourceItems(
 	sourceSchema *dynamoTableSchema,
 	readTS uint64,
 ) error {
+	readPin := d.pinReadTS(readTS)
+	defer readPin.Release()
+
 	prefix := dynamoItemPrefixForTable(targetSchema.TableName, sourceSchema.Generation)
 	upper := prefixScanEnd(prefix)
 	cursor := bytes.Clone(prefix)
@@ -7741,6 +7762,9 @@ func (d *DynamoDBServer) scanAllByPrefix(ctx context.Context, prefix []byte) ([]
 }
 
 func (d *DynamoDBServer) scanAllByPrefixAt(ctx context.Context, prefix []byte, readTS uint64) ([]*store.KVPair, error) {
+	readPin := d.pinReadTS(readTS)
+	defer readPin.Release()
+
 	end := prefixScanEnd(prefix)
 	start := bytes.Clone(prefix)
 

adapter/dynamodb_retention_test.go

@@ -0,0 +1,114 @@
+package adapter
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/bootjp/elastickv/kv"
+	"github.com/bootjp/elastickv/store"
+	"github.com/stretchr/testify/require"
+)
+
+type blockingDynamoScanStore struct {
+	store.MVCCStore
+	started chan uint64
+	release chan struct{}
+	once    sync.Once
+}
+
+func (s *blockingDynamoScanStore) ScanAt(ctx context.Context, start []byte, end []byte, limit int, ts uint64) ([]*store.KVPair, error) {
+	s.once.Do(func() {
+		s.started <- ts
+		<-s.release
+	})
+	return s.MVCCStore.ScanAt(ctx, start, end, limit, ts)
+}
+
+func requireReadTSPinned(t *testing.T, tracker *kv.ActiveTimestampTracker, want uint64) {
+	t.Helper()
+	require.Eventually(t, func() bool {
+		return tracker.Oldest() == want
+	}, time.Second, 10*time.Millisecond)
+}
+
+func waitForScanStart(t *testing.T, started <-chan uint64) uint64 {
+	t.Helper()
+	select {
+	case ts := <-started:
+		return ts
+	case <-time.After(time.Second):
+		t.Fatal("timed out waiting for scan")
+		return 0
+	}
+}
+
+func TestDynamoDBScanAllByPrefixAtPinsReadTS(t *testing.T) {
+	t.Parallel()
+
+	tracker := kv.NewActiveTimestampTracker()
+	backing := store.NewMVCCStore()
+	require.NoError(t, backing.PutAt(context.Background(), []byte("prefix/item"), []byte("value"), 10, 0))
+
+	scanStore := &blockingDynamoScanStore{
+		MVCCStore: backing,
+		started:   make(chan uint64, 1),
+		release:   make(chan struct{}),
+	}
+	server := NewDynamoDBServer(nil, scanStore, &stubAdapterCoordinator{}, WithDynamoDBActiveTimestampTracker(tracker))
+
+	const readTS uint64 = 42
+	done := make(chan error, 1)
+	go func() {
+		_, err := server.scanAllByPrefixAt(context.Background(), []byte("prefix/"), readTS)
+		done <- err
+	}()
+
+	require.Equal(t, readTS, waitForScanStart(t, scanStore.started))
+	requireReadTSPinned(t, tracker, readTS)
+
+	close(scanStore.release)
+
+	select {
+	case err := <-done:
+		require.NoError(t, err)
+	case <-time.After(time.Second):
+		t.Fatal("timed out waiting for scanAllByPrefixAt")
+	}
+	require.Zero(t, tracker.Oldest())
+}
+
+func TestDynamoDBMigrateLegacySourceItemsPinsReadTS(t *testing.T) {
+	t.Parallel()
+
+	tracker := kv.NewActiveTimestampTracker()
+	scanStore := &blockingDynamoScanStore{
+		MVCCStore: store.NewMVCCStore(),
+		started:   make(chan uint64, 1),
+		release:   make(chan struct{}),
+	}
+	server := NewDynamoDBServer(nil, scanStore, &stubAdapterCoordinator{}, WithDynamoDBActiveTimestampTracker(tracker))
+
+	targetSchema := &dynamoTableSchema{TableName: "tbl", Generation: 2}
+	sourceSchema := &dynamoTableSchema{TableName: "tbl", Generation: 1}
+
+	const readTS uint64 = 55
+	done := make(chan error, 1)
+	go func() {
+		done <- server.migrateLegacySourceItems(context.Background(), targetSchema, sourceSchema, readTS)
+	}()
+
+	require.Equal(t, readTS, waitForScanStart(t, scanStore.started))
+	requireReadTSPinned(t, tracker, readTS)
+
+	close(scanStore.release)
+
+	select {
+	case err := <-done:
+		require.NoError(t, err)
+	case <-time.After(time.Second):
+		t.Fatal("timed out waiting for migrateLegacySourceItems")
+	}
+	require.Zero(t, tracker.Oldest())
+}

kv/compactor.go

@@ -11,7 +11,7 @@ import (
 )
 
 const (
-	defaultFSMCompactorInterval        = 30 * time.Second
+	defaultFSMCompactorInterval        = 5 * time.Minute
 	defaultFSMCompactorRetentionWindow = 30 * time.Minute
 	defaultFSMCompactorTimeout         = 5 * time.Second
 )
@@ -172,7 +172,6 @@ func (c *FSMCompactor) compactRuntime(ctx context.Context, runtime FSMCompactRun
 	if err := runtime.Store.Compact(compactCtx, safeMinTS); err != nil {
 		return errors.Wrapf(err, "compact group %d", runtime.GroupID)
 	}
-	retention.SetMinRetainedTS(safeMinTS)
 	c.logger.InfoContext(compactCtx, "fsm compacted",
 		"group_id", runtime.GroupID,
 		"min_retained_ts", safeMinTS,

kv/compactor_test.go

@@ -2,6 +2,7 @@ package kv
 
 import (
 	"context"
+	"os"
 	"testing"
 	"time"
 
@@ -110,3 +111,42 @@ func TestFSMCompactorSkipsLaggingRuntime(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, []byte("v10"), val)
 }
+
+func TestFSMCompactorCompactsEligiblePebbleRuntime(t *testing.T) {
+	dir, err := os.MkdirTemp("", "fsm-compactor-pebble-*")
+	require.NoError(t, err)
+	defer os.RemoveAll(dir)
+
+	st, err := store.NewPebbleStore(dir)
+	require.NoError(t, err)
+	defer st.Close()
+
+	ctx := context.Background()
+	require.NoError(t, st.PutAt(ctx, []byte("k"), []byte("v10"), 10, 0))
+	require.NoError(t, st.PutAt(ctx, []byte("k"), []byte("v20"), 20, 0))
+	require.NoError(t, st.PutAt(ctx, []byte("k"), []byte("v30"), 30, 0))
+
+	compactor := NewFSMCompactor(
+		[]FSMCompactRuntime{{
+			GroupID: 1,
+			Raft: fakeRaftStats{stats: map[string]string{
+				"state":         "Follower",
+				"fsm_pending":   "0",
+				"applied_index": "10",
+				"commit_index":  "10",
+			}},
+			Store: st,
+		}},
+		WithFSMCompactorInterval(time.Hour),
+		WithFSMCompactorRetentionWindow(time.Millisecond),
+	)
+
+	require.NoError(t, compactor.SyncOnce(ctx))
+
+	_, err = st.GetAt(ctx, []byte("k"), 20)
+	require.ErrorIs(t, err, store.ErrReadTSCompacted)
+
+	val, err := st.GetAt(ctx, []byte("k"), 30)
+	require.NoError(t, err)
+	require.Equal(t, []byte("v30"), val)
+}

kv/shard_store.go

@@ -187,40 +187,59 @@ func (s *ShardStore) reverseScanRoutesAt(
 	seenGroups := make(map[uint64]struct{})
 	for i := len(routes) - 1; i >= 0; i-- {
 		route := routes[i]
-		scanStart := start
-		scanEnd := end
 		if clampToRoutes {
-			if len(out) >= limit {
-				break
-			}
-			scanStart = clampScanStart(start, route.Start)
-			scanEnd = clampScanEnd(end, route.End)
-			kvs, err := s.scanRouteAtDirection(ctx, route, scanStart, scanEnd, limit-len(out), ts, true)
+			kvs, done, err := s.clampedReverseScanRouteAt(ctx, route, start, end, limit, len(out), ts)
 			if err != nil {
 				return nil, err
 			}
-			out = append(out, kvs...)
-		} else {
-			// When clampToRoutes is false (e.g. S3 manifest scans spanning multiple
-			// shards), keys from different routes may interleave in descending order.
-			// Fetch up to limit from every route and merge+sort descending so the
-			// result honours the ReverseScanAt contract.
-			// De-duplicate by GroupID: after a range split both halves share the same
-			// GroupID (same backing shard store), so only scan each group once.
-			if _, seen := seenGroups[route.GroupID]; seen {
-				continue
-			}
-			seenGroups[route.GroupID] = struct{}{}
-			kvs, err := s.scanRouteAtDirection(ctx, route, scanStart, scanEnd, limit, ts, true)
-			if err != nil {
-				return nil, err
+			if done {
+				break
 			}
-			out = mergeAndTrimReverseScanResults(out, kvs, limit)
+			out = append(out, kvs...)
+			continue
 		}
+
+		// When clampToRoutes is false (e.g. S3 manifest scans spanning multiple
+		// shards), keys from different routes may interleave in descending order.
+		// Fetch up to limit from every route and merge+sort descending so the
+		// result honours the ReverseScanAt contract.
+		// De-duplicate by GroupID: after a range split both halves share the same
+		// GroupID (same backing shard store), so only scan each group once.
+		if _, seen := seenGroups[route.GroupID]; seen {
+			continue
+		}
+		seenGroups[route.GroupID] = struct{}{}
+		kvs, err := s.scanRouteAtDirection(ctx, route, start, end, limit, ts, true)
+		if err != nil {
+			return nil, err
+		}
+		out = mergeAndTrimReverseScanResults(out, kvs, limit)
 	}
 	return out, nil
 }
 
+func (s *ShardStore) clampedReverseScanRouteAt(
+	ctx context.Context,
+	route distribution.Route,
+	start []byte,
+	end []byte,
+	limit int,
+	currentLen int,
+	ts uint64,
+) ([]*store.KVPair, bool, error) {
+	if currentLen >= limit {
+		return nil, true, nil
+	}
+
+	scanStart := clampScanStart(start, route.Start)
+	scanEnd := clampScanEnd(end, route.End)
+	kvs, err := s.scanRouteAtDirection(ctx, route, scanStart, scanEnd, limit-currentLen, ts, true)
+	if err != nil {
+		return nil, false, err
+	}
+	return kvs, false, nil
+}
+
 func (s *ShardStore) scanRouteAtDirection(
 	ctx context.Context,
 	route distribution.Route,

main.go

@@ -398,7 +398,7 @@ func startRedisServer(ctx context.Context, lc *net.ListenConfig, eg *errgroup.Gr
 	return nil
 }
 
-func startDynamoDBServer(ctx context.Context, lc *net.ListenConfig, eg *errgroup.Group, dynamoAddr string, shardStore *kv.ShardStore, coordinate kv.Coordinator, metricsRegistry *monitoring.Registry) error {
+func startDynamoDBServer(ctx context.Context, lc *net.ListenConfig, eg *errgroup.Group, dynamoAddr string, shardStore *kv.ShardStore, coordinate kv.Coordinator, metricsRegistry *monitoring.Registry, readTracker *kv.ActiveTimestampTracker) error {
 	dynamoL, err := lc.Listen(ctx, "tcp", dynamoAddr)
 	if err != nil {
 		return errors.Wrapf(err, "failed to listen on %s", dynamoAddr)
@@ -407,6 +407,7 @@ func startDynamoDBServer(ctx context.Context, lc *net.ListenConfig, eg *errgroup
 		dynamoL,
 		shardStore,
 		coordinate,
+		adapter.WithDynamoDBActiveTimestampTracker(readTracker),
 		adapter.WithDynamoDBRequestObserver(metricsRegistry.DynamoDBObserver()),
 	)
 	eg.Go(func() error {
@@ -563,7 +564,7 @@ func (r runtimeServerRunner) start() error {
 	if err := startRaftServers(r.ctx, r.lc, r.eg, r.runtimes, r.shardStore, r.coordinate, r.distServer, r.pubsubRelay); err != nil {
 		return waitErrgroupAfterStartupFailure(r.cancel, r.eg, err)
 	}
-	if err := startDynamoDBServer(r.ctx, r.lc, r.eg, r.dynamoAddress, r.shardStore, r.coordinate, r.metricsRegistry); err != nil {
+	if err := startDynamoDBServer(r.ctx, r.lc, r.eg, r.dynamoAddress, r.shardStore, r.coordinate, r.metricsRegistry, r.readTracker); err != nil {
 		return waitErrgroupAfterStartupFailure(r.cancel, r.eg, err)
 	}
 	if err := startMetricsServer(r.ctx, r.lc, r.eg, r.metricsAddress, r.metricsToken, r.metricsRegistry.Handler()); err != nil {