diff --git a/adapter/internal.go b/adapter/internal.go index 30029da94..03986377d 100644 --- a/adapter/internal.go +++ b/adapter/internal.go @@ -254,6 +254,28 @@ func (i *Internal) PromoteStagedVersions(ctx context.Context, req *pb.PromoteSta }, nil } +func (i *Internal) ApplyTargetStagedReadiness(ctx context.Context, req *pb.TargetStagedReadinessRequest) (*pb.TargetStagedReadinessResponse, error) { + if req == nil { + return nil, errors.WithStack(status.Error(codes.InvalidArgument, "target staged readiness request is nil")) + } + if req.GetJobId() == 0 { + return nil, errors.WithStack(status.Error(codes.InvalidArgument, "target staged readiness job_id is required")) + } + if req.GetMigrationJobId() == 0 { + return nil, errors.WithStack(status.Error(codes.InvalidArgument, "target staged readiness migration_job_id is required")) + } + if i.migrationProposer == nil { + return nil, errors.WithStack(status.Error(codes.FailedPrecondition, "target staged readiness proposer is not configured")) + } + if err := i.verifyInternalLeader(ctx); err != nil { + return nil, err + } + if err := i.proposeTargetStagedReadiness(ctx, req); err != nil { + return nil, errors.WithStack(err) + } + return &pb.TargetStagedReadinessResponse{}, nil +} + func (i *Internal) verifyMigrationPromoteEnabled(ctx context.Context) error { if i.migrationPromoteGate == nil { return nil @@ -349,6 +371,24 @@ func (i *Internal) proposeMigrationPromote(ctx context.Context, req *pb.PromoteS } } +func (i *Internal) proposeTargetStagedReadiness(ctx context.Context, req *pb.TargetStagedReadinessRequest) error { + cmd, err := kv.MarshalTargetStagedReadinessCommand(req) + if err != nil { + return errors.WithStack(err) + } + resp, err := i.proposeMigrationCommand(ctx, cmd, "target staged readiness") + if err != nil { + return errors.WithStack(err) + } + if resp == nil { + return nil + } + if err, ok := resp.(error); ok { + return errors.WithStack(err) + } + return errors.WithStack(errors.Newf("unexpected target readiness apply response type %T", resp)) +} + func (i *Internal) proposeMigrationCommand(ctx context.Context, cmd []byte, label string) (any, error) { result, err := i.migrationProposer.Propose(ctx, cmd) if err != nil { diff --git a/adapter/internal_migration_test.go b/adapter/internal_migration_test.go index 519b8647b..13f9a2524 100644 --- a/adapter/internal_migration_test.go +++ b/adapter/internal_migration_test.go @@ -498,3 +498,43 @@ func TestInternalPromoteStagedVersionsAppliesStoreBatch(t *testing.T) { require.ErrorIs(t, err, store.ErrKeyNotFound) require.GreaterOrEqual(t, clock.Current(), uint64(30)) } + +func TestInternalApplyTargetStagedReadinessProposesThroughRaft(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + proposer := &applyingMigrationProposer{ + fsm: kv.NewKvFSMWithHLC(st, nil), + } + internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, + WithInternalStore(st), + WithInternalMigrationProposer(proposer), + ) + + _, err := internal.ApplyTargetStagedReadiness(ctx, &pb.TargetStagedReadinessRequest{ + JobId: 9, + RouteStart: []byte("a"), + RouteEnd: []byte("z"), + ExpectedCutoverVersion: 3, + MigrationJobId: 7, + MinWriteTsExclusive: 100, + Armed: true, + }) + require.NoError(t, err) + require.Equal(t, uint64(1), proposer.calls) + + reader, ok := st.(store.MigrationTargetReadinessReader) + require.True(t, ok) + states, err := reader.MigrationTargetReadinessStates(ctx) + require.NoError(t, err) + require.Equal(t, []store.TargetStagedReadinessState{{ + JobID: 9, + RouteStart: []byte("a"), + RouteEnd: []byte("z"), + ExpectedCutoverVersion: 3, + MigrationJobID: 7, + MinWriteTSExclusive: 100, + Armed: true, + }}, states) +} diff --git a/adapter/redis_compat_commands_stream_test.go b/adapter/redis_compat_commands_stream_test.go index ff605d3c1..cebd67e55 100644 --- a/adapter/redis_compat_commands_stream_test.go +++ b/adapter/redis_compat_commands_stream_test.go @@ -279,7 +279,7 @@ func TestRedis_StreamXReadLatencyIsConstant(t *testing.T) { rdb := redis.NewClient(&redis.Options{Addr: nodes[0].redisAddress}) defer func() { _ = rdb.Close() }() - ctx := context.Background() + ctx := t.Context() const ( total = 10_000 @@ -291,13 +291,21 @@ func TestRedis_StreamXReadLatencyIsConstant(t *testing.T) { } measure := func() time.Duration { - start := time.Now() - streams, err := rdb.XRead(ctx, &redis.XReadArgs{ - Streams: []string{"stream-lat", lastID}, - Count: 10, - Block: 10 * time.Millisecond, - }).Result() - elapsed := time.Since(start) + var ( + streams []redis.XStream + elapsed time.Duration + ) + err := retryNotLeader(ctx, func() error { + start := time.Now() + var xerr error + streams, xerr = rdb.XRead(ctx, &redis.XReadArgs{ + Streams: []string{"stream-lat", lastID}, + Count: 10, + Block: 10 * time.Millisecond, + }).Result() + elapsed = time.Since(start) + return xerr + }) require.True(t, errors.Is(err, redis.Nil) || err == nil) require.Empty(t, streams) return elapsed diff --git a/kv/fsm.go b/kv/fsm.go index bd2edd712..f5214652c 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -122,9 +122,12 @@ type RouteSnapshot interface { // OwnerOf returns the Raft group ID that owned key at this // snapshot's version. (0, false) when no route covered key. OwnerOf(key []byte) (uint64, bool) - // RouteOf returns the complete route descriptor covering key. + // RouteOf returns the complete route descriptor covering key at this + // snapshot's version. Used by apply-time target-readiness checks to + // prove staged/cleared descriptor state before mutating the store. RouteOf(key []byte) (distribution.Route, bool) - // IntersectingRoutes returns every route intersecting [start, end). + // IntersectingRoutes returns the route descriptors that intersect + // [start, end) at this snapshot's version. A nil end denotes +infinity. IntersectingRoutes(start, end []byte) []distribution.Route // WriteFencedForKey reports whether key is currently inside a // WriteFenced route in this snapshot. @@ -132,6 +135,12 @@ type RouteSnapshot interface { // WriteFencedIntersects reports whether [start, end) intersects // any WriteFenced route in this snapshot. WriteFencedIntersects(start, end []byte) bool + // WriteFloorForKey returns the post-migration write timestamp floor + // for key, when the current route retains one. + WriteFloorForKey(key []byte) (uint64, bool) + // WriteFloorIntersects returns the maximum post-migration write + // timestamp floor across routes intersecting [start, end). + WriteFloorIntersects(start, end []byte) (uint64, bool) } // SetApplyIndex implements raftengine.ApplyIndexAware. The engine @@ -373,6 +382,8 @@ func (f *kvFSM) applyReservedOpcode(ctx context.Context, data []byte) (any, bool return f.applyMigrationImport(ctx, data[1:]), true case data[0] == raftEncodeMigrationPromote: return f.applyMigrationPromote(ctx, data[1:]), true + case data[0] == raftEncodeTargetReadiness: + return f.applyTargetStagedReadiness(ctx, data[1:]), true case data[0] >= fsmwire.OpEncryptionMin && data[0] <= fsmwire.OpEncryptionMax: return f.applyEncryption(f.pendingApplyIdx, data[0], data[1:]), true default: @@ -412,6 +423,10 @@ const ( // data promotion chunk. Every target voter atomically copies staged MVCC // versions into the live keyspace and removes the promoted staged rows. raftEncodeMigrationPromote byte = 0x0b + // raftEncodeTargetReadiness carries the target-local staged-readiness + // guard. It must be replicated through the target Raft group before the + // migration controller can treat a target as fail-closed for cutover. + raftEncodeTargetReadiness byte = 0x0c ) func decodeRaftRequests(data []byte) ([]*pb.Request, error) { @@ -507,7 +522,7 @@ func (f *kvFSM) handleRawRequest(ctx context.Context, r *pb.Request, commitTS ui return f.handleDelPrefix(ctx, prefix, commitTS) } - if err := f.validateRawMutationsForApply(ctx, r.Mutations); err != nil { + if err := f.validateRawMutationsForApply(ctx, r.Mutations, commitTS); err != nil { return err } @@ -524,16 +539,16 @@ func (f *kvFSM) handleRawRequest(ctx context.Context, r *pb.Request, commitTS ui return nil } -func (f *kvFSM) validateRawMutationsForApply(ctx context.Context, muts []*pb.Mutation) error { +func (f *kvFSM) validateRawMutationsForApply(ctx context.Context, muts []*pb.Mutation, commitTS uint64) error { for _, mut := range muts { - if err := f.validateRawMutationForApply(ctx, mut); err != nil { + if err := f.validateRawMutationForApply(ctx, mut, commitTS); err != nil { return err } } return nil } -func (f *kvFSM) validateRawMutationForApply(ctx context.Context, mut *pb.Mutation) error { +func (f *kvFSM) validateRawMutationForApply(ctx context.Context, mut *pb.Mutation, commitTS uint64) error { if mut == nil || len(mut.Key) == 0 { return errors.WithStack(ErrInvalidRequest) } @@ -544,6 +559,12 @@ func (f *kvFSM) validateRawMutationForApply(ctx context.Context, mut *pb.Mutatio if err := f.verifyRouteNotFencedForKey(mut.Key); err != nil { return err } + if err := f.verifyTargetReadinessForRange(ctx, mut.Key, nextScanCursor(mut.Key)); err != nil { + return err + } + if err := f.verifyRouteWriteFloorForKey(mut.Key, commitTS); err != nil { + return err + } if err := f.assertNoConflictingTxnLock(ctx, mut.Key, nil, 0); err != nil { return err } @@ -567,6 +588,19 @@ func (f *kvFSM) handleDelPrefix(ctx context.Context, prefix []byte, commitTS uin if err := f.verifyRouteNotFencedForPrefix(prefix); err != nil { return err } + if err := f.verifyTargetReadinessForPrefix(ctx, prefix); err != nil { + return err + } + if err := f.verifyRouteWriteFloorForPrefix(prefix, commitTS); err != nil { + return err + } + routes := f.stagedVisibilityRoutesForPrefix(prefix) + deleteStagedPrefix := func(stagedPrefix []byte, stagedExcludePrefix []byte) error { + return f.store.DeletePrefixAtRaftAt(ctx, stagedPrefix, stagedExcludePrefix, commitTS, 0) + } + if err := deleteStagedVisibilityPrefixes(routes, prefix, txnCommonPrefix, deleteStagedPrefix); err != nil { + return err + } if err := f.store.DeletePrefixAtRaftAt(ctx, prefix, txnCommonPrefix, commitTS, f.pendingApplyIdx); err != nil { return errors.WithStack(err) } @@ -574,6 +608,25 @@ func (f *kvFSM) handleDelPrefix(ctx context.Context, prefix []byte, commitTS uin return nil } +func (f *kvFSM) stagedVisibilityRoutesForPrefix(prefix []byte) []distribution.Route { + if f.routes == nil { + return nil + } + snap, ok := f.routes.Current() + if !ok { + return nil + } + start, end := routePrefixRange(prefix) + routes := snap.IntersectingRoutes(start, end) + out := make([]distribution.Route, 0, len(routes)) + for _, route := range routes { + if route.GroupID == f.shardGroupID && routeHasStagedVisibility(route) { + out = append(out, route) + } + } + return out +} + func (f *kvFSM) verifyRouteNotFencedForMutations(muts []*pb.Mutation) error { for _, mut := range muts { if mut == nil || len(mut.Key) == 0 || isTxnInternalKey(mut.Key) { @@ -619,6 +672,180 @@ func (f *kvFSM) verifyRouteNotFencedForPrefix(prefix []byte) error { return errors.Wrapf(ErrRouteWriteFenced, "prefix %q route range [%q,%q)", prefix, start, end) } +func (f *kvFSM) verifyRouteWriteFloorForMutations(muts []*pb.Mutation, commitTS uint64) error { + for _, mut := range muts { + if mut == nil || len(mut.Key) == 0 || isTxnInternalKey(mut.Key) { + continue + } + if err := f.verifyRouteWriteFloorForKey(mut.Key, commitTS); err != nil { + return err + } + } + return nil +} + +func (f *kvFSM) verifyTargetReadinessForMutations(ctx context.Context, muts []*pb.Mutation) error { + for _, mut := range muts { + if mut == nil || len(mut.Key) == 0 || isTxnInternalKey(mut.Key) { + continue + } + if err := f.verifyTargetReadinessForRange(ctx, mut.Key, nextScanCursor(mut.Key)); err != nil { + return err + } + } + return nil +} + +func (f *kvFSM) verifyTargetReadinessForReadKeys(ctx context.Context, keys [][]byte) error { + for _, key := range keys { + if isTxnInternalKey(key) { + continue + } + if err := f.verifyTargetReadinessForRange(ctx, key, nextScanCursor(key)); err != nil { + return err + } + } + return nil +} + +func (f *kvFSM) verifyTargetReadinessForTxnFootprint(ctx context.Context, muts []*pb.Mutation, readKeys [][]byte, primaryKey []byte) error { + if len(primaryKey) != 0 && !isTxnInternalKey(primaryKey) { + if err := f.verifyTargetReadinessForRange(ctx, primaryKey, nextScanCursor(primaryKey)); err != nil { + return err + } + } + if err := f.verifyTargetReadinessForReadKeys(ctx, readKeys); err != nil { + return err + } + return f.verifyTargetReadinessForMutations(ctx, muts) +} + +func (f *kvFSM) verifyTargetReadinessForPrefix(ctx context.Context, prefix []byte) error { + start, end := routePrefixRange(prefix) + return f.verifyTargetReadinessForRouteRange(ctx, start, end) +} + +func (f *kvFSM) verifyTargetReadinessForRange(ctx context.Context, start []byte, end []byte) error { + routeStart, routeEnd := readinessRouteRange(start, end) + return f.verifyTargetReadinessForRouteRange(ctx, routeStart, routeEnd) +} + +func (f *kvFSM) verifyTargetReadinessForRouteRange(ctx context.Context, routeStart []byte, routeEnd []byte) error { + _, err := f.targetReadyRoutesForRouteRange(ctx, routeStart, routeEnd) + return err +} + +func (f *kvFSM) targetReadyRoutesForRange(ctx context.Context, start []byte, end []byte) ([]distribution.Route, error) { + routeStart, routeEnd := readinessRouteRange(start, end) + return f.targetReadyRoutesForRouteRange(ctx, routeStart, routeEnd) +} + +func (f *kvFSM) targetReadyRoutesForRouteRange(ctx context.Context, routeStart []byte, routeEnd []byte) ([]distribution.Route, error) { + routes, catalogVersion, proof := f.currentShardRoutesForRouteRange(routeStart, routeEnd) + + reader, ok := f.store.(store.MigrationTargetReadinessReader) + if !ok { + return routes, nil + } + states, err := reader.MigrationTargetReadinessStates(ctx) + if err != nil { + return nil, errors.WithStack(err) + } + if len(states) == 0 { + return routes, nil + } + if targetReadinessStatesSatisfied(states, routes, routeStart, routeEnd, f.shardGroupID, catalogVersion, proof) { + return routes, nil + } + return nil, errors.WithStack(ErrRouteCutoverPending) +} + +func (f *kvFSM) currentShardRoutesForRouteRange(routeStart []byte, routeEnd []byte) ([]distribution.Route, uint64, bool) { + if f.routes == nil { + return nil, 0, false + } + snap, ok := f.routes.Current() + if !ok { + return nil, 0, false + } + routes := make([]distribution.Route, 0) + for _, route := range snap.IntersectingRoutes(routeStart, routeEnd) { + if route.GroupID == f.shardGroupID { + routes = append(routes, route) + } + } + return routes, snap.Version(), true +} + +func targetReadinessStatesSatisfied( + states []store.TargetStagedReadinessState, + routes []distribution.Route, + routeStart []byte, + routeEnd []byte, + groupID uint64, + catalogVersion uint64, + proof bool, +) bool { + for _, ready := range states { + if !ready.Armed || !routeRangeIntersects(routeStart, routeEnd, ready.RouteStart, ready.RouteEnd) { + continue + } + if !proof || !routesSatisfyTargetReadiness(routes, ready, groupID, catalogVersion) { + return false + } + } + return true +} + +func routesSatisfyTargetReadiness(routes []distribution.Route, ready store.TargetStagedReadinessState, groupID uint64, catalogVersion uint64) bool { + matched := false + for _, route := range routes { + if route.GroupID != groupID { + continue + } + if !routeRangeIntersects(route.Start, route.End, ready.RouteStart, ready.RouteEnd) { + continue + } + matched = true + if !routeSatisfiesTargetReadiness(route, ready, catalogVersion) { + return false + } + } + return matched +} + +func (f *kvFSM) verifyRouteWriteFloorForKey(key []byte, commitTS uint64) error { + if f.routes == nil { + return nil + } + snap, ok := f.routes.Current() + if !ok { + return nil + } + rkey := routeKey(key) + floor, ok := snap.WriteFloorForKey(rkey) + if !ok || commitTS > floor { + return nil + } + return errors.Wrapf(ErrRouteWriteBelowFloor, "commit_ts %d <= floor %d for key %q routeKey %q", commitTS, floor, key, rkey) +} + +func (f *kvFSM) verifyRouteWriteFloorForPrefix(prefix []byte, commitTS uint64) error { + if f.routes == nil { + return nil + } + snap, ok := f.routes.Current() + if !ok { + return nil + } + start, end := routePrefixRange(prefix) + floor, ok := snap.WriteFloorIntersects(start, end) + if !ok || commitTS > floor { + return nil + } + return errors.Wrapf(ErrRouteWriteBelowFloor, "commit_ts %d <= floor %d for prefix %q route range [%q,%q)", commitTS, floor, prefix, start, end) +} + func routePrefixRange(prefix []byte) ([]byte, []byte) { if len(prefix) == 0 { return []byte(""), nil @@ -940,7 +1167,7 @@ func (f *kvFSM) validateConflicts(ctx context.Context, muts []*pb.Mutation, star } seen[keyStr] = struct{}{} - latest, exists, err := f.store.LatestCommitTS(ctx, mut.Key) + latest, exists, err := f.latestCommitTSForTargetReadyKey(ctx, mut.Key) if err != nil { return errors.WithStack(err) } @@ -951,6 +1178,64 @@ func (f *kvFSM) validateConflicts(ctx context.Context, muts []*pb.Mutation, star return nil } +func (f *kvFSM) validateReadConflicts(ctx context.Context, keys [][]byte, startTS uint64) error { + seen := make(map[string]struct{}, len(keys)) + for _, key := range keys { + if len(key) == 0 || isTxnInternalKey(key) { + continue + } + keyStr := string(key) + if _, ok := seen[keyStr]; ok { + continue + } + seen[keyStr] = struct{}{} + + latest, exists, err := f.latestCommitTSForTargetReadyKey(ctx, key) + if err != nil { + return errors.WithStack(err) + } + if exists && latest > startTS { + return errors.WithStack(store.NewWriteConflictError(key)) + } + } + return nil +} + +func (f *kvFSM) latestCommitTSForTargetReadyKey(ctx context.Context, key []byte) (uint64, bool, error) { + routes, err := f.targetReadyRoutesForRange(ctx, key, nextScanCursor(key)) + if err != nil { + return 0, false, err + } + liveTS, liveExists, err := f.store.LatestCommitTS(ctx, key) + if err != nil { + return 0, false, errors.WithStack(err) + } + latest, exists := liveTS, liveExists + for _, route := range routes { + if !routeHasStagedVisibility(route) { + continue + } + stagedTS, stagedExists, err := f.store.LatestCommitTS(ctx, distribution.MigrationStagedDataKey(route.MigrationJobID, key)) + if err != nil { + return 0, false, errors.WithStack(err) + } + if stagedExists && (!exists || stagedTS > latest) { + latest, exists = stagedTS, true + } + } + return latest, exists, nil +} + +func (f *kvFSM) validateTxnConflicts(ctx context.Context, muts []*pb.Mutation, readKeys [][]byte, startTS uint64) error { + if err := f.validateConflicts(ctx, muts, startTS); err != nil { + return errors.WithStack(err) + } + if err := f.validateReadConflicts(ctx, readKeys, startTS); err != nil { + return errors.WithStack(err) + } + return nil +} + func uniqueMutations(muts []*pb.Mutation) ([]*pb.Mutation, error) { if len(muts) == 0 { return []*pb.Mutation{}, nil @@ -992,15 +1277,19 @@ func (f *kvFSM) handlePrepareRequest(ctx context.Context, r *pb.Request) error { } startTS := r.Ts - uniq, err := uniqueMutations(muts) - if err != nil { + floorTS := startTS + if meta.CommitTS != 0 { + floorTS = meta.CommitTS + } + if err := f.verifyTargetReadinessForTxnFootprint(ctx, muts, r.ReadKeys, meta.PrimaryKey); err != nil { return err } - if err := f.verifyRouteNotFencedForMutations(uniq); err != nil { + uniq, err := f.uniqueMutationsNotFenced(ctx, muts, floorTS) + if err != nil { return err } - if err := f.validateConflicts(ctx, uniq, startTS); err != nil { - return errors.WithStack(err) + if err := f.validateTxnConflicts(ctx, uniq, r.ReadKeys, startTS); err != nil { + return err } expireAt := txnLockExpireAt(meta.LockTTLms) @@ -1057,7 +1346,7 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com // applying this log entry. The retention-window > max-retry-latency // invariant prevents the rare case where a real never-landed retry // arrives with PrevCommitTS below pebble's compacted floor. - dedup, err := f.dedupProbeOnePhase(ctx, meta) + dedup, err := f.dedupProbeOnePhase(ctx, meta, muts, r.ReadKeys) if err != nil { return err } @@ -1065,12 +1354,7 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com return nil } - uniq, err := f.uniqueMutationsNotFenced(muts) - if err != nil { - return err - } - - storeMuts, err := f.buildOnePhaseStoreMutations(ctx, uniq) + uniq, storeMuts, err := f.onePhaseStoreMutations(ctx, muts, r.ReadKeys, startTS, commitTS) if err != nil { return err } @@ -1081,34 +1365,57 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com return nil } -func (f *kvFSM) uniqueMutationsNotFenced(muts []*pb.Mutation) ([]*pb.Mutation, error) { - uniq, err := uniqueMutations(muts) +func (f *kvFSM) onePhaseStoreMutations( + ctx context.Context, + muts []*pb.Mutation, + readKeys [][]byte, + startTS uint64, + commitTS uint64, +) ([]*pb.Mutation, []*store.KVPairMutation, error) { + uniq, err := f.uniqueMutationsNotFenced(ctx, muts, commitTS) if err != nil { - return nil, err + return nil, nil, err } - if err := f.verifyRouteNotFencedForMutations(uniq); err != nil { - return nil, err + if err := f.validateTxnConflicts(ctx, uniq, readKeys, startTS); err != nil { + return nil, nil, err } - return uniq, nil + storeMuts, err := f.buildOnePhaseStoreMutations(ctx, uniq) + if err != nil { + return nil, nil, err + } + return uniq, storeMuts, nil } -func uniqueTxnMutations(muts []*pb.Mutation) ([]*pb.Mutation, error) { +func (f *kvFSM) uniqueMutationsNotFenced(ctx context.Context, muts []*pb.Mutation, commitTS uint64) ([]*pb.Mutation, error) { uniq, err := uniqueMutations(muts) if err != nil { return nil, err } + if err := f.verifyRouteNotFencedForMutations(uniq); err != nil { + return nil, err + } + if err := f.verifyTargetReadinessForMutations(ctx, uniq); err != nil { + return nil, err + } + if err := f.verifyRouteWriteFloorForMutations(uniq, commitTS); err != nil { + return nil, err + } return uniq, nil } -// dedupProbeOnePhase decides whether handleOnePhaseTxnRequest should no-op -// because the entry is a retry whose prior attempt already landed. Extracted -// to keep handleOnePhaseTxnRequest under the cyclop budget; the determinism -// rationale lives at the call site. +// dedupProbeOnePhase first fail-closes the target readiness footprint, then +// decides whether handleOnePhaseTxnRequest should no-op because the entry is a +// retry whose prior attempt already landed. Extracted to keep +// handleOnePhaseTxnRequest under the cyclop budget; the determinism rationale +// lives at the call site. // // Returns (true, nil) → the entry must no-op (prior attempt landed). // Returns (false, nil) → fall through to normal apply. // Returns (false, err) → propagate err; apply must not proceed. -func (f *kvFSM) dedupProbeOnePhase(ctx context.Context, meta TxnMeta) (bool, error) { +func (f *kvFSM) dedupProbeOnePhase(ctx context.Context, meta TxnMeta, muts []*pb.Mutation, readKeys [][]byte) (bool, error) { + if err := f.verifyTargetReadinessForTxnFootprint(ctx, muts, readKeys, meta.PrimaryKey); err != nil { + return false, err + } if meta.PrevCommitTS == 0 { return false, nil } @@ -1169,7 +1476,7 @@ func (f *kvFSM) handleCommitRequest(ctx context.Context, r *pb.Request) error { if err != nil { return err } - uniq, err := uniqueTxnMutations(muts) + uniq, err := f.uniqueMutationsAboveWriteFloor(ctx, muts, commitTS) if err != nil { return err } @@ -1286,7 +1593,7 @@ func (f *kvFSM) handleAbortRequest(ctx context.Context, r *pb.Request, abortTS u // shouldClearAbortKey (lock-missing ⇒ nothing to do) and for the // rollback-marker Put in appendRollbackRecord. - uniq, err := uniqueMutations(muts) + uniq, err := f.uniqueAbortCleanupMutations(ctx, muts) if err != nil { return err } @@ -1306,6 +1613,24 @@ func (f *kvFSM) handleAbortRequest(ctx context.Context, r *pb.Request, abortTS u return errors.WithStack(f.store.ApplyMutationsRaftAt(ctx, storeMuts, nil, startTS, abortTS, f.pendingApplyIdx)) } +func (f *kvFSM) uniqueMutationsAboveWriteFloor(ctx context.Context, muts []*pb.Mutation, commitTS uint64) ([]*pb.Mutation, error) { + uniq, err := uniqueMutations(muts) + if err != nil { + return nil, err + } + if err := f.verifyTargetReadinessForMutations(ctx, uniq); err != nil { + return nil, err + } + if err := f.verifyRouteWriteFloorForMutations(uniq, commitTS); err != nil { + return nil, err + } + return uniq, nil +} + +func (f *kvFSM) uniqueAbortCleanupMutations(_ context.Context, muts []*pb.Mutation) ([]*pb.Mutation, error) { + return uniqueMutations(muts) +} + func (f *kvFSM) buildPrepareStoreMutations(ctx context.Context, muts []*pb.Mutation, primaryKey []byte, startTS, expireAt uint64) ([]*store.KVPairMutation, error) { storeMuts := make([]*store.KVPairMutation, 0, len(muts)*txnPrepareStoreMutationFactor) for _, mut := range muts { diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index 278c6800c..f11cf5a4d 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -12,6 +12,35 @@ import ( "github.com/stretchr/testify/require" ) +type deletePrefixIndexRecordingStore struct { + store.MVCCStore + + deletePrefixIndexes []uint64 + deletePrefixPrefixes [][]byte +} + +func (s *deletePrefixIndexRecordingStore) DeletePrefixAtRaftAt(ctx context.Context, prefix []byte, excludePrefix []byte, commitTS, appliedIndex uint64) error { + s.deletePrefixIndexes = append(s.deletePrefixIndexes, appliedIndex) + s.deletePrefixPrefixes = append(s.deletePrefixPrefixes, append([]byte(nil), prefix...)) + return s.MVCCStore.DeletePrefixAtRaftAt(ctx, prefix, excludePrefix, commitTS, appliedIndex) +} + +func (s *deletePrefixIndexRecordingStore) MigrationTargetReadinessStates(ctx context.Context) ([]store.TargetStagedReadinessState, error) { + reader, ok := s.MVCCStore.(store.MigrationTargetReadinessReader) + if !ok { + return nil, nil + } + return reader.MigrationTargetReadinessStates(ctx) +} + +func (s *deletePrefixIndexRecordingStore) ApplyTargetStagedReadiness(ctx context.Context, state store.TargetStagedReadinessState) error { + writer, ok := s.MVCCStore.(store.MigrationTargetReadinessWriter) + if !ok { + return store.ErrNotSupported + } + return writer.ApplyTargetStagedReadiness(ctx, state) +} + func newWriteFencedFSM(t *testing.T) *kvFSM { t.Helper() @@ -28,11 +57,67 @@ func newWriteFloorFSM(t *testing.T) *kvFSM { engine := distribution.NewEngine() applyComposed1Snapshot(t, engine, 1, []distribution.RouteDescriptor{ - {RouteID: 1, Start: []byte(""), End: nil, GroupID: 1, State: distribution.RouteStateActive, MinWriteTSExclusive: 100}, + { + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + MinWriteTSExclusive: 100, + }, }) return newComposed1FSM(t, engine, 1) } +func newTargetReadinessFSM(t *testing.T, route distribution.RouteDescriptor) *kvFSM { + t.Helper() + + engine := distribution.NewEngine() + applyComposed1Snapshot(t, engine, 2, []distribution.RouteDescriptor{route}) + fsm := newComposed1FSM(t, engine, route.GroupID) + writer, ok := fsm.store.(store.MigrationTargetReadinessWriter) + require.True(t, ok) + require.NoError(t, writer.ApplyTargetStagedReadiness(context.Background(), store.TargetStagedReadinessState{ + JobID: 9, + RouteStart: []byte("a"), + RouteEnd: []byte("z"), + ExpectedCutoverVersion: 2, + MigrationJobID: 9, + MinWriteTSExclusive: 100, + Armed: true, + })) + return fsm +} + +func applyTargetReadinessToFSM(t *testing.T, fsm *kvFSM, state store.TargetStagedReadinessState) { + t.Helper() + + writer, ok := fsm.store.(store.MigrationTargetReadinessWriter) + require.True(t, ok) + require.NoError(t, writer.ApplyTargetStagedReadiness(context.Background(), state)) +} + +func newReadinessReadKeyFSM(t *testing.T) *kvFSM { + t.Helper() + + engine := distribution.NewEngine() + applyComposed1Snapshot(t, engine, 2, []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte("a"), End: []byte("m"), GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: []byte("m"), End: []byte("z"), GroupID: 1, State: distribution.RouteStateActive}, + }) + fsm := newComposed1FSM(t, engine, 1) + applyTargetReadinessToFSM(t, fsm, store.TargetStagedReadinessState{ + JobID: 9, + RouteStart: []byte("m"), + RouteEnd: []byte("z"), + ExpectedCutoverVersion: 2, + MigrationJobID: 9, + MinWriteTSExclusive: 100, + Armed: true, + }) + return fsm +} + func s3BucketAuxiliaryFenceRoutes(bucket string, rawGroupID, fencedGroupID uint64) []distribution.RouteDescriptor { start := s3keys.RoutePrefixForBucketAnyGeneration(bucket) end := prefixScanEnd(start) @@ -64,6 +149,79 @@ func TestFSMRejectsRawPointWriteOnWriteFencedRoute(t *testing.T) { require.ErrorIs(t, getErr, store.ErrKeyNotFound) } +func TestFSMRejectsRawPointWriteWithoutTargetReadinessProof(t *testing.T) { + t.Parallel() + + fsm := newTargetReadinessFSM(t, distribution.RouteDescriptor{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + }) + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte("b"), Value: []byte("v")}}, + }, 120) + require.ErrorIs(t, err, ErrRouteCutoverPending) + + _, getErr := fsm.store.GetAt(context.Background(), []byte("b"), ^uint64(0)) + require.ErrorIs(t, getErr, store.ErrKeyNotFound) +} + +func TestFSMAllowsRawPointWriteWithClearedTargetReadinessProof(t *testing.T) { + t.Parallel() + + fsm := newTargetReadinessFSM(t, distribution.RouteDescriptor{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + MinWriteTSExclusive: 100, + }) + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte("b"), Value: []byte("v")}}, + }, 101) + require.NoError(t, err) + + got, getErr := fsm.store.GetAt(context.Background(), []byte("b"), ^uint64(0)) + require.NoError(t, getErr) + require.Equal(t, []byte("v"), got) +} + +func TestFSMRejectsTargetReadinessProofFromAnotherGroup(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + applyComposed1Snapshot(t, engine, 1, []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 2, + State: distribution.RouteStateActive, + MinWriteTSExclusive: 100, + }, + }) + fsm := newComposed1FSM(t, engine, 1) + writer, ok := fsm.store.(store.MigrationTargetReadinessWriter) + require.True(t, ok) + require.NoError(t, writer.ApplyTargetStagedReadiness(context.Background(), store.TargetStagedReadinessState{ + JobID: 9, + RouteStart: []byte("a"), + RouteEnd: []byte("z"), + ExpectedCutoverVersion: 2, + MigrationJobID: 9, + MinWriteTSExclusive: 100, + Armed: true, + })) + + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte("b"), Value: []byte("v")}}, + }, 101) + require.ErrorIs(t, err, ErrRouteCutoverPending) +} + func TestFSMRejectsS3BucketAuxiliaryPointWriteOnWriteFencedRoute(t *testing.T) { t.Parallel() @@ -101,6 +259,87 @@ func TestFSMRejectsDelPrefixIntersectingWriteFencedRoute(t *testing.T) { require.Equal(t, []byte("v"), got) } +func TestFSMRejectsDelPrefixWithoutTargetReadinessProof(t *testing.T) { + t.Parallel() + + fsm := newTargetReadinessFSM(t, distribution.RouteDescriptor{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + }) + require.NoError(t, fsm.store.PutAt(context.Background(), []byte("b"), []byte("v"), 1, 0)) + + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_DEL_PREFIX, Key: []byte("b")}}, + }, 120) + require.ErrorIs(t, err, ErrRouteCutoverPending) + + got, getErr := fsm.store.GetAt(context.Background(), []byte("b"), ^uint64(0)) + require.NoError(t, getErr) + require.Equal(t, []byte("v"), got) +} + +func TestFSMDelPrefixTombstonesStagedVisibilityRows(t *testing.T) { + t.Parallel() + + ctx := context.Background() + fsm := newTargetReadinessFSM(t, distribution.RouteDescriptor{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + MinWriteTSExclusive: 100, + }) + rawKey := []byte("b") + stagedKey := distribution.MigrationStagedDataKey(9, rawKey) + require.NoError(t, fsm.store.PutAt(ctx, stagedKey, []byte("staged"), 110, 0)) + + err := fsm.handleRawRequest(ctx, &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_DEL_PREFIX, Key: rawKey}}, + }, 120) + require.NoError(t, err) + + _, err = fsm.store.GetAt(ctx, stagedKey, 130) + require.ErrorIs(t, err, store.ErrKeyNotFound) +} + +func TestFSMDelPrefixAdvancesApplyIndexOnlyOnLiveDelete(t *testing.T) { + t.Parallel() + + ctx := context.Background() + fsm := newTargetReadinessFSM(t, distribution.RouteDescriptor{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + MinWriteTSExclusive: 100, + }) + recording := &deletePrefixIndexRecordingStore{MVCCStore: fsm.store} + fsm.store = recording + fsm.SetApplyIndex(55) + + rawKey := []byte("b") + stagedKey := distribution.MigrationStagedDataKey(9, rawKey) + require.NoError(t, fsm.store.PutAt(ctx, stagedKey, []byte("staged"), 110, 0)) + + err := fsm.handleRawRequest(ctx, &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_DEL_PREFIX, Key: rawKey}}, + }, 120) + require.NoError(t, err) + + require.Equal(t, []uint64{0, 55}, recording.deletePrefixIndexes) + require.Equal(t, stagedKey, recording.deletePrefixPrefixes[0]) + require.Equal(t, rawKey, recording.deletePrefixPrefixes[1]) +} + func TestFSMRejectsFullRangeDelPrefixWhenRouteIsWriteFenced(t *testing.T) { t.Parallel() @@ -117,6 +356,35 @@ func TestFSMRejectsFullRangeDelPrefixWhenRouteIsWriteFenced(t *testing.T) { require.Equal(t, []byte("v"), got) } +func TestFSMRejectsRawPointWriteBelowRouteFloor(t *testing.T) { + t.Parallel() + + fsm := newWriteFloorFSM(t) + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte("b"), Value: []byte("v")}}, + }, 100) + require.ErrorIs(t, err, ErrRouteWriteBelowFloor) + + _, getErr := fsm.store.GetAt(context.Background(), []byte("b"), ^uint64(0)) + require.ErrorIs(t, getErr, store.ErrKeyNotFound) +} + +func TestFSMRejectsDelPrefixBelowRouteFloor(t *testing.T) { + t.Parallel() + + fsm := newWriteFloorFSM(t) + require.NoError(t, fsm.store.PutAt(context.Background(), []byte("b"), []byte("v"), 1, 0)) + + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_DEL_PREFIX, Key: []byte("b")}}, + }, 100) + require.ErrorIs(t, err, ErrRouteWriteBelowFloor) + + got, getErr := fsm.store.GetAt(context.Background(), []byte("b"), ^uint64(0)) + require.NoError(t, getErr) + require.Equal(t, []byte("v"), got) +} + func TestFSMRejectsBroadInternalDelPrefixWhenRouteIsWriteFenced(t *testing.T) { t.Parallel() @@ -134,6 +402,188 @@ func TestFSMRejectsBroadInternalDelPrefixWhenRouteIsWriteFenced(t *testing.T) { require.Equal(t, []byte("v"), got) } +func TestFSMRejectsOnePhaseTxnBelowRouteFloor(t *testing.T) { + t.Parallel() + + fsm := newWriteFloorFSM(t) + err := fsm.handleTxnRequest(context.Background(), onePhaseReq(10, 100, 0, []byte("b"), []byte("v")), 100) + require.ErrorIs(t, err, ErrRouteWriteBelowFloor) + + _, getErr := fsm.store.GetAt(context.Background(), []byte("b"), ^uint64(0)) + require.ErrorIs(t, getErr, store.ErrKeyNotFound) +} + +func TestFSMOnePhaseDedupChecksTargetReadinessBeforeNoOp(t *testing.T) { + t.Parallel() + + ctx := context.Background() + fsm := newTargetReadinessFSM(t, distribution.RouteDescriptor{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + }) + require.NoError(t, fsm.store.PutAt(ctx, []byte("b"), []byte("landed"), 20, 0)) + + err := fsm.handleTxnRequest(ctx, onePhaseReq(30, 40, 20, []byte("b"), []byte("retry")), 40) + require.ErrorIs(t, err, ErrRouteCutoverPending) + + landedAt40, probeErr := fsm.store.CommittedVersionAt(ctx, []byte("b"), 40) + require.NoError(t, probeErr) + require.False(t, landedAt40) +} + +func TestFSMOnePhaseTxnChecksReadKeysForTargetReadiness(t *testing.T) { + t.Parallel() + + ctx := context.Background() + fsm := newReadinessReadKeyFSM(t) + req := onePhaseReq(10, 20, 0, []byte("b"), []byte("v")) + req.ReadKeys = [][]byte{[]byte("n")} + + err := fsm.handleTxnRequest(ctx, req, 20) + require.ErrorIs(t, err, ErrRouteCutoverPending) + + _, getErr := fsm.store.GetAt(ctx, []byte("b"), ^uint64(0)) + require.ErrorIs(t, getErr, store.ErrKeyNotFound) +} + +func TestFSMPrepareTxnChecksReadKeysForTargetReadiness(t *testing.T) { + t.Parallel() + + ctx := context.Background() + fsm := newReadinessReadKeyFSM(t) + req := &pb.Request{ + IsTxn: true, + Phase: pb.Phase_PREPARE, + Ts: 10, + ReadKeys: [][]byte{[]byte("n")}, + Mutations: []*pb.Mutation{ + { + Op: pb.Op_PUT, + Key: []byte(txnMetaPrefix), + Value: EncodeTxnMeta(TxnMeta{ + PrimaryKey: []byte("b"), + LockTTLms: defaultTxnLockTTLms, + }), + }, + {Op: pb.Op_PUT, Key: []byte("b"), Value: []byte("v")}, + }, + } + + err := fsm.handleTxnRequest(ctx, req, 10) + require.ErrorIs(t, err, ErrRouteCutoverPending) + + _, getErr := fsm.store.GetAt(ctx, txnLockKey([]byte("b")), ^uint64(0)) + require.ErrorIs(t, getErr, store.ErrKeyNotFound) +} + +func TestFSMPrepareTxnChecksStagedVisibilityWriteConflicts(t *testing.T) { + t.Parallel() + + ctx := context.Background() + fsm := newTargetReadinessFSM(t, distribution.RouteDescriptor{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + MinWriteTSExclusive: 100, + }) + require.NoError(t, fsm.store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("b")), []byte("staged"), 20, 0)) + + err := fsm.handleTxnRequest(ctx, &pb.Request{ + IsTxn: true, + Phase: pb.Phase_PREPARE, + Ts: 10, + Mutations: []*pb.Mutation{ + { + Op: pb.Op_PUT, + Key: []byte(txnMetaPrefix), + Value: EncodeTxnMeta(TxnMeta{ + PrimaryKey: []byte("b"), + CommitTS: 120, + LockTTLms: defaultTxnLockTTLms, + }), + }, + {Op: pb.Op_PUT, Key: []byte("b"), Value: []byte("v")}, + }, + }, 10) + require.ErrorIs(t, err, store.ErrWriteConflict) + + _, getErr := fsm.store.GetAt(ctx, txnLockKey([]byte("b")), ^uint64(0)) + require.ErrorIs(t, getErr, store.ErrKeyNotFound) +} + +func TestFSMOnePhaseTxnChecksStagedVisibilityReadConflicts(t *testing.T) { + t.Parallel() + + ctx := context.Background() + fsm := newTargetReadinessFSM(t, distribution.RouteDescriptor{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + MinWriteTSExclusive: 100, + }) + require.NoError(t, fsm.store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("n")), []byte("staged"), 20, 0)) + + req := onePhaseReq(10, 120, 0, []byte("b"), []byte("v")) + req.ReadKeys = [][]byte{[]byte("n")} + + err := fsm.handleTxnRequest(ctx, req, 120) + require.ErrorIs(t, err, store.ErrWriteConflict) + + _, getErr := fsm.store.GetAt(ctx, []byte("b"), ^uint64(0)) + require.ErrorIs(t, getErr, store.ErrKeyNotFound) +} + +func TestFSMPrepareUsesCommitTSForRouteFloorWhenPresent(t *testing.T) { + t.Parallel() + + fsm := newWriteFloorFSM(t) + err := fsm.handleTxnRequest(context.Background(), &pb.Request{ + IsTxn: true, + Phase: pb.Phase_PREPARE, + Ts: 50, + Mutations: []*pb.Mutation{ + { + Op: pb.Op_PUT, + Key: []byte(txnMetaPrefix), + Value: EncodeTxnMeta(TxnMeta{PrimaryKey: []byte("b"), LockTTLms: defaultTxnLockTTLms, CommitTS: 101}), + }, + {Op: pb.Op_PUT, Key: []byte("b"), Value: []byte("v")}, + }, + }, 50) + require.NoError(t, err) +} + +func TestFSMPrepareWithoutCommitTSUsesStartTSForRouteFloor(t *testing.T) { + t.Parallel() + + fsm := newWriteFloorFSM(t) + err := fsm.handleTxnRequest(context.Background(), &pb.Request{ + IsTxn: true, + Phase: pb.Phase_PREPARE, + Ts: 100, + Mutations: []*pb.Mutation{ + { + Op: pb.Op_PUT, + Key: []byte(txnMetaPrefix), + Value: EncodeTxnMeta(TxnMeta{PrimaryKey: []byte("b"), LockTTLms: defaultTxnLockTTLms}), + }, + {Op: pb.Op_PUT, Key: []byte("b"), Value: []byte("v")}, + }, + }, 100) + require.ErrorIs(t, err, ErrRouteWriteBelowFloor) +} + func TestFSMRejectsPrepareOnWriteFencedRouteButAllowsAbort(t *testing.T) { t.Parallel() @@ -163,85 +613,80 @@ func TestFSMRejectsPrepareOnWriteFencedRouteButAllowsAbort(t *testing.T) { require.False(t, errors.Is(err, ErrRouteWriteFenced), "ABORT must keep the narrow cleanup lane open") } -func TestFSMAllowsRawPointWriteAtMigrationTimestampFloorDuringReplay(t *testing.T) { +func TestFSMAbortCleanupBypassesRetainedWriteFloor(t *testing.T) { t.Parallel() ctx := context.Background() fsm := newWriteFloorFSM(t) - err := fsm.handleRawRequest(ctx, &pb.Request{ - Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("replayed")}}, - }, 100) - require.NoError(t, err) - got, getErr := fsm.store.GetAt(ctx, []byte("z"), ^uint64(0)) - require.NoError(t, getErr) - require.Equal(t, []byte("replayed"), got) -} - -func TestFSMAllowsDelPrefixAtMigrationTimestampFloorDuringReplay(t *testing.T) { - t.Parallel() - - ctx := context.Background() - fsm := newWriteFloorFSM(t) - require.NoError(t, fsm.store.PutAt(ctx, []byte("z"), []byte("v"), 10, 0)) + startTS := uint64(10) + abortTS := uint64(11) + primaryKey := []byte("b") + require.NoError(t, fsm.store.PutAt(ctx, txnLockKey(primaryKey), encodeTxnLock(txnLock{ + StartTS: startTS, + PrimaryKey: primaryKey, + IsPrimaryKey: true, + }), startTS, 0)) + require.NoError(t, fsm.store.PutAt(ctx, txnIntentKey(primaryKey), encodeTxnIntent(txnIntent{ + StartTS: startTS, + Op: txnIntentOpPut, + Value: []byte("v"), + }), startTS, 0)) - err := fsm.handleRawRequest(ctx, &pb.Request{ - Mutations: []*pb.Mutation{{Op: pb.Op_DEL_PREFIX, Key: []byte("z")}}, - }, 100) - require.NoError(t, err) - - _, getErr := fsm.store.GetAt(ctx, []byte("z"), ^uint64(0)) - require.ErrorIs(t, getErr, store.ErrKeyNotFound) -} - -func TestFSMAllowsOnePhaseTxnAtMigrationTimestampFloorDuringReplay(t *testing.T) { - t.Parallel() - - ctx := context.Background() - fsm := newWriteFloorFSM(t) - req := &pb.Request{ + abort := &pb.Request{ IsTxn: true, - Phase: pb.Phase_NONE, - Ts: 90, + Phase: pb.Phase_ABORT, + Ts: startTS, Mutations: []*pb.Mutation{ - {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: []byte("z"), CommitTS: 100})}, - {Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("low")}, + {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: primaryKey, CommitTS: abortTS})}, + {Op: pb.Op_PUT, Key: primaryKey}, }, } - err := fsm.handleTxnRequest(ctx, req, 100) - require.NoError(t, err) - got, getErr := fsm.store.GetAt(ctx, []byte("z"), ^uint64(0)) - require.NoError(t, getErr) - require.Equal(t, []byte("low"), got) + require.NoError(t, fsm.handleTxnRequest(ctx, abort, abortTS)) + + _, lockErr := fsm.store.GetAt(ctx, txnLockKey(primaryKey), ^uint64(0)) + require.ErrorIs(t, lockErr, store.ErrKeyNotFound) + _, intentErr := fsm.store.GetAt(ctx, txnIntentKey(primaryKey), ^uint64(0)) + require.ErrorIs(t, intentErr, store.ErrKeyNotFound) } -func TestFSMAllowsCommitAtMigrationTimestampFloorDuringReplay(t *testing.T) { +func TestFSMAbortCleanupBypassesTargetReadiness(t *testing.T) { t.Parallel() ctx := context.Background() - fsm := newWriteFloorFSM(t) - prepare := &pb.Request{ - IsTxn: true, - Phase: pb.Phase_PREPARE, - Ts: 90, - Mutations: []*pb.Mutation{ - {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: []byte("z"), LockTTLms: defaultTxnLockTTLms})}, - {Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("v")}, - }, - } - require.NoError(t, fsm.handleTxnRequest(ctx, prepare, 90)) + fsm := newTargetReadinessFSM(t, distribution.RouteDescriptor{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + }) + startTS := uint64(10) + abortTS := uint64(11) + primaryKey := []byte("b") + require.NoError(t, fsm.store.PutAt(ctx, txnLockKey(primaryKey), encodeTxnLock(txnLock{ + StartTS: startTS, + PrimaryKey: primaryKey, + IsPrimaryKey: true, + }), startTS, 0)) + require.NoError(t, fsm.store.PutAt(ctx, txnIntentKey(primaryKey), encodeTxnIntent(txnIntent{ + StartTS: startTS, + Op: txnIntentOpPut, + Value: []byte("v"), + }), startTS, 0)) - commit := &pb.Request{ + abort := &pb.Request{ IsTxn: true, - Phase: pb.Phase_COMMIT, - Ts: 90, + Phase: pb.Phase_ABORT, + Ts: startTS, Mutations: []*pb.Mutation{ - {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: []byte("z"), CommitTS: 100})}, - {Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("low")}, + {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: primaryKey, CommitTS: abortTS})}, + {Op: pb.Op_PUT, Key: primaryKey}, }, } - err := fsm.handleTxnRequest(ctx, commit, 100) - require.NoError(t, err) - got, getErr := fsm.store.GetAt(ctx, []byte("z"), ^uint64(0)) - require.NoError(t, getErr) - require.Equal(t, []byte("v"), got) + require.NoError(t, fsm.handleTxnRequest(ctx, abort, abortTS)) + + _, lockErr := fsm.store.GetAt(ctx, txnLockKey(primaryKey), ^uint64(0)) + require.ErrorIs(t, lockErr, store.ErrKeyNotFound) + _, intentErr := fsm.store.GetAt(ctx, txnIntentKey(primaryKey), ^uint64(0)) + require.ErrorIs(t, intentErr, store.ErrKeyNotFound) } diff --git a/kv/fsm_migration_readiness.go b/kv/fsm_migration_readiness.go new file mode 100644 index 000000000..91ea21cd8 --- /dev/null +++ b/kv/fsm_migration_readiness.go @@ -0,0 +1,59 @@ +package kv + +import ( + "bytes" + "context" + + pb "github.com/bootjp/elastickv/proto" + "github.com/bootjp/elastickv/store" + "github.com/cockroachdb/errors" + "google.golang.org/protobuf/proto" +) + +// MarshalTargetStagedReadinessCommand encodes a target-group readiness guard +// as a Raft FSM command. The target Internal RPC handler uses this instead of +// mutating its local store directly so an acknowledged guard has been applied +// by the target group's voters. +func MarshalTargetStagedReadinessCommand(req *pb.TargetStagedReadinessRequest) ([]byte, error) { + if req == nil { + return nil, errors.WithStack(ErrInvalidRequest) + } + b, err := proto.Marshal(req) + if err != nil { + return nil, errors.WithStack(err) + } + if len(b) >= maxMarshaledCommandSize { + return nil, errors.New("marshaled target readiness request too large") + } + return prependByte(raftEncodeTargetReadiness, b), nil +} + +func (f *kvFSM) applyTargetStagedReadiness(ctx context.Context, data []byte) any { + req := &pb.TargetStagedReadinessRequest{} + if err := proto.Unmarshal(data, req); err != nil { + return errors.WithStack(err) + } + writer, ok := f.store.(store.MigrationTargetReadinessWriter) + if !ok { + return errors.WithStack(store.ErrNotSupported) + } + if err := writer.ApplyTargetStagedReadiness(ctx, targetStagedReadinessStateFromProto(req)); err != nil { + return errors.WithStack(err) + } + return nil +} + +func targetStagedReadinessStateFromProto(req *pb.TargetStagedReadinessRequest) store.TargetStagedReadinessState { + if req == nil { + return store.TargetStagedReadinessState{} + } + return store.TargetStagedReadinessState{ + JobID: req.GetJobId(), + RouteStart: bytes.Clone(req.GetRouteStart()), + RouteEnd: bytes.Clone(req.GetRouteEnd()), + ExpectedCutoverVersion: req.GetExpectedCutoverVersion(), + MigrationJobID: req.GetMigrationJobId(), + MinWriteTSExclusive: req.GetMinWriteTsExclusive(), + Armed: req.GetArmed(), + } +} diff --git a/kv/fsm_migration_readiness_test.go b/kv/fsm_migration_readiness_test.go new file mode 100644 index 000000000..8d5939b84 --- /dev/null +++ b/kv/fsm_migration_readiness_test.go @@ -0,0 +1,44 @@ +package kv + +import ( + "context" + "testing" + + pb "github.com/bootjp/elastickv/proto" + "github.com/bootjp/elastickv/store" + "github.com/stretchr/testify/require" +) + +func TestApplyTargetStagedReadinessCommandPersistsGuard(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + fsm := &kvFSM{store: st} + + cmd, err := MarshalTargetStagedReadinessCommand(&pb.TargetStagedReadinessRequest{ + JobId: 9, + RouteStart: []byte("a"), + RouteEnd: []byte("z"), + ExpectedCutoverVersion: 3, + MigrationJobId: 7, + MinWriteTsExclusive: 100, + Armed: true, + }) + require.NoError(t, err) + require.Nil(t, fsm.Apply(cmd)) + + reader, ok := st.(store.MigrationTargetReadinessReader) + require.True(t, ok) + states, err := reader.MigrationTargetReadinessStates(ctx) + require.NoError(t, err) + require.Equal(t, []store.TargetStagedReadinessState{{ + JobID: 9, + RouteStart: []byte("a"), + RouteEnd: []byte("z"), + ExpectedCutoverVersion: 3, + MigrationJobID: 7, + MinWriteTSExclusive: 100, + Armed: true, + }}, states) +} diff --git a/kv/route_history.go b/kv/route_history.go index 8bc50c58e..94eec1076 100644 --- a/kv/route_history.go +++ b/kv/route_history.go @@ -82,3 +82,21 @@ func (s distributionRouteSnapshot) WriteFencedIntersects(start, end []byte) bool } return false } + +func (s distributionRouteSnapshot) WriteFloorForKey(key []byte) (uint64, bool) { + route, ok := s.snap.RouteOf(key) + if !ok || route.MinWriteTSExclusive == 0 { + return 0, false + } + return route.MinWriteTSExclusive, true +} + +func (s distributionRouteSnapshot) WriteFloorIntersects(start, end []byte) (uint64, bool) { + var maxFloor uint64 + for _, route := range s.snap.IntersectingRoutes(start, end) { + if route.MinWriteTSExclusive > maxFloor { + maxFloor = route.MinWriteTSExclusive + } + } + return maxFloor, maxFloor > 0 +} diff --git a/kv/shard_store.go b/kv/shard_store.go index 750e4b08b..e12a73659 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -26,7 +26,9 @@ type ShardStore struct { } var ErrCrossShardMutationBatchNotSupported = errors.New("cross-shard mutation batches are not supported") +var ErrRouteCutoverPending = errors.New("route cutover pending") var ErrExplicitGroupStagedVisibilityUnresolved = errors.New("explicit group read cannot resolve staged visibility route") +var ErrRouteWriteBelowFloor = ErrRouteWriteTimestampTooLow // NewShardStore creates a sharded MVCC store wrapper. func NewShardStore(engine *distribution.Engine, groups map[uint64]*ShardGroup) *ShardStore { @@ -112,6 +114,11 @@ func isLinearizableRaftLeader(ctx context.Context, engine raftengine.LeaderView) } func (s *ShardStore) leaderGetAt(ctx context.Context, g *ShardGroup, route distribution.Route, key []byte, ts uint64) ([]byte, error) { + var err error + route, err = s.targetReadyRouteForRange(ctx, g, route, key, nextScanCursor(key)) + if err != nil { + return nil, err + } if !isTxnInternalKey(key) { if err := s.maybeResolveTxnLock(ctx, g, key, ts); err != nil { return nil, err @@ -121,6 +128,11 @@ func (s *ShardStore) leaderGetAt(ctx context.Context, g *ShardGroup, route distr } func (s *ShardStore) localGetAt(ctx context.Context, g *ShardGroup, route distribution.Route, key []byte, ts uint64) ([]byte, error) { + var err error + route, err = s.targetReadyRouteForRange(ctx, g, route, key, nextScanCursor(key)) + if err != nil { + return nil, err + } if routeHasStagedVisibility(route) { return s.getAtWithStagedVisibility(ctx, g, route, key, ts) } @@ -135,6 +147,173 @@ func routeHasStagedVisibility(route distribution.Route) bool { return route.StagedVisibilityActive && route.MigrationJobID != 0 } +func routeSatisfiesTargetReadiness(route distribution.Route, ready store.TargetStagedReadinessState, catalogVersion uint64) bool { + if !routeRangeIntersects(route.Start, route.End, ready.RouteStart, ready.RouteEnd) { + return false + } + if route.MinWriteTSExclusive < ready.MinWriteTSExclusive { + return false + } + if route.StagedVisibilityActive { + return route.MigrationJobID == ready.MigrationJobID + } + if route.MigrationJobID != 0 { + return false + } + return ready.ExpectedCutoverVersion == 0 || catalogVersion >= ready.ExpectedCutoverVersion +} + +func (s *ShardStore) targetReadyRouteForRange(ctx context.Context, g *ShardGroup, route distribution.Route, start []byte, end []byte) (distribution.Route, error) { + routeStart, routeEnd := readinessRouteRangeForScan(start, end) + return s.targetReadyRouteForRouteRange(ctx, g, route, routeStart, routeEnd) +} + +func (s *ShardStore) verifyTargetReadinessForRange(ctx context.Context, g *ShardGroup, route distribution.Route, start []byte, end []byte) error { + _, err := s.targetReadyRouteForRange(ctx, g, route, start, end) + return err +} + +func (s *ShardStore) targetReadyRouteForRouteRange(ctx context.Context, g *ShardGroup, route distribution.Route, routeStart []byte, routeEnd []byte) (distribution.Route, error) { + routes, err := s.targetReadyRoutesForRouteRange(ctx, g, route, routeStart, routeEnd) + if err != nil { + return route, err + } + if len(routes) == 1 { + return routes[0], nil + } + return route, nil +} + +func (s *ShardStore) targetReadyRoutesForRouteRange(ctx context.Context, g *ShardGroup, route distribution.Route, routeStart []byte, routeEnd []byte) ([]distribution.Route, error) { + if g == nil || g.Store == nil { + return []distribution.Route{route}, nil + } + reader, ok := g.Store.(store.MigrationTargetReadinessReader) + if !ok { + return []distribution.Route{route}, nil + } + states, err := reader.MigrationTargetReadinessStates(ctx) + if err != nil { + return nil, errors.WithStack(err) + } + applicable := targetReadinessApplicableStates(states, route, routeStart, routeEnd) + if len(applicable) == 0 { + return []distribution.Route{route}, nil + } + + proofRoutes, catalogVersion, ok := s.readinessProofRoutes(route, routeStart, routeEnd) + if !ok { + return nil, errors.WithStack(ErrRouteCutoverPending) + } + if !readinessProofSatisfiesStates(applicable, proofRoutes, route.GroupID, catalogVersion) { + return nil, errors.WithStack(ErrRouteCutoverPending) + } + return proofRoutes, nil +} + +func targetReadinessApplicableStates( + states []store.TargetStagedReadinessState, + route distribution.Route, + routeStart []byte, + routeEnd []byte, +) []store.TargetStagedReadinessState { + applicable := make([]store.TargetStagedReadinessState, 0, len(states)) + for _, ready := range states { + if targetReadinessAppliesToRoute(route, routeStart, routeEnd, ready) { + applicable = append(applicable, ready) + } + } + return applicable +} + +func readinessProofSatisfiesStates( + states []store.TargetStagedReadinessState, + routes []distribution.Route, + groupID uint64, + catalogVersion uint64, +) bool { + for _, ready := range states { + if !routesSatisfyTargetReadiness(routes, ready, groupID, catalogVersion) { + return false + } + } + return true +} + +func (s *ShardStore) verifyTargetReadinessForRouteRange(ctx context.Context, g *ShardGroup, route distribution.Route, routeStart []byte, routeEnd []byte) error { + _, err := s.targetReadyRouteForRouteRange(ctx, g, route, routeStart, routeEnd) + return err +} + +func (s *ShardStore) readinessProofRoutes(route distribution.Route, routeStart []byte, routeEnd []byte) ([]distribution.Route, uint64, bool) { + if s == nil || s.engine == nil { + return []distribution.Route{route}, 0, true + } + snap, ok := s.engine.Current() + if !ok { + return nil, 0, false + } + routes := snap.IntersectingRoutes(routeStart, routeEnd) + proof := routes[:0] + for _, candidate := range routes { + if candidate.GroupID != route.GroupID { + continue + } + if (route.Start != nil || route.End != nil || route.RouteID != 0) && + !routeRangeIntersects(candidate.Start, candidate.End, route.Start, route.End) { + continue + } + proof = append(proof, candidate) + } + return proof, snap.Version(), len(proof) > 0 +} + +func targetReadinessAppliesToRoute(route distribution.Route, routeStart []byte, routeEnd []byte, ready store.TargetStagedReadinessState) bool { + if !ready.Armed || !routeRangeIntersects(routeStart, routeEnd, ready.RouteStart, ready.RouteEnd) { + return false + } + if route.RouteID != 0 && !routeRangeIntersects(route.Start, route.End, ready.RouteStart, ready.RouteEnd) { + return false + } + return true +} + +func readinessRouteRange(start []byte, end []byte) ([]byte, []byte) { + routeStart := routeKey(start) + if end == nil { + return routeStart, nil + } + routeEnd := routeKey(end) + if bytes.Compare(routeEnd, routeStart) <= 0 { + routeEnd = nextScanCursor(routeStart) + } + return routeStart, routeEnd +} + +func readinessRouteRangeForScan(start []byte, end []byte) ([]byte, []byte) { + if routeStart, routeEnd, ok := s3keys.ManifestScanRouteBounds(start, end); ok { + return routeStart, routeEnd + } + return readinessRouteRange(start, end) +} + +func verifyRouteWriteFloor(route distribution.Route, commitTS uint64) error { + if route.MinWriteTSExclusive == 0 || commitTS == 0 || commitTS > route.MinWriteTSExclusive { + return nil + } + return errors.Wrapf(ErrRouteWriteBelowFloor, "commit_ts %d <= floor %d", commitTS, route.MinWriteTSExclusive) +} + +func routeRangeIntersects(aStart, aEnd, bStart, bEnd []byte) bool { + if aEnd != nil && bytes.Compare(aEnd, bStart) <= 0 { + return false + } + if bEnd != nil && bytes.Compare(bEnd, aStart) <= 0 { + return false + } + return true +} + func (s *ShardStore) routeForExplicitGroupKey(groupID uint64, key []byte) (distribution.Route, error) { fallback := distribution.Route{GroupID: groupID} if s == nil || s.engine == nil { @@ -185,6 +364,7 @@ func latestMVCCVersionAt(ctx context.Context, st store.MVCCStore, key []byte, ts StartKey: key, EndKey: prefixScanEnd(key), MaxCommitTSInclusive: ts, + ReadTS: ts, MaxVersions: 1, MaxScannedBytes: 0, MinCommitTSExclusive: 0, @@ -263,19 +443,20 @@ func (s *ShardStore) ExistsAt(ctx context.Context, key []byte, ts uint64) (bool, // pending.length fast-path during churn. Mirrors LeaderRoutedStore's fix // for codex P1 #796. func (s *ShardStore) CommittedVersionAt(ctx context.Context, key []byte, commitTS uint64) (bool, error) { - g, ok := s.groupForKey(key) + route, g, ok := s.routeAndGroupForKey(key) if !ok || g.Store == nil { return false, nil } + var err error + route, err = s.targetReadyRouteForRange(ctx, g, route, key, nextScanCursor(key)) + if err != nil { + return false, err + } // engineForGroup may be nil in test fixtures that wire ShardStore // without raft; preserve the existing local-only fallback there. engine := engineForGroup(g) if engine == nil { - exists, err := g.Store.CommittedVersionAt(ctx, key, commitTS) - if err != nil { - return false, errors.WithStack(err) - } - return exists, nil + return committedVersionAtForRoute(ctx, g.Store, route, key, commitTS) } if !isLinearizableRaftLeader(ctx, engine) && !tryEngineLinearizableFence(ctx, engine) { // Not the linearizable leader for this group AND the ReadIndex @@ -285,7 +466,23 @@ func (s *ShardStore) CommittedVersionAt(ctx context.Context, key []byte, commitT // serialization. return false, nil } - exists, err := g.Store.CommittedVersionAt(ctx, key, commitTS) + route, err = s.targetReadyRouteForRange(ctx, g, route, key, nextScanCursor(key)) + if err != nil { + return false, err + } + return committedVersionAtForRoute(ctx, g.Store, route, key, commitTS) +} + +func committedVersionAtForRoute(ctx context.Context, st store.MVCCStore, route distribution.Route, key []byte, commitTS uint64) (bool, error) { + exists, err := st.CommittedVersionAt(ctx, key, commitTS) + if err != nil { + return false, errors.WithStack(err) + } + if exists || !routeHasStagedVisibility(route) { + return exists, nil + } + stagedKey := distribution.MigrationStagedDataKey(route.MigrationJobID, key) + exists, err = st.CommittedVersionAt(ctx, stagedKey, commitTS) if err != nil { return false, errors.WithStack(err) } @@ -385,27 +582,49 @@ func (s *ShardStore) ScanGroupAt(ctx context.Context, groupID uint64, start []by if err != nil { return nil, err } - out := make([]*store.KVPair, 0) + routeStart, routeEnd := readinessRouteRangeForScan(start, end) + if err := s.verifyExplicitGroupRoutesForRange(ctx, groupID, routes, routeStart, routeEnd); err != nil { + return nil, err + } + return s.scanExplicitGroupRoutesAt(ctx, routes, start, end, limit, ts, clampToRoutes) +} + +func (s *ShardStore) scanExplicitGroupRoutesAt(ctx context.Context, routes []distribution.Route, start []byte, end []byte, limit int, ts uint64, clampToRoutes bool) ([]*store.KVPair, error) { + if len(routes) == 1 { + return s.scanRouteAtDirection(ctx, routes[0], start, end, limit, ts, false, true) + } + out := make([]*store.KVPair, 0, limit) for _, route := range routes { - scanStart := start - scanEnd := end + scanStart, scanEnd := start, end if clampToRoutes { scanStart = clampScanStart(start, route.Start) scanEnd = clampScanEnd(end, route.End) } - kvs, err := s.scanRouteAtDirection(ctx, route, scanStart, scanEnd, limit, ts, false, true) + kvs, err := s.scanRouteAtDirection(ctx, route, scanStart, scanEnd, limit-len(out), ts, false, true) if err != nil { return nil, err } - out = append(out, kvs...) - if len(out) >= limit { - clear(out[limit:]) - return out[:limit], nil + for _, kvp := range kvs { + if !kvBelongsToRoute(kvp, route) { + continue + } + out = append(out, kvp) + if len(out) == limit { + return out, nil + } } } return out, nil } +func kvBelongsToRoute(kvp *store.KVPair, route distribution.Route) bool { + if kvp == nil || route.RouteID == 0 { + return true + } + key := routeKey(kvp.Key) + return routeRangeIntersects(key, nextScanCursor(key), route.Start, route.End) +} + func (s *ShardStore) ReverseScanAt(ctx context.Context, start []byte, end []byte, limit int, ts uint64) ([]*store.KVPair, error) { return s.reverseScanAtWithReadFence(ctx, start, end, limit, ts, 0, nil, nil) } @@ -978,6 +1197,11 @@ func (s *ShardStore) scanRouteAtDirectionPhysicalLimit( if !ok || g == nil || g.Store == nil { return nil, false, nil } + var err error + route, err = s.targetReadyRouteForRange(ctx, g, route, start, end) + if err != nil { + return nil, false, err + } if engineForGroup(g) == nil { if routeHasStagedVisibility(route) { @@ -991,10 +1215,7 @@ func (s *ShardStore) scanRouteAtDirectionPhysicalLimit( } if isLinearizableRaftLeader(ctx, engineForGroup(g)) { - if routeHasStagedVisibility(route) { - return nil, true, nil - } - return s.scanRouteAtLeaderPhysicalLimit(ctx, g, route, start, end, visibleLimit, physicalLimit, ts, reverse) + return s.scanReadyLeaderPhysicalLimit(ctx, g, route, start, end, visibleLimit, physicalLimit, ts, reverse) } // RawScanAt cannot enforce physicalLimit, so report truncation and let @@ -1002,6 +1223,27 @@ func (s *ShardStore) scanRouteAtDirectionPhysicalLimit( return nil, true, nil } +func (s *ShardStore) scanReadyLeaderPhysicalLimit( + ctx context.Context, + g *ShardGroup, + route distribution.Route, + start []byte, + end []byte, + visibleLimit int, + physicalLimit int, + ts uint64, + reverse bool, +) ([]*store.KVPair, bool, error) { + route, err := s.targetReadyRouteForRange(ctx, g, route, start, end) + if err != nil { + return nil, false, err + } + if routeHasStagedVisibility(route) { + return nil, true, nil + } + return s.scanRouteAtLeaderPhysicalLimit(ctx, g, route, start, end, visibleLimit, physicalLimit, ts, reverse) +} + func scanLocalPhysicalLimit( ctx context.Context, st store.MVCCStore, @@ -1052,6 +1294,11 @@ func (s *ShardStore) scanRouteLocal( ts uint64, reverse bool, ) ([]*store.KVPair, error) { + var err error + route, err = s.targetReadyRouteForRange(ctx, g, route, start, end) + if err != nil { + return nil, err + } if routeHasStagedVisibility(route) { return s.scanRouteWithStagedVisibility(ctx, g, route, start, end, limit, ts, reverse) } @@ -1074,6 +1321,11 @@ func (s *ShardStore) scanRouteAtLeaderPhysicalLimit( ts uint64, reverse bool, ) ([]*store.KVPair, bool, error) { + var err error + route, err = s.targetReadyRouteForRange(ctx, g, route, start, end) + if err != nil { + return nil, false, err + } kvs, limitReached, err := scanLocalPhysicalLimit(ctx, g.Store, start, end, visibleLimit, physicalLimit, ts, reverse) if err != nil { return nil, limitReached, errors.WithStack(err) @@ -1097,10 +1349,12 @@ func (s *ShardStore) scanRouteAtLeader( ts uint64, reverse bool, ) ([]*store.KVPair, error) { - var ( - kvs []*store.KVPair - err error - ) + var err error + route, err = s.targetReadyRouteForRange(ctx, g, route, start, end) + if err != nil { + return nil, err + } + var kvs []*store.KVPair switch { case routeHasStagedVisibility(route): kvs, err = s.scanRouteWithStagedVisibility(ctx, g, route, start, end, limit, ts, reverse) @@ -1625,9 +1879,14 @@ func clampScanEnd(end []byte, routeEnd []byte) []byte { func (s *ShardStore) PutAt(ctx context.Context, key []byte, value []byte, commitTS uint64, expireAt uint64) error { route, g, ok := s.routeAndGroupForKey(key) - if !ok || g.Store == nil { + if !ok || g == nil || g.Store == nil { return store.ErrNotSupported } + var err error + route, err = s.targetReadyRouteForRange(ctx, g, route, key, nextScanCursor(key)) + if err != nil { + return err + } if err := ensureRouteWriteTimestampFloor(route, key, commitTS); err != nil { return err } @@ -1639,9 +1898,14 @@ func (s *ShardStore) PutAt(ctx context.Context, key []byte, value []byte, commit func (s *ShardStore) DeleteAt(ctx context.Context, key []byte, commitTS uint64) error { route, g, ok := s.routeAndGroupForKey(key) - if !ok || g.Store == nil { + if !ok || g == nil || g.Store == nil { return store.ErrNotSupported } + var err error + route, err = s.targetReadyRouteForRange(ctx, g, route, key, nextScanCursor(key)) + if err != nil { + return err + } if err := ensureRouteWriteTimestampFloor(route, key, commitTS); err != nil { return err } @@ -1653,9 +1917,14 @@ func (s *ShardStore) DeleteAt(ctx context.Context, key []byte, commitTS uint64) func (s *ShardStore) PutWithTTLAt(ctx context.Context, key []byte, value []byte, commitTS uint64, expireAt uint64) error { route, g, ok := s.routeAndGroupForKey(key) - if !ok || g.Store == nil { + if !ok || g == nil || g.Store == nil { return store.ErrNotSupported } + var err error + route, err = s.targetReadyRouteForRange(ctx, g, route, key, nextScanCursor(key)) + if err != nil { + return err + } if err := ensureRouteWriteTimestampFloor(route, key, commitTS); err != nil { return err } @@ -1667,9 +1936,14 @@ func (s *ShardStore) PutWithTTLAt(ctx context.Context, key []byte, value []byte, func (s *ShardStore) ExpireAt(ctx context.Context, key []byte, expireAt uint64, commitTS uint64) error { route, g, ok := s.routeAndGroupForKey(key) - if !ok || g.Store == nil { + if !ok || g == nil || g.Store == nil { return store.ErrNotSupported } + var err error + route, err = s.targetReadyRouteForRange(ctx, g, route, key, nextScanCursor(key)) + if err != nil { + return err + } if err := ensureRouteWriteTimestampFloor(route, key, commitTS); err != nil { return err } @@ -1714,6 +1988,11 @@ func (s *ShardStore) LatestCommitTSWithReadFence(ctx context.Context, key []byte } func (s *ShardStore) localLatestCommitTS(ctx context.Context, g *ShardGroup, route distribution.Route, key []byte) (uint64, bool, error) { + var err error + route, err = s.targetReadyRouteForRange(ctx, g, route, key, nextScanCursor(key)) + if err != nil { + return 0, false, err + } liveTS, liveExists, err := g.Store.LatestCommitTS(ctx, key) if err != nil { return 0, false, errors.WithStack(err) @@ -2394,7 +2673,7 @@ func (s *ShardStore) ApplyMutations(ctx context.Context, mutations []*store.KVPa if err != nil || group == nil { return err } - if err := s.ensureMutationWriteTimestampFloors(mutations, commitTS); err != nil { + if err := s.verifyMutationRoutes(ctx, mutations, readKeys, commitTS); err != nil { return err } readKeys = s.readKeysWithStagedVisibilityAliases(group, readKeys) @@ -2402,6 +2681,40 @@ func (s *ShardStore) ApplyMutations(ctx context.Context, mutations []*store.KVPa return errors.WithStack(group.Store.ApplyMutations(ctx, mutations, readKeys, startTS, commitTS)) } +func (s *ShardStore) verifyMutationRoutes(ctx context.Context, mutations []*store.KVPairMutation, readKeys [][]byte, commitTS uint64) error { + for _, mut := range mutations { + if err := s.verifyMutationWriteRoute(ctx, mut.Key, commitTS); err != nil { + return err + } + } + for _, key := range readKeys { + route, g, ok := s.routeAndGroupForKey(key) + if !ok || g == nil || g.Store == nil { + return store.ErrNotSupported + } + if err := s.verifyTargetReadinessForRange(ctx, g, route, key, nextScanCursor(key)); err != nil { + return err + } + } + return nil +} + +func (s *ShardStore) verifyMutationWriteRoute(ctx context.Context, key []byte, commitTS uint64) error { + route, g, ok := s.routeAndGroupForKey(key) + if !ok || g == nil || g.Store == nil { + return store.ErrNotSupported + } + var err error + route, err = s.targetReadyRouteForRange(ctx, g, route, key, nextScanCursor(key)) + if err != nil { + return err + } + if err := ensureRouteWriteTimestampFloor(route, key, commitTS); err != nil { + return err + } + return s.ensureS3BucketAuxiliaryWriteTimestampFloor(key, commitTS) +} + // ApplyMutationsRaft is the raft-apply variant; see store.MVCCStore for the // durability contract. Only the FSM may call this method. func (s *ShardStore) ApplyMutationsRaft(ctx context.Context, mutations []*store.KVPairMutation, readKeys [][]byte, startTS, commitTS uint64) error { @@ -2434,28 +2747,6 @@ func ensureRouteWriteTimestampFloor(route distribution.Route, key []byte, commit return errors.Wrapf(ErrRouteWriteTimestampTooLow, "key %q routeKey %q commit_ts=%d floor=%d", key, routeKey(key), commitTS, route.MinWriteTSExclusive) } -func (s *ShardStore) ensureMutationWriteTimestampFloors(mutations []*store.KVPairMutation, commitTS uint64) error { - if commitTS == 0 { - return nil - } - for _, mut := range mutations { - if mut == nil || len(mut.Key) == 0 { - continue - } - route, _, ok := s.routeAndGroupForKey(mut.Key) - if !ok { - return store.ErrNotSupported - } - if err := ensureRouteWriteTimestampFloor(route, mut.Key, commitTS); err != nil { - return err - } - if err := s.ensureS3BucketAuxiliaryWriteTimestampFloor(mut.Key, commitTS); err != nil { - return err - } - } - return nil -} - func (s *ShardStore) ensureS3BucketAuxiliaryWriteTimestampFloor(key []byte, commitTS uint64) error { if s == nil || s.engine == nil || commitTS == 0 { return nil @@ -2544,13 +2835,20 @@ func (s *ShardStore) resolveSingleShardGroup(mutations []*store.KVPairMutation) // DeletePrefixAt applies a prefix delete to every shard in the store. func (s *ShardStore) DeletePrefixAt(ctx context.Context, prefix []byte, excludePrefix []byte, commitTS uint64) error { - if err := s.ensurePrefixWriteTimestampFloors(prefix, commitTS); err != nil { + routes, err := s.verifyPrefixDeleteRoutes(ctx, prefix, commitTS) + if err != nil { return err } - for _, g := range s.groups { + for groupID, g := range s.groups { if g == nil || g.Store == nil { continue } + deleteStagedPrefix := func(stagedPrefix []byte, stagedExcludePrefix []byte) error { + return g.Store.DeletePrefixAt(ctx, stagedPrefix, stagedExcludePrefix, commitTS) + } + if err := deleteStagedVisibilityPrefixes(routesForGroupID(routes, groupID), prefix, excludePrefix, deleteStagedPrefix); err != nil { + return err + } if err := g.Store.DeletePrefixAt(ctx, prefix, excludePrefix, commitTS); err != nil { return errors.WithStack(err) } @@ -2558,25 +2856,19 @@ func (s *ShardStore) DeletePrefixAt(ctx context.Context, prefix []byte, excludeP return nil } -func (s *ShardStore) ensurePrefixWriteTimestampFloors(prefix []byte, commitTS uint64) error { - if s == nil || s.engine == nil || commitTS == 0 { - return nil - } - start, end := routePrefixRange(prefix) - for _, route := range s.engine.GetIntersectingRoutes(start, end) { - if route.MinWriteTSExclusive != 0 && commitTS <= route.MinWriteTSExclusive { - return errors.Wrapf(ErrRouteWriteTimestampTooLow, "prefix %q route range [%q,%q) commit_ts=%d floor=%d", prefix, start, end, commitTS, route.MinWriteTSExclusive) - } - } - return nil -} - // DeletePrefixAtRaft is the raft-apply variant of DeletePrefixAt. func (s *ShardStore) DeletePrefixAtRaft(ctx context.Context, prefix []byte, excludePrefix []byte, commitTS uint64) error { - for _, g := range s.groups { + routes := s.stagedVisibilityRoutesForPrefix(prefix) + for groupID, g := range s.groups { if g == nil || g.Store == nil { continue } + deleteStagedPrefix := func(stagedPrefix []byte, stagedExcludePrefix []byte) error { + return g.Store.DeletePrefixAtRaft(ctx, stagedPrefix, stagedExcludePrefix, commitTS) + } + if err := deleteStagedVisibilityPrefixes(routesForGroupID(routes, groupID), prefix, excludePrefix, deleteStagedPrefix); err != nil { + return err + } if err := g.Store.DeletePrefixAtRaft(ctx, prefix, excludePrefix, commitTS); err != nil { return errors.WithStack(err) } @@ -2599,10 +2891,17 @@ func (s *ShardStore) DeletePrefixAtRaft(ctx context.Context, prefix []byte, excl // is the receiver only when an aggregate (admin / coordinator) path // is replaying a global FLUSHALL, which is not raft-applied. func (s *ShardStore) DeletePrefixAtRaftAt(ctx context.Context, prefix []byte, excludePrefix []byte, commitTS, appliedIndex uint64) error { - for _, g := range s.groups { + routes := s.stagedVisibilityRoutesForPrefix(prefix) + for groupID, g := range s.groups { if g == nil || g.Store == nil { continue } + deleteStagedPrefix := func(stagedPrefix []byte, stagedExcludePrefix []byte) error { + return g.Store.DeletePrefixAtRaftAt(ctx, stagedPrefix, stagedExcludePrefix, commitTS, 0) + } + if err := deleteStagedVisibilityPrefixes(routesForGroupID(routes, groupID), prefix, excludePrefix, deleteStagedPrefix); err != nil { + return err + } // Pass appliedIndex through to every group. In the // single-group call-path (the production raft-apply case) // this is correct: appliedIndex IS that group's raft entry @@ -2620,6 +2919,90 @@ func (s *ShardStore) DeletePrefixAtRaftAt(ctx context.Context, prefix []byte, ex return nil } +func (s *ShardStore) stagedVisibilityRoutesForPrefix(prefix []byte) []distribution.Route { + if s == nil || s.engine == nil { + return nil + } + start, end := routePrefixRange(prefix) + routes := s.engine.GetIntersectingRoutes(start, end) + out := make([]distribution.Route, 0, len(routes)) + for _, route := range routes { + if routeHasStagedVisibility(route) { + out = append(out, route) + } + } + return out +} + +func (s *ShardStore) verifyPrefixDeleteRoutes(ctx context.Context, prefix []byte, commitTS uint64) ([]distribution.Route, error) { + if s == nil || s.engine == nil { + return nil, nil + } + routeStart, routeEnd := routePrefixRange(prefix) + routes := s.engine.GetIntersectingRoutes(routeStart, routeEnd) + if len(routes) == 0 { + return nil, errors.WithStack(ErrRouteCutoverPending) + } + verified := make([]distribution.Route, 0, len(routes)) + for _, route := range routes { + g, ok := s.groupForID(route.GroupID) + if !ok || g == nil || g.Store == nil { + return nil, store.ErrNotSupported + } + proofRoutes, err := s.verifyPrefixDeleteRoute(ctx, g, route, routeStart, routeEnd, commitTS) + if err != nil { + return nil, err + } + verified = append(verified, proofRoutes...) + } + return verified, nil +} + +func (s *ShardStore) verifyPrefixDeleteRoute(ctx context.Context, g *ShardGroup, route distribution.Route, routeStart []byte, routeEnd []byte, commitTS uint64) ([]distribution.Route, error) { + proofRoutes, err := s.targetReadyRoutesForRouteRange(ctx, g, route, routeStart, routeEnd) + if err != nil { + return nil, err + } + for _, proofRoute := range proofRoutes { + if err := verifyRouteWriteFloor(proofRoute, commitTS); err != nil { + return nil, err + } + } + return proofRoutes, nil +} + +func routesForGroupID(routes []distribution.Route, groupID uint64) []distribution.Route { + out := make([]distribution.Route, 0, len(routes)) + for _, route := range routes { + if route.GroupID == groupID { + out = append(out, route) + } + } + return out +} + +func deleteStagedVisibilityPrefixes(routes []distribution.Route, prefix []byte, excludePrefix []byte, deletePrefix func([]byte, []byte) error) error { + seen := make(map[uint64]struct{}) + for _, route := range routes { + if !routeHasStagedVisibility(route) { + continue + } + if _, ok := seen[route.MigrationJobID]; ok { + continue + } + seen[route.MigrationJobID] = struct{}{} + stagedPrefix := distribution.MigrationStagedDataKey(route.MigrationJobID, prefix) + var stagedExcludePrefix []byte + if len(excludePrefix) > 0 { + stagedExcludePrefix = distribution.MigrationStagedDataKey(route.MigrationJobID, excludePrefix) + } + if err := deletePrefix(stagedPrefix, stagedExcludePrefix); err != nil { + return errors.WithStack(err) + } + } + return nil +} + func (s *ShardStore) LastCommitTS() uint64 { var max uint64 for _, g := range s.groups { @@ -2805,6 +3188,20 @@ func (s *ShardStore) groupForKey(key []byte) (*ShardGroup, bool) { return g, ok } +func (s *ShardStore) verifyExplicitGroupRoutesForRange(ctx context.Context, groupID uint64, routes []distribution.Route, start []byte, end []byte) error { + g, ok := s.groupForID(groupID) + if !ok || g == nil || g.Store == nil { + return store.ErrNotSupported + } + routeStart, routeEnd := readinessRouteRangeForScan(start, end) + for _, route := range routes { + if err := s.verifyTargetReadinessForRouteRange(ctx, g, route, routeStart, routeEnd); err != nil { + return err + } + } + return nil +} + func (s *ShardStore) routeAndGroupForKey(key []byte) (distribution.Route, *ShardGroup, bool) { if route, ok := s.stagedVisibilityRouteForS3BucketAuxiliaryKey(key); ok { g, ok := s.groups[route.GroupID] diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 99a82eda0..a10427273 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -2,10 +2,12 @@ package kv import ( "context" + "errors" "fmt" "testing" "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/internal/raftengine" "github.com/bootjp/elastickv/internal/s3keys" pb "github.com/bootjp/elastickv/proto" "github.com/bootjp/elastickv/store" @@ -35,6 +37,84 @@ func newStagedVisibilityShardStore(t *testing.T) (*ShardStore, *ShardGroup) { return NewShardStore(engine, map[uint64]*ShardGroup{1: group}), group } +func applyTargetReadiness(t *testing.T, group *ShardGroup) { + t.Helper() + applyTargetReadinessState(t, group, store.TargetStagedReadinessState{ + JobID: 9, + RouteStart: []byte("a"), + RouteEnd: []byte("z"), + ExpectedCutoverVersion: 2, + MigrationJobID: 9, + MinWriteTSExclusive: 100, + Armed: true, + }) +} + +func applyTargetReadinessState(t *testing.T, group *ShardGroup, state store.TargetStagedReadinessState) { + t.Helper() + writer, ok := group.Store.(store.MigrationTargetReadinessWriter) + require.True(t, ok) + require.NoError(t, writer.ApplyTargetStagedReadiness(context.Background(), state)) +} + +func newReadinessShardStore(t *testing.T, route distribution.RouteDescriptor) (*ShardStore, *ShardGroup) { + t.Helper() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 2, + Routes: []distribution.RouteDescriptor{route}, + })) + group := &ShardGroup{Store: store.NewMVCCStore()} + return NewShardStore(engine, map[uint64]*ShardGroup{route.GroupID: group}), group +} + +type readinessFenceEngine struct { + state raftengine.State + onLinearizableRead func() +} + +func (e *readinessFenceEngine) State() raftengine.State { + if e.state != "" { + return e.state + } + return raftengine.StateFollower +} + +func (e *readinessFenceEngine) Leader() raftengine.LeaderInfo { + return raftengine.LeaderInfo{ID: "leader"} +} + +func (e *readinessFenceEngine) VerifyLeader(context.Context) error { + return nil +} + +func (e *readinessFenceEngine) LinearizableRead(context.Context) (uint64, error) { + if e.onLinearizableRead != nil { + e.onLinearizableRead() + } + return 1, nil +} + +func (e *readinessFenceEngine) Propose(context.Context, []byte) (*raftengine.ProposalResult, error) { + return nil, errors.New("unexpected propose") +} + +func (e *readinessFenceEngine) ProposeAdmin(context.Context, []byte) (*raftengine.ProposalResult, error) { + return nil, errors.New("unexpected propose admin") +} + +func (e *readinessFenceEngine) Status() raftengine.Status { + return raftengine.Status{State: e.State()} +} + +func (e *readinessFenceEngine) Configuration(context.Context) (raftengine.Configuration, error) { + return raftengine.Configuration{}, nil +} + +func (e *readinessFenceEngine) Close() error { + return nil +} + func newStagedVisibilityPebbleShardStore(t *testing.T) (*ShardStore, *ShardGroup) { t.Helper() @@ -85,6 +165,224 @@ func TestShardStoreGetAt_MergesStagedVisibility(t *testing.T) { require.ErrorIs(t, err, store.ErrKeyNotFound) } +func TestShardStoreCommittedVersionAtChecksStagedVisibilityKey(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + rawKey := []byte("k") + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, rawKey), []byte("staged"), 77, 0)) + + landed, err := st.CommittedVersionAt(ctx, rawKey, 77) + require.NoError(t, err) + require.True(t, landed) +} + +func TestShardStoreCommittedVersionAtChecksTargetReadiness(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newReadinessShardStore(t, distribution.RouteDescriptor{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + }) + applyTargetReadiness(t, group) + require.NoError(t, group.Store.PutAt(ctx, []byte("k"), []byte("live"), 77, 0)) + + landed, err := st.CommittedVersionAt(ctx, []byte("k"), 77) + require.False(t, landed) + require.ErrorIs(t, err, ErrRouteCutoverPending) +} + +func TestShardStoreCommittedVersionAtRechecksTargetReadinessAfterFence(t *testing.T) { + ctx := context.Background() + st, group := newReadinessShardStore(t, distribution.RouteDescriptor{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + }) + group.Engine = &readinessFenceEngine{ + onLinearizableRead: func() { + applyTargetReadiness(t, group) + }, + } + require.NoError(t, group.Store.PutAt(ctx, []byte("k"), []byte("live"), 77, 0)) + + landed, err := st.CommittedVersionAt(ctx, []byte("k"), 77) + require.False(t, landed) + require.ErrorIs(t, err, ErrRouteCutoverPending) +} + +func TestShardStoreTargetReadinessFailsClosedWithoutDescriptorProof(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newReadinessShardStore(t, distribution.RouteDescriptor{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + }) + applyTargetReadiness(t, group) + + _, err := st.GetAt(ctx, []byte("k"), 120) + require.ErrorIs(t, err, ErrRouteCutoverPending) + + err = st.PutAt(ctx, []byte("k"), []byte("v"), 120, 0) + require.ErrorIs(t, err, ErrRouteCutoverPending) +} + +func TestShardStoreTargetReadinessNormalizesInternalKeyRange(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newReadinessShardStore(t, distribution.RouteDescriptor{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + }) + applyTargetReadiness(t, group) + + itemKey := store.ListItemKey([]byte("b"), 0) + require.NoError(t, group.Store.PutAt(ctx, itemKey, []byte("item"), 120, 0)) + + _, err := st.GetAt(ctx, itemKey, 120) + require.ErrorIs(t, err, ErrRouteCutoverPending) +} + +func TestShardStoreTargetReadinessAcceptsStagedDescriptor(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + applyTargetReadiness(t, group) + rawKey := []byte("k") + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, rawKey), []byte("staged"), 120, 0)) + + got, err := st.GetAt(ctx, rawKey, 130) + require.NoError(t, err) + require.Equal(t, []byte("staged"), got) +} + +func TestShardStoreTargetReadinessAcceptsClearedDescriptorAndRetainsFloor(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newReadinessShardStore(t, distribution.RouteDescriptor{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + MinWriteTSExclusive: 100, + }) + applyTargetReadiness(t, group) + require.NoError(t, group.Store.PutAt(ctx, []byte("k"), []byte("live"), 120, 0)) + + got, err := st.GetAt(ctx, []byte("k"), 130) + require.NoError(t, err) + require.Equal(t, []byte("live"), got) + + err = st.PutAt(ctx, []byte("k"), []byte("low"), 100, 0) + require.ErrorIs(t, err, ErrRouteWriteBelowFloor) + + err = st.ExpireAt(ctx, []byte("k"), 200, 100) + require.ErrorIs(t, err, ErrRouteWriteBelowFloor) + + err = st.ApplyMutations(ctx, []*store.KVPairMutation{{ + Op: store.OpTypePut, + Key: []byte("n"), + Value: []byte("low"), + }}, nil, 0, 100) + require.ErrorIs(t, err, ErrRouteWriteBelowFloor) + + err = st.PutAt(ctx, []byte("k"), []byte("ok"), 101, 0) + require.NoError(t, err) + + err = st.ApplyMutations(ctx, []*store.KVPairMutation{{ + Op: store.OpTypePut, + Key: []byte("n"), + Value: []byte("ok"), + }}, nil, 0, 101) + require.NoError(t, err) +} + +func TestShardStoreTargetReadinessRejectsClearedDescriptorBeforeCutoverVersion(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + MinWriteTSExclusive: 100, + }}, + })) + group := &ShardGroup{Store: store.NewMVCCStore()} + st := NewShardStore(engine, map[uint64]*ShardGroup{1: group}) + applyTargetReadiness(t, group) + require.NoError(t, group.Store.PutAt(ctx, []byte("k"), []byte("live"), 120, 0)) + + _, err := st.GetAt(ctx, []byte("k"), 130) + require.ErrorIs(t, err, ErrRouteCutoverPending) +} + +func TestShardStoreTargetReadinessUsesSingleCatalogSnapshotProof(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + MinWriteTSExclusive: 100, + }}, + })) + staleRoute, ok := engine.GetRoute([]byte("b")) + require.True(t, ok) + group := &ShardGroup{Store: store.NewMVCCStore()} + shards := NewShardStore(engine, map[uint64]*ShardGroup{1: group}) + applyTargetReadiness(t, group) + require.NoError(t, group.Store.PutAt(ctx, []byte("b"), []byte("live-old"), 80, 0)) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("b")), []byte("staged-new"), 120, 0)) + + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 2, + Routes: []distribution.RouteDescriptor{{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + MinWriteTSExclusive: 100, + }}, + })) + + got, err := shards.localGetAt(ctx, group, staleRoute, []byte("b"), 130) + require.NoError(t, err) + require.Equal(t, []byte("staged-new"), got) +} + func TestShardStoreGetAt_MergesStagedVisibilityPebbleExactKey(t *testing.T) { t.Parallel() @@ -101,6 +399,199 @@ func TestShardStoreGetAt_MergesStagedVisibilityPebbleExactKey(t *testing.T) { require.Equal(t, []byte("staged-new"), got) } +func TestShardStoreDeletePrefixAtUsesReadinessProofRouteFloor(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + }}, + })) + group := &ShardGroup{Store: store.NewMVCCStore()} + shards := NewShardStore(engine, map[uint64]*ShardGroup{1: group}) + applyTargetReadiness(t, group) + require.NoError(t, group.Store.PutAt(ctx, []byte("b"), []byte("live"), 80, 0)) + + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 2, + Routes: []distribution.RouteDescriptor{{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + MinWriteTSExclusive: 100, + }}, + })) + + err := shards.DeletePrefixAt(ctx, []byte("b"), nil, 100) + require.ErrorIs(t, err, ErrRouteWriteBelowFloor) + + got, getErr := group.Store.GetAt(ctx, []byte("b"), ^uint64(0)) + require.NoError(t, getErr) + require.Equal(t, []byte("live"), got) +} + +func TestShardStoreDeletePrefixAtUsesReadinessProofRouteForStagedCleanup(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + MinWriteTSExclusive: 100, + }}, + })) + group := &ShardGroup{Store: store.NewMVCCStore()} + shards := NewShardStore(engine, map[uint64]*ShardGroup{1: group}) + applyTargetReadiness(t, group) + stagedKey := distribution.MigrationStagedDataKey(9, []byte("b")) + require.NoError(t, group.Store.PutAt(ctx, stagedKey, []byte("staged"), 120, 0)) + + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 2, + Routes: []distribution.RouteDescriptor{{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + MinWriteTSExclusive: 100, + }}, + })) + + require.NoError(t, shards.DeletePrefixAt(ctx, []byte("b"), nil, 130)) + + _, err := group.Store.GetAt(ctx, stagedKey, 140) + require.ErrorIs(t, err, store.ErrKeyNotFound) + _, err = shards.GetAt(ctx, []byte("b"), 140) + require.ErrorIs(t, err, store.ErrKeyNotFound) +} + +func TestShardStoreExplicitGroupReadUsesRouteProofForReadiness(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newReadinessShardStore(t, distribution.RouteDescriptor{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 42, + State: distribution.RouteStateActive, + MinWriteTSExclusive: 100, + }) + applyTargetReadiness(t, group) + require.NoError(t, group.Store.PutAt(ctx, []byte("b"), []byte("live"), 120, 0)) + require.NoError(t, group.Store.PutAt(ctx, []byte("c"), []byte("scan"), 121, 0)) + + got, err := st.GetGroupAt(ctx, 42, []byte("b"), 130) + require.NoError(t, err) + require.Equal(t, []byte("live"), got) + + kvs, err := st.ScanGroupAt(ctx, 42, []byte("a"), []byte("z"), 10, 130) + require.NoError(t, err) + require.Len(t, kvs, 2) +} + +func TestShardStoreScanGroupAtChecksEveryCoveredRoute(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 2, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("m"), + GroupID: 42, + State: distribution.RouteStateActive, + MinWriteTSExclusive: 100, + }, + { + RouteID: 2, + Start: []byte("m"), + End: []byte("z"), + GroupID: 42, + State: distribution.RouteStateActive, + }, + }, + })) + group := &ShardGroup{Store: store.NewMVCCStore()} + st := NewShardStore(engine, map[uint64]*ShardGroup{42: group}) + applyTargetReadinessState(t, group, store.TargetStagedReadinessState{ + JobID: 9, + RouteStart: []byte("m"), + RouteEnd: []byte("z"), + ExpectedCutoverVersion: 2, + MigrationJobID: 9, + MinWriteTSExclusive: 100, + Armed: true, + }) + require.NoError(t, group.Store.PutAt(ctx, []byte("b"), []byte("left"), 120, 0)) + require.NoError(t, group.Store.PutAt(ctx, []byte("x"), []byte("right"), 120, 0)) + + _, err := st.ScanGroupAt(ctx, 42, []byte("a"), []byte("z"), 10, 130) + require.ErrorIs(t, err, ErrRouteCutoverPending) +} + +func TestShardStoreScanGroupAtSplitsCoveredRoutesForStagedVisibility(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 2, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("m"), + GroupID: 42, + State: distribution.RouteStateActive, + }, + { + RouteID: 2, + Start: []byte("m"), + End: []byte("z"), + GroupID: 42, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 10, + }, + }, + })) + group := &ShardGroup{Store: store.NewMVCCStore()} + st := NewShardStore(engine, map[uint64]*ShardGroup{42: group}) + require.NoError(t, group.Store.PutAt(ctx, []byte("b"), []byte("left-live"), 110, 0)) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(10, []byte("x")), []byte("right-staged"), 120, 0)) + + kvs, err := st.ScanGroupAt(ctx, 42, []byte("a"), []byte("z"), 10, 130) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: []byte("b"), Value: []byte("left-live")}, + {Key: []byte("x"), Value: []byte("right-staged")}, + }, kvs) +} + func TestShardStoreGetAt_MergesStagedVisibilityForS3BucketAuxiliary(t *testing.T) { t.Parallel() @@ -144,6 +635,146 @@ func TestShardStoreGetAt_MergesStagedVisibilityForS3BucketAuxiliary(t *testing.T } } +func TestShardStorePhysicalLimitScanChecksTargetReadiness(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newReadinessShardStore(t, distribution.RouteDescriptor{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + }) + applyTargetReadiness(t, group) + + userKey := []byte("b") + start := store.ListItemKey(userKey, 0) + end := store.ListItemKey(userKey, 2) + require.NoError(t, group.Store.PutAt(ctx, start, []byte("v"), 120, 0)) + + _, _, err := st.ScanAtPhysicalLimit(ctx, start, end, 10, 10, 130) + require.ErrorIs(t, err, ErrRouteCutoverPending) +} + +func TestShardStorePhysicalLimitScanRechecksTargetReadinessAfterFence(t *testing.T) { + ctx := context.Background() + st, group := newReadinessShardStore(t, distribution.RouteDescriptor{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + }) + group.Engine = &readinessFenceEngine{ + state: raftengine.StateLeader, + onLinearizableRead: func() { + applyTargetReadiness(t, group) + }, + } + + userKey := []byte("b") + start := store.ListItemKey(userKey, 0) + end := store.ListItemKey(userKey, 2) + require.NoError(t, group.Store.PutAt(ctx, start, []byte("v"), 120, 0)) + + _, _, err := st.ScanAtPhysicalLimit(ctx, start, end, 10, 10, 130) + require.ErrorIs(t, err, ErrRouteCutoverPending) +} + +func TestShardStoreS3ManifestScanChecksTargetReadinessInRouteSpace(t *testing.T) { + t.Parallel() + + ctx := context.Background() + start := s3keys.ObjectManifestScanStart("bucket-a", 1, "z/") + end := prefixScanEnd(start) + routeStart, routeEnd, ok := s3keys.ManifestScanRouteBounds(start, end) + require.True(t, ok) + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 2, + Routes: []distribution.RouteDescriptor{{ + RouteID: 1, + Start: routeStart, + End: routeEnd, + GroupID: 1, + State: distribution.RouteStateActive, + }}, + })) + group := &ShardGroup{Store: store.NewMVCCStore()} + st := NewShardStore(engine, map[uint64]*ShardGroup{1: group}) + applyTargetReadinessState(t, group, store.TargetStagedReadinessState{ + JobID: 9, + RouteStart: routeStart, + RouteEnd: routeEnd, + ExpectedCutoverVersion: 2, + MigrationJobID: 9, + MinWriteTSExclusive: 100, + Armed: true, + }) + + key := s3keys.ObjectManifestKey("bucket-a", 1, "z/object-0") + require.NoError(t, group.Store.PutAt(ctx, key, []byte("manifest"), 120, 0)) + + _, err := st.ScanAt(ctx, start, end, 10, 130) + require.ErrorIs(t, err, ErrRouteCutoverPending) + _, _, err = st.ScanAtPhysicalLimit(ctx, start, end, 10, 10, 130) + require.ErrorIs(t, err, ErrRouteCutoverPending) +} + +func TestShardStoreTargetReadinessSkipsNonOverlappingGuardForScannedRoute(t *testing.T) { + t.Parallel() + + ctx := context.Background() + start := s3keys.ObjectManifestScanStart("bucket-a", 1, "a/") + end := s3keys.ObjectManifestScanStart("bucket-a", 1, "z/") + routeStart, routeEnd, ok := s3keys.ManifestScanRouteBounds(start, end) + require.True(t, ok) + routeBoundary, _, ok := s3keys.ManifestScanRouteBounds( + s3keys.ObjectManifestScanStart("bucket-a", 1, "m/"), + end, + ) + require.True(t, ok) + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 2, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: routeStart, + End: routeBoundary, + GroupID: 1, + State: distribution.RouteStateActive, + MinWriteTSExclusive: 100, + }, + { + RouteID: 2, + Start: routeBoundary, + End: routeEnd, + GroupID: 1, + State: distribution.RouteStateActive, + }, + }, + })) + group := &ShardGroup{Store: store.NewMVCCStore()} + st := NewShardStore(engine, map[uint64]*ShardGroup{1: group}) + applyTargetReadinessState(t, group, store.TargetStagedReadinessState{ + JobID: 9, + RouteStart: routeStart, + RouteEnd: routeBoundary, + ExpectedCutoverVersion: 2, + MigrationJobID: 9, + MinWriteTSExclusive: 100, + Armed: true, + }) + + kvs, err := st.ScanAt(ctx, start, end, 10, 130) + require.NoError(t, err) + require.Empty(t, kvs) +} + func TestShardStoreRejectsS3BucketAuxiliaryWriteAtMigrationTimestampFloor(t *testing.T) { t.Parallel() @@ -195,6 +826,132 @@ func TestShardStoreRejectsS3BucketAuxiliaryWriteAtMigrationTimestampFloor(t *tes } } +func TestShardStoreExplicitGroupS3ManifestScanUsesRouteProofForReadiness(t *testing.T) { + t.Parallel() + + ctx := context.Background() + start := s3keys.ObjectManifestScanStart("bucket-a", 1, "z/") + end := prefixScanEnd(start) + routeStart, routeEnd, ok := s3keys.ManifestScanRouteBounds(start, end) + require.True(t, ok) + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 2, + Routes: []distribution.RouteDescriptor{{ + RouteID: 1, + Start: routeStart, + End: routeEnd, + GroupID: 42, + State: distribution.RouteStateActive, + MinWriteTSExclusive: 100, + }}, + })) + group := &ShardGroup{Store: store.NewMVCCStore()} + st := NewShardStore(engine, map[uint64]*ShardGroup{42: group}) + applyTargetReadinessState(t, group, store.TargetStagedReadinessState{ + JobID: 9, + RouteStart: routeStart, + RouteEnd: routeEnd, + ExpectedCutoverVersion: 2, + MigrationJobID: 9, + MinWriteTSExclusive: 100, + Armed: true, + }) + + key := s3keys.ObjectManifestKey("bucket-a", 1, "z/object-0") + require.NoError(t, group.Store.PutAt(ctx, key, []byte("manifest"), 120, 0)) + + kvs, err := st.ScanGroupAt(ctx, 42, start, end, 10, 130) + require.NoError(t, err) + require.Len(t, kvs, 1) + require.Equal(t, key, kvs[0].Key) +} + +func TestShardStoreDeletePrefixAtChecksTargetReadinessBeforeDelete(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newReadinessShardStore(t, distribution.RouteDescriptor{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + }) + applyTargetReadiness(t, group) + require.NoError(t, group.Store.PutAt(ctx, []byte("b"), []byte("v"), 1, 0)) + + err := st.DeletePrefixAt(ctx, []byte("b"), nil, 120) + require.ErrorIs(t, err, ErrRouteCutoverPending) + + got, getErr := group.Store.GetAt(ctx, []byte("b"), ^uint64(0)) + require.NoError(t, getErr) + require.Equal(t, []byte("v"), got) +} + +func TestShardStoreDeletePrefixAtFailsClosedWithoutRouteProof(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 2, + Routes: []distribution.RouteDescriptor{{ + RouteID: 1, + Start: []byte("m"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + }}, + })) + group := &ShardGroup{Store: store.NewMVCCStore()} + st := NewShardStore(engine, map[uint64]*ShardGroup{1: group}) + require.NoError(t, group.Store.PutAt(ctx, []byte("b"), []byte("v"), 1, 0)) + + err := st.DeletePrefixAt(ctx, []byte("b"), nil, 120) + require.ErrorIs(t, err, ErrRouteCutoverPending) + + got, getErr := group.Store.GetAt(ctx, []byte("b"), ^uint64(0)) + require.NoError(t, getErr) + require.Equal(t, []byte("v"), got) +} + +func TestShardStoreDeletePrefixAtTombstonesStagedVisibilityRows(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + rawKey := []byte("b") + stagedKey := distribution.MigrationStagedDataKey(9, rawKey) + require.NoError(t, group.Store.PutAt(ctx, stagedKey, []byte("staged"), 120, 0)) + + require.NoError(t, st.DeletePrefixAt(ctx, rawKey, nil, 130)) + + _, err := st.GetAt(ctx, rawKey, 140) + require.ErrorIs(t, err, store.ErrKeyNotFound) + _, err = group.Store.GetAt(ctx, stagedKey, 140) + require.ErrorIs(t, err, store.ErrKeyNotFound) +} + +func TestShardStoreDeletePrefixAtRaftAtAdvancesApplyIndexOnlyOnLiveDelete(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + recording := &deletePrefixIndexRecordingStore{MVCCStore: group.Store} + group.Store = recording + rawKey := []byte("b") + stagedKey := distribution.MigrationStagedDataKey(9, rawKey) + require.NoError(t, group.Store.PutAt(ctx, stagedKey, []byte("staged"), 120, 0)) + + require.NoError(t, st.DeletePrefixAtRaftAt(ctx, rawKey, nil, 130, 88)) + + require.Equal(t, []uint64{0, 88}, recording.deletePrefixIndexes) + require.Equal(t, stagedKey, recording.deletePrefixPrefixes[0]) + require.Equal(t, rawKey, recording.deletePrefixPrefixes[1]) +} + func TestShardStoreStagedVisibilityReadTSCompacted(t *testing.T) { t.Parallel() @@ -255,6 +1012,23 @@ func TestShardStoreScanAndLatestCommitTS_MergeStagedVisibility(t *testing.T) { require.Equal(t, uint64(40), ts) } +func TestShardStoreStagedVisibilityPreservesCompactionErrors(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + retention, ok := group.Store.(store.RetentionController) + require.True(t, ok) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("b")), []byte("staged"), 10, 0)) + retention.SetMinRetainedTS(20) + + _, err := st.GetAt(ctx, []byte("b"), 15) + require.ErrorIs(t, err, store.ErrReadTSCompacted) + + _, err = st.ScanAt(ctx, []byte("a"), []byte("z"), 10, 15) + require.ErrorIs(t, err, store.ErrReadTSCompacted) +} + func TestShardStoreRouteFilteredLeaderScanUsesStagedVisibility(t *testing.T) { t.Parallel() diff --git a/kv/shard_store_txn_lock_test.go b/kv/shard_store_txn_lock_test.go index d30491c64..dbed0364c 100644 --- a/kv/shard_store_txn_lock_test.go +++ b/kv/shard_store_txn_lock_test.go @@ -114,6 +114,36 @@ func TestShardStoreGetAt_ReturnsTxnLockedForPendingLock(t *testing.T) { require.True(t, errors.Is(err, ErrTxnLocked), "expected ErrTxnLocked, got %v", err) } +func TestShardStoreLeaderGetAtChecksReadinessBeforeResolvingPendingLock(t *testing.T) { + t.Parallel() + + ctx := context.Background() + route := distribution.RouteDescriptor{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + } + shardStore, group := newReadinessShardStore(t, route) + applyTargetReadiness(t, group) + + lockTS := uint64(1) + key := []byte("b") + require.NoError(t, group.Store.PutAt(ctx, txnLockKey(key), encodeTxnLock(txnLock{ + StartTS: lockTS, + PrimaryKey: key, + }), lockTS, 0)) + + resolvedRoute, _, ok := shardStore.routeAndGroupForKey(key) + require.True(t, ok) + _, err := shardStore.leaderGetAt(ctx, group, resolvedRoute, key, ^uint64(0)) + require.ErrorIs(t, err, ErrRouteCutoverPending) + + _, lockErr := group.Store.GetAt(ctx, txnLockKey(key), ^uint64(0)) + require.NoError(t, lockErr) +} + func TestShardStoreGetAt_ReturnsTxnLockedForPendingCrossShardTxn(t *testing.T) { t.Parallel() @@ -588,6 +618,104 @@ func TestShardStoreScanAt_ResolvesCommittedSecondaryLockWithoutCommittedValue(t require.Equal(t, []byte("v2"), kvs[0].Value) } +func TestShardStorePhysicalLimitScanPreservesRouteProofDuringLockResolution(t *testing.T) { + t.Parallel() + + ctx := context.Background() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 2, + Routes: []distribution.RouteDescriptor{{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + MinWriteTSExclusive: 100, + }}, + })) + + st1 := store.NewMVCCStore() + fsm := NewKvFSMWithHLC(st1, NewHLC(), WithRouteHistory(WrapDistributionEngine(engine), 1)) + applyFSM, ok := fsm.(*kvFSM) + require.True(t, ok) + txn := &localApplyTransactional{fsm: applyFSM} + + groups := map[uint64]*ShardGroup{ + 1: {Store: st1, Txn: txn}, + } + shardStore := NewShardStore(engine, groups) + applyTargetReadiness(t, groups[1]) + + startTS := uint64(101) + commitTS := uint64(120) + primaryKey := []byte("b") + secondaryKey := []byte("c") + require.NoError(t, st1.PutAt(ctx, secondaryKey, []byte("old"), 1, 0)) + prepare := &pb.Request{ + IsTxn: true, + Phase: pb.Phase_PREPARE, + Ts: startTS, + Mutations: []*pb.Mutation{ + {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: primaryKey, LockTTLms: defaultTxnLockTTLms})}, + {Op: pb.Op_PUT, Key: primaryKey, Value: []byte("vp")}, + {Op: pb.Op_PUT, Key: secondaryKey, Value: []byte("vs")}, + }, + } + _, err := groups[1].Txn.Commit(ctx, []*pb.Request{prepare}) + require.NoError(t, err) + commitPrimary := &pb.Request{ + IsTxn: true, + Phase: pb.Phase_COMMIT, + Ts: startTS, + Mutations: []*pb.Mutation{ + {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: primaryKey, CommitTS: commitTS})}, + {Op: pb.Op_PUT, Key: primaryKey}, + }, + } + _, err = groups[1].Txn.Commit(ctx, []*pb.Request{commitPrimary}) + require.NoError(t, err) + + route, ok := engine.GetRoute(secondaryKey) + require.True(t, ok) + kvs, _, err := shardStore.scanRouteAtLeaderPhysicalLimit(ctx, groups[1], route, []byte(""), nil, 10, 10, commitTS, false) + require.NoError(t, err) + got := map[string]string{} + for _, kvp := range kvs { + got[string(kvp.Key)] = string(kvp.Value) + } + require.Equal(t, "vs", got[string(secondaryKey)]) + + _, lockErr := st1.GetAt(ctx, txnLockKey(secondaryKey), ^uint64(0)) + require.ErrorIs(t, lockErr, store.ErrKeyNotFound) +} + +type localApplyTransactional struct { + fsm *kvFSM +} + +func (t *localApplyTransactional) Commit(ctx context.Context, reqs []*pb.Request) (*TransactionResponse, error) { + for _, req := range reqs { + if err := t.fsm.handleTxnRequest(ctx, req, txnRequestResolveTS(req)); err != nil { + return nil, err + } + } + return &TransactionResponse{}, nil +} + +func (t *localApplyTransactional) Abort(ctx context.Context, reqs []*pb.Request) (*TransactionResponse, error) { + return t.Commit(ctx, reqs) +} + +func txnRequestResolveTS(req *pb.Request) uint64 { + meta, _, err := extractTxnMeta(req.GetMutations()) + if err == nil && meta.CommitTS != 0 { + return meta.CommitTS + } + return req.GetTs() +} + func TestShardStoreScanAt_FiltersTxnInternalKeysWithoutRaft(t *testing.T) { t.Parallel() diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index d943302ad..cb2868d6f 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -1039,6 +1039,9 @@ func (c *ShardedCoordinator) dispatchDelPrefixBroadcast(ctx context.Context, isT if err := c.rejectWriteTimestampFloorDelPrefixes(elems, ts); err != nil { return nil, err } + if err := c.rejectDelPrefixesWithoutTargetReadinessProof(ctx, elems, ts); err != nil { + return nil, err + } requests := make([]*pb.Request, 0, len(elems)) for _, elem := range elems { requests = append(requests, &pb.Request{ @@ -1102,6 +1105,22 @@ func (c *ShardedCoordinator) rejectWriteFencedDelPrefixes(elems []*Elem[OP]) err return nil } +func (c *ShardedCoordinator) rejectDelPrefixesWithoutTargetReadinessProof(ctx context.Context, elems []*Elem[OP], commitTS uint64) error { + shards, ok := c.store.(*ShardStore) + if !ok || shards == nil { + return nil + } + for _, elem := range elems { + if elem == nil { + continue + } + if _, err := shards.verifyPrefixDeleteRoutes(ctx, elem.Key, commitTS); err != nil { + return errors.Wrapf(err, "prefix %q", elem.Key) + } + } + return nil +} + func (c *ShardedCoordinator) rejectWriteTimestampFloorPointElems(elems []*Elem[OP], commitTS uint64) error { if c == nil || c.engine == nil || commitTS == 0 { return nil @@ -1399,7 +1418,7 @@ type preparedGroup struct { } func (c *ShardedCoordinator) prewriteTxn(ctx context.Context, startTS, commitTS uint64, primaryKey []byte, grouped map[uint64][]*pb.Mutation, gids []uint64, groupedReadKeys map[uint64][][]byte, observedRouteVersion uint64) ([]preparedGroup, error) { - prepareMeta := txnMetaMutation(primaryKey, defaultTxnLockTTLms, 0) + prepareMeta := txnMetaMutation(primaryKey, defaultTxnLockTTLms, commitTS) prepared := make([]preparedGroup, 0, len(gids)) for _, gid := range gids { @@ -2140,6 +2159,9 @@ func (c *ShardedCoordinator) validateReadKeysOnShard(ctx context.Context, gid ui return errors.WithStack(err) } for _, key := range keys { + if err := c.verifyTargetReadinessForReadKeyOnShard(ctx, gid, g, key); err != nil { + return err + } ts, exists, err := c.latestCommitTSForReadKeyOnShard(ctx, gid, g, key) if err != nil { return errors.WithStack(err) @@ -2151,6 +2173,43 @@ func (c *ShardedCoordinator) validateReadKeysOnShard(ctx context.Context, gid ui return nil } +func (c *ShardedCoordinator) verifyTargetReadinessForReadKeyOnShard(ctx context.Context, gid uint64, g *ShardGroup, key []byte) error { + reader, ok := g.Store.(store.MigrationTargetReadinessReader) + if !ok { + return nil + } + states, err := reader.MigrationTargetReadinessStates(ctx) + if err != nil { + return errors.WithStack(err) + } + if len(states) == 0 { + return nil + } + routeStart, routeEnd := readinessRouteRange(key, nextScanCursor(key)) + routes, catalogVersion, proof := c.currentShardRoutesForRouteRange(gid, routeStart, routeEnd) + if targetReadinessStatesSatisfied(states, routes, routeStart, routeEnd, gid, catalogVersion, proof) { + return nil + } + return errors.WithStack(ErrRouteCutoverPending) +} + +func (c *ShardedCoordinator) currentShardRoutesForRouteRange(gid uint64, routeStart []byte, routeEnd []byte) ([]distribution.Route, uint64, bool) { + if c == nil || c.engine == nil { + return nil, 0, false + } + snap, ok := c.engine.Current() + if !ok { + return nil, 0, false + } + routes := make([]distribution.Route, 0) + for _, route := range snap.IntersectingRoutes(routeStart, routeEnd) { + if route.GroupID == gid { + routes = append(routes, route) + } + } + return routes, snap.Version(), true +} + func (c *ShardedCoordinator) latestCommitTSForReadKeyOnShard(ctx context.Context, gid uint64, g *ShardGroup, key []byte) (uint64, bool, error) { liveTS, liveExists, err := g.Store.LatestCommitTS(ctx, key) if err != nil { diff --git a/kv/sharded_coordinator_del_prefix_test.go b/kv/sharded_coordinator_del_prefix_test.go index 26408e36b..28c8f6779 100644 --- a/kv/sharded_coordinator_del_prefix_test.go +++ b/kv/sharded_coordinator_del_prefix_test.go @@ -329,6 +329,70 @@ func TestShardedCoordinatorRejectsDelPrefixIntersectingWriteFencedRoute(t *testi require.Empty(t, g2Txn.requests) } +func TestShardedCoordinatorRejectsDelPrefixBelowRouteFloorBeforeBroadcast(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: []byte("m"), GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: []byte("m"), End: nil, GroupID: 2, State: distribution.RouteStateActive, MinWriteTSExclusive: ^uint64(0)}, + }, + })) + + g1Txn := &recordingTransactional{} + g2Txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: DelPrefix, Key: []byte("z")}}, + }) + require.ErrorIs(t, err, ErrRouteWriteBelowFloor) + require.Empty(t, g1Txn.requests) + require.Empty(t, g2Txn.requests) +} + +func TestShardedCoordinatorRejectsDelPrefixWithoutTargetReadinessProofBeforeBroadcast(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 2, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: []byte("m"), GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: []byte("m"), End: nil, GroupID: 2, State: distribution.RouteStateActive}, + }, + })) + + g1Txn := &recordingTransactional{} + g2Txn := &recordingTransactional{} + g1 := &ShardGroup{Txn: g1Txn, Store: store.NewMVCCStore()} + g2 := &ShardGroup{Txn: g2Txn, Store: store.NewMVCCStore()} + applyTargetReadinessState(t, g2, store.TargetStagedReadinessState{ + JobID: 9, + RouteStart: []byte("m"), + RouteEnd: nil, + ExpectedCutoverVersion: 2, + MigrationJobID: 9, + MinWriteTSExclusive: 100, + Armed: true, + }) + groups := map[uint64]*ShardGroup{1: g1, 2: g2} + shardStore := NewShardStore(engine, groups) + coord := NewShardedCoordinator(engine, groups, 1, NewHLC(), shardStore) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: DelPrefix, Key: []byte("z")}}, + }) + require.ErrorIs(t, err, ErrRouteCutoverPending) + require.Empty(t, g1Txn.requests) + require.Empty(t, g2Txn.requests) +} + func TestShardedCoordinatorRejectsFullRangeDelPrefixWhenRouteIsWriteFenced(t *testing.T) { t.Parallel() diff --git a/kv/sharded_coordinator_txn_test.go b/kv/sharded_coordinator_txn_test.go index ee8cce672..221cbfb61 100644 --- a/kv/sharded_coordinator_txn_test.go +++ b/kv/sharded_coordinator_txn_test.go @@ -297,8 +297,8 @@ func TestShardedCoordinatorDispatchTxn_CrossShardPhasesAndCommitIndex(t *testing require.Equal(t, []byte("b"), prepareMeta2.PrimaryKey) require.Equal(t, defaultTxnLockTTLms, prepareMeta1.LockTTLms) require.Equal(t, defaultTxnLockTTLms, prepareMeta2.LockTTLms) - require.Zero(t, prepareMeta1.CommitTS) - require.Zero(t, prepareMeta2.CommitTS) + require.Greater(t, prepareMeta1.CommitTS, startTS) + require.Equal(t, prepareMeta1.CommitTS, prepareMeta2.CommitTS) require.Equal(t, pb.Phase_COMMIT, g1Commit.Phase) require.Equal(t, pb.Phase_COMMIT, g2Commit.Phase) @@ -317,7 +317,7 @@ func TestShardedCoordinatorDispatchTxn_CrossShardPhasesAndCommitIndex(t *testing require.Equal(t, []byte("b"), commitMeta2.PrimaryKey) require.Zero(t, commitMeta1.LockTTLms) require.Zero(t, commitMeta2.LockTTLms) - require.Greater(t, commitMeta1.CommitTS, startTS) + require.Equal(t, prepareMeta1.CommitTS, commitMeta1.CommitTS) require.Equal(t, commitMeta1.CommitTS, commitMeta2.CommitTS) } @@ -397,6 +397,10 @@ func TestShardedCoordinatorDispatchTxn_UsesProvidedCommitTS(t *testing.T) { commitMeta2 := requestTxnMeta(t, g2Txn.requests[1]) require.Equal(t, commitTS, commitMeta1.CommitTS) require.Equal(t, commitTS, commitMeta2.CommitTS) + prepareMeta1 := requestTxnMeta(t, g1Txn.requests[0]) + prepareMeta2 := requestTxnMeta(t, g2Txn.requests[0]) + require.Equal(t, commitTS, prepareMeta1.CommitTS) + require.Equal(t, commitTS, prepareMeta2.CommitTS) } func TestShardedCoordinatorDispatchTxn_RejectsMigrationTimestampFloor(t *testing.T) { @@ -580,6 +584,7 @@ func TestGroupReadKeysByShardID_RoutesS3BucketAuxiliaryToStagedOwner(t *testing. type stubMVCCStore struct { store.MVCCStore latestTS map[string]uint64 + readiness []store.TargetStagedReadinessState returnErr error } @@ -591,6 +596,10 @@ func (s *stubMVCCStore) LatestCommitTS(_ context.Context, key []byte) (uint64, b return ts, ok, nil } +func (s *stubMVCCStore) MigrationTargetReadinessStates(_ context.Context) ([]store.TargetStagedReadinessState, error) { + return s.readiness, nil +} + // noopEngine satisfies raftengine.Engine for unit tests. // LinearizableRead returns immediately (simulates an already-up-to-date FSM). type noopEngine struct{} @@ -677,6 +686,111 @@ func TestValidateReadOnlyShards_NoConflictWhenKeyUnchanged(t *testing.T) { require.NoError(t, err) } +func TestValidateReadOnlyShards_FailsClosedOnTargetReadiness(t *testing.T) { + t.Parallel() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("m"), + GroupID: 1, + State: distribution.RouteStateActive, + }, + { + RouteID: 2, + Start: []byte("m"), + End: nil, + GroupID: 2, + State: distribution.RouteStateActive, + }, + }, + })) + + readOnlyStore := &stubMVCCStore{ + latestTS: map[string]uint64{ + "x": 5, + }, + readiness: []store.TargetStagedReadinessState{ + { + JobID: 7, + RouteStart: []byte("m"), + RouteEnd: nil, + ExpectedCutoverVersion: 2, + MigrationJobID: 7, + MinWriteTSExclusive: 100, + Armed: true, + }, + }, + } + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {}, + 2: {Store: readOnlyStore, Engine: noopEngine{}}, + }, 1, NewHLC(), nil) + + groupedReadKeys := map[uint64][][]byte{ + 2: {[]byte("x")}, + } + err := coord.validateReadOnlyShards(context.Background(), groupedReadKeys, []uint64{1}, 10) + require.Error(t, err) + require.ErrorIs(t, err, ErrRouteCutoverPending) +} + +func TestValidateReadOnlyShards_ChecksStagedVisibilityLatestCommitTS(t *testing.T) { + t.Parallel() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 2, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("m"), + GroupID: 1, + State: distribution.RouteStateActive, + }, + { + RouteID: 2, + Start: []byte("m"), + End: []byte("z"), + GroupID: 2, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 7, + }, + }, + })) + + readOnlyStore := &stubMVCCStore{ + latestTS: map[string]uint64{ + "x": 5, + string(distribution.MigrationStagedDataKey(7, []byte("x"))): 20, + }, + readiness: []store.TargetStagedReadinessState{ + { + JobID: 7, + RouteStart: []byte("m"), + RouteEnd: []byte("z"), + ExpectedCutoverVersion: 2, + MigrationJobID: 7, + Armed: true, + }, + }, + } + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {}, + 2: {Store: readOnlyStore, Engine: noopEngine{}}, + }, 1, NewHLC(), nil) + + groupedReadKeys := map[uint64][][]byte{ + 2: {[]byte("x")}, + } + err := coord.validateReadOnlyShards(context.Background(), groupedReadKeys, []uint64{1}, 10) + require.ErrorIs(t, err, store.ErrWriteConflict) +} + func TestValidateReadOnlyShards_PropagatesStoreError(t *testing.T) { t.Parallel() engine := distribution.NewEngine() diff --git a/main.go b/main.go index c77369312..ef663e72e 100644 --- a/main.go +++ b/main.go @@ -1829,6 +1829,7 @@ func (g *startupPublicKVGate) blocked() bool { func startupRotationGatedMethod(fullMethod string) bool { switch fullMethod { case pb.Internal_Forward_FullMethodName, + pb.Internal_ApplyTargetStagedReadiness_FullMethodName, pb.AdminForward_Forward_FullMethodName, pb.Distribution_SplitRange_FullMethodName, pb.RaftAdmin_AddVoter_FullMethodName, diff --git a/main_encryption_rotate_on_startup_test.go b/main_encryption_rotate_on_startup_test.go index 5e16348ad..750dc673c 100644 --- a/main_encryption_rotate_on_startup_test.go +++ b/main_encryption_rotate_on_startup_test.go @@ -419,6 +419,7 @@ func TestStartupPublicKVGate_BlocksMutatorsUntilReady(t *testing.T) { pb.RawKV_RawGet_FullMethodName, pb.TransactionalKV_Get_FullMethodName, pb.Internal_Forward_FullMethodName, + pb.Internal_ApplyTargetStagedReadiness_FullMethodName, pb.AdminForward_Forward_FullMethodName, pb.Distribution_SplitRange_FullMethodName, pb.RaftAdmin_AddVoter_FullMethodName, diff --git a/proto/internal.pb.go b/proto/internal.pb.go index 91bcff36b..bfdd7db15 100644 --- a/proto/internal.pb.go +++ b/proto/internal.pb.go @@ -1075,6 +1075,134 @@ func (x *PromoteStagedVersionsResponse) GetMaxPromotedTs() uint64 { return 0 } +type TargetStagedReadinessRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + JobId uint64 `protobuf:"varint,1,opt,name=job_id,json=jobId,proto3" json:"job_id,omitempty"` + RouteStart []byte `protobuf:"bytes,2,opt,name=route_start,json=routeStart,proto3" json:"route_start,omitempty"` + RouteEnd []byte `protobuf:"bytes,3,opt,name=route_end,json=routeEnd,proto3" json:"route_end,omitempty"` + ExpectedCutoverVersion uint64 `protobuf:"varint,4,opt,name=expected_cutover_version,json=expectedCutoverVersion,proto3" json:"expected_cutover_version,omitempty"` + MigrationJobId uint64 `protobuf:"varint,5,opt,name=migration_job_id,json=migrationJobId,proto3" json:"migration_job_id,omitempty"` + MinWriteTsExclusive uint64 `protobuf:"varint,6,opt,name=min_write_ts_exclusive,json=minWriteTsExclusive,proto3" json:"min_write_ts_exclusive,omitempty"` + Armed bool `protobuf:"varint,7,opt,name=armed,proto3" json:"armed,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *TargetStagedReadinessRequest) Reset() { + *x = TargetStagedReadinessRequest{} + mi := &file_internal_proto_msgTypes[14] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *TargetStagedReadinessRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*TargetStagedReadinessRequest) ProtoMessage() {} + +func (x *TargetStagedReadinessRequest) ProtoReflect() protoreflect.Message { + mi := &file_internal_proto_msgTypes[14] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use TargetStagedReadinessRequest.ProtoReflect.Descriptor instead. +func (*TargetStagedReadinessRequest) Descriptor() ([]byte, []int) { + return file_internal_proto_rawDescGZIP(), []int{14} +} + +func (x *TargetStagedReadinessRequest) GetJobId() uint64 { + if x != nil { + return x.JobId + } + return 0 +} + +func (x *TargetStagedReadinessRequest) GetRouteStart() []byte { + if x != nil { + return x.RouteStart + } + return nil +} + +func (x *TargetStagedReadinessRequest) GetRouteEnd() []byte { + if x != nil { + return x.RouteEnd + } + return nil +} + +func (x *TargetStagedReadinessRequest) GetExpectedCutoverVersion() uint64 { + if x != nil { + return x.ExpectedCutoverVersion + } + return 0 +} + +func (x *TargetStagedReadinessRequest) GetMigrationJobId() uint64 { + if x != nil { + return x.MigrationJobId + } + return 0 +} + +func (x *TargetStagedReadinessRequest) GetMinWriteTsExclusive() uint64 { + if x != nil { + return x.MinWriteTsExclusive + } + return 0 +} + +func (x *TargetStagedReadinessRequest) GetArmed() bool { + if x != nil { + return x.Armed + } + return false +} + +type TargetStagedReadinessResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *TargetStagedReadinessResponse) Reset() { + *x = TargetStagedReadinessResponse{} + mi := &file_internal_proto_msgTypes[15] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *TargetStagedReadinessResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*TargetStagedReadinessResponse) ProtoMessage() {} + +func (x *TargetStagedReadinessResponse) ProtoReflect() protoreflect.Message { + mi := &file_internal_proto_msgTypes[15] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use TargetStagedReadinessResponse.ProtoReflect.Descriptor instead. +func (*TargetStagedReadinessResponse) Descriptor() ([]byte, []int) { + return file_internal_proto_rawDescGZIP(), []int{15} +} + var File_internal_proto protoreflect.FileDescriptor const file_internal_proto_rawDesc = "" + @@ -1155,7 +1283,17 @@ const file_internal_proto_rawDesc = "" + "nextCursor\x12\x12\n" + "\x04done\x18\x02 \x01(\bR\x04done\x12#\n" + "\rpromoted_rows\x18\x03 \x01(\x04R\fpromotedRows\x12&\n" + - "\x0fmax_promoted_ts\x18\x04 \x01(\x04R\rmaxPromotedTs*&\n" + + "\x0fmax_promoted_ts\x18\x04 \x01(\x04R\rmaxPromotedTs\"\xa2\x02\n" + + "\x1cTargetStagedReadinessRequest\x12\x15\n" + + "\x06job_id\x18\x01 \x01(\x04R\x05jobId\x12\x1f\n" + + "\vroute_start\x18\x02 \x01(\fR\n" + + "routeStart\x12\x1b\n" + + "\troute_end\x18\x03 \x01(\fR\brouteEnd\x128\n" + + "\x18expected_cutover_version\x18\x04 \x01(\x04R\x16expectedCutoverVersion\x12(\n" + + "\x10migration_job_id\x18\x05 \x01(\x04R\x0emigrationJobId\x123\n" + + "\x16min_write_ts_exclusive\x18\x06 \x01(\x04R\x13minWriteTsExclusive\x12\x14\n" + + "\x05armed\x18\a \x01(\bR\x05armed\"\x1f\n" + + "\x1dTargetStagedReadinessResponse*&\n" + "\x02Op\x12\a\n" + "\x03PUT\x10\x00\x12\a\n" + "\x03DEL\x10\x01\x12\x0e\n" + @@ -1166,13 +1304,14 @@ const file_internal_proto_rawDesc = "" + "\aPREPARE\x10\x01\x12\n" + "\n" + "\x06COMMIT\x10\x02\x12\t\n" + - "\x05ABORT\x10\x032\xfd\x02\n" + + "\x05ABORT\x10\x032\xdc\x03\n" + "\bInternal\x12.\n" + "\aForward\x12\x0f.ForwardRequest\x1a\x10.ForwardResponse\"\x00\x12=\n" + "\fRelayPublish\x12\x14.RelayPublishRequest\x1a\x15.RelayPublishResponse\"\x00\x12T\n" + "\x13ExportRangeVersions\x12\x1b.ExportRangeVersionsRequest\x1a\x1c.ExportRangeVersionsResponse\"\x000\x01\x12R\n" + "\x13ImportRangeVersions\x12\x1b.ImportRangeVersionsRequest\x1a\x1c.ImportRangeVersionsResponse\"\x00\x12X\n" + - "\x15PromoteStagedVersions\x12\x1d.PromoteStagedVersionsRequest\x1a\x1e.PromoteStagedVersionsResponse\"\x00B#Z!github.com/bootjp/elastickv/protob\x06proto3" + "\x15PromoteStagedVersions\x12\x1d.PromoteStagedVersionsRequest\x1a\x1e.PromoteStagedVersionsResponse\"\x00\x12]\n" + + "\x1aApplyTargetStagedReadiness\x12\x1d.TargetStagedReadinessRequest\x1a\x1e.TargetStagedReadinessResponse\"\x00B#Z!github.com/bootjp/elastickv/protob\x06proto3" var ( file_internal_proto_rawDescOnce sync.Once @@ -1187,7 +1326,7 @@ func file_internal_proto_rawDescGZIP() []byte { } var file_internal_proto_enumTypes = make([]protoimpl.EnumInfo, 2) -var file_internal_proto_msgTypes = make([]protoimpl.MessageInfo, 14) +var file_internal_proto_msgTypes = make([]protoimpl.MessageInfo, 16) var file_internal_proto_goTypes = []any{ (Op)(0), // 0: Op (Phase)(0), // 1: Phase @@ -1205,6 +1344,8 @@ var file_internal_proto_goTypes = []any{ (*ImportRangeVersionsResponse)(nil), // 13: ImportRangeVersionsResponse (*PromoteStagedVersionsRequest)(nil), // 14: PromoteStagedVersionsRequest (*PromoteStagedVersionsResponse)(nil), // 15: PromoteStagedVersionsResponse + (*TargetStagedReadinessRequest)(nil), // 16: TargetStagedReadinessRequest + (*TargetStagedReadinessResponse)(nil), // 17: TargetStagedReadinessResponse } var file_internal_proto_depIdxs = []int32{ 0, // 0: Mutation.op:type_name -> Op @@ -1219,13 +1360,15 @@ var file_internal_proto_depIdxs = []int32{ 9, // 9: Internal.ExportRangeVersions:input_type -> ExportRangeVersionsRequest 12, // 10: Internal.ImportRangeVersions:input_type -> ImportRangeVersionsRequest 14, // 11: Internal.PromoteStagedVersions:input_type -> PromoteStagedVersionsRequest - 6, // 12: Internal.Forward:output_type -> ForwardResponse - 8, // 13: Internal.RelayPublish:output_type -> RelayPublishResponse - 10, // 14: Internal.ExportRangeVersions:output_type -> ExportRangeVersionsResponse - 13, // 15: Internal.ImportRangeVersions:output_type -> ImportRangeVersionsResponse - 15, // 16: Internal.PromoteStagedVersions:output_type -> PromoteStagedVersionsResponse - 12, // [12:17] is the sub-list for method output_type - 7, // [7:12] is the sub-list for method input_type + 16, // 12: Internal.ApplyTargetStagedReadiness:input_type -> TargetStagedReadinessRequest + 6, // 13: Internal.Forward:output_type -> ForwardResponse + 8, // 14: Internal.RelayPublish:output_type -> RelayPublishResponse + 10, // 15: Internal.ExportRangeVersions:output_type -> ExportRangeVersionsResponse + 13, // 16: Internal.ImportRangeVersions:output_type -> ImportRangeVersionsResponse + 15, // 17: Internal.PromoteStagedVersions:output_type -> PromoteStagedVersionsResponse + 17, // 18: Internal.ApplyTargetStagedReadiness:output_type -> TargetStagedReadinessResponse + 13, // [13:19] is the sub-list for method output_type + 7, // [7:13] is the sub-list for method input_type 7, // [7:7] is the sub-list for extension type_name 7, // [7:7] is the sub-list for extension extendee 0, // [0:7] is the sub-list for field type_name @@ -1242,7 +1385,7 @@ func file_internal_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_internal_proto_rawDesc), len(file_internal_proto_rawDesc)), NumEnums: 2, - NumMessages: 14, + NumMessages: 16, NumExtensions: 0, NumServices: 1, }, diff --git a/proto/internal.proto b/proto/internal.proto index b45bf8fb5..1b168812a 100644 --- a/proto/internal.proto +++ b/proto/internal.proto @@ -10,6 +10,7 @@ service Internal { rpc ExportRangeVersions(ExportRangeVersionsRequest) returns (stream ExportRangeVersionsResponse) {} rpc ImportRangeVersions(ImportRangeVersionsRequest) returns (ImportRangeVersionsResponse) {} rpc PromoteStagedVersions(PromoteStagedVersionsRequest) returns (PromoteStagedVersionsResponse) {} + rpc ApplyTargetStagedReadiness(TargetStagedReadinessRequest) returns (TargetStagedReadinessResponse) {} } // internal.proto is node to node communication message in raft replication. @@ -144,3 +145,15 @@ message PromoteStagedVersionsResponse { uint64 promoted_rows = 3; uint64 max_promoted_ts = 4; } + +message TargetStagedReadinessRequest { + uint64 job_id = 1; + bytes route_start = 2; + bytes route_end = 3; + uint64 expected_cutover_version = 4; + uint64 migration_job_id = 5; + uint64 min_write_ts_exclusive = 6; + bool armed = 7; +} + +message TargetStagedReadinessResponse {} diff --git a/proto/internal_grpc.pb.go b/proto/internal_grpc.pb.go index ba679ed80..6af68894d 100644 --- a/proto/internal_grpc.pb.go +++ b/proto/internal_grpc.pb.go @@ -19,11 +19,12 @@ import ( const _ = grpc.SupportPackageIsVersion9 const ( - Internal_Forward_FullMethodName = "/Internal/Forward" - Internal_RelayPublish_FullMethodName = "/Internal/RelayPublish" - Internal_ExportRangeVersions_FullMethodName = "/Internal/ExportRangeVersions" - Internal_ImportRangeVersions_FullMethodName = "/Internal/ImportRangeVersions" - Internal_PromoteStagedVersions_FullMethodName = "/Internal/PromoteStagedVersions" + Internal_Forward_FullMethodName = "/Internal/Forward" + Internal_RelayPublish_FullMethodName = "/Internal/RelayPublish" + Internal_ExportRangeVersions_FullMethodName = "/Internal/ExportRangeVersions" + Internal_ImportRangeVersions_FullMethodName = "/Internal/ImportRangeVersions" + Internal_PromoteStagedVersions_FullMethodName = "/Internal/PromoteStagedVersions" + Internal_ApplyTargetStagedReadiness_FullMethodName = "/Internal/ApplyTargetStagedReadiness" ) // InternalClient is the client API for Internal service. @@ -36,6 +37,7 @@ type InternalClient interface { ExportRangeVersions(ctx context.Context, in *ExportRangeVersionsRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[ExportRangeVersionsResponse], error) ImportRangeVersions(ctx context.Context, in *ImportRangeVersionsRequest, opts ...grpc.CallOption) (*ImportRangeVersionsResponse, error) PromoteStagedVersions(ctx context.Context, in *PromoteStagedVersionsRequest, opts ...grpc.CallOption) (*PromoteStagedVersionsResponse, error) + ApplyTargetStagedReadiness(ctx context.Context, in *TargetStagedReadinessRequest, opts ...grpc.CallOption) (*TargetStagedReadinessResponse, error) } type internalClient struct { @@ -105,6 +107,16 @@ func (c *internalClient) PromoteStagedVersions(ctx context.Context, in *PromoteS return out, nil } +func (c *internalClient) ApplyTargetStagedReadiness(ctx context.Context, in *TargetStagedReadinessRequest, opts ...grpc.CallOption) (*TargetStagedReadinessResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(TargetStagedReadinessResponse) + err := c.cc.Invoke(ctx, Internal_ApplyTargetStagedReadiness_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + // InternalServer is the server API for Internal service. // All implementations must embed UnimplementedInternalServer // for forward compatibility. @@ -115,6 +127,7 @@ type InternalServer interface { ExportRangeVersions(*ExportRangeVersionsRequest, grpc.ServerStreamingServer[ExportRangeVersionsResponse]) error ImportRangeVersions(context.Context, *ImportRangeVersionsRequest) (*ImportRangeVersionsResponse, error) PromoteStagedVersions(context.Context, *PromoteStagedVersionsRequest) (*PromoteStagedVersionsResponse, error) + ApplyTargetStagedReadiness(context.Context, *TargetStagedReadinessRequest) (*TargetStagedReadinessResponse, error) mustEmbedUnimplementedInternalServer() } @@ -140,6 +153,9 @@ func (UnimplementedInternalServer) ImportRangeVersions(context.Context, *ImportR func (UnimplementedInternalServer) PromoteStagedVersions(context.Context, *PromoteStagedVersionsRequest) (*PromoteStagedVersionsResponse, error) { return nil, status.Error(codes.Unimplemented, "method PromoteStagedVersions not implemented") } +func (UnimplementedInternalServer) ApplyTargetStagedReadiness(context.Context, *TargetStagedReadinessRequest) (*TargetStagedReadinessResponse, error) { + return nil, status.Error(codes.Unimplemented, "method ApplyTargetStagedReadiness not implemented") +} func (UnimplementedInternalServer) mustEmbedUnimplementedInternalServer() {} func (UnimplementedInternalServer) testEmbeddedByValue() {} @@ -244,6 +260,24 @@ func _Internal_PromoteStagedVersions_Handler(srv interface{}, ctx context.Contex return interceptor(ctx, in, info, handler) } +func _Internal_ApplyTargetStagedReadiness_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(TargetStagedReadinessRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(InternalServer).ApplyTargetStagedReadiness(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: Internal_ApplyTargetStagedReadiness_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(InternalServer).ApplyTargetStagedReadiness(ctx, req.(*TargetStagedReadinessRequest)) + } + return interceptor(ctx, in, info, handler) +} + // Internal_ServiceDesc is the grpc.ServiceDesc for Internal service. // It's only intended for direct use with grpc.RegisterService, // and not to be introspected or modified (even as a copy) @@ -267,6 +301,10 @@ var Internal_ServiceDesc = grpc.ServiceDesc{ MethodName: "PromoteStagedVersions", Handler: _Internal_PromoteStagedVersions_Handler, }, + { + MethodName: "ApplyTargetStagedReadiness", + Handler: _Internal_ApplyTargetStagedReadiness_Handler, + }, }, Streams: []grpc.StreamDesc{ { diff --git a/store/lsm_migration.go b/store/lsm_migration.go index c114b24ce..e744354eb 100644 --- a/store/lsm_migration.go +++ b/store/lsm_migration.go @@ -25,6 +25,9 @@ func (s *pebbleStore) exportVersionsLocked(ctx context.Context, opts ExportVersi if opts.MaxVersions <= 0 { return ExportVersionsResult{Done: true}, nil } + if readTSCompacted(exportRetentionReadTS(opts), s.effectiveMinRetainedTS()) { + return ExportVersionsResult{}, ErrReadTSCompacted + } iter, err := s.db.NewIter(pebbleExportIterOptions(opts)) if err != nil { diff --git a/store/lsm_store.go b/store/lsm_store.go index 96bb16203..369ab55db 100644 --- a/store/lsm_store.go +++ b/store/lsm_store.go @@ -182,19 +182,20 @@ var metaAppliedIndexBytes = []byte(metaAppliedIndex) // (PR #915 round-3) so the snapshot-persist checkpoint cannot rewind // metaAppliedIndex below a concurrent per-Apply value. // 4. mtx – guards the in-memory metadata fields -// (lastCommitTS, minRetainedTS, pendingMinRetainedTS). +// (lastCommitTS, minRetainedTS, pendingMinRetainedTS, migrationReadinessCache). type pebbleStore struct { - db *pebble.DB - cache *pebble.Cache // owns one ref; Unref on Close / reopen. - dbMu sync.RWMutex // guards s.db pointer – see lock ordering above - log *slog.Logger - lastCommitTS uint64 - minRetainedTS uint64 - pendingMinRetainedTS uint64 - mtx sync.RWMutex - applyMu sync.Mutex // serializes ApplyMutations: conflict check → commit - maintenanceMu sync.Mutex - dir string + db *pebble.DB + cache *pebble.Cache // owns one ref; Unref on Close / reopen. + dbMu sync.RWMutex // guards s.db pointer – see lock ordering above + log *slog.Logger + lastCommitTS uint64 + minRetainedTS uint64 + pendingMinRetainedTS uint64 + migrationReadinessCache []TargetStagedReadinessState + mtx sync.RWMutex + applyMu sync.Mutex // serializes ApplyMutations: conflict check → commit + maintenanceMu sync.Mutex + dir string // writeConflicts tracks per-(kind, key_prefix) OCC conflict counts // detected inside ApplyMutations. Polled by the monitoring // WriteConflictCollector; not part of the authoritative OCC path. @@ -340,9 +341,15 @@ func NewPebbleStore(dir string, opts ...PebbleStoreOption) (MVCCStore, error) { cleanupOnInitFail() return nil, err } + readinessStates, err := s.loadPebbleTargetReadinessStatesFromDB() + if err != nil { + cleanupOnInitFail() + return nil, err + } s.lastCommitTS = maxTS s.minRetainedTS = minRetainedTS s.pendingMinRetainedTS = pendingMinRetainedTS + s.migrationReadinessCache = readinessStates return s, nil } @@ -2340,6 +2347,7 @@ func (s *pebbleStore) handleRestorePeekError(err error, header []byte) error { s.lastCommitTS = 0 s.minRetainedTS = 0 s.pendingMinRetainedTS = 0 + s.migrationReadinessCache = nil if setErr := writePebbleUint64(s.db, metaLastCommitTSBytes, 0, pebble.NoSync); setErr != nil { return errors.WithStack(setErr) } @@ -2510,22 +2518,34 @@ func writeNativeSnapshotToTempDir(r io.Reader, tmpDir string, ts uint64) error { // Entries are written to a temporary Pebble directory and only swapped into // place after the CRC32 checksum is verified, preserving the existing store // on failure. -func readStreamingMVCCRestoreHeader(r io.Reader) (io.Reader, hash.Hash32, uint32, uint64, uint64, error) { - expectedChecksum, err := readMVCCSnapshotHeader(r) +func readStreamingMVCCRestoreHeader(r io.Reader) (io.Reader, hash.Hash32, uint32, uint64, uint64, []TargetStagedReadinessState, error) { + expectedChecksum, version, err := readMVCCSnapshotHeader(r) if err != nil { - return nil, nil, 0, 0, 0, err + return nil, nil, 0, 0, 0, nil, err } hash := crc32.NewIEEE() body := io.TeeReader(r, hash) lastCommitTS, minRetainedTS, err := readMVCCSnapshotMetadata(body) if err != nil { - return nil, nil, 0, 0, 0, err + return nil, nil, 0, 0, 0, nil, err + } + readinessStates, err := readMVCCSnapshotReadinessStates(body, version) + if err != nil { + return nil, nil, 0, 0, 0, nil, err } - return body, hash, expectedChecksum, lastCommitTS, minRetainedTS, nil + return body, hash, expectedChecksum, lastCommitTS, minRetainedTS, readinessStates, nil } -func writeStreamingMVCCRestoreTempDB(dir string, body io.Reader, hash hash.Hash32, expectedChecksum uint32, lastCommitTS uint64, minRetainedTS uint64) (string, error) { +func writeStreamingMVCCRestoreTempDB( + dir string, + body io.Reader, + hash hash.Hash32, + expectedChecksum uint32, + lastCommitTS uint64, + minRetainedTS uint64, + readinessStates []TargetStagedReadinessState, +) (string, error) { tmpDir := filepath.Clean(dir) + ".restore-tmp" if err := os.RemoveAll(tmpDir); err != nil { return "", errors.WithStack(err) @@ -2545,6 +2565,10 @@ func writeStreamingMVCCRestoreTempDB(dir string, body io.Reader, hash hash.Hash3 _ = os.RemoveAll(tmpDir) } + if err := writePebbleTargetReadinessStates(tmpDB, readinessStates); err != nil { + cleanupTmp() + return "", err + } if err := writeMVCCEntriesToDB(body, tmpDB); err != nil { cleanupTmp() return "", err @@ -2564,13 +2588,22 @@ func writeStreamingMVCCRestoreTempDB(dir string, body io.Reader, hash hash.Hash3 return tmpDir, nil } +func writePebbleTargetReadinessStates(db *pebble.DB, states []TargetStagedReadinessState) error { + for _, state := range states { + if err := db.Set(migrationReadyKey(state.JobID), encodeTargetStagedReadinessState(state), pebble.NoSync); err != nil { + return errors.WithStack(err) + } + } + return nil +} + func (s *pebbleStore) restoreFromStreamingMVCC(r io.Reader) error { - body, hash, expectedChecksum, lastCommitTS, minRetainedTS, err := readStreamingMVCCRestoreHeader(r) + body, hash, expectedChecksum, lastCommitTS, minRetainedTS, readinessStates, err := readStreamingMVCCRestoreHeader(r) if err != nil { return err } - tmpDir, err := writeStreamingMVCCRestoreTempDB(s.dir, body, hash, expectedChecksum, lastCommitTS, minRetainedTS) + tmpDir, err := writeStreamingMVCCRestoreTempDB(s.dir, body, hash, expectedChecksum, lastCommitTS, minRetainedTS, readinessStates) if err != nil { return err } @@ -2657,9 +2690,15 @@ func (s *pebbleStore) swapInTempDB(tmpDir string) error { cleanupOnSwapFail() return err } + readinessStates, err := s.loadPebbleTargetReadinessStatesFromDB() + if err != nil { + cleanupOnSwapFail() + return err + } s.lastCommitTS = lastCommitTS s.minRetainedTS = minRetainedTS s.pendingMinRetainedTS = pendingMinRetainedTS + s.migrationReadinessCache = readinessStates return nil } diff --git a/store/lsm_store_test.go b/store/lsm_store_test.go index 27b0fe306..5e5a8a366 100644 --- a/store/lsm_store_test.go +++ b/store/lsm_store_test.go @@ -690,6 +690,18 @@ func TestPebbleStore_RestoreFromStreamingMVCC(t *testing.T) { require.NoError(t, src.PutAt(ctx, []byte("key1"), []byte("val1-updated"), 20, 0)) require.NoError(t, src.DeleteAt(ctx, []byte("key2"), 15)) require.NoError(t, src.PutWithTTLAt(ctx, []byte("key3"), []byte("val3"), 30, 9999)) + readyState := TargetStagedReadinessState{ + JobID: 9, + RouteStart: []byte("a"), + RouteEnd: []byte("z"), + ExpectedCutoverVersion: 12, + MigrationJobID: 9, + MinWriteTSExclusive: 100, + Armed: true, + } + readyWriter, ok := src.(MigrationTargetReadinessWriter) + require.True(t, ok) + require.NoError(t, readyWriter.ApplyTargetStagedReadiness(ctx, readyState)) snap, err := src.Snapshot() require.NoError(t, err) @@ -727,6 +739,11 @@ func TestPebbleStore_RestoreFromStreamingMVCC(t *testing.T) { assert.Equal(t, []byte("val3"), val) assert.Equal(t, src.LastCommitTS(), dst.LastCommitTS()) + readyReader, ok := dst.(MigrationTargetReadinessReader) + require.True(t, ok) + states, err := readyReader.MigrationTargetReadinessStates(ctx) + require.NoError(t, err) + require.Equal(t, []TargetStagedReadinessState{readyState}, states) } func TestPebbleStore_RestoreFromStreamingMVCCPreservesMinRetainedTS(t *testing.T) { @@ -775,6 +792,17 @@ func TestPebbleStore_Restore_EmptySnapshot(t *testing.T) { // Pre-populate so we can verify the restore clears the data. require.NoError(t, s.PutAt(ctx, []byte("k"), []byte("v"), 42, 0)) + readyWriter, ok := s.(MigrationTargetReadinessWriter) + require.True(t, ok) + require.NoError(t, readyWriter.ApplyTargetStagedReadiness(ctx, TargetStagedReadinessState{ + JobID: 9, + RouteStart: []byte("a"), + RouteEnd: []byte("z"), + ExpectedCutoverVersion: 12, + MigrationJobID: 9, + MinWriteTSExclusive: 100, + Armed: true, + })) requirePebbleRetentionController(t, s).SetMinRetainedTS(21) require.Equal(t, uint64(42), s.LastCommitTS()) @@ -786,6 +814,11 @@ func TestPebbleStore_Restore_EmptySnapshot(t *testing.T) { assert.ErrorIs(t, err, ErrKeyNotFound) assert.Equal(t, uint64(0), s.LastCommitTS()) assert.Equal(t, uint64(0), requirePebbleRetentionController(t, s).MinRetainedTS()) + readyReader, ok := s.(MigrationTargetReadinessReader) + require.True(t, ok) + states, err := readyReader.MigrationTargetReadinessStates(ctx) + require.NoError(t, err) + assert.Empty(t, states) } // TestPebbleStore_Restore_TruncatedHeader verifies that a partial magic diff --git a/store/migration_promote.go b/store/migration_promote.go index bfcd1af44..403cd9d1a 100644 --- a/store/migration_promote.go +++ b/store/migration_promote.go @@ -101,7 +101,7 @@ func (s *mvccStore) planMemoryPromotionLocked( if err != nil { return ExportVersionsResult{}, nil, PromoteVersionsResult{}, err } - exported, err := s.exportMemoryVersionsLocked(ctx, exportOpts, pos) + exported, err := s.exportVersionsLocked(ctx, exportOpts, pos) if err != nil { return ExportVersionsResult{}, nil, PromoteVersionsResult{}, err } @@ -139,36 +139,6 @@ func (s *mvccStore) MigrationPromotionState(_ context.Context, jobID uint64) (Pr return clonePromotionState(state), ok, nil } -func (s *mvccStore) exportMemoryVersionsLocked(ctx context.Context, opts ExportVersionsOptions, pos exportCursorPosition) (ExportVersionsResult, error) { - result := newExportVersionsResult(opts.MaxVersions) - it := s.tree.Iterator() - if !s.seekMemoryExportStart(&it, opts.StartKey, pos) { - result.Done = true - return result, nil - } - for ok := true; ok; ok = it.Next() { - key, keyOK := it.Key().([]byte) - if err := checkExportKey(ctx, key, keyOK, opts.EndKey); err != nil { - if errors.Is(err, errExportReachedEnd) { - result.Done = true - result.NextCursor = nil - return result, nil - } - return ExportVersionsResult{}, err - } - if !keyOK { - continue - } - done, err := exportMemoryIteratorKey(ctx, opts, pos, key, it.Value(), &result) - if err != nil || !done { - return result, err - } - } - result.Done = true - result.NextCursor = nil - return result, nil -} - func (s *mvccStore) removeVersionLocked(key []byte, commitTS uint64) bool { existing, ok := s.tree.Get(key) if !ok { diff --git a/store/migration_readiness.go b/store/migration_readiness.go new file mode 100644 index 000000000..d26d25d36 --- /dev/null +++ b/store/migration_readiness.go @@ -0,0 +1,246 @@ +package store + +import ( + "bytes" + "context" + "encoding/binary" + "sort" + + "github.com/cockroachdb/errors" + "github.com/cockroachdb/pebble/v2" +) + +const ( + targetReadinessCodecVersion byte = 1 + targetReadinessArmedFlag byte = 1 +) + +func validateTargetStagedReadinessState(state TargetStagedReadinessState) error { + switch { + case state.JobID == 0: + return errors.New("target staged readiness job_id is required") + case state.MigrationJobID == 0: + return errors.New("target staged readiness migration_job_id is required") + case state.RouteEnd != nil && bytes.Compare(state.RouteStart, state.RouteEnd) >= 0: + return errors.New("target staged readiness route range is invalid") + default: + return nil + } +} + +func (s *mvccStore) ApplyTargetStagedReadiness(_ context.Context, state TargetStagedReadinessState) error { + if err := validateTargetStagedReadinessState(state); err != nil { + return err + } + s.mtx.Lock() + defer s.mtx.Unlock() + cloned := cloneTargetStagedReadinessState(state) + s.migrationReadiness[state.JobID] = cloned + s.migrationReadinessCache = upsertTargetStagedReadinessStateCache(s.migrationReadinessCache, cloned) + return nil +} + +func (s *mvccStore) MigrationTargetReadinessStates(_ context.Context) ([]TargetStagedReadinessState, error) { + s.mtx.RLock() + defer s.mtx.RUnlock() + return cloneTargetStagedReadinessStates(s.migrationReadinessCache), nil +} + +func (s *pebbleStore) ApplyTargetStagedReadiness(_ context.Context, state TargetStagedReadinessState) error { + if err := validateTargetStagedReadinessState(state); err != nil { + return err + } + s.dbMu.RLock() + defer s.dbMu.RUnlock() + s.mtx.Lock() + defer s.mtx.Unlock() + if err := s.db.Set(migrationReadyKey(state.JobID), encodeTargetStagedReadinessState(state), s.directApplyWriteOpts()); err != nil { + return errors.WithStack(err) + } + s.migrationReadinessCache = upsertTargetStagedReadinessStateCache(s.migrationReadinessCache, state) + return nil +} + +func (s *pebbleStore) MigrationTargetReadinessStates(_ context.Context) ([]TargetStagedReadinessState, error) { + s.mtx.RLock() + defer s.mtx.RUnlock() + return cloneTargetStagedReadinessStates(s.migrationReadinessCache), nil +} + +func (s *pebbleStore) loadPebbleTargetReadinessStatesFromDB() ([]TargetStagedReadinessState, error) { + iter, err := s.db.NewIter(&pebble.IterOptions{ + LowerBound: []byte(migrationReadyPrefix), + UpperBound: PrefixScanEnd([]byte(migrationReadyPrefix)), + }) + if err != nil { + return nil, errors.WithStack(err) + } + defer iter.Close() + + var out []TargetStagedReadinessState + for valid := iter.First(); valid; valid = iter.Next() { + if !isTargetReadinessRecordKey(iter.Key()) { + continue + } + state, ok := decodeTargetStagedReadinessState(iter.Value()) + if !ok { + return nil, errors.New("corrupt target staged readiness state") + } + out = append(out, state) + } + if err := iter.Error(); err != nil { + return nil, errors.WithStack(err) + } + sortTargetStagedReadinessStates(out) + return out, nil +} + +func isTargetReadinessRecordKey(rawKey []byte) bool { + return len(rawKey) == len(migrationReadyPrefix)+migrationUint64Bytes && + bytes.HasPrefix(rawKey, []byte(migrationReadyPrefix)) +} + +func encodeTargetStagedReadinessState(state TargetStagedReadinessState) []byte { + buf := make([]byte, 0, targetReadinessEncodedSize(state)) + buf = append(buf, targetReadinessCodecVersion) + if state.Armed { + buf = append(buf, targetReadinessArmedFlag) + } else { + buf = append(buf, 0) + } + buf = binary.BigEndian.AppendUint64(buf, state.JobID) + buf = binary.BigEndian.AppendUint64(buf, state.ExpectedCutoverVersion) + buf = binary.BigEndian.AppendUint64(buf, state.MigrationJobID) + buf = binary.BigEndian.AppendUint64(buf, state.MinWriteTSExclusive) + buf = appendVarBytes(buf, state.RouteStart) + if state.RouteEnd == nil { + buf = append(buf, 0) + return buf + } + buf = append(buf, 1) + return appendVarBytes(buf, state.RouteEnd) +} + +func decodeTargetStagedReadinessState(data []byte) (TargetStagedReadinessState, bool) { + state, rest, ok := decodeTargetReadinessHeader(data) + if !ok { + return TargetStagedReadinessState{}, false + } + if !decodeTargetReadinessRange(&state, rest) { + return TargetStagedReadinessState{}, false + } + if err := validateTargetStagedReadinessState(state); err != nil { + return TargetStagedReadinessState{}, false + } + return state, true +} + +func decodeTargetReadinessHeader(data []byte) (TargetStagedReadinessState, []byte, bool) { + if len(data) < 2+4*migrationUint64Bytes || data[0] != targetReadinessCodecVersion { + return TargetStagedReadinessState{}, nil, false + } + if data[1] != 0 && data[1] != targetReadinessArmedFlag { + return TargetStagedReadinessState{}, nil, false + } + state := TargetStagedReadinessState{Armed: data[1] == targetReadinessArmedFlag} + rest := data[2:] + state.JobID = binary.BigEndian.Uint64(rest[:migrationUint64Bytes]) + rest = rest[migrationUint64Bytes:] + state.ExpectedCutoverVersion = binary.BigEndian.Uint64(rest[:migrationUint64Bytes]) + rest = rest[migrationUint64Bytes:] + state.MigrationJobID = binary.BigEndian.Uint64(rest[:migrationUint64Bytes]) + rest = rest[migrationUint64Bytes:] + state.MinWriteTSExclusive = binary.BigEndian.Uint64(rest[:migrationUint64Bytes]) + rest = rest[migrationUint64Bytes:] + return state, rest, true +} + +func decodeTargetReadinessRange(state *TargetStagedReadinessState, data []byte) bool { + routeStart, rest, ok := readVarBytes(data) + if !ok || len(rest) == 0 { + return false + } + state.RouteStart = routeStart + endPresent := rest[0] + rest = rest[1:] + if endPresent == 0 { + return len(rest) == 0 + } + if endPresent != 1 { + return false + } + state.RouteEnd, rest, ok = readVarBytes(rest) + return ok && len(rest) == 0 +} + +func targetReadinessEncodedSize(state TargetStagedReadinessState) int { + size := 2 + 4*migrationUint64Bytes + binary.MaxVarintLen64 + len(state.RouteStart) + 1 + if state.RouteEnd != nil { + size += binary.MaxVarintLen64 + len(state.RouteEnd) + } + return size +} + +func appendVarBytes(dst []byte, value []byte) []byte { + dst = binary.AppendUvarint(dst, lenAsUint64(len(value))) + return append(dst, value...) +} + +func readVarBytes(data []byte) ([]byte, []byte, bool) { + l, n := binary.Uvarint(data) + if n <= 0 { + return nil, nil, false + } + rest := data[n:] + if l > lenAsUint64(len(rest)) { + return nil, nil, false + } + return bytes.Clone(rest[:l]), rest[l:], true +} + +func cloneTargetStagedReadinessState(state TargetStagedReadinessState) TargetStagedReadinessState { + return TargetStagedReadinessState{ + JobID: state.JobID, + RouteStart: bytes.Clone(state.RouteStart), + RouteEnd: bytes.Clone(state.RouteEnd), + ExpectedCutoverVersion: state.ExpectedCutoverVersion, + MigrationJobID: state.MigrationJobID, + MinWriteTSExclusive: state.MinWriteTSExclusive, + Armed: state.Armed, + } +} + +func cloneTargetStagedReadinessStates(states []TargetStagedReadinessState) []TargetStagedReadinessState { + if len(states) == 0 { + return nil + } + out := make([]TargetStagedReadinessState, 0, len(states)) + for _, state := range states { + out = append(out, cloneTargetStagedReadinessState(state)) + } + return out +} + +func upsertTargetStagedReadinessStateCache(cache []TargetStagedReadinessState, state TargetStagedReadinessState) []TargetStagedReadinessState { + cloned := cloneTargetStagedReadinessState(state) + for i := range cache { + if cache[i].JobID == cloned.JobID { + cache[i] = cloned + return cache + } + } + cache = append(cache, cloned) + sortTargetStagedReadinessStates(cache) + return cache +} + +func sortTargetStagedReadinessStates(states []TargetStagedReadinessState) { + sort.Slice(states, func(i, j int) bool { + return states[i].JobID < states[j].JobID + }) +} + +var _ MigrationTargetReadinessReader = (*mvccStore)(nil) +var _ MigrationTargetReadinessReader = (*pebbleStore)(nil) +var _ MigrationTargetReadinessWriter = (*mvccStore)(nil) +var _ MigrationTargetReadinessWriter = (*pebbleStore)(nil) diff --git a/store/migration_versions.go b/store/migration_versions.go index ef3ef8254..9e29297bd 100644 --- a/store/migration_versions.go +++ b/store/migration_versions.go @@ -22,7 +22,9 @@ const ( migrationMetadataVersion = 1 migrationAckPrefix = "!migstage|ack|" + migrationReadyPrefix = "!migstage|ready|" migrationUint64Bytes = 8 + migrationAckKeyIDBytes = 2 * migrationUint64Bytes exportVersionSizeOverhead = 24 defaultSparseExportMaxScannedBytes = 1 << 20 ) @@ -192,10 +194,18 @@ func exportUsesSparseScanBudget(opts ExportVersionsOptions) bool { opts.EndKey != nil } +func migrationReadyKey(jobID uint64) []byte { + key := make([]byte, len(migrationReadyPrefix)+migrationUint64Bytes) + copy(key, migrationReadyPrefix) + binary.BigEndian.PutUint64(key[len(migrationReadyPrefix):], jobID) + return key +} + func isMigrationMetadataKey(rawKey []byte) bool { return bytes.Equal(rawKey, migrationAckMetaKeyBytes) || bytes.Equal(rawKey, migrationHLCFloorMetaKeyBytes) || - bytes.Equal(rawKey, migrationPromoteMetaKeyBytes) + bytes.Equal(rawKey, migrationPromoteMetaKeyBytes) || + (len(rawKey) == len(migrationReadyPrefix)+migrationUint64Bytes && bytes.HasPrefix(rawKey, []byte(migrationReadyPrefix))) } func encodeMigrationImportAcks(acks map[migrationAckID]migrationImportAck) []byte { @@ -420,7 +430,14 @@ func (s *mvccStore) ExportVersions(ctx context.Context, opts ExportVersionsOptio s.mtx.RLock() defer s.mtx.RUnlock() + if err := s.checkExportReadTSLocked(opts); err != nil { + return ExportVersionsResult{}, err + } + return s.exportVersionsLocked(ctx, opts, pos) +} + +func (s *mvccStore) exportVersionsLocked(ctx context.Context, opts ExportVersionsOptions, pos exportCursorPosition) (ExportVersionsResult, error) { result := newExportVersionsResult(opts.MaxVersions) it := s.tree.Iterator() if !s.seekMemoryExportStart(&it, opts.StartKey, pos) { @@ -451,6 +468,20 @@ func (s *mvccStore) ExportVersions(ctx context.Context, opts ExportVersionsOptio return result, nil } +func (s *mvccStore) checkExportReadTSLocked(opts ExportVersionsOptions) error { + if readTSCompacted(exportRetentionReadTS(opts), s.minRetainedTS) { + return ErrReadTSCompacted + } + return nil +} + +func exportRetentionReadTS(opts ExportVersionsOptions) uint64 { + if opts.ReadTS != 0 { + return opts.ReadTS + } + return opts.MaxCommitTSInclusive +} + var errExportReachedEnd = errors.New("export reached end") var errExportChunkFull = errors.New("export chunk full") diff --git a/store/migration_versions_test.go b/store/migration_versions_test.go index c83ab6ca9..1aab19768 100644 --- a/store/migration_versions_test.go +++ b/store/migration_versions_test.go @@ -57,6 +57,42 @@ func TestExportVersionsPreservesRawVersionMetadata(t *testing.T) { }) } +func TestExportVersionsReadTSPreservesCompactionErrors(t *testing.T) { + runMigrationStoreSuite(t, func(t *testing.T, st MVCCStore) { + ctx := context.Background() + require.NoError(t, st.PutAt(ctx, []byte("k"), []byte("v10"), 10, 0)) + retention, ok := st.(RetentionController) + require.True(t, ok) + retention.SetMinRetainedTS(20) + + _, err := st.ExportVersions(ctx, ExportVersionsOptions{ + StartKey: []byte("k"), + EndKey: []byte("l"), + MaxCommitTSInclusive: 15, + ReadTS: 15, + MaxVersions: 10, + }) + require.ErrorIs(t, err, ErrReadTSCompacted) + + _, err = st.ExportVersions(ctx, ExportVersionsOptions{ + StartKey: []byte("k"), + EndKey: []byte("l"), + MaxCommitTSInclusive: 15, + MaxVersions: 10, + }) + require.ErrorIs(t, err, ErrReadTSCompacted) + + result, err := st.ExportVersions(ctx, ExportVersionsOptions{ + StartKey: []byte("k"), + EndKey: []byte("l"), + MaxCommitTSInclusive: 25, + MaxVersions: 10, + }) + require.NoError(t, err) + require.Len(t, result.Versions, 1) + }) +} + func TestExportVersionsExcludesTxnLocks(t *testing.T) { runMigrationStoreSuite(t, func(t *testing.T, st MVCCStore) { ctx := context.Background() @@ -934,6 +970,121 @@ func TestPromoteVersionsMovesStagedVersionsAndDeletesStagedRows(t *testing.T) { }) } +func TestTargetStagedReadinessStatePersistsAndIsCloned(t *testing.T) { + runMigrationStoreSuite(t, func(t *testing.T, st MVCCStore) { + ctx := context.Background() + writer, ok := st.(MigrationTargetReadinessWriter) + require.True(t, ok) + reader, ok := st.(MigrationTargetReadinessReader) + require.True(t, ok) + + state := TargetStagedReadinessState{ + JobID: 9, + RouteStart: []byte("a"), + RouteEnd: []byte("z"), + ExpectedCutoverVersion: 12, + MigrationJobID: 9, + MinWriteTSExclusive: 100, + Armed: true, + } + require.NoError(t, writer.ApplyTargetStagedReadiness(ctx, state)) + + states, err := reader.MigrationTargetReadinessStates(ctx) + require.NoError(t, err) + require.Equal(t, []TargetStagedReadinessState{state}, states) + + states[0].RouteStart[0] = 'x' + states, err = reader.MigrationTargetReadinessStates(ctx) + require.NoError(t, err) + require.Equal(t, []byte("a"), states[0].RouteStart) + + updated := state + updated.MinWriteTSExclusive = 101 + updated.RouteStart = []byte("b") + require.NoError(t, writer.ApplyTargetStagedReadiness(ctx, updated)) + + states, err = reader.MigrationTargetReadinessStates(ctx) + require.NoError(t, err) + require.Equal(t, []TargetStagedReadinessState{updated}, states) + + exported, err := st.ExportVersions(ctx, ExportVersionsOptions{ + StartKey: []byte("!"), + EndKey: []byte("~"), + MaxVersions: 100, + }) + require.NoError(t, err) + require.Empty(t, exported.Versions) + }) +} + +func TestPebbleTargetStagedReadinessPersistsAcrossReopen(t *testing.T) { + ctx := context.Background() + dir, err := os.MkdirTemp("", "migration-ready-persist-*") + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, os.RemoveAll(dir)) }) + + st, err := NewPebbleStore(dir) + require.NoError(t, err) + state := TargetStagedReadinessState{ + JobID: 11, + RouteStart: []byte("m"), + RouteEnd: nil, + ExpectedCutoverVersion: 22, + MigrationJobID: 11, + MinWriteTSExclusive: 333, + Armed: true, + } + writer, ok := st.(MigrationTargetReadinessWriter) + require.True(t, ok) + require.NoError(t, writer.ApplyTargetStagedReadiness(ctx, state)) + require.NoError(t, st.Close()) + + reopened, err := NewPebbleStore(dir) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, reopened.Close()) }) + reader, ok := reopened.(MigrationTargetReadinessReader) + require.True(t, ok) + states, err := reader.MigrationTargetReadinessStates(ctx) + require.NoError(t, err) + require.Equal(t, []TargetStagedReadinessState{state}, states) +} + +func TestPebbleTargetStagedReadinessIgnoresNonRecordPrefixKeys(t *testing.T) { + ctx := context.Background() + dir, err := os.MkdirTemp("", "migration-ready-prefix-*") + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, os.RemoveAll(dir)) }) + + st, err := NewPebbleStore(dir) + require.NoError(t, err) + state := TargetStagedReadinessState{ + JobID: 11, + RouteStart: []byte("m"), + RouteEnd: nil, + ExpectedCutoverVersion: 22, + MigrationJobID: 11, + MinWriteTSExclusive: 333, + Armed: true, + } + writer, ok := st.(MigrationTargetReadinessWriter) + require.True(t, ok) + require.NoError(t, writer.ApplyTargetStagedReadiness(ctx, state)) + pebbleStore, ok := st.(*pebbleStore) + require.True(t, ok) + junkKey := append([]byte(migrationReadyPrefix), []byte("side-record")...) + require.NoError(t, pebbleStore.db.Set(junkKey, []byte("not-readiness-state"), pebbleStore.directApplyWriteOpts())) + require.NoError(t, st.Close()) + + reopened, err := NewPebbleStore(dir) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, reopened.Close()) }) + reader, ok := reopened.(MigrationTargetReadinessReader) + require.True(t, ok) + states, err := reader.MigrationTargetReadinessStates(ctx) + require.NoError(t, err) + require.Equal(t, []TargetStagedReadinessState{state}, states) +} + func TestPromoteVersionsIgnoresClientCursorWhenStateMissing(t *testing.T) { runMigrationStoreSuite(t, func(t *testing.T, st MVCCStore) { ctx := context.Background() diff --git a/store/mvcc_store.go b/store/mvcc_store.go index 88817d94d..7c414b083 100644 --- a/store/mvcc_store.go +++ b/store/mvcc_store.go @@ -25,9 +25,12 @@ type VersionedValue struct { } const ( - mvccSnapshotVersion = uint32(1) - maxSnapshotKeySize = 1 << 20 // 1 MiB per key - maxSnapshotVersionCount = 1 << 20 // 1M versions per key + mvccSnapshotVersion = uint32(2) + mvccSnapshotLegacyVersion = uint32(1) + maxSnapshotKeySize = 1 << 20 // 1 MiB per key + maxSnapshotVersionCount = 1 << 20 // 1M versions per key + maxSnapshotReadinessStateCount = 1 << 20 + maxSnapshotReadinessEncodedBytes = 1 << 20 ) // maxSnapshotValueSize caps the allowed size of a single value during streaming @@ -60,14 +63,16 @@ func byteSliceComparator(a, b any) int { // mvccStore is an in-memory MVCC implementation backed by a treemap for // deterministic iteration order and range scans. type mvccStore struct { - tree *treemap.Map // key []byte -> []VersionedValue - mtx sync.RWMutex - log *slog.Logger - lastCommitTS uint64 - minRetainedTS uint64 - migrationAcks map[migrationAckID]migrationImportAck - migrationHLCFloors map[uint64]uint64 - migrationPromotions map[uint64]PromotionState + tree *treemap.Map // key []byte -> []VersionedValue + mtx sync.RWMutex + log *slog.Logger + lastCommitTS uint64 + minRetainedTS uint64 + migrationAcks map[migrationAckID]migrationImportAck + migrationHLCFloors map[uint64]uint64 + migrationPromotions map[uint64]PromotionState + migrationReadiness map[uint64]TargetStagedReadinessState + migrationReadinessCache []TargetStagedReadinessState // writeConflicts mirrors the per-(kind, key_prefix) counter from // the pebble-backed store so the in-memory implementation shows up // in the same Prometheus series (even if the counts are usually @@ -117,6 +122,7 @@ func NewMVCCStore(opts ...MVCCStoreOption) MVCCStore { migrationAcks: make(map[migrationAckID]migrationImportAck), migrationHLCFloors: make(map[uint64]uint64), migrationPromotions: make(map[uint64]PromotionState), + migrationReadiness: make(map[uint64]TargetStagedReadinessState), writeConflicts: newWriteConflictCounter(), } for _, opt := range opts { @@ -840,6 +846,9 @@ func (s *mvccStore) writeSnapshotBody(f *os.File) (uint32, error) { if err := binary.Write(w, binary.LittleEndian, s.minRetainedTS); err != nil { return 0, errors.WithStack(err) } + if err := writeMVCCSnapshotReadinessStates(w, s.migrationReadinessCache); err != nil { + return 0, err + } iter := s.tree.Iterator() for iter.Next() { key, ok := iter.Key().([]byte) @@ -910,6 +919,22 @@ func writeMVCCSnapshotVersion(w io.Writer, version VersionedValue) error { return nil } +func writeMVCCSnapshotReadinessStates(w io.Writer, states []TargetStagedReadinessState) error { + if err := binary.Write(w, binary.LittleEndian, uint64(len(states))); err != nil { + return errors.WithStack(err) + } + for _, state := range states { + encoded := encodeTargetStagedReadinessState(state) + if err := binary.Write(w, binary.LittleEndian, uint64(len(encoded))); err != nil { + return errors.WithStack(err) + } + if _, err := w.Write(encoded); err != nil { + return errors.WithStack(err) + } + } + return nil +} + func mvccSnapshotTombstoneByte(tombstone bool) byte { if tombstone { return 1 @@ -918,12 +943,12 @@ func mvccSnapshotTombstoneByte(tombstone bool) byte { } func (s *mvccStore) restoreStreamingSnapshot(r io.Reader) error { - expected, err := readMVCCSnapshotHeader(r) + expected, version, err := readMVCCSnapshotHeader(r) if err != nil { return err } - tree, lastCommitTS, minRetainedTS, actual, err := restoreStreamingMVCCSnapshotBody(r) + tree, lastCommitTS, minRetainedTS, readinessStates, actual, err := restoreStreamingMVCCSnapshotBody(r, version) if err != nil { return err } @@ -939,48 +964,54 @@ func (s *mvccStore) restoreStreamingSnapshot(r io.Reader) error { s.migrationAcks = make(map[migrationAckID]migrationImportAck) s.migrationHLCFloors = make(map[uint64]uint64) s.migrationPromotions = make(map[uint64]PromotionState) + s.migrationReadiness = targetReadinessStateMap(readinessStates) + s.migrationReadinessCache = cloneTargetStagedReadinessStates(readinessStates) return nil } -func readMVCCSnapshotHeader(r io.Reader) (uint32, error) { +func readMVCCSnapshotHeader(r io.Reader) (uint32, uint32, error) { var magic [8]byte if _, err := io.ReadFull(r, magic[:]); err != nil { - return 0, errors.WithStack(err) + return 0, 0, errors.WithStack(err) } if magic != mvccSnapshotMagic { - return 0, errors.WithStack(ErrInvalidChecksum) + return 0, 0, errors.WithStack(ErrInvalidChecksum) } var version uint32 if err := binary.Read(r, binary.LittleEndian, &version); err != nil { - return 0, errors.WithStack(err) + return 0, 0, errors.WithStack(err) } - if version != mvccSnapshotVersion { - return 0, errors.WithStack(errors.Newf("unsupported mvcc snapshot version %d", version)) + if version != mvccSnapshotLegacyVersion && version != mvccSnapshotVersion { + return 0, 0, errors.WithStack(errors.Newf("unsupported mvcc snapshot version %d", version)) } var expected uint32 if err := binary.Read(r, binary.LittleEndian, &expected); err != nil { - return 0, errors.WithStack(err) + return 0, 0, errors.WithStack(err) } - return expected, nil + return expected, version, nil } -func restoreStreamingMVCCSnapshotBody(r io.Reader) (*treemap.Map, uint64, uint64, uint32, error) { +func restoreStreamingMVCCSnapshotBody(r io.Reader, version uint32) (*treemap.Map, uint64, uint64, []TargetStagedReadinessState, uint32, error) { hash := crc32.NewIEEE() body := io.TeeReader(r, hash) lastCommitTS, minRetainedTS, err := readMVCCSnapshotMetadata(body) if err != nil { - return nil, 0, 0, 0, err + return nil, 0, 0, nil, 0, err + } + readinessStates, err := readMVCCSnapshotReadinessStates(body, version) + if err != nil { + return nil, 0, 0, nil, 0, err } tree, err := readMVCCSnapshotTree(body) if err != nil { - return nil, 0, 0, 0, err + return nil, 0, 0, nil, 0, err } - return tree, lastCommitTS, minRetainedTS, hash.Sum32(), nil + return tree, lastCommitTS, minRetainedTS, readinessStates, hash.Sum32(), nil } func readMVCCSnapshotMetadata(r io.Reader) (uint64, uint64, error) { @@ -995,6 +1026,49 @@ func readMVCCSnapshotMetadata(r io.Reader) (uint64, uint64, error) { return lastCommitTS, minRetainedTS, nil } +func readMVCCSnapshotReadinessStates(r io.Reader, version uint32) ([]TargetStagedReadinessState, error) { + if version < mvccSnapshotVersion { + return nil, nil + } + var count uint64 + if err := binary.Read(r, binary.LittleEndian, &count); err != nil { + return nil, errors.WithStack(err) + } + if count > maxSnapshotReadinessStateCount { + return nil, errors.Wrapf(ErrSnapshotVersionCountTooLarge, "readiness states %d > %d", count, maxSnapshotReadinessStateCount) + } + states := make([]TargetStagedReadinessState, 0, count) + for range count { + var encodedLen uint64 + if err := binary.Read(r, binary.LittleEndian, &encodedLen); err != nil { + return nil, errors.WithStack(err) + } + if encodedLen > maxSnapshotReadinessEncodedBytes { + return nil, errors.Wrapf(ErrValueTooLarge, "readiness state %d > %d", encodedLen, maxSnapshotReadinessEncodedBytes) + } + encoded := make([]byte, encodedLen) + if _, err := io.ReadFull(r, encoded); err != nil { + return nil, errors.WithStack(err) + } + state, ok := decodeTargetStagedReadinessState(encoded) + if !ok { + return nil, errors.New("corrupt target staged readiness state") + } + states = append(states, state) + } + sortTargetStagedReadinessStates(states) + return states, nil +} + +func targetReadinessStateMap(states []TargetStagedReadinessState) map[uint64]TargetStagedReadinessState { + out := make(map[uint64]TargetStagedReadinessState, len(states)) + for _, state := range states { + cloned := cloneTargetStagedReadinessState(state) + out[cloned.JobID] = cloned + } + return out +} + func readMVCCSnapshotTree(r io.Reader) (*treemap.Map, error) { tree := treemap.NewWith(byteSliceComparator) for { diff --git a/store/mvcc_store_snapshot_test.go b/store/mvcc_store_snapshot_test.go index 3032a1cc3..5024a6c26 100644 --- a/store/mvcc_store_snapshot_test.go +++ b/store/mvcc_store_snapshot_test.go @@ -38,6 +38,61 @@ func TestMVCCStore_SnapshotRestoreRoundTrip(t *testing.T) { require.Equal(t, []byte("v2"), v) } +func TestMVCCStore_SnapshotRestorePreservesTargetReadiness(t *testing.T) { + t.Parallel() + + ctx := context.Background() + src := newTestMVCCStore(t) + + stale := TargetStagedReadinessState{ + JobID: 8, + RouteStart: []byte("old"), + RouteEnd: []byte("oldz"), + ExpectedCutoverVersion: 1, + MigrationJobID: 8, + MinWriteTSExclusive: 80, + Armed: true, + } + require.NoError(t, src.ApplyTargetStagedReadiness(ctx, stale)) + + snapState := TargetStagedReadinessState{ + JobID: 9, + RouteStart: []byte("a"), + RouteEnd: []byte("z"), + ExpectedCutoverVersion: 12, + MigrationJobID: 9, + MinWriteTSExclusive: 100, + Armed: true, + } + require.NoError(t, src.ApplyTargetStagedReadiness(ctx, snapState)) + + snap, err := src.Snapshot() + require.NoError(t, err) + defer snap.Close() + raw := snapshotBytes(t, snap) + + dst := newTestMVCCStore(t) + require.NoError(t, dst.ApplyTargetStagedReadiness(ctx, TargetStagedReadinessState{ + JobID: 77, + RouteStart: []byte("dst-only"), + RouteEnd: []byte("dst-z"), + ExpectedCutoverVersion: 2, + MigrationJobID: 77, + MinWriteTSExclusive: 700, + Armed: true, + })) + + require.NoError(t, dst.Restore(bytes.NewReader(raw))) + states, err := dst.MigrationTargetReadinessStates(ctx) + require.NoError(t, err) + require.Equal(t, []TargetStagedReadinessState{stale, snapState}, states) + + states[0].RouteStart[0] = 'x' + states, err = dst.MigrationTargetReadinessStates(ctx) + require.NoError(t, err) + require.Equal(t, []byte("old"), states[0].RouteStart) +} + func TestMVCCStore_RestoreRejectsInvalidChecksum(t *testing.T) { t.Parallel() diff --git a/store/store.go b/store/store.go index 700550f6b..ca80528b3 100644 --- a/store/store.go +++ b/store/store.go @@ -86,12 +86,14 @@ type ExportVersionsOptions struct { EndKey []byte MinCommitTSExclusive uint64 MaxCommitTSInclusive uint64 - Cursor []byte - MaxVersions int - MaxBytes uint64 - MaxScannedBytes uint64 - KeyFamily uint32 - AcceptKey func([]byte) bool + // ReadTS asks export to enforce the same retention watermark as GetAt/ScanAt. + ReadTS uint64 + Cursor []byte + MaxVersions int + MaxBytes uint64 + MaxScannedBytes uint64 + KeyFamily uint32 + AcceptKey func([]byte) bool } // ExportVersionsResult is one resumable chunk of raw MVCC versions. @@ -155,6 +157,19 @@ type PromotionState struct { LastError string } +// TargetStagedReadinessState is a target-local guard that makes a target voter +// fail closed for a moving route until it has either the staged descriptor or +// the cleared descriptor with the retained write timestamp floor. +type TargetStagedReadinessState struct { + JobID uint64 + RouteStart []byte + RouteEnd []byte + ExpectedCutoverVersion uint64 + MigrationJobID uint64 + MinWriteTSExclusive uint64 + Armed bool +} + // MigrationPromoter is implemented by stores that can promote staged range // migration data into the live keyspace. type MigrationPromoter interface { @@ -166,6 +181,17 @@ type MigrationPromotionStateReader interface { MigrationPromotionState(ctx context.Context, jobID uint64) (PromotionState, bool, error) } +// MigrationTargetReadinessWriter persists a target-local staged-readiness +// guard through the target Raft group. +type MigrationTargetReadinessWriter interface { + ApplyTargetStagedReadiness(ctx context.Context, state TargetStagedReadinessState) error +} + +// MigrationTargetReadinessReader reads retained target-local readiness guards. +type MigrationTargetReadinessReader interface { + MigrationTargetReadinessStates(ctx context.Context) ([]TargetStagedReadinessState, error) +} + // OpType describes a mutation kind. type OpType int