From 2ea0d381f212cd9dabe51080f87e155ea119aa7f Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 01:27:31 +0900 Subject: [PATCH 01/29] migration: add target readiness guard --- kv/shard_store.go | 139 ++++++++++++++++-- kv/shard_store_test.go | 103 +++++++++++++ store/lsm_store.go | 37 +++-- store/migration_readiness.go | 238 +++++++++++++++++++++++++++++++ store/migration_versions.go | 11 +- store/migration_versions_test.go | 79 ++++++++++ store/mvcc_store.go | 19 +-- store/store.go | 24 ++++ 8 files changed, 621 insertions(+), 29 deletions(-) create mode 100644 store/migration_readiness.go diff --git a/kv/shard_store.go b/kv/shard_store.go index ba5ef6ac9..aaf18b57c 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -26,6 +26,8 @@ type ShardStore struct { } var ErrCrossShardMutationBatchNotSupported = errors.New("cross-shard mutation batches are not supported") +var ErrRouteCutoverPending = errors.New("route cutover pending") +var ErrRouteWriteBelowFloor = errors.New("route write timestamp is below route floor") // NewShardStore creates a sharded MVCC store wrapper. func NewShardStore(engine *distribution.Engine, groups map[uint64]*ShardGroup) *ShardStore { @@ -98,6 +100,9 @@ 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) { + if err := s.verifyTargetReadinessForRange(ctx, g, route, key, nextScanCursor(key)); err != nil { + return nil, err + } if routeHasStagedVisibility(route) { return s.getAtWithStagedVisibility(ctx, g, route, key, ts) } @@ -112,6 +117,59 @@ func routeHasStagedVisibility(route distribution.Route) bool { return route.StagedVisibilityActive && route.MigrationJobID != 0 } +func routeSatisfiesTargetReadiness(route distribution.Route, ready store.TargetStagedReadinessState) 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 + } + return route.MigrationJobID == 0 +} + +func (s *ShardStore) verifyTargetReadinessForRange(ctx context.Context, g *ShardGroup, route distribution.Route, start []byte, end []byte) error { + if g == nil || g.Store == nil { + return nil + } + reader, ok := g.Store.(store.MigrationTargetReadinessReader) + if !ok { + return nil + } + states, err := reader.MigrationTargetReadinessStates(ctx) + if err != nil { + return errors.WithStack(err) + } + for _, ready := range states { + if !ready.Armed || !routeRangeIntersects(start, end, ready.RouteStart, ready.RouteEnd) { + continue + } + if !routeSatisfiesTargetReadiness(route, ready) { + return errors.WithStack(ErrRouteCutoverPending) + } + } + return nil +} + +func verifyRouteWriteFloor(route distribution.Route, commitTS uint64) error { + if route.MinWriteTSExclusive == 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) getAtWithStagedVisibility(ctx context.Context, g *ShardGroup, route distribution.Route, key []byte, ts uint64) ([]byte, error) { live, liveOK, err := latestMVCCVersionAt(ctx, g.Store, key, ts) if err != nil { @@ -574,6 +632,9 @@ func (s *ShardStore) scanRouteLocal( ts uint64, reverse bool, ) ([]*store.KVPair, error) { + if err := s.verifyTargetReadinessForRange(ctx, g, route, start, end); err != nil { + return nil, err + } if routeHasStagedVisibility(route) { return s.scanRouteWithStagedVisibility(ctx, g, route, start, end, limit, ts, reverse) } @@ -618,6 +679,9 @@ func (s *ShardStore) scanRouteAtLeader( ts uint64, reverse bool, ) ([]*store.KVPair, error) { + if err := s.verifyTargetReadinessForRange(ctx, g, route, start, end); err != nil { + return nil, err + } var ( kvs []*store.KVPair err error @@ -905,34 +969,58 @@ func clampScanEnd(end []byte, routeEnd []byte) []byte { } func (s *ShardStore) PutAt(ctx context.Context, key []byte, value []byte, commitTS uint64, expireAt uint64) error { - g, ok := s.groupForKey(key) - if !ok || g.Store == nil { + 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 + } + if err := verifyRouteWriteFloor(route, commitTS); err != nil { + return err + } return errors.WithStack(g.Store.PutAt(ctx, key, value, commitTS, expireAt)) } func (s *ShardStore) DeleteAt(ctx context.Context, key []byte, commitTS uint64) error { - g, ok := s.groupForKey(key) - if !ok || g.Store == nil { + 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 + } + if err := verifyRouteWriteFloor(route, commitTS); err != nil { + return err + } return errors.WithStack(g.Store.DeleteAt(ctx, key, commitTS)) } func (s *ShardStore) PutWithTTLAt(ctx context.Context, key []byte, value []byte, commitTS uint64, expireAt uint64) error { - g, ok := s.groupForKey(key) - if !ok || g.Store == nil { + 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 + } + if err := verifyRouteWriteFloor(route, commitTS); err != nil { + return err + } return errors.WithStack(g.Store.PutWithTTLAt(ctx, key, value, commitTS, expireAt)) } func (s *ShardStore) ExpireAt(ctx context.Context, key []byte, expireAt uint64, commitTS uint64) error { - g, ok := s.groupForKey(key) - if !ok || g.Store == nil { + 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 + } + if err := verifyRouteWriteFloor(route, commitTS); err != nil { + return err + } return errors.WithStack(g.Store.ExpireAt(ctx, key, expireAt, commitTS)) } @@ -967,6 +1055,9 @@ func (s *ShardStore) LatestCommitTS(ctx context.Context, key []byte) (uint64, bo } func (s *ShardStore) localLatestCommitTS(ctx context.Context, g *ShardGroup, route distribution.Route, key []byte) (uint64, bool, error) { + if err := s.verifyTargetReadinessForRange(ctx, g, route, key, nextScanCursor(key)); err != nil { + return 0, false, err + } liveTS, liveExists, err := g.Store.LatestCommitTS(ctx, key) if err != nil { return 0, false, errors.WithStack(err) @@ -1606,9 +1697,41 @@ func (s *ShardStore) ApplyMutations(ctx context.Context, mutations []*store.KVPa if err != nil || group == nil { return err } + if err := s.verifyMutationRoutes(ctx, mutations, readKeys, commitTS); err != nil { + return err + } 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 + } + if err := s.verifyTargetReadinessForRange(ctx, g, route, key, nextScanCursor(key)); err != nil { + return err + } + return verifyRouteWriteFloor(route, 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 { diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 100b70e27..cfe9aec59 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -33,6 +33,32 @@ func newStagedVisibilityShardStore(t *testing.T) (*ShardStore, *ShardGroup) { return NewShardStore(engine, map[uint64]*ShardGroup{1: group}), group } +func applyTargetReadiness(t *testing.T, group *ShardGroup, floor uint64) { + t.Helper() + writer, ok := group.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: floor, + Armed: true, + })) +} + +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 +} + func TestShardStoreGetAt_MergesStagedVisibility(t *testing.T) { t.Parallel() @@ -57,6 +83,83 @@ func TestShardStoreGetAt_MergesStagedVisibility(t *testing.T) { require.ErrorIs(t, err, store.ErrKeyNotFound) } +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, 100) + + _, 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 TestShardStoreTargetReadinessAcceptsStagedDescriptor(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + applyTargetReadiness(t, group, 100) + 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, 100) + 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 TestShardStoreScanAndLatestCommitTS_MergeStagedVisibility(t *testing.T) { t.Parallel() diff --git a/store/lsm_store.go b/store/lsm_store.go index 18913cddc..a1a2ac350 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 } @@ -2583,9 +2590,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/migration_readiness.go b/store/migration_readiness.go new file mode 100644 index 000000000..7aa160ca3 --- /dev/null +++ b/store/migration_readiness.go @@ -0,0 +1,238 @@ +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() + if err := s.db.Set(migrationReadyKey(state.JobID), encodeTargetStagedReadinessState(state), s.directApplyWriteOpts()); err != nil { + return errors.WithStack(err) + } + s.mtx.Lock() + defer s.mtx.Unlock() + 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() { + 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 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 28fd2f73d..b1e275ecb 100644 --- a/store/migration_versions.go +++ b/store/migration_versions.go @@ -16,6 +16,7 @@ const ( migrationAckPrefix = "!migstage|ack|" migrationHLCFloorPrefix = "!migstage|hlc_floor|" migrationPromotePrefix = "!migstage|promote|" + migrationReadyPrefix = "!migstage|ready|" migrationUint64Bytes = 8 migrationAckKeyIDBytes = 2 * migrationUint64Bytes exportVersionSizeOverhead = 24 @@ -92,10 +93,18 @@ func migrationPromoteKey(jobID uint64) []byte { return key } +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 (len(rawKey) == len(migrationAckPrefix)+migrationAckKeyIDBytes && bytes.HasPrefix(rawKey, []byte(migrationAckPrefix))) || (len(rawKey) == len(migrationHLCFloorPrefix)+migrationUint64Bytes && bytes.HasPrefix(rawKey, []byte(migrationHLCFloorPrefix))) || - (len(rawKey) == len(migrationPromotePrefix)+migrationUint64Bytes && bytes.HasPrefix(rawKey, []byte(migrationPromotePrefix))) + (len(rawKey) == len(migrationPromotePrefix)+migrationUint64Bytes && bytes.HasPrefix(rawKey, []byte(migrationPromotePrefix))) || + (len(rawKey) == len(migrationReadyPrefix)+migrationUint64Bytes && bytes.HasPrefix(rawKey, []byte(migrationReadyPrefix))) } func encodeMigrationImportAck(ack migrationImportAck) []byte { diff --git a/store/migration_versions_test.go b/store/migration_versions_test.go index 91548bb00..7332f318b 100644 --- a/store/migration_versions_test.go +++ b/store/migration_versions_test.go @@ -279,6 +279,85 @@ 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 TestPebbleImportMetadataPersistsAcrossReopen(t *testing.T) { ctx := context.Background() dir, err := os.MkdirTemp("", "migration-import-persist-*") diff --git a/store/mvcc_store.go b/store/mvcc_store.go index b20c35254..1bea2a9f7 100644 --- a/store/mvcc_store.go +++ b/store/mvcc_store.go @@ -60,14 +60,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[string]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[string]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 +119,7 @@ func NewMVCCStore(opts ...MVCCStoreOption) MVCCStore { migrationAcks: make(map[string]migrationImportAck), migrationHLCFloors: make(map[uint64]uint64), migrationPromotions: make(map[uint64]PromotionState), + migrationReadiness: make(map[uint64]TargetStagedReadinessState), writeConflicts: newWriteConflictCounter(), } for _, opt := range opts { diff --git a/store/store.go b/store/store.go index a645a95c5..ae53203f9 100644 --- a/store/store.go +++ b/store/store.go @@ -148,6 +148,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 { @@ -159,6 +172,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 From b93f5b9d473a3b88d1069794c665f479911d543b Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 02:03:50 +0900 Subject: [PATCH 02/29] distribution: add split job management RPCs --- adapter/distribution_server.go | 195 +++++++++++++++++++++++++ adapter/distribution_server_test.go | 152 ++++++++++++++++++- distribution/split_job_catalog.go | 5 + distribution/split_job_catalog_test.go | 92 ++++++++++++ distribution/split_job_lifecycle.go | 145 ++++++++++++++++++ 5 files changed, 586 insertions(+), 3 deletions(-) create mode 100644 distribution/split_job_lifecycle.go diff --git a/adapter/distribution_server.go b/adapter/distribution_server.go index 9703c828b..7052a8561 100644 --- a/adapter/distribution_server.go +++ b/adapter/distribution_server.go @@ -4,12 +4,14 @@ import ( "bytes" "context" "math" + "strings" "sync" "time" "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/kv" pb "github.com/bootjp/elastickv/proto" + "github.com/bootjp/elastickv/store" "github.com/cockroachdb/errors" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" @@ -78,6 +80,7 @@ var ( errDistributionCoordinatorRequired = errors.New("distribution coordinator is not configured") errDistributionEngineNotConfigured = errors.New("distribution engine is not configured") errDistributionCatalogVersionNotFound = errors.New("route catalog version not found") + errDistributionClusterNotReady = errors.New("cluster is not ready for split migration") ) // NewDistributionServer creates a new server. @@ -172,6 +175,97 @@ func (s *DistributionServer) GetIntersectingRoutes(ctx context.Context, req *pb. }, nil } +func (s *DistributionServer) StartSplitMigration(ctx context.Context, req *pb.StartSplitMigrationRequest) (*pb.StartSplitMigrationResponse, error) { + if req.GetTargetGroupId() == 0 { + return nil, grpcStatusError(codes.InvalidArgument, distribution.ErrCatalogSplitJobTargetGroupRequired.Error()) + } + if err := s.verifyCatalogLeader(ctx); err != nil { + return nil, err + } + if s.catalog == nil { + return nil, grpcStatusError(codes.FailedPrecondition, errDistributionCatalogNotConfigured.Error()) + } + return nil, grpcStatusError(codes.FailedPrecondition, errDistributionClusterNotReady.Error()) +} + +func (s *DistributionServer) GetSplitJob(ctx context.Context, req *pb.GetSplitJobRequest) (*pb.GetSplitJobResponse, error) { + if req.GetJobId() == 0 { + return nil, grpcStatusError(codes.InvalidArgument, distribution.ErrCatalogSplitJobIDRequired.Error()) + } + if s.catalog == nil { + return nil, grpcStatusError(codes.FailedPrecondition, errDistributionCatalogNotConfigured.Error()) + } + job, found, err := s.catalog.SplitJob(ctx, req.GetJobId()) + if err != nil { + return nil, splitJobCatalogStatusError(err) + } + if !found { + return nil, grpcStatusError(codes.NotFound, distribution.ErrCatalogSplitJobNotFound.Error()) + } + return &pb.GetSplitJobResponse{Job: distribution.SplitJobToProto(job)}, nil +} + +func (s *DistributionServer) ListSplitJobs(ctx context.Context, req *pb.ListSplitJobsRequest) (*pb.ListSplitJobsResponse, error) { + if s.catalog == nil { + return nil, grpcStatusError(codes.FailedPrecondition, errDistributionCatalogNotConfigured.Error()) + } + if req == nil { + req = &pb.ListSplitJobsRequest{} + } + phaseFilter := normalizeSplitJobPhaseFilter(req.GetPhase()) + if phaseFilter != "" && !validSplitJobPhaseFilter(phaseFilter) { + return nil, grpcStatusErrorf(codes.InvalidArgument, "unknown split job phase %q", req.GetPhase()) + } + jobs, err := s.catalog.ListSplitJobs(ctx) + if err != nil { + return nil, splitJobCatalogStatusError(err) + } + resp := &pb.ListSplitJobsResponse{} + for _, job := range jobs { + if !splitJobPassesListFilter(job, req.GetSinceTerminalAtMs(), phaseFilter) { + continue + } + resp.Jobs = append(resp.Jobs, distribution.SplitJobToProto(job)) + } + return resp, nil +} + +func (s *DistributionServer) AbandonSplitJob(ctx context.Context, req *pb.AbandonSplitJobRequest) (*pb.AbandonSplitJobResponse, error) { + if req.GetJobId() == 0 { + return nil, grpcStatusError(codes.InvalidArgument, distribution.ErrCatalogSplitJobIDRequired.Error()) + } + if err := s.verifyCatalogLeader(ctx); err != nil { + return nil, err + } + if s.catalog == nil { + return nil, grpcStatusError(codes.FailedPrecondition, errDistributionCatalogNotConfigured.Error()) + } + if err := s.updateSplitJobViaCoordinator(ctx, req.GetJobId(), func(job distribution.SplitJob) (distribution.SplitJob, error) { + return distribution.BeginSplitJobAbandon(job, time.Now().UnixMilli()) + }); err != nil { + return nil, err + } + return &pb.AbandonSplitJobResponse{}, nil +} + +func (s *DistributionServer) RetrySplitJob(ctx context.Context, req *pb.RetrySplitJobRequest) (*pb.RetrySplitJobResponse, error) { + if req.GetJobId() == 0 { + return nil, grpcStatusError(codes.InvalidArgument, distribution.ErrCatalogSplitJobIDRequired.Error()) + } + if err := s.verifyCatalogLeader(ctx); err != nil { + return nil, err + } + if s.catalog == nil { + return nil, grpcStatusError(codes.FailedPrecondition, errDistributionCatalogNotConfigured.Error()) + } + if err := s.updateSplitJobViaCoordinator(ctx, req.GetJobId(), func(job distribution.SplitJob) (distribution.SplitJob, error) { + return distribution.RetrySplitJobState(job, time.Now().UnixMilli()) + }); err != nil { + return nil, err + } + return &pb.RetrySplitJobResponse{}, nil +} + func (s *DistributionServer) routeSnapshotAt(version uint64) (distribution.RouteHistorySnapshot, error) { if s.engine == nil { return distribution.RouteHistorySnapshot{}, grpcStatusError(codes.FailedPrecondition, errDistributionEngineNotConfigured.Error()) @@ -401,6 +495,42 @@ func (s *DistributionServer) applyEngineSnapshot(snapshot distribution.CatalogSn return nil } +func (s *DistributionServer) updateSplitJobViaCoordinator( + ctx context.Context, + jobID uint64, + transition func(distribution.SplitJob) (distribution.SplitJob, error), +) error { + expected, readTS, err := s.catalog.LiveSplitJobForUpdate(ctx, jobID) + if err != nil { + return splitJobCatalogStatusError(err) + } + next, err := transition(expected) + if err != nil { + return splitJobCatalogStatusError(err) + } + if splitJobsEqualByEncoding(expected, next) { + return nil + } + encoded, err := distribution.EncodeSplitJob(next) + if err != nil { + return splitJobCatalogStatusError(err) + } + key := distribution.CatalogSplitJobKey(jobID) + if _, err := s.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + Elems: []*kv.Elem[kv.OP]{{ + Op: kv.Put, + Key: key, + Value: encoded, + }}, + IsTxn: true, + StartTS: readTS, + ReadKeys: [][]byte{key}, + }); err != nil { + return splitJobCoordinatorStatusError(err) + } + return nil +} + func validateExpectedCatalogVersion(currentVersion, expectedVersion uint64) error { if currentVersion != expectedVersion { return grpcStatusError(codes.Aborted, errDistributionCatalogConflict.Error()) @@ -493,6 +623,71 @@ func routeRangeIntersects(aStart, aEnd, bStart, bEnd []byte) bool { return true } +func splitJobPassesListFilter(job distribution.SplitJob, sinceTerminalAtMs uint64, phaseFilter string) bool { + if sinceTerminalAtMs > 0 && job.TerminalAtMs > 0 && uint64(job.TerminalAtMs) < sinceTerminalAtMs { + return false + } + if phaseFilter == "" { + return true + } + return normalizeSplitJobPhaseFilter(pb.SplitJobPhase(job.Phase).String()) == phaseFilter || + normalizeSplitJobPhaseFilter(strings.TrimPrefix(pb.SplitJobPhase(job.Phase).String(), "SPLIT_JOB_PHASE_")) == phaseFilter +} + +func normalizeSplitJobPhaseFilter(phase string) string { + phase = strings.TrimSpace(phase) + if phase == "" { + return "" + } + phase = strings.ReplaceAll(phase, "-", "_") + phase = strings.ReplaceAll(phase, " ", "_") + return strings.ToUpper(phase) +} + +func validSplitJobPhaseFilter(phase string) bool { + for _, candidate := range pb.SplitJobPhase_name { + full := normalizeSplitJobPhaseFilter(candidate) + short := normalizeSplitJobPhaseFilter(strings.TrimPrefix(candidate, "SPLIT_JOB_PHASE_")) + if phase == full || phase == short { + return true + } + } + return false +} + +func splitJobsEqualByEncoding(left, right distribution.SplitJob) bool { + leftRaw, leftErr := distribution.EncodeSplitJob(left) + rightRaw, rightErr := distribution.EncodeSplitJob(right) + return leftErr == nil && rightErr == nil && bytes.Equal(leftRaw, rightRaw) +} + +func splitJobCatalogStatusError(err error) error { + switch { + case errors.Is(err, distribution.ErrCatalogSplitJobIDRequired): + return grpcStatusError(codes.InvalidArgument, distribution.ErrCatalogSplitJobIDRequired.Error()) + case errors.Is(err, distribution.ErrCatalogSplitJobNotFound): + return grpcStatusError(codes.NotFound, distribution.ErrCatalogSplitJobNotFound.Error()) + case errors.Is(err, distribution.ErrCatalogSplitJobCannotRetry), + errors.Is(err, distribution.ErrCatalogSplitJobCannotAbandon): + return grpcStatusError(codes.FailedPrecondition, err.Error()) + case errors.Is(err, distribution.ErrCatalogSplitJobConflict): + return grpcStatusError(codes.Aborted, distribution.ErrCatalogSplitJobConflict.Error()) + default: + return grpcStatusErrorf(codes.Internal, "split job catalog: %v", err) + } +} + +func splitJobCoordinatorStatusError(err error) error { + switch { + case errors.Is(err, store.ErrWriteConflict): + return grpcStatusError(codes.Aborted, distribution.ErrCatalogSplitJobConflict.Error()) + case errors.Is(err, kv.ErrLeaderNotFound): + return grpcStatusError(codes.FailedPrecondition, errDistributionNotLeader.Error()) + default: + return grpcStatusErrorf(codes.Internal, "commit split job mutation: %v", err) + } +} + func splitCatalogRoutes( parent distribution.RouteDescriptor, splitKey []byte, diff --git a/adapter/distribution_server_test.go b/adapter/distribution_server_test.go index f717dfb56..434180c0a 100644 --- a/adapter/distribution_server_test.go +++ b/adapter/distribution_server_test.go @@ -129,6 +129,137 @@ func TestDistributionServerListRoutes_RequiresCatalog(t *testing.T) { require.ErrorContains(t, err, errDistributionCatalogNotConfigured.Error()) } +func TestDistributionServerStartSplitMigration_FailsClosedUntilCapabilityGate(t *testing.T) { + t.Parallel() + + ctx := context.Background() + baseStore := store.NewMVCCStore() + catalog := distribution.NewCatalogStore(baseStore) + coordinator := newDistributionCoordinatorStub(baseStore, true) + s := NewDistributionServer(distribution.NewEngine(), catalog, WithDistributionCoordinator(coordinator)) + + _, err := s.StartSplitMigration(ctx, &pb.StartSplitMigrationRequest{ + ExpectedCatalogVersion: 1, + RouteId: 1, + SplitKey: []byte("g"), + TargetGroupId: 2, + }) + require.Error(t, err) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) + require.ErrorContains(t, err, errDistributionClusterNotReady.Error()) + require.Zero(t, coordinator.dispatchCalls) + + jobs, listErr := catalog.ListSplitJobs(ctx) + require.NoError(t, listErr) + require.Empty(t, jobs) +} + +func TestDistributionServerSplitJobRPCs_ReadAndListCatalogJobs(t *testing.T) { + t.Parallel() + + ctx := context.Background() + catalog := distribution.NewCatalogStore(store.NewMVCCStore()) + live := sampleDistributionSplitJob(10) + live.Phase = distribution.SplitJobPhaseDeltaCopy + require.NoError(t, catalog.CreateSplitJob(ctx, live)) + history := sampleDistributionSplitJob(11) + history.Phase = distribution.SplitJobPhaseDone + history.TerminalAtMs = 2000 + require.NoError(t, catalog.CreateSplitJob(ctx, history)) + require.NoError(t, catalog.MoveSplitJobToHistory(ctx, history, history)) + + s := NewDistributionServer(distribution.NewEngine(), catalog) + got, err := s.GetSplitJob(ctx, &pb.GetSplitJobRequest{JobId: live.JobID}) + require.NoError(t, err) + require.Equal(t, live.JobID, got.Job.JobId) + require.Equal(t, pb.SplitJobPhase_SPLIT_JOB_PHASE_DELTA_COPY, got.Job.Phase) + + listAll, err := s.ListSplitJobs(ctx, &pb.ListSplitJobsRequest{}) + require.NoError(t, err) + require.Len(t, listAll.Jobs, 2) + require.Equal(t, []uint64{live.JobID, history.JobID}, []uint64{listAll.Jobs[0].JobId, listAll.Jobs[1].JobId}) + + listDone, err := s.ListSplitJobs(ctx, &pb.ListSplitJobsRequest{Phase: "done"}) + require.NoError(t, err) + require.Len(t, listDone.Jobs, 1) + require.Equal(t, history.JobID, listDone.Jobs[0].JobId) + + _, err = s.ListSplitJobs(ctx, &pb.ListSplitJobsRequest{Phase: "not-a-phase"}) + require.Error(t, err) + require.Equal(t, codes.InvalidArgument, status.Code(err)) +} + +func TestDistributionServerRetrySplitJob_UsesCoordinatorCAS(t *testing.T) { + t.Parallel() + + ctx := context.Background() + baseStore := store.NewMVCCStore() + catalog := distribution.NewCatalogStore(baseStore) + job := sampleDistributionSplitJob(12) + job.Phase = distribution.SplitJobPhaseFailed + job.RetryPhase = distribution.SplitJobPhaseFence + job.LastError = "retry me" + require.NoError(t, catalog.CreateSplitJob(ctx, job)) + + coordinator := newDistributionCoordinatorStub(baseStore, true) + s := NewDistributionServer(distribution.NewEngine(), catalog, WithDistributionCoordinator(coordinator)) + _, err := s.RetrySplitJob(ctx, &pb.RetrySplitJobRequest{JobId: job.JobID}) + require.NoError(t, err) + require.Equal(t, 1, coordinator.dispatchCalls) + + got, found, err := catalog.SplitJob(ctx, job.JobID) + require.NoError(t, err) + require.True(t, found) + require.Equal(t, distribution.SplitJobPhaseFence, got.Phase) + require.Equal(t, distribution.SplitJobPhaseNone, got.RetryPhase) + require.Empty(t, got.LastError) +} + +func TestDistributionServerAbandonSplitJob_RecordsAbandoningViaCoordinator(t *testing.T) { + t.Parallel() + + ctx := context.Background() + baseStore := store.NewMVCCStore() + catalog := distribution.NewCatalogStore(baseStore) + job := sampleDistributionSplitJob(13) + job.Phase = distribution.SplitJobPhaseFailed + job.RetryPhase = distribution.SplitJobPhaseBackfill + job.AbandonFromPhase = distribution.SplitJobPhaseNone + require.NoError(t, catalog.CreateSplitJob(ctx, job)) + + coordinator := newDistributionCoordinatorStub(baseStore, true) + s := NewDistributionServer(distribution.NewEngine(), catalog, WithDistributionCoordinator(coordinator)) + _, err := s.AbandonSplitJob(ctx, &pb.AbandonSplitJobRequest{JobId: job.JobID}) + require.NoError(t, err) + require.Equal(t, 1, coordinator.dispatchCalls) + + got, found, err := catalog.SplitJob(ctx, job.JobID) + require.NoError(t, err) + require.True(t, found) + require.Equal(t, distribution.SplitJobPhaseAbandoning, got.Phase) + require.Equal(t, distribution.SplitJobPhaseNone, got.RetryPhase) + require.Equal(t, distribution.SplitJobPhaseBackfill, got.AbandonFromPhase) +} + +func TestDistributionServerRetrySplitJob_RequiresCatalogLeader(t *testing.T) { + t.Parallel() + + ctx := context.Background() + baseStore := store.NewMVCCStore() + catalog := distribution.NewCatalogStore(baseStore) + job := sampleDistributionSplitJob(14) + job.Phase = distribution.SplitJobPhaseFailed + job.RetryPhase = distribution.SplitJobPhaseBackfill + require.NoError(t, catalog.CreateSplitJob(ctx, job)) + + coordinator := newDistributionCoordinatorStub(baseStore, false) + s := NewDistributionServer(distribution.NewEngine(), catalog, WithDistributionCoordinator(coordinator)) + _, err := s.RetrySplitJob(ctx, &pb.RetrySplitJobRequest{JobId: job.JobID}) + require.Error(t, err) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) + require.Zero(t, coordinator.dispatchCalls) +} + func TestDistributionServerGetRouteOwnership_UsesExactVersionSnapshot(t *testing.T) { t.Parallel() @@ -781,6 +912,19 @@ func seededDistributionServerWithoutCoordinator(t *testing.T) (*DistributionServ return NewDistributionServer(distribution.NewEngine(), catalog), saved.Version } +func sampleDistributionSplitJob(jobID uint64) distribution.SplitJob { + return distribution.SplitJob{ + JobID: jobID, + SourceRouteID: 1, + SplitKey: []byte("g"), + TargetGroupID: 2, + Phase: distribution.SplitJobPhaseBackfill, + RetryPhase: distribution.SplitJobPhaseNone, + StartedAtMs: 1000, + UpdatedAtMs: 1000, + } +} + type distributionCoordinatorStub struct { store store.MVCCStore leader bool @@ -814,16 +958,17 @@ func (s *distributionCoordinatorStub) Dispatch(ctx context.Context, reqs *kv.Ope if s.asyncApplyDelay > 0 { done := s.asyncApplyDone delay := s.asyncApplyDelay + readKeys := cloneByteSlices(reqs.ReadKeys) go func() { time.Sleep(delay) - err := s.applyDispatch(ctx, mutations, startTS, commitTS) + err := s.applyDispatch(ctx, mutations, readKeys, startTS, commitTS) if done != nil { done <- err } }() return &kv.CoordinateResponse{CommitIndex: commitTS}, nil } - if err := s.applyDispatch(ctx, mutations, startTS, commitTS); err != nil { + if err := s.applyDispatch(ctx, mutations, reqs.ReadKeys, startTS, commitTS); err != nil { return nil, err } return &kv.CoordinateResponse{CommitIndex: commitTS}, nil @@ -854,10 +999,11 @@ func (s *distributionCoordinatorStub) nextTimestamps(startTS uint64) (uint64, ui func (s *distributionCoordinatorStub) applyDispatch( ctx context.Context, mutations []*store.KVPairMutation, + readKeys [][]byte, startTS uint64, commitTS uint64, ) error { - if err := s.store.ApplyMutations(ctx, mutations, nil, startTS, commitTS); err != nil { + if err := s.store.ApplyMutations(ctx, mutations, readKeys, startTS, commitTS); err != nil { return err } if s.afterDispatch != nil { diff --git a/distribution/split_job_catalog.go b/distribution/split_job_catalog.go index 808ad8498..082f9aaf6 100644 --- a/distribution/split_job_catalog.go +++ b/distribution/split_job_catalog.go @@ -844,6 +844,11 @@ func CloneSplitJob(job SplitJob) SplitJob { } } +// SplitJobToProto converts a catalog SplitJob into its wire representation. +func SplitJobToProto(job SplitJob) *pb.SplitJob { + return splitJobToProto(job) +} + func splitJobToProto(job SplitJob) *pb.SplitJob { job = CloneSplitJob(job) return &pb.SplitJob{ diff --git a/distribution/split_job_catalog_test.go b/distribution/split_job_catalog_test.go index 0ac91a59b..cdeef3c3a 100644 --- a/distribution/split_job_catalog_test.go +++ b/distribution/split_job_catalog_test.go @@ -242,6 +242,98 @@ func TestCatalogStoreSaveSplitJobRejectsStaleExpectedJob(t *testing.T) { assertSplitJobEqual(t, advanced, got) } +func TestCatalogStoreRetrySplitJobUsesDurableRetryPhase(t *testing.T) { + cs := NewCatalogStore(store.NewMVCCStore()) + ctx := context.Background() + job := sampleSplitJob(18) + job.Phase = SplitJobPhaseFailed + job.RetryPhase = SplitJobPhaseDeltaCopy + job.LastError = "transient" + + if err := cs.CreateSplitJob(ctx, job); err != nil { + t.Fatalf("create split job: %v", err) + } + retried, err := cs.RetrySplitJob(ctx, job.JobID, 1200) + if err != nil { + t.Fatalf("retry split job: %v", err) + } + if retried.Phase != SplitJobPhaseDeltaCopy || retried.RetryPhase != SplitJobPhaseNone { + t.Fatalf("unexpected retry transition: %+v", retried) + } + if retried.AbandonFromPhase != SplitJobPhaseNone || retried.LastError != "" || retried.UpdatedAtMs != 1200 { + t.Fatalf("unexpected retry witness cleanup: %+v", retried) + } + got, found, err := cs.SplitJob(ctx, job.JobID) + if err != nil { + t.Fatalf("load retried split job: %v", err) + } + if !found { + t.Fatal("expected retried split job") + } + assertSplitJobEqual(t, retried, got) +} + +func TestCatalogStoreRetrySplitJobRejectsMissingRetryWitness(t *testing.T) { + cs := NewCatalogStore(store.NewMVCCStore()) + ctx := context.Background() + job := sampleSplitJob(19) + job.Phase = SplitJobPhaseFailed + job.RetryPhase = SplitJobPhaseNone + + if err := cs.CreateSplitJob(ctx, job); err != nil { + t.Fatalf("create split job: %v", err) + } + if _, err := cs.RetrySplitJob(ctx, job.JobID, 1200); !errors.Is(err, ErrCatalogSplitJobCannotRetry) { + t.Fatalf("expected ErrCatalogSplitJobCannotRetry, got %v", err) + } +} + +func TestCatalogStoreBeginSplitJobAbandonRecordsPreCutoverPhase(t *testing.T) { + cs := NewCatalogStore(store.NewMVCCStore()) + ctx := context.Background() + job := sampleSplitJob(20) + job.Phase = SplitJobPhaseFailed + job.RetryPhase = SplitJobPhaseFence + job.AbandonFromPhase = SplitJobPhaseNone + + if err := cs.CreateSplitJob(ctx, job); err != nil { + t.Fatalf("create split job: %v", err) + } + abandoning, err := cs.BeginSplitJobAbandon(ctx, job.JobID, 1300) + if err != nil { + t.Fatalf("begin split job abandon: %v", err) + } + if abandoning.Phase != SplitJobPhaseAbandoning || + abandoning.RetryPhase != SplitJobPhaseNone || + abandoning.AbandonFromPhase != SplitJobPhaseFence || + abandoning.UpdatedAtMs != 1300 { + t.Fatalf("unexpected abandon transition: %+v", abandoning) + } + got, found, err := cs.SplitJob(ctx, job.JobID) + if err != nil { + t.Fatalf("load abandoning split job: %v", err) + } + if !found { + t.Fatal("expected abandoning split job") + } + assertSplitJobEqual(t, abandoning, got) +} + +func TestCatalogStoreBeginSplitJobAbandonRejectsPostCutoverPhases(t *testing.T) { + cs := NewCatalogStore(store.NewMVCCStore()) + ctx := context.Background() + job := sampleSplitJob(21) + job.Phase = SplitJobPhaseCleanup + job.RetryPhase = SplitJobPhaseNone + + if err := cs.CreateSplitJob(ctx, job); err != nil { + t.Fatalf("create split job: %v", err) + } + if _, err := cs.BeginSplitJobAbandon(ctx, job.JobID, 1300); !errors.Is(err, ErrCatalogSplitJobCannotAbandon) { + t.Fatalf("expected ErrCatalogSplitJobCannotAbandon, got %v", err) + } +} + func TestCatalogStoreListSplitJobsIncludesLiveAndHistory(t *testing.T) { cs := NewCatalogStore(store.NewMVCCStore()) ctx := context.Background() diff --git a/distribution/split_job_lifecycle.go b/distribution/split_job_lifecycle.go new file mode 100644 index 000000000..2f6170c1c --- /dev/null +++ b/distribution/split_job_lifecycle.go @@ -0,0 +1,145 @@ +package distribution + +import ( + "bytes" + "context" + + "github.com/cockroachdb/errors" +) + +var ( + ErrCatalogSplitJobNotFound = errors.New("catalog split job is not found") + ErrCatalogSplitJobCannotRetry = errors.New("catalog split job cannot be retried") + ErrCatalogSplitJobCannotAbandon = errors.New("catalog split job cannot be abandoned") +) + +// RetrySplitJobState returns the durable state transition for RetrySplitJob. +// The retry target must come from the recorded RetryPhase witness; callers must +// not infer a phase from cursors, timestamps, or diagnostics after restart. +func RetrySplitJobState(job SplitJob, nowMs int64) (SplitJob, error) { + out := CloneSplitJob(job) + if out.Phase != SplitJobPhaseFailed || !splitJobRetryPhaseAllowed(out.RetryPhase) { + return SplitJob{}, errors.WithStack(ErrCatalogSplitJobCannotRetry) + } + out.Phase = out.RetryPhase + out.RetryPhase = SplitJobPhaseNone + out.AbandonFromPhase = SplitJobPhaseNone + out.LastError = "" + out.UpdatedAtMs = nowMs + return out, nil +} + +// BeginSplitJobAbandon records the durable ABANDONING witness before any +// pre-CUTOVER cleanup side effect is removed. +func BeginSplitJobAbandon(job SplitJob, nowMs int64) (SplitJob, error) { + out := CloneSplitJob(job) + from, ok := splitJobAbandonFromPhase(out) + if !ok { + return SplitJob{}, errors.WithStack(ErrCatalogSplitJobCannotAbandon) + } + if out.Phase == SplitJobPhaseAbandoning { + return out, nil + } + out.Phase = SplitJobPhaseAbandoning + out.RetryPhase = SplitJobPhaseNone + out.AbandonFromPhase = from + out.UpdatedAtMs = nowMs + return out, nil +} + +// RetrySplitJob CASes a FAILED live job back to its recorded retry phase. +func (s *CatalogStore) RetrySplitJob(ctx context.Context, jobID uint64, nowMs int64) (SplitJob, error) { + expected, _, err := s.LiveSplitJobForUpdate(ctx, jobID) + if err != nil { + return SplitJob{}, err + } + next, err := RetrySplitJobState(expected, nowMs) + if err != nil { + return SplitJob{}, err + } + return next, s.SaveSplitJob(ctx, expected, next) +} + +// BeginSplitJobAbandon CASes a live pre-CUTOVER job into ABANDONING. +func (s *CatalogStore) BeginSplitJobAbandon(ctx context.Context, jobID uint64, nowMs int64) (SplitJob, error) { + expected, _, err := s.LiveSplitJobForUpdate(ctx, jobID) + if err != nil { + return SplitJob{}, err + } + next, err := BeginSplitJobAbandon(expected, nowMs) + if err != nil { + return SplitJob{}, err + } + if splitJobsEquivalent(expected, next) { + return next, nil + } + return next, s.SaveSplitJob(ctx, expected, next) +} + +// LiveSplitJobForUpdate reads the latest live split job and the snapshot +// timestamp used for the read. RPC callers use readTS as a CAS StartTS when +// proposing the update through Raft. +func (s *CatalogStore) LiveSplitJobForUpdate(ctx context.Context, jobID uint64) (SplitJob, uint64, error) { + if err := ensureCatalogStore(s); err != nil { + return SplitJob{}, 0, err + } + if jobID == 0 { + return SplitJob{}, 0, errors.WithStack(ErrCatalogSplitJobIDRequired) + } + ctx = contextOrBackground(ctx) + readTS := s.store.LastCommitTS() + job, found, err := s.liveSplitJobAt(ctx, jobID, readTS) + if err != nil { + return SplitJob{}, 0, err + } + if !found { + return SplitJob{}, 0, errors.WithStack(ErrCatalogSplitJobNotFound) + } + return job, readTS, nil +} + +func splitJobRetryPhaseAllowed(phase SplitJobPhase) bool { + switch phase { + case SplitJobPhaseBackfill, SplitJobPhaseFence, SplitJobPhaseDeltaCopy, SplitJobPhaseCutover, SplitJobPhaseCleanup: + return true + case SplitJobPhaseNone, SplitJobPhasePlanned, SplitJobPhaseDone, SplitJobPhaseFailed, SplitJobPhaseAbandoning, SplitJobPhaseAbandoned: + return false + } + return false +} + +func splitJobAbandonFromPhase(job SplitJob) (SplitJobPhase, bool) { + switch job.Phase { + case SplitJobPhasePlanned, SplitJobPhaseBackfill, SplitJobPhaseFence, SplitJobPhaseDeltaCopy: + return job.Phase, true + case SplitJobPhaseFailed: + if splitJobAbandonRetryPhaseAllowed(job.RetryPhase) { + return job.RetryPhase, true + } + case SplitJobPhaseAbandoning: + if splitJobAbandonRetryPhaseAllowed(job.AbandonFromPhase) { + return job.AbandonFromPhase, true + } + case SplitJobPhaseNone, SplitJobPhaseCutover, SplitJobPhaseCleanup, SplitJobPhaseDone, SplitJobPhaseAbandoned: + } + return SplitJobPhaseNone, false +} + +func splitJobAbandonRetryPhaseAllowed(phase SplitJobPhase) bool { + switch phase { + case SplitJobPhasePlanned, SplitJobPhaseBackfill, SplitJobPhaseFence, SplitJobPhaseDeltaCopy: + return true + case SplitJobPhaseNone, SplitJobPhaseCutover, SplitJobPhaseCleanup, SplitJobPhaseDone, SplitJobPhaseFailed, SplitJobPhaseAbandoning, SplitJobPhaseAbandoned: + return false + } + return false +} + +func splitJobsEquivalent(left, right SplitJob) bool { + leftRaw, leftErr := EncodeSplitJob(left) + rightRaw, rightErr := EncodeSplitJob(right) + if leftErr != nil || rightErr != nil { + return false + } + return bytes.Equal(leftRaw, rightRaw) +} From 58c212801aca3dcc87478ca51707fdf2a3f58818 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 02:57:33 +0900 Subject: [PATCH 03/29] distribution: page split job listing --- adapter/distribution_server.go | 157 +++++++++++++++++++++++++--- adapter/distribution_server_test.go | 102 ++++++++++++++++++ distribution/split_job_lifecycle.go | 6 +- 3 files changed, 247 insertions(+), 18 deletions(-) diff --git a/adapter/distribution_server.go b/adapter/distribution_server.go index 7052a8561..78106c6bb 100644 --- a/adapter/distribution_server.go +++ b/adapter/distribution_server.go @@ -3,12 +3,15 @@ package adapter import ( "bytes" "context" + "encoding/binary" "math" + "sort" "strings" "sync" "time" "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/internal/raftengine" "github.com/bootjp/elastickv/kv" pb "github.com/bootjp/elastickv/proto" "github.com/bootjp/elastickv/store" @@ -62,8 +65,13 @@ func WithCatalogReloadRetryPolicy(attempts int, interval time.Duration) Distribu } const ( - childRouteCount = 2 - splitMutationOpCount = childRouteCount + 3 + childRouteCount = 2 + splitMutationOpCount = childRouteCount + 3 + listSplitJobsDefaultPageSize = 200 + splitJobListCursorVersion = byte(1) + splitJobListCursorTerminalOff = 1 + splitJobListCursorJobIDOff = splitJobListCursorTerminalOff + 8 + splitJobListCursorEncodedBytes = splitJobListCursorJobIDOff + 8 ) var ( @@ -220,12 +228,9 @@ func (s *DistributionServer) ListSplitJobs(ctx context.Context, req *pb.ListSpli if err != nil { return nil, splitJobCatalogStatusError(err) } - resp := &pb.ListSplitJobsResponse{} - for _, job := range jobs { - if !splitJobPassesListFilter(job, req.GetSinceTerminalAtMs(), phaseFilter) { - continue - } - resp.Jobs = append(resp.Jobs, distribution.SplitJobToProto(job)) + resp, err := splitJobListPage(jobs, req, phaseFilter) + if err != nil { + return nil, err } return resp, nil } @@ -508,7 +513,7 @@ func (s *DistributionServer) updateSplitJobViaCoordinator( if err != nil { return splitJobCatalogStatusError(err) } - if splitJobsEqualByEncoding(expected, next) { + if distribution.SplitJobsEquivalent(expected, next) { return nil } encoded, err := distribution.EncodeSplitJob(next) @@ -655,12 +660,6 @@ func validSplitJobPhaseFilter(phase string) bool { return false } -func splitJobsEqualByEncoding(left, right distribution.SplitJob) bool { - leftRaw, leftErr := distribution.EncodeSplitJob(left) - rightRaw, rightErr := distribution.EncodeSplitJob(right) - return leftErr == nil && rightErr == nil && bytes.Equal(leftRaw, rightRaw) -} - func splitJobCatalogStatusError(err error) error { switch { case errors.Is(err, distribution.ErrCatalogSplitJobIDRequired): @@ -681,13 +680,139 @@ func splitJobCoordinatorStatusError(err error) error { switch { case errors.Is(err, store.ErrWriteConflict): return grpcStatusError(codes.Aborted, distribution.ErrCatalogSplitJobConflict.Error()) - case errors.Is(err, kv.ErrLeaderNotFound): + case splitJobCoordinatorLeadershipError(err): return grpcStatusError(codes.FailedPrecondition, errDistributionNotLeader.Error()) default: return grpcStatusErrorf(codes.Internal, "commit split job mutation: %v", err) } } +func splitJobListPage(jobs []distribution.SplitJob, req *pb.ListSplitJobsRequest, phaseFilter string) (*pb.ListSplitJobsResponse, error) { + filtered := make([]distribution.SplitJob, 0, len(jobs)) + for _, job := range jobs { + if splitJobPassesListFilter(job, req.GetSinceTerminalAtMs(), phaseFilter) { + filtered = append(filtered, job) + } + } + sortSplitJobsForList(filtered) + + start, err := splitJobListStartIndex(filtered, req.GetPageCursor()) + if err != nil { + return nil, err + } + end := start + listSplitJobsDefaultPageSize + if end > len(filtered) { + end = len(filtered) + } + + resp := &pb.ListSplitJobsResponse{ + Jobs: make([]*pb.SplitJob, 0, end-start), + } + for _, job := range filtered[start:end] { + resp.Jobs = append(resp.Jobs, distribution.SplitJobToProto(job)) + } + if end < len(filtered) { + resp.NextPageCursor = encodeSplitJobListCursor(filtered[end-1]) + } + return resp, nil +} + +func sortSplitJobsForList(jobs []distribution.SplitJob) { + sort.Slice(jobs, func(i, j int) bool { + left, right := jobs[i], jobs[j] + leftLive := left.TerminalAtMs <= 0 + rightLive := right.TerminalAtMs <= 0 + if leftLive != rightLive { + return leftLive + } + if leftLive { + if left.UpdatedAtMs != right.UpdatedAtMs { + return left.UpdatedAtMs > right.UpdatedAtMs + } + return left.JobID > right.JobID + } + if left.TerminalAtMs != right.TerminalAtMs { + return left.TerminalAtMs > right.TerminalAtMs + } + return left.JobID > right.JobID + }) +} + +func splitJobListStartIndex(jobs []distribution.SplitJob, cursor []byte) (int, error) { + if len(cursor) == 0 { + return 0, nil + } + terminalAtMs, jobID, err := decodeSplitJobListCursor(cursor) + if err != nil { + return 0, err + } + for i, job := range jobs { + if job.JobID == jobID && splitJobListCursorTerminalAtMs(job) == terminalAtMs { + return i + 1, nil + } + } + return 0, grpcStatusError(codes.InvalidArgument, "split job page cursor does not match the filtered result set") +} + +func encodeSplitJobListCursor(job distribution.SplitJob) []byte { + cursor := make([]byte, splitJobListCursorEncodedBytes) + cursor[0] = splitJobListCursorVersion + binary.BigEndian.PutUint64( + cursor[splitJobListCursorTerminalOff:splitJobListCursorJobIDOff], + splitJobListCursorTerminalAtMs(job), + ) + binary.BigEndian.PutUint64(cursor[splitJobListCursorJobIDOff:], job.JobID) + return cursor +} + +func decodeSplitJobListCursor(cursor []byte) (uint64, uint64, error) { + if len(cursor) != splitJobListCursorEncodedBytes || cursor[0] != splitJobListCursorVersion { + return 0, 0, grpcStatusError(codes.InvalidArgument, "invalid split job page cursor") + } + terminalAtMs := binary.BigEndian.Uint64(cursor[splitJobListCursorTerminalOff:splitJobListCursorJobIDOff]) + jobID := binary.BigEndian.Uint64(cursor[splitJobListCursorJobIDOff:]) + if jobID == 0 { + return 0, 0, grpcStatusError(codes.InvalidArgument, "invalid split job page cursor") + } + return terminalAtMs, jobID, nil +} + +func splitJobListCursorTerminalAtMs(job distribution.SplitJob) uint64 { + if job.TerminalAtMs <= 0 { + return 0 + } + return uint64(job.TerminalAtMs) +} + +func splitJobCoordinatorLeadershipError(err error) bool { + if err == nil { + return false + } + if errors.Is(err, kv.ErrLeaderNotFound) || + errors.Is(err, raftengine.ErrNotLeader) || + errors.Is(err, raftengine.ErrLeadershipLost) || + errors.Is(err, raftengine.ErrLeadershipTransferInProgress) { + return true + } + return hasSplitJobLeaderErrorSuffix(err.Error()) +} + +var splitJobLeaderErrorPhrases = []string{ + "not leader", + "leader not found", + "leadership lost", + "leadership transfer in progress", +} + +func hasSplitJobLeaderErrorSuffix(msg string) bool { + for _, phrase := range splitJobLeaderErrorPhrases { + if len(msg) >= len(phrase) && strings.EqualFold(msg[len(msg)-len(phrase):], phrase) { + return true + } + } + return false +} + func splitCatalogRoutes( parent distribution.RouteDescriptor, splitKey []byte, diff --git a/adapter/distribution_server_test.go b/adapter/distribution_server_test.go index 434180c0a..109f21d62 100644 --- a/adapter/distribution_server_test.go +++ b/adapter/distribution_server_test.go @@ -2,10 +2,12 @@ package adapter import ( "context" + "errors" "testing" "time" "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/internal/raftengine" "github.com/bootjp/elastickv/kv" pb "github.com/bootjp/elastickv/proto" "github.com/bootjp/elastickv/store" @@ -189,6 +191,57 @@ func TestDistributionServerSplitJobRPCs_ReadAndListCatalogJobs(t *testing.T) { require.Equal(t, codes.InvalidArgument, status.Code(err)) } +func TestDistributionServerListSplitJobs_PaginatesNewestHistory(t *testing.T) { + t.Parallel() + + ctx := context.Background() + catalog := distribution.NewCatalogStore(store.NewMVCCStore()) + for jobID, terminalAtMs := uint64(1), int64(1001); jobID <= 205; jobID, terminalAtMs = jobID+1, terminalAtMs+1 { + job := sampleDistributionSplitJob(jobID) + job.Phase = distribution.SplitJobPhaseDone + job.TerminalAtMs = terminalAtMs + job.UpdatedAtMs = job.TerminalAtMs + require.NoError(t, catalog.CreateSplitJob(ctx, job)) + require.NoError(t, catalog.MoveSplitJobToHistory(ctx, job, job)) + } + + s := NewDistributionServer(distribution.NewEngine(), catalog) + first, err := s.ListSplitJobs(ctx, &pb.ListSplitJobsRequest{}) + require.NoError(t, err) + require.Len(t, first.Jobs, listSplitJobsDefaultPageSize) + require.NotEmpty(t, first.NextPageCursor) + require.Equal(t, uint64(205), first.Jobs[0].JobId) + require.Equal(t, uint64(6), first.Jobs[len(first.Jobs)-1].JobId) + + second, err := s.ListSplitJobs(ctx, &pb.ListSplitJobsRequest{PageCursor: first.NextPageCursor}) + require.NoError(t, err) + require.Empty(t, second.NextPageCursor) + require.Equal(t, []uint64{5, 4, 3, 2, 1}, splitJobIDs(second.Jobs)) +} + +func TestDistributionServerListSplitJobs_RejectsInvalidCursor(t *testing.T) { + t.Parallel() + + ctx := context.Background() + catalog := distribution.NewCatalogStore(store.NewMVCCStore()) + job := sampleDistributionSplitJob(1) + job.Phase = distribution.SplitJobPhaseDone + job.TerminalAtMs = 1000 + require.NoError(t, catalog.CreateSplitJob(ctx, job)) + require.NoError(t, catalog.MoveSplitJobToHistory(ctx, job, job)) + + s := NewDistributionServer(distribution.NewEngine(), catalog) + _, err := s.ListSplitJobs(ctx, &pb.ListSplitJobsRequest{PageCursor: []byte("bad")}) + require.Error(t, err) + require.Equal(t, codes.InvalidArgument, status.Code(err)) + + missing := sampleDistributionSplitJob(999) + missing.TerminalAtMs = 9999 + _, err = s.ListSplitJobs(ctx, &pb.ListSplitJobsRequest{PageCursor: encodeSplitJobListCursor(missing)}) + require.Error(t, err) + require.Equal(t, codes.InvalidArgument, status.Code(err)) +} + func TestDistributionServerRetrySplitJob_UsesCoordinatorCAS(t *testing.T) { t.Parallel() @@ -215,6 +268,43 @@ func TestDistributionServerRetrySplitJob_UsesCoordinatorCAS(t *testing.T) { require.Empty(t, got.LastError) } +func TestDistributionServerRetrySplitJob_MapsDispatchLeadershipLoss(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + err error + }{ + {name: "leader not found", err: kv.ErrLeaderNotFound}, + {name: "raft not leader", err: raftengine.ErrNotLeader}, + {name: "leadership lost", err: raftengine.ErrLeadershipLost}, + {name: "transfer in progress", err: raftengine.ErrLeadershipTransferInProgress}, + {name: "wrapped grpc detail", err: errors.New("rpc error: code = Unknown desc = raft engine: leadership lost")}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + ctx := context.Background() + baseStore := store.NewMVCCStore() + catalog := distribution.NewCatalogStore(baseStore) + job := sampleDistributionSplitJob(15) + job.Phase = distribution.SplitJobPhaseFailed + job.RetryPhase = distribution.SplitJobPhaseFence + require.NoError(t, catalog.CreateSplitJob(ctx, job)) + + coordinator := newDistributionCoordinatorStub(baseStore, true) + coordinator.dispatchErr = tc.err + s := NewDistributionServer(distribution.NewEngine(), catalog, WithDistributionCoordinator(coordinator)) + _, err := s.RetrySplitJob(ctx, &pb.RetrySplitJobRequest{JobId: job.JobID}) + require.Error(t, err) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) + require.ErrorContains(t, err, errDistributionNotLeader.Error()) + require.Equal(t, 1, coordinator.dispatchCalls) + }) + } +} + func TestDistributionServerAbandonSplitJob_RecordsAbandoningViaCoordinator(t *testing.T) { t.Parallel() @@ -925,9 +1015,18 @@ func sampleDistributionSplitJob(jobID uint64) distribution.SplitJob { } } +func splitJobIDs(jobs []*pb.SplitJob) []uint64 { + ids := make([]uint64, 0, len(jobs)) + for _, job := range jobs { + ids = append(ids, job.GetJobId()) + } + return ids +} + type distributionCoordinatorStub struct { store store.MVCCStore leader bool + dispatchErr error nextTS uint64 lastStartTS uint64 afterDispatch func(context.Context, store.MVCCStore, uint64) error @@ -948,6 +1047,9 @@ func (s *distributionCoordinatorStub) Dispatch(ctx context.Context, reqs *kv.Ope return nil, err } s.dispatchCalls++ + if s.dispatchErr != nil { + return nil, s.dispatchErr + } startTS, commitTS := s.nextTimestamps(reqs.StartTS) s.lastStartTS = startTS diff --git a/distribution/split_job_lifecycle.go b/distribution/split_job_lifecycle.go index 2f6170c1c..f2427a4e0 100644 --- a/distribution/split_job_lifecycle.go +++ b/distribution/split_job_lifecycle.go @@ -70,7 +70,7 @@ func (s *CatalogStore) BeginSplitJobAbandon(ctx context.Context, jobID uint64, n if err != nil { return SplitJob{}, err } - if splitJobsEquivalent(expected, next) { + if SplitJobsEquivalent(expected, next) { return next, nil } return next, s.SaveSplitJob(ctx, expected, next) @@ -135,7 +135,9 @@ func splitJobAbandonRetryPhaseAllowed(phase SplitJobPhase) bool { return false } -func splitJobsEquivalent(left, right SplitJob) bool { +// SplitJobsEquivalent reports whether two split jobs encode to the same +// durable catalog value. +func SplitJobsEquivalent(left, right SplitJob) bool { leftRaw, leftErr := EncodeSplitJob(left) rightRaw, rightErr := EncodeSplitJob(right) if leftErr != nil || rightErr != nil { From 5ba3e82db274ca12a076e70d4092048e811d565e Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 03:18:10 +0900 Subject: [PATCH 04/29] kv: enforce target readiness guards --- kv/fsm.go | 110 ++++++++++++++++++++++++++---- kv/fsm_migration_fence_test.go | 57 ++++++++++++++++ kv/route_history.go | 18 +++++ kv/shard_store.go | 57 ++++++++++++++-- kv/shard_store_test.go | 98 ++++++++++++++++++++++++-- store/lsm_store.go | 42 +++++++++--- store/lsm_store_test.go | 33 +++++++++ store/mvcc_store.go | 105 +++++++++++++++++++++++----- store/mvcc_store_snapshot_test.go | 55 +++++++++++++++ 9 files changed, 526 insertions(+), 49 deletions(-) diff --git a/kv/fsm.go b/kv/fsm.go index 269b13954..a4bc98dfe 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -127,6 +127,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 @@ -501,17 +507,7 @@ func (f *kvFSM) handleRawRequest(ctx context.Context, r *pb.Request, commitTS ui } for _, mut := range r.Mutations { - if mut == nil || len(mut.Key) == 0 { - return errors.WithStack(ErrInvalidRequest) - } - // Raw requests should not mutate txn-internal keys. - if isTxnInternalKey(mut.Key) { - return errors.WithStack(ErrInvalidRequest) - } - if err := f.verifyRouteNotFencedForKey(mut.Key); err != nil { - return err - } - if err := f.assertNoConflictingTxnLock(ctx, mut.Key, nil, 0); err != nil { + if err := f.verifyRawMutationCanApply(ctx, mut, commitTS); err != nil { return err } } @@ -529,6 +525,26 @@ func (f *kvFSM) handleRawRequest(ctx context.Context, r *pb.Request, commitTS ui return nil } +func (f *kvFSM) verifyRawMutationCanApply(ctx context.Context, mut *pb.Mutation, commitTS uint64) error { + if mut == nil || len(mut.Key) == 0 { + return errors.WithStack(ErrInvalidRequest) + } + // Raw requests should not mutate txn-internal keys. + if isTxnInternalKey(mut.Key) { + return errors.WithStack(ErrInvalidRequest) + } + if err := f.verifyRouteNotFencedForKey(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 + } + return nil +} + // extractDelPrefix checks if the mutations contain a DEL_PREFIX operation. // If found, it validates that no other operation types are mixed in. func extractDelPrefix(muts []*pb.Mutation) (bool, []byte) { @@ -546,6 +562,9 @@ func (f *kvFSM) handleDelPrefix(ctx context.Context, prefix []byte, commitTS uin if err := f.verifyRouteNotFencedForPrefix(prefix); err != nil { return err } + if err := f.verifyRouteWriteFloorForPrefix(prefix, commitTS); err != nil { + return err + } if err := f.store.DeletePrefixAtRaftAt(ctx, prefix, txnCommonPrefix, commitTS, f.pendingApplyIdx); err != nil { return errors.WithStack(err) } @@ -595,6 +614,50 @@ 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) 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 @@ -970,6 +1033,9 @@ func (f *kvFSM) handlePrepareRequest(ctx context.Context, r *pb.Request) error { if err := f.verifyRouteNotFencedForMutations(uniq); err != nil { return err } + if err := f.verifyRouteWriteFloorForMutations(uniq, startTS); err != nil { + return err + } if err := f.validateConflicts(ctx, uniq, startTS); err != nil { return errors.WithStack(err) } @@ -1036,7 +1102,7 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com return nil } - uniq, err := f.uniqueMutationsNotFenced(muts) + uniq, err := f.uniqueMutationsNotFenced(muts, commitTS) if err != nil { return err } @@ -1052,7 +1118,7 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com return nil } -func (f *kvFSM) uniqueMutationsNotFenced(muts []*pb.Mutation) ([]*pb.Mutation, error) { +func (f *kvFSM) uniqueMutationsNotFenced(muts []*pb.Mutation, commitTS uint64) ([]*pb.Mutation, error) { uniq, err := uniqueMutations(muts) if err != nil { return nil, err @@ -1060,6 +1126,9 @@ func (f *kvFSM) uniqueMutationsNotFenced(muts []*pb.Mutation) ([]*pb.Mutation, e if err := f.verifyRouteNotFencedForMutations(uniq); err != nil { return nil, err } + if err := f.verifyRouteWriteFloorForMutations(uniq, commitTS); err != nil { + return nil, err + } return uniq, nil } @@ -1102,7 +1171,7 @@ func (f *kvFSM) handleCommitRequest(ctx context.Context, r *pb.Request) error { if err != nil { return err } - uniq, err := uniqueMutations(muts) + uniq, err := f.uniqueMutationsAboveWriteFloor(muts, commitTS) if err != nil { return err } @@ -1219,7 +1288,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.uniqueMutationsAboveWriteFloor(muts, abortTS) if err != nil { return err } @@ -1239,6 +1308,17 @@ 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(muts []*pb.Mutation, commitTS uint64) ([]*pb.Mutation, error) { + uniq, err := uniqueMutations(muts) + if err != nil { + return nil, err + } + if err := f.verifyRouteWriteFloorForMutations(uniq, commitTS); err != nil { + return nil, err + } + return uniq, nil +} + 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 e47e71f95..00b02e31b 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -22,6 +22,23 @@ func newWriteFencedFSM(t *testing.T) *kvFSM { return newComposed1FSM(t, engine, 1) } +func newWriteFloorFSM(t *testing.T) *kvFSM { + t.Helper() + + engine := distribution.NewEngine() + applyComposed1Snapshot(t, engine, 1, []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + MinWriteTSExclusive: 100, + }, + }) + return newComposed1FSM(t, engine, 1) +} + func TestFSMRejectsRawPointWriteOnWriteFencedRoute(t *testing.T) { t.Parallel() @@ -67,6 +84,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() @@ -84,6 +130,17 @@ 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 TestFSMRejectsPrepareOnWriteFencedRouteButAllowsAbort(t *testing.T) { t.Parallel() diff --git a/kv/route_history.go b/kv/route_history.go index c6ac85608..bd6989513 100644 --- a/kv/route_history.go +++ b/kv/route_history.go @@ -74,3 +74,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 aaf18b57c..f2e1beba6 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -63,12 +63,13 @@ func (s *ShardStore) GetGroupAt(ctx context.Context, groupID uint64, key []byte, if !ok || g.Store == nil { return nil, store.ErrKeyNotFound } + route := s.explicitGroupRouteForKey(groupID, key) if engineForGroup(g) == nil { - return s.localGetAt(ctx, g, distribution.Route{}, key, ts) + return s.localGetAt(ctx, g, route, key, ts) } if isLinearizableRaftLeader(ctx, engineForGroup(g)) { - return s.leaderGetAt(ctx, g, distribution.Route{}, key, ts) + return s.leaderGetAt(ctx, g, route, key, ts) } return s.proxyRawGet(ctx, g, key, ts, groupID) } @@ -142,8 +143,9 @@ func (s *ShardStore) verifyTargetReadinessForRange(ctx context.Context, g *Shard if err != nil { return errors.WithStack(err) } + routeStart, routeEnd := readinessRouteRange(start, end) for _, ready := range states { - if !ready.Armed || !routeRangeIntersects(start, end, ready.RouteStart, ready.RouteEnd) { + if !ready.Armed || !routeRangeIntersects(routeStart, routeEnd, ready.RouteStart, ready.RouteEnd) { continue } if !routeSatisfiesTargetReadiness(route, ready) { @@ -153,6 +155,18 @@ func (s *ShardStore) verifyTargetReadinessForRange(ctx context.Context, g *Shard return nil } +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 verifyRouteWriteFloor(route distribution.Route, commitTS uint64) error { if route.MinWriteTSExclusive == 0 || commitTS > route.MinWriteTSExclusive { return nil @@ -350,7 +364,7 @@ func (s *ShardStore) ScanGroupAt(ctx context.Context, groupID uint64, start []by if limit <= 0 { return []*store.KVPair{}, nil } - return s.scanRouteAtDirection(ctx, distribution.Route{GroupID: groupID}, start, end, limit, ts, false, true) + return s.scanRouteAtDirection(ctx, s.explicitGroupRouteForRange(groupID, start, end), start, end, limit, ts, false, true) } func (s *ShardStore) ReverseScanAt(ctx context.Context, start []byte, end []byte, limit int, ts uint64) ([]*store.KVPair, error) { @@ -558,6 +572,9 @@ func (s *ShardStore) scanRouteAtDirectionPhysicalLimit( if !ok || g == nil || g.Store == nil { return nil, false, nil } + if err := s.verifyTargetReadinessForRange(ctx, g, route, start, end); err != nil { + return nil, false, err + } if engineForGroup(g) == nil { if routeHasStagedVisibility(route) { @@ -1739,6 +1756,9 @@ func (s *ShardStore) ApplyMutationsRaft(ctx context.Context, mutations []*store. if err != nil || group == nil { return err } + if err := s.verifyMutationRoutes(ctx, mutations, readKeys, commitTS); err != nil { + return err + } return errors.WithStack(group.Store.ApplyMutationsRaft(ctx, mutations, readKeys, startTS, commitTS)) } @@ -1750,6 +1770,9 @@ func (s *ShardStore) ApplyMutationsRaftAt(ctx context.Context, mutations []*stor if err != nil || group == nil { return err } + if err := s.verifyMutationRoutes(ctx, mutations, readKeys, commitTS); err != nil { + return err + } return errors.WithStack(group.Store.ApplyMutationsRaftAt(ctx, mutations, readKeys, startTS, commitTS, appliedIndex)) } @@ -2020,6 +2043,32 @@ func (s *ShardStore) groupForKey(key []byte) (*ShardGroup, bool) { return g, ok } +func (s *ShardStore) explicitGroupRouteForKey(groupID uint64, key []byte) distribution.Route { + fallback := distribution.Route{GroupID: groupID} + if s == nil || s.engine == nil { + return fallback + } + route, ok := s.engine.GetRoute(routeKey(key)) + if !ok || route.GroupID != groupID { + return fallback + } + return route +} + +func (s *ShardStore) explicitGroupRouteForRange(groupID uint64, start []byte, end []byte) distribution.Route { + fallback := distribution.Route{GroupID: groupID} + if s == nil || s.engine == nil { + return fallback + } + routeStart, routeEnd := readinessRouteRange(start, end) + for _, route := range s.engine.GetIntersectingRoutes(routeStart, routeEnd) { + if route.GroupID == groupID { + return route + } + } + return fallback +} + func (s *ShardStore) routeAndGroupForKey(key []byte) (distribution.Route, *ShardGroup, bool) { route, ok := s.engine.GetRoute(routeKey(key)) if !ok { diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index cfe9aec59..c1fe77814 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -33,7 +33,7 @@ func newStagedVisibilityShardStore(t *testing.T) (*ShardStore, *ShardGroup) { return NewShardStore(engine, map[uint64]*ShardGroup{1: group}), group } -func applyTargetReadiness(t *testing.T, group *ShardGroup, floor uint64) { +func applyTargetReadiness(t *testing.T, group *ShardGroup) { t.Helper() writer, ok := group.Store.(store.MigrationTargetReadinessWriter) require.True(t, ok) @@ -43,7 +43,7 @@ func applyTargetReadiness(t *testing.T, group *ShardGroup, floor uint64) { RouteEnd: []byte("z"), ExpectedCutoverVersion: 2, MigrationJobID: 9, - MinWriteTSExclusive: floor, + MinWriteTSExclusive: 100, Armed: true, })) } @@ -94,7 +94,7 @@ func TestShardStoreTargetReadinessFailsClosedWithoutDescriptorProof(t *testing.T GroupID: 1, State: distribution.RouteStateActive, }) - applyTargetReadiness(t, group, 100) + applyTargetReadiness(t, group) _, err := st.GetAt(ctx, []byte("k"), 120) require.ErrorIs(t, err, ErrRouteCutoverPending) @@ -103,12 +103,32 @@ func TestShardStoreTargetReadinessFailsClosedWithoutDescriptorProof(t *testing.T 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, 100) + applyTargetReadiness(t, group) rawKey := []byte("k") require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, rawKey), []byte("staged"), 120, 0)) @@ -129,7 +149,7 @@ func TestShardStoreTargetReadinessAcceptsClearedDescriptorAndRetainsFloor(t *tes State: distribution.RouteStateActive, MinWriteTSExclusive: 100, }) - applyTargetReadiness(t, group, 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) @@ -160,6 +180,74 @@ func TestShardStoreTargetReadinessAcceptsClearedDescriptorAndRetainsFloor(t *tes require.NoError(t, err) } +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 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 TestShardStoreApplyMutationsRaftAtEnforcesRouteFloor(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, _ := newReadinessShardStore(t, distribution.RouteDescriptor{ + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + MinWriteTSExclusive: 100, + }) + + err := st.ApplyMutationsRaftAt(ctx, []*store.KVPairMutation{{ + Op: store.OpTypePut, + Key: []byte("b"), + Value: []byte("low"), + }}, nil, 0, 100, 7) + require.ErrorIs(t, err, ErrRouteWriteBelowFloor) +} + func TestShardStoreScanAndLatestCommitTS_MergeStagedVisibility(t *testing.T) { t.Parallel() diff --git a/store/lsm_store.go b/store/lsm_store.go index a1a2ac350..dbd38e3a1 100644 --- a/store/lsm_store.go +++ b/store/lsm_store.go @@ -2273,6 +2273,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) } @@ -2443,22 +2444,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 } - return body, hash, expectedChecksum, lastCommitTS, minRetainedTS, nil + readinessStates, err := readMVCCSnapshotReadinessStates(body, version) + if err != nil { + return nil, nil, 0, 0, 0, nil, err + } + 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) @@ -2478,6 +2491,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 @@ -2497,13 +2514,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 } 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/mvcc_store.go b/store/mvcc_store.go index 1bea2a9f7..926032757 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 @@ -843,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) @@ -913,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 @@ -921,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 +961,54 @@ func (s *mvccStore) restoreStreamingSnapshot(r io.Reader) error { s.tree = tree s.lastCommitTS = lastCommitTS s.minRetainedTS = minRetainedTS + 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 +1023,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 3ebdbcac8..d7dfb93c4 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() From a2ee041227b34923434c2ada1f931c3c4593d128 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 04:19:48 +0900 Subject: [PATCH 05/29] Guard target readiness apply paths --- kv/fsm.go | 105 ++++++++++++++++-- kv/fsm_migration_fence_test.go | 82 ++++++++++++++ kv/route_history.go | 8 ++ kv/shard_store.go | 43 +++++++- kv/shard_store_test.go | 44 ++++++++ kv/shard_store_txn_lock_test.go | 128 ++++++++++++++++++++++ kv/sharded_coordinator.go | 21 ++++ kv/sharded_coordinator_del_prefix_test.go | 27 +++++ 8 files changed, 443 insertions(+), 15 deletions(-) diff --git a/kv/fsm.go b/kv/fsm.go index a4bc98dfe..58fdccf32 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -9,6 +9,7 @@ import ( "log/slog" "os" + "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/internal/encryption/fsmwire" "github.com/bootjp/elastickv/internal/raftengine" "github.com/bootjp/elastickv/internal/s3keys" @@ -121,6 +122,13 @@ 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 route descriptor that covered 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 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. WriteFencedForKey(key []byte) bool @@ -536,6 +544,9 @@ func (f *kvFSM) verifyRawMutationCanApply(ctx context.Context, mut *pb.Mutation, 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 } @@ -562,6 +573,9 @@ 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 } @@ -626,6 +640,73 @@ func (f *kvFSM) verifyRouteWriteFloorForMutations(muts []*pb.Mutation, commitTS 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) 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 { + reader, ok := f.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 + } + + var ( + snap RouteSnapshot + proof bool + ) + if f.routes != nil { + snap, proof = f.routes.Current() + } + for _, ready := range states { + if !ready.Armed || !routeRangeIntersects(routeStart, routeEnd, ready.RouteStart, ready.RouteEnd) { + continue + } + if !proof || !routesSatisfyTargetReadiness(snap.IntersectingRoutes(routeStart, routeEnd), ready) { + return errors.WithStack(ErrRouteCutoverPending) + } + } + return nil +} + +func routesSatisfyTargetReadiness(routes []distribution.Route, ready store.TargetStagedReadinessState) bool { + matched := false + for _, route := range routes { + if !routeRangeIntersects(route.Start, route.End, ready.RouteStart, ready.RouteEnd) { + continue + } + matched = true + if !routeSatisfiesTargetReadiness(route, ready) { + return false + } + } + return matched +} + func (f *kvFSM) verifyRouteWriteFloorForKey(key []byte, commitTS uint64) error { if f.routes == nil { return nil @@ -1026,16 +1107,10 @@ func (f *kvFSM) handlePrepareRequest(ctx context.Context, r *pb.Request) error { } startTS := r.Ts - uniq, err := uniqueMutations(muts) + uniq, err := f.uniqueMutationsNotFenced(ctx, muts, startTS) if err != nil { return err } - if err := f.verifyRouteNotFencedForMutations(uniq); err != nil { - return err - } - if err := f.verifyRouteWriteFloorForMutations(uniq, startTS); err != nil { - return err - } if err := f.validateConflicts(ctx, uniq, startTS); err != nil { return errors.WithStack(err) } @@ -1102,7 +1177,7 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com return nil } - uniq, err := f.uniqueMutationsNotFenced(muts, commitTS) + uniq, err := f.uniqueMutationsNotFenced(ctx, muts, commitTS) if err != nil { return err } @@ -1118,7 +1193,7 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com return nil } -func (f *kvFSM) uniqueMutationsNotFenced(muts []*pb.Mutation, commitTS uint64) ([]*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 @@ -1126,6 +1201,9 @@ func (f *kvFSM) uniqueMutationsNotFenced(muts []*pb.Mutation, commitTS uint64) ( 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 } @@ -1171,7 +1249,7 @@ func (f *kvFSM) handleCommitRequest(ctx context.Context, r *pb.Request) error { if err != nil { return err } - uniq, err := f.uniqueMutationsAboveWriteFloor(muts, commitTS) + uniq, err := f.uniqueMutationsAboveWriteFloor(ctx, muts, commitTS) if err != nil { return err } @@ -1288,7 +1366,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 := f.uniqueMutationsAboveWriteFloor(muts, abortTS) + uniq, err := f.uniqueMutationsAboveWriteFloor(ctx, muts, abortTS) if err != nil { return err } @@ -1308,11 +1386,14 @@ 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(muts []*pb.Mutation, commitTS uint64) ([]*pb.Mutation, error) { +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 } diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index 00b02e31b..0f593df46 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -39,6 +39,26 @@ func newWriteFloorFSM(t *testing.T) *kvFSM { return newComposed1FSM(t, engine, 1) } +func newTargetReadinessFSM(t *testing.T, route distribution.RouteDescriptor) *kvFSM { + t.Helper() + + engine := distribution.NewEngine() + applyComposed1Snapshot(t, engine, 1, []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 TestFSMRejectsRawPointWriteOnWriteFencedRoute(t *testing.T) { t.Parallel() @@ -52,6 +72,46 @@ 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 TestFSMRejectsDelPrefixIntersectingWriteFencedRoute(t *testing.T) { t.Parallel() @@ -68,6 +128,28 @@ 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 TestFSMRejectsFullRangeDelPrefixWhenRouteIsWriteFenced(t *testing.T) { t.Parallel() diff --git a/kv/route_history.go b/kv/route_history.go index bd6989513..94eec1076 100644 --- a/kv/route_history.go +++ b/kv/route_history.go @@ -61,6 +61,14 @@ func (s distributionRouteSnapshot) OwnerOf(key []byte) (uint64, bool) { return s.snap.OwnerOf(key) } +func (s distributionRouteSnapshot) RouteOf(key []byte) (distribution.Route, bool) { + return s.snap.RouteOf(key) +} + +func (s distributionRouteSnapshot) IntersectingRoutes(start, end []byte) []distribution.Route { + return s.snap.IntersectingRoutes(start, end) +} + func (s distributionRouteSnapshot) WriteFencedForKey(key []byte) bool { route, ok := s.snap.RouteOf(key) return ok && route.State == distribution.RouteStateWriteFenced diff --git a/kv/shard_store.go b/kv/shard_store.go index f2e1beba6..246fac205 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -92,6 +92,9 @@ 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) { + if err := s.verifyTargetReadinessForRange(ctx, g, route, key, nextScanCursor(key)); err != nil { + return nil, err + } if !isTxnInternalKey(key) { if err := s.maybeResolveTxnLock(ctx, g, key, ts); err != nil { return nil, err @@ -132,6 +135,11 @@ func routeSatisfiesTargetReadiness(route distribution.Route, ready store.TargetS } func (s *ShardStore) verifyTargetReadinessForRange(ctx context.Context, g *ShardGroup, route distribution.Route, start []byte, end []byte) error { + routeStart, routeEnd := readinessRouteRange(start, end) + return s.verifyTargetReadinessForRouteRange(ctx, g, route, routeStart, routeEnd) +} + +func (s *ShardStore) verifyTargetReadinessForRouteRange(ctx context.Context, g *ShardGroup, route distribution.Route, routeStart []byte, routeEnd []byte) error { if g == nil || g.Store == nil { return nil } @@ -143,7 +151,6 @@ func (s *ShardStore) verifyTargetReadinessForRange(ctx context.Context, g *Shard if err != nil { return errors.WithStack(err) } - routeStart, routeEnd := readinessRouteRange(start, end) for _, ready := range states { if !ready.Armed || !routeRangeIntersects(routeStart, routeEnd, ready.RouteStart, ready.RouteEnd) { continue @@ -591,7 +598,7 @@ func (s *ShardStore) scanRouteAtDirectionPhysicalLimit( if routeHasStagedVisibility(route) { return nil, true, nil } - return s.scanRouteAtLeaderPhysicalLimit(ctx, g, start, end, visibleLimit, physicalLimit, ts, reverse) + return s.scanRouteAtLeaderPhysicalLimit(ctx, g, route, start, end, visibleLimit, physicalLimit, ts, reverse) } // RawScanAt cannot enforce physicalLimit, so report truncation and let @@ -666,6 +673,7 @@ func (s *ShardStore) scanRouteLocal( func (s *ShardStore) scanRouteAtLeaderPhysicalLimit( ctx context.Context, g *ShardGroup, + route distribution.Route, start []byte, end []byte, visibleLimit int, @@ -682,7 +690,7 @@ func (s *ShardStore) scanRouteAtLeaderPhysicalLimit( if err != nil { return nil, limitReached, err } - resolved, err := s.resolveScanLocks(ctx, g, distribution.Route{}, kvs, lockKVs, ts) + resolved, err := s.resolveScanLocks(ctx, g, route, kvs, lockKVs, ts) return resolved, limitReached, err } @@ -1802,6 +1810,9 @@ 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.verifyPrefixDeleteRoutes(ctx, prefix, commitTS); err != nil { + return err + } for _, g := range s.groups { if g == nil || g.Store == nil { continue @@ -1815,6 +1826,9 @@ func (s *ShardStore) DeletePrefixAt(ctx context.Context, prefix []byte, excludeP // DeletePrefixAtRaft is the raft-apply variant of DeletePrefixAt. func (s *ShardStore) DeletePrefixAtRaft(ctx context.Context, prefix []byte, excludePrefix []byte, commitTS uint64) error { + if err := s.verifyPrefixDeleteRoutes(ctx, prefix, commitTS); err != nil { + return err + } for _, g := range s.groups { if g == nil || g.Store == nil { continue @@ -1841,6 +1855,9 @@ 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 { + if err := s.verifyPrefixDeleteRoutes(ctx, prefix, commitTS); err != nil { + return err + } for _, g := range s.groups { if g == nil || g.Store == nil { continue @@ -1862,6 +1879,26 @@ func (s *ShardStore) DeletePrefixAtRaftAt(ctx context.Context, prefix []byte, ex return nil } +func (s *ShardStore) verifyPrefixDeleteRoutes(ctx context.Context, prefix []byte, commitTS uint64) error { + if s == nil || s.engine == nil { + return nil + } + routeStart, routeEnd := routePrefixRange(prefix) + for _, route := range s.engine.GetIntersectingRoutes(routeStart, routeEnd) { + g, ok := s.groupForID(route.GroupID) + if !ok || g == nil || g.Store == nil { + return store.ErrNotSupported + } + if err := s.verifyTargetReadinessForRouteRange(ctx, g, route, routeStart, routeEnd); err != nil { + return err + } + if err := verifyRouteWriteFloor(route, commitTS); err != nil { + return err + } + } + return nil +} + func (s *ShardStore) LastCommitTS() uint64 { var max uint64 for _, g := range s.groups { diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index c1fe77814..7088c01ec 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -248,6 +248,50 @@ func TestShardStoreApplyMutationsRaftAtEnforcesRouteFloor(t *testing.T) { require.ErrorIs(t, err, ErrRouteWriteBelowFloor) } +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 TestShardStoreDeletePrefixAtRaftEnforcesRouteFloorBeforeDelete(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, + }) + require.NoError(t, group.Store.PutAt(ctx, []byte("b"), []byte("v"), 1, 0)) + + err := st.DeletePrefixAtRaft(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("v"), got) +} + func TestShardStoreScanAndLatestCommitTS_MergeStagedVisibility(t *testing.T) { t.Parallel() diff --git a/kv/shard_store_txn_lock_test.go b/kv/shard_store_txn_lock_test.go index 8c50ff3a4..c395b3dc9 100644 --- a/kv/shard_store_txn_lock_test.go +++ b/kv/shard_store_txn_lock_test.go @@ -113,6 +113,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() @@ -441,6 +471,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: 1, + 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 7ed2ccb7c..285036d3e 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -1036,6 +1036,9 @@ func (c *ShardedCoordinator) dispatchDelPrefixBroadcast(ctx context.Context, isT if err != nil { return nil, err } + if err := c.rejectDelPrefixesBelowWriteFloor(elems, ts); err != nil { + return nil, err + } requests := make([]*pb.Request, 0, len(elems)) for _, elem := range elems { requests = append(requests, &pb.Request{ @@ -1084,6 +1087,24 @@ func (c *ShardedCoordinator) rejectWriteFencedDelPrefixes(elems []*Elem[OP]) err return nil } +func (c *ShardedCoordinator) rejectDelPrefixesBelowWriteFloor(elems []*Elem[OP], commitTS uint64) error { + if c == nil || c.engine == nil { + return nil + } + for _, elem := range elems { + if elem == nil { + continue + } + start, end := routePrefixRange(elem.Key) + for _, route := range c.engine.GetIntersectingRoutes(start, end) { + if err := verifyRouteWriteFloor(route, commitTS); err != nil { + return errors.Wrapf(err, "prefix %q route range [%q,%q)", elem.Key, start, end) + } + } + } + return nil +} + // broadcastToAllGroups sends the same set of requests to every shard group in // parallel and returns the maximum commit index. func (c *ShardedCoordinator) broadcastToAllGroups(ctx context.Context, requests []*pb.Request) (*CoordinateResponse, error) { diff --git a/kv/sharded_coordinator_del_prefix_test.go b/kv/sharded_coordinator_del_prefix_test.go index 4b83dc131..5b712db58 100644 --- a/kv/sharded_coordinator_del_prefix_test.go +++ b/kv/sharded_coordinator_del_prefix_test.go @@ -173,6 +173,33 @@ 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 TestShardedCoordinatorRejectsFullRangeDelPrefixWhenRouteIsWriteFenced(t *testing.T) { t.Parallel() From 6ac85a1ecdc7ec247b14b8a6542cb11089f2fadc Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 05:03:15 +0900 Subject: [PATCH 06/29] Fix target readiness guard gaps --- adapter/redis_compat_commands_stream_test.go | 39 +++++++---- kv/fsm.go | 20 +++++- kv/fsm_migration_fence_test.go | 69 ++++++++++++++++++++ 3 files changed, 112 insertions(+), 16 deletions(-) diff --git a/adapter/redis_compat_commands_stream_test.go b/adapter/redis_compat_commands_stream_test.go index bc78c846f..cf666f497 100644 --- a/adapter/redis_compat_commands_stream_test.go +++ b/adapter/redis_compat_commands_stream_test.go @@ -278,7 +278,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 @@ -286,23 +286,36 @@ func TestRedis_StreamXReadLatencyIsConstant(t *testing.T) { ) lastID := "" for i := range total { - id, err := rdb.XAdd(ctx, &redis.XAddArgs{ - Stream: "stream-lat", - ID: "*", - Values: []string{"i", fmt.Sprint(i)}, - }).Result() + var id string + err := retryNotLeader(ctx, func() error { + var xerr error + id, xerr = rdb.XAdd(ctx, &redis.XAddArgs{ + Stream: "stream-lat", + ID: "*", + Values: []string{"i", fmt.Sprint(i)}, + }).Result() + return xerr + }) require.NoError(t, err) lastID = id } 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 58fdccf32..31f38b0b5 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -686,16 +686,19 @@ func (f *kvFSM) verifyTargetReadinessForRouteRange(ctx context.Context, routeSta if !ready.Armed || !routeRangeIntersects(routeStart, routeEnd, ready.RouteStart, ready.RouteEnd) { continue } - if !proof || !routesSatisfyTargetReadiness(snap.IntersectingRoutes(routeStart, routeEnd), ready) { + if !proof || !routesSatisfyTargetReadiness(snap.IntersectingRoutes(routeStart, routeEnd), ready, f.shardGroupID) { return errors.WithStack(ErrRouteCutoverPending) } } return nil } -func routesSatisfyTargetReadiness(routes []distribution.Route, ready store.TargetStagedReadinessState) bool { +func routesSatisfyTargetReadiness(routes []distribution.Route, ready store.TargetStagedReadinessState, groupID uint64) bool { matched := false for _, route := range routes { + if route.GroupID != groupID { + continue + } if !routeRangeIntersects(route.Start, route.End, ready.RouteStart, ready.RouteEnd) { continue } @@ -1366,7 +1369,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 := f.uniqueMutationsAboveWriteFloor(ctx, muts, abortTS) + uniq, err := f.uniqueAbortCleanupMutations(ctx, muts) if err != nil { return err } @@ -1400,6 +1403,17 @@ func (f *kvFSM) uniqueMutationsAboveWriteFloor(ctx context.Context, muts []*pb.M return uniq, nil } +func (f *kvFSM) uniqueAbortCleanupMutations(ctx context.Context, muts []*pb.Mutation) ([]*pb.Mutation, error) { + uniq, err := uniqueMutations(muts) + if err != nil { + return nil, err + } + if err := f.verifyTargetReadinessForMutations(ctx, uniq); err != nil { + return nil, err + } + return uniq, nil +} + 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 0f593df46..d0d856c39 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -112,6 +112,39 @@ func TestFSMAllowsRawPointWriteWithClearedTargetReadinessProof(t *testing.T) { 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 TestFSMRejectsDelPrefixIntersectingWriteFencedRoute(t *testing.T) { t.Parallel() @@ -251,3 +284,39 @@ func TestFSMRejectsPrepareOnWriteFencedRouteButAllowsAbort(t *testing.T) { err := fsm.handleTxnRequest(ctx, abort, 11) require.False(t, errors.Is(err, ErrRouteWriteFenced), "ABORT must keep the narrow cleanup lane open") } + +func TestFSMAbortCleanupBypassesRetainedWriteFloor(t *testing.T) { + t.Parallel() + + ctx := context.Background() + fsm := newWriteFloorFSM(t) + 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)) + + abort := &pb.Request{ + IsTxn: true, + Phase: pb.Phase_ABORT, + Ts: startTS, + Mutations: []*pb.Mutation{ + {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: primaryKey, CommitTS: abortTS})}, + {Op: pb.Op_PUT, Key: primaryKey}, + }, + } + 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) +} From ce91b665196bcbe9a1817c2017b14dbdb052a0b4 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 05:05:57 +0900 Subject: [PATCH 07/29] Guard manifest scan readiness routes --- kv/shard_store.go | 9 ++++++- kv/shard_store_test.go | 54 ++++++++++++++++++++++++++++++++++++++---- 2 files changed, 58 insertions(+), 5 deletions(-) diff --git a/kv/shard_store.go b/kv/shard_store.go index 246fac205..eaacf33a7 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -135,7 +135,7 @@ func routeSatisfiesTargetReadiness(route distribution.Route, ready store.TargetS } func (s *ShardStore) verifyTargetReadinessForRange(ctx context.Context, g *ShardGroup, route distribution.Route, start []byte, end []byte) error { - routeStart, routeEnd := readinessRouteRange(start, end) + routeStart, routeEnd := readinessRouteRangeForScan(start, end) return s.verifyTargetReadinessForRouteRange(ctx, g, route, routeStart, routeEnd) } @@ -174,6 +174,13 @@ func readinessRouteRange(start []byte, end []byte) ([]byte, []byte) { 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 > route.MinWriteTSExclusive { return nil diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 7088c01ec..baa5ef041 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -35,9 +35,7 @@ func newStagedVisibilityShardStore(t *testing.T) (*ShardStore, *ShardGroup) { func applyTargetReadiness(t *testing.T, group *ShardGroup) { t.Helper() - writer, ok := group.Store.(store.MigrationTargetReadinessWriter) - require.True(t, ok) - require.NoError(t, writer.ApplyTargetStagedReadiness(context.Background(), store.TargetStagedReadinessState{ + applyTargetReadinessState(t, group, store.TargetStagedReadinessState{ JobID: 9, RouteStart: []byte("a"), RouteEnd: []byte("z"), @@ -45,7 +43,14 @@ func applyTargetReadiness(t *testing.T, group *ShardGroup) { 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) { @@ -227,6 +232,47 @@ func TestShardStorePhysicalLimitScanChecksTargetReadiness(t *testing.T) { 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 TestShardStoreApplyMutationsRaftAtEnforcesRouteFloor(t *testing.T) { t.Parallel() From f2bcd65bccee6f69e78ee7ff0673861db8d88172 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 05:36:59 +0900 Subject: [PATCH 08/29] Fix target readiness route proofs --- kv/fsm.go | 6 ++- kv/fsm_migration_fence_test.go | 40 +++++++++++++++ kv/shard_store.go | 14 ++++- kv/shard_store_test.go | 94 ++++++++++++++++++++++++++++++++++ 4 files changed, 151 insertions(+), 3 deletions(-) diff --git a/kv/fsm.go b/kv/fsm.go index 31f38b0b5..b3a8d6d42 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -1110,7 +1110,11 @@ func (f *kvFSM) handlePrepareRequest(ctx context.Context, r *pb.Request) error { } startTS := r.Ts - uniq, err := f.uniqueMutationsNotFenced(ctx, muts, startTS) + floorTS := startTS + if meta.CommitTS != 0 { + floorTS = meta.CommitTS + } + uniq, err := f.uniqueMutationsNotFenced(ctx, muts, floorTS) if err != nil { return err } diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index d0d856c39..35cfee44d 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -256,6 +256,46 @@ func TestFSMRejectsOnePhaseTxnBelowRouteFloor(t *testing.T) { 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() diff --git a/kv/shard_store.go b/kv/shard_store.go index eaacf33a7..06349ae36 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -152,7 +152,7 @@ func (s *ShardStore) verifyTargetReadinessForRouteRange(ctx context.Context, g * return errors.WithStack(err) } for _, ready := range states { - if !ready.Armed || !routeRangeIntersects(routeStart, routeEnd, ready.RouteStart, ready.RouteEnd) { + if !targetReadinessAppliesToRoute(route, routeStart, routeEnd, ready) { continue } if !routeSatisfiesTargetReadiness(route, ready) { @@ -162,6 +162,16 @@ func (s *ShardStore) verifyTargetReadinessForRouteRange(ctx context.Context, g * return nil } +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 { @@ -2104,7 +2114,7 @@ func (s *ShardStore) explicitGroupRouteForRange(groupID uint64, start []byte, en if s == nil || s.engine == nil { return fallback } - routeStart, routeEnd := readinessRouteRange(start, end) + routeStart, routeEnd := readinessRouteRangeForScan(start, end) for _, route := range s.engine.GetIntersectingRoutes(routeStart, routeEnd) { if route.GroupID == groupID { return route diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index baa5ef041..5748cdfc0 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -273,6 +273,100 @@ func TestShardStoreS3ManifestScanChecksTargetReadinessInRouteSpace(t *testing.T) 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 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 TestShardStoreApplyMutationsRaftAtEnforcesRouteFloor(t *testing.T) { t.Parallel() From f727f3079a06bd8cb8c772ba0e6331f0194aa230 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 06:11:36 +0900 Subject: [PATCH 09/29] Enable split migration job creation --- adapter/distribution_server.go | 200 +++++++++++++++++++++++++--- adapter/distribution_server_test.go | 168 +++++++++++++++++++++++ distribution/split_job_catalog.go | 1 + 3 files changed, 354 insertions(+), 15 deletions(-) diff --git a/adapter/distribution_server.go b/adapter/distribution_server.go index 78106c6bb..c1b1fa7f4 100644 --- a/adapter/distribution_server.go +++ b/adapter/distribution_server.go @@ -22,12 +22,13 @@ import ( // DistributionServer serves distribution related gRPC APIs. type DistributionServer struct { - mu sync.Mutex - engine *distribution.Engine - catalog *distribution.CatalogStore - coordinator kv.Coordinator - readTracker *kv.ActiveTimestampTracker - reloadRetry struct { + mu sync.Mutex + engine *distribution.Engine + catalog *distribution.CatalogStore + coordinator kv.Coordinator + readTracker *kv.ActiveTimestampTracker + migrationCapabilityGate SplitMigrationCapabilityGate + reloadRetry struct { attempts int interval time.Duration } @@ -37,6 +38,10 @@ type DistributionServer struct { // DistributionServerOption configures DistributionServer behavior. type DistributionServerOption func(*DistributionServer) +// SplitMigrationCapabilityGate reports whether this node can safely create +// migration-only side effects. A nil gate keeps StartSplitMigration fail-closed. +type SplitMigrationCapabilityGate func(context.Context) error + // WithDistributionCoordinator configures the coordinator used for Raft-backed // catalog mutations in SplitRange. func WithDistributionCoordinator(coordinator kv.Coordinator) DistributionServerOption { @@ -51,6 +56,12 @@ func WithDistributionActiveTimestampTracker(tracker *kv.ActiveTimestampTracker) } } +func WithSplitMigrationCapabilityGate(gate SplitMigrationCapabilityGate) DistributionServerOption { + return func(s *DistributionServer) { + s.migrationCapabilityGate = gate + } +} + // WithCatalogReloadRetryPolicy configures the retry policy used after split // commit when waiting for the local catalog snapshot to become visible. func WithCatalogReloadRetryPolicy(attempts int, interval time.Duration) DistributionServerOption { @@ -184,16 +195,97 @@ func (s *DistributionServer) GetIntersectingRoutes(ctx context.Context, req *pb. } func (s *DistributionServer) StartSplitMigration(ctx context.Context, req *pb.StartSplitMigrationRequest) (*pb.StartSplitMigrationResponse, error) { + s.mu.Lock() + defer s.mu.Unlock() + + if err := s.verifyStartSplitMigrationReady(ctx, req); err != nil { + return nil, err + } + snapshot, err := s.loadCatalogSnapshot(ctx) + if err != nil { + return nil, err + } + readPin := s.pinReadTS(snapshot.ReadTS) + defer readPin.Release() + + parent, err := s.startSplitMigrationParent(ctx, snapshot, req) + if err != nil { + return nil, err + } + job, err := s.newSplitMigrationJob(ctx, snapshot, parent, req) + if err != nil { + return nil, err + } + if err := s.createSplitJobViaCoordinator(ctx, snapshot.ReadTS, job); err != nil { + return nil, err + } + return &pb.StartSplitMigrationResponse{ + CatalogVersion: snapshot.Version, + JobId: job.JobID, + }, nil +} + +func (s *DistributionServer) verifyStartSplitMigrationReady(ctx context.Context, req *pb.StartSplitMigrationRequest) error { if req.GetTargetGroupId() == 0 { - return nil, grpcStatusError(codes.InvalidArgument, distribution.ErrCatalogSplitJobTargetGroupRequired.Error()) + return grpcStatusError(codes.InvalidArgument, distribution.ErrCatalogSplitJobTargetGroupRequired.Error()) } if err := s.verifyCatalogLeader(ctx); err != nil { - return nil, err + return err } if s.catalog == nil { - return nil, grpcStatusError(codes.FailedPrecondition, errDistributionCatalogNotConfigured.Error()) + return grpcStatusError(codes.FailedPrecondition, errDistributionCatalogNotConfigured.Error()) + } + return s.verifySplitMigrationCapability(ctx) +} + +func (s *DistributionServer) startSplitMigrationParent( + ctx context.Context, + snapshot distribution.CatalogSnapshot, + req *pb.StartSplitMigrationRequest, +) (distribution.RouteDescriptor, error) { + if err := validateExpectedCatalogVersion(snapshot.Version, req.GetExpectedCatalogVersion()); err != nil { + return distribution.RouteDescriptor{}, err + } + parent, found := findRouteByID(snapshot.Routes, req.GetRouteId()) + if !found { + return distribution.RouteDescriptor{}, grpcStatusError(codes.NotFound, errDistributionUnknownRoute.Error()) + } + splitKey := distribution.CloneBytes(req.GetSplitKey()) + if err := validateSplitKey(parent, splitKey); err != nil { + return distribution.RouteDescriptor{}, err + } + if err := distribution.ValidateMigrationRouteRange(splitKey, parent.End); err != nil { + return distribution.RouteDescriptor{}, splitMigrationRangeStatusError(err) + } + if parent.GroupID == req.GetTargetGroupId() { + return distribution.RouteDescriptor{}, grpcStatusError(codes.InvalidArgument, "target group must differ from source route group") + } + if err := s.rejectLiveSplitJob(ctx, snapshot, parent); err != nil { + return distribution.RouteDescriptor{}, err + } + return parent, nil +} + +func (s *DistributionServer) newSplitMigrationJob( + ctx context.Context, + snapshot distribution.CatalogSnapshot, + parent distribution.RouteDescriptor, + req *pb.StartSplitMigrationRequest, +) (distribution.SplitJob, error) { + jobID, err := s.catalog.NextSplitJobIDAt(ctx, snapshot.ReadTS) + if err != nil { + return distribution.SplitJob{}, splitJobCatalogStatusError(err) + } + job, err := distribution.InitializeSplitJobPlan(distribution.SplitJob{ + JobID: jobID, + SourceRouteID: parent.RouteID, + SplitKey: distribution.CloneBytes(req.GetSplitKey()), + TargetGroupID: req.GetTargetGroupId(), + }, parent, time.Now().UnixMilli()) + if err != nil { + return distribution.SplitJob{}, splitJobCatalogStatusError(err) } - return nil, grpcStatusError(codes.FailedPrecondition, errDistributionClusterNotReady.Error()) + return job, nil } func (s *DistributionServer) GetSplitJob(ctx context.Context, req *pb.GetSplitJobRequest) (*pb.GetSplitJobResponse, error) { @@ -303,7 +395,7 @@ func (s *DistributionServer) SplitRange(ctx context.Context, req *pb.SplitRangeR return nil, err } - parent, _, found := findRouteByID(snapshot.Routes, req.GetRouteId()) + parent, found := findRouteByID(snapshot.Routes, req.GetRouteId()) if !found { return nil, grpcStatusError(codes.NotFound, errDistributionUnknownRoute.Error()) } @@ -536,6 +628,51 @@ func (s *DistributionServer) updateSplitJobViaCoordinator( return nil } +func (s *DistributionServer) createSplitJobViaCoordinator(ctx context.Context, readTS uint64, job distribution.SplitJob) error { + if job.JobID == math.MaxUint64 { + return splitJobCatalogStatusError(distribution.ErrCatalogSplitJobIDOverflow) + } + encoded, err := distribution.EncodeSplitJob(job) + if err != nil { + return splitJobCatalogStatusError(err) + } + jobKey := distribution.CatalogSplitJobKey(job.JobID) + nextIDKey := distribution.CatalogNextSplitJobIDKey() + if _, err := s.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + Elems: []*kv.Elem[kv.OP]{ + { + Op: kv.Put, + Key: jobKey, + Value: encoded, + }, + { + Op: kv.Put, + Key: nextIDKey, + Value: distribution.EncodeCatalogNextSplitJobID(job.JobID + 1), + }, + }, + IsTxn: true, + StartTS: readTS, + ReadKeys: [][]byte{jobKey, nextIDKey}, + }); err != nil { + return splitJobCoordinatorStatusError(err) + } + return nil +} + +func (s *DistributionServer) verifySplitMigrationCapability(ctx context.Context) error { + if s.migrationCapabilityGate == nil { + return grpcStatusError(codes.FailedPrecondition, errDistributionClusterNotReady.Error()) + } + if err := s.migrationCapabilityGate(ctx); err != nil { + if status.Code(err) != codes.OK { + return err + } + return grpcStatusError(codes.FailedPrecondition, errDistributionClusterNotReady.Error()) + } + return nil +} + func validateExpectedCatalogVersion(currentVersion, expectedVersion uint64) error { if currentVersion != expectedVersion { return grpcStatusError(codes.Aborted, errDistributionCatalogConflict.Error()) @@ -583,6 +720,29 @@ func (s *DistributionServer) rejectLiveSplitJobOverlap(ctx context.Context, snap return nil } +func (s *DistributionServer) rejectLiveSplitJob(ctx context.Context, snapshot distribution.CatalogSnapshot, parent distribution.RouteDescriptor) error { + jobs, err := s.catalog.ListSplitJobsAt(ctx, snapshot.ReadTS) + if err != nil { + return grpcStatusErrorf(codes.Internal, "load split jobs: %v", err) + } + liveJobs := 0 + for _, job := range jobs { + if !splitJobIsLive(job) { + continue + } + liveJobs++ + for _, interval := range liveSplitJobIntervals(job, snapshot.Routes) { + if routeRangeIntersects(parent.Start, parent.End, interval.start, interval.end) { + return grpcStatusError(codes.Aborted, distribution.ErrSplitJobOverlap.Error()) + } + } + } + if liveJobs > 0 { + return grpcStatusError(codes.ResourceExhausted, distribution.ErrTooManyInFlightSplitJobs.Error()) + } + return nil +} + func splitJobIsLive(job distribution.SplitJob) bool { return job.Phase != distribution.SplitJobPhaseDone && job.Phase != distribution.SplitJobPhaseAbandoned } @@ -676,6 +836,16 @@ func splitJobCatalogStatusError(err error) error { } } +func splitMigrationRangeStatusError(err error) error { + switch { + case errors.Is(err, distribution.ErrMigrationReservedRange), + errors.Is(err, distribution.ErrMigrationInvalidRoute): + return grpcStatusError(codes.InvalidArgument, err.Error()) + default: + return grpcStatusErrorf(codes.Internal, "validate split migration range: %v", err) + } +} + func splitJobCoordinatorStatusError(err error) error { switch { case errors.Is(err, store.ErrWriteConflict): @@ -865,13 +1035,13 @@ func (s *DistributionServer) allocateChildRouteIDs(ctx context.Context, readTS u return leftID, rightID, nil } -func findRouteByID(routes []distribution.RouteDescriptor, routeID uint64) (distribution.RouteDescriptor, int, bool) { - for i, route := range routes { +func findRouteByID(routes []distribution.RouteDescriptor, routeID uint64) (distribution.RouteDescriptor, bool) { + for _, route := range routes { if route.RouteID == routeID { - return distribution.CloneRouteDescriptor(route), i, true + return distribution.CloneRouteDescriptor(route), true } } - return distribution.RouteDescriptor{}, -1, false + return distribution.RouteDescriptor{}, false } func toProtoRouteDescriptors(routes []distribution.RouteDescriptor) []*pb.RouteDescriptor { diff --git a/adapter/distribution_server_test.go b/adapter/distribution_server_test.go index 109f21d62..8e51c4f99 100644 --- a/adapter/distribution_server_test.go +++ b/adapter/distribution_server_test.go @@ -156,6 +156,174 @@ func TestDistributionServerStartSplitMigration_FailsClosedUntilCapabilityGate(t require.Empty(t, jobs) } +func TestDistributionServerStartSplitMigration_CreatesPlannedJobWhenGateOpen(t *testing.T) { + t.Parallel() + + ctx := context.Background() + baseStore := store.NewMVCCStore() + catalog := distribution.NewCatalogStore(baseStore) + saved, err := catalog.Save(ctx, 0, []distribution.RouteDescriptor{{ + RouteID: 1, + Start: []byte("a"), + End: []byte("m"), + GroupID: 1, + State: distribution.RouteStateActive, + ParentRouteID: 0, + }}) + require.NoError(t, err) + + coordinator := newDistributionCoordinatorStub(baseStore, true) + gateCalls := 0 + s := NewDistributionServer( + distribution.NewEngine(), + catalog, + WithDistributionCoordinator(coordinator), + WithSplitMigrationCapabilityGate(func(context.Context) error { + gateCalls++ + return nil + }), + ) + + resp, err := s.StartSplitMigration(ctx, &pb.StartSplitMigrationRequest{ + ExpectedCatalogVersion: saved.Version, + RouteId: 1, + SplitKey: []byte("g"), + TargetGroupId: 2, + }) + require.NoError(t, err) + require.Equal(t, saved.Version, resp.CatalogVersion) + require.Equal(t, uint64(1), resp.JobId) + require.Equal(t, 1, gateCalls) + require.Equal(t, 1, coordinator.dispatchCalls) + + jobs, err := catalog.ListSplitJobs(ctx) + require.NoError(t, err) + require.Len(t, jobs, 1) + job := jobs[0] + require.Equal(t, uint64(1), job.JobID) + require.Equal(t, uint64(1), job.SourceRouteID) + require.Equal(t, []byte("g"), job.SplitKey) + require.Equal(t, uint64(2), job.TargetGroupID) + require.Equal(t, distribution.SplitJobPhasePlanned, job.Phase) + require.NotZero(t, job.StartedAtMs) + require.NotZero(t, job.UpdatedAtMs) + require.NotEmpty(t, job.BracketProgress) + next, err := catalog.NextSplitJobID(ctx) + require.NoError(t, err) + require.Equal(t, uint64(2), next) +} + +func TestDistributionServerStartSplitMigration_RejectsSecondLiveJob(t *testing.T) { + t.Parallel() + + ctx := context.Background() + baseStore := store.NewMVCCStore() + catalog := distribution.NewCatalogStore(baseStore) + saved, err := catalog.Save(ctx, 0, []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("m"), + GroupID: 1, + State: distribution.RouteStateActive, + ParentRouteID: 0, + }, + { + RouteID: 2, + Start: []byte("m"), + End: nil, + GroupID: 2, + State: distribution.RouteStateActive, + ParentRouteID: 0, + }, + }) + require.NoError(t, err) + require.NoError(t, catalog.CreateSplitJob(ctx, distribution.SplitJob{ + JobID: 1, + SourceRouteID: 2, + SplitKey: []byte("t"), + TargetGroupID: 3, + Phase: distribution.SplitJobPhaseBackfill, + StartedAtMs: 1000, + UpdatedAtMs: 1000, + })) + + coordinator := newDistributionCoordinatorStub(baseStore, true) + s := NewDistributionServer( + distribution.NewEngine(), + catalog, + WithDistributionCoordinator(coordinator), + WithSplitMigrationCapabilityGate(func(context.Context) error { return nil }), + ) + + _, err = s.StartSplitMigration(ctx, &pb.StartSplitMigrationRequest{ + ExpectedCatalogVersion: saved.Version, + RouteId: 1, + SplitKey: []byte("g"), + TargetGroupId: 4, + }) + require.Error(t, err) + require.Equal(t, codes.ResourceExhausted, status.Code(err)) + require.ErrorContains(t, err, distribution.ErrTooManyInFlightSplitJobs.Error()) + require.Zero(t, coordinator.dispatchCalls) +} + +func TestDistributionServerStartSplitMigration_RejectsReservedMigrationRange(t *testing.T) { + t.Parallel() + + ctx := context.Background() + for _, tc := range []struct { + name string + splitKey []byte + end []byte + }{ + {name: "dist catalog", splitKey: []byte("!dist|"), end: nil}, + {name: "dist staged catalog", splitKey: []byte("!dist|migstage|"), end: nil}, + {name: "migration staged", splitKey: []byte("!migstage|ready|1"), end: nil}, + {name: "migration tracker", splitKey: []byte("!migwrite|"), end: nil}, + {name: "migration fence", splitKey: []byte("!migfence|"), end: nil}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + baseStore := store.NewMVCCStore() + catalog := distribution.NewCatalogStore(baseStore) + saved, err := catalog.Save(ctx, 0, []distribution.RouteDescriptor{{ + RouteID: 1, + Start: []byte(""), + End: tc.end, + GroupID: 1, + State: distribution.RouteStateActive, + ParentRouteID: 0, + }}) + require.NoError(t, err) + + coordinator := newDistributionCoordinatorStub(baseStore, true) + s := NewDistributionServer( + distribution.NewEngine(), + catalog, + WithDistributionCoordinator(coordinator), + WithSplitMigrationCapabilityGate(func(context.Context) error { return nil }), + ) + + _, err = s.StartSplitMigration(ctx, &pb.StartSplitMigrationRequest{ + ExpectedCatalogVersion: saved.Version, + RouteId: 1, + SplitKey: tc.splitKey, + TargetGroupId: 2, + }) + require.Error(t, err) + require.Equal(t, codes.InvalidArgument, status.Code(err)) + require.ErrorContains(t, err, distribution.ErrMigrationReservedRange.Error()) + require.Zero(t, coordinator.dispatchCalls) + + jobs, listErr := catalog.ListSplitJobs(ctx) + require.NoError(t, listErr) + require.Empty(t, jobs) + }) + } +} + func TestDistributionServerSplitJobRPCs_ReadAndListCatalogJobs(t *testing.T) { t.Parallel() diff --git a/distribution/split_job_catalog.go b/distribution/split_job_catalog.go index 082f9aaf6..1f7fc241b 100644 --- a/distribution/split_job_catalog.go +++ b/distribution/split_job_catalog.go @@ -36,6 +36,7 @@ var ( ErrCatalogSplitJobConflict = errors.New("catalog split job conflict") ErrCatalogSplitJobTerminalRequired = errors.New("catalog split job terminal state is required") ErrSplitJobOverlap = errors.New("split job overlaps requested route") + ErrTooManyInFlightSplitJobs = errors.New("too many in-flight split jobs") ) // SplitJobPhase is the durable phase of a split migration job. From 53d38e31d0e8f9fa32a76a970218c0c32f7c44af Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 06:21:34 +0900 Subject: [PATCH 10/29] Harden split migration readiness checks --- adapter/distribution_server.go | 7 ++--- adapter/distribution_server_test.go | 46 +++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 5 deletions(-) diff --git a/adapter/distribution_server.go b/adapter/distribution_server.go index c1b1fa7f4..389559b4e 100644 --- a/adapter/distribution_server.go +++ b/adapter/distribution_server.go @@ -431,7 +431,7 @@ func (s *DistributionServer) SplitRange(ctx context.Context, req *pb.SplitRangeR func (s *DistributionServer) pinReadTS(ts uint64) *kv.ActiveTimestampToken { if s == nil || s.readTracker == nil { - return nil + return &kv.ActiveTimestampToken{} } return s.readTracker.Pin(ts) } @@ -665,10 +665,7 @@ func (s *DistributionServer) verifySplitMigrationCapability(ctx context.Context) return grpcStatusError(codes.FailedPrecondition, errDistributionClusterNotReady.Error()) } if err := s.migrationCapabilityGate(ctx); err != nil { - if status.Code(err) != codes.OK { - return err - } - return grpcStatusError(codes.FailedPrecondition, errDistributionClusterNotReady.Error()) + return err } return nil } diff --git a/adapter/distribution_server_test.go b/adapter/distribution_server_test.go index 8e51c4f99..ceb73fe36 100644 --- a/adapter/distribution_server_test.go +++ b/adapter/distribution_server_test.go @@ -74,6 +74,24 @@ func TestWithCatalogReloadRetryPolicy_OverridesDefaults(t *testing.T) { require.Equal(t, 5*time.Millisecond, s.reloadRetry.interval) } +func TestDistributionServerPinReadTSWithoutTrackerReturnsReleasableToken(t *testing.T) { + t.Parallel() + + s := NewDistributionServer(distribution.NewEngine(), nil) + token := s.pinReadTS(10) + require.NotNil(t, token) + require.NotPanics(t, func() { + token.Release() + }) + + var nilServer *DistributionServer + nilToken := nilServer.pinReadTS(10) + require.NotNil(t, nilToken) + require.NotPanics(t, func() { + nilToken.Release() + }) +} + func TestDistributionServerListRoutes_ReadsDurableCatalog(t *testing.T) { t.Parallel() @@ -156,6 +174,34 @@ func TestDistributionServerStartSplitMigration_FailsClosedUntilCapabilityGate(t require.Empty(t, jobs) } +func TestDistributionServerStartSplitMigrationReturnsCapabilityGateError(t *testing.T) { + t.Parallel() + + ctx := context.Background() + baseStore := store.NewMVCCStore() + catalog := distribution.NewCatalogStore(baseStore) + coordinator := newDistributionCoordinatorStub(baseStore, true) + s := NewDistributionServer( + distribution.NewEngine(), + catalog, + WithDistributionCoordinator(coordinator), + WithSplitMigrationCapabilityGate(func(context.Context) error { + return status.Error(codes.Unavailable, "split migration capability not ready") + }), + ) + + _, err := s.StartSplitMigration(ctx, &pb.StartSplitMigrationRequest{ + ExpectedCatalogVersion: 1, + RouteId: 1, + SplitKey: []byte("g"), + TargetGroupId: 2, + }) + require.Error(t, err) + require.Equal(t, codes.Unavailable, status.Code(err)) + require.ErrorContains(t, err, "split migration capability not ready") + require.Zero(t, coordinator.dispatchCalls) +} + func TestDistributionServerStartSplitMigration_CreatesPlannedJobWhenGateOpen(t *testing.T) { t.Parallel() From a624a2169ed8ab7d63a224274a022c86236cc8f9 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 06:32:03 +0900 Subject: [PATCH 11/29] Harden target readiness route proofs --- kv/fsm.go | 6 +- kv/fsm_migration_fence_test.go | 2 +- kv/shard_store.go | 57 +++++++++++++++---- kv/shard_store_test.go | 95 ++++++++++++++++++++++++++++++++ kv/shard_store_txn_lock_test.go | 2 +- store/migration_readiness.go | 8 +++ store/migration_versions_test.go | 36 ++++++++++++ 7 files changed, 190 insertions(+), 16 deletions(-) diff --git a/kv/fsm.go b/kv/fsm.go index b3a8d6d42..85dd67590 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -686,14 +686,14 @@ func (f *kvFSM) verifyTargetReadinessForRouteRange(ctx context.Context, routeSta if !ready.Armed || !routeRangeIntersects(routeStart, routeEnd, ready.RouteStart, ready.RouteEnd) { continue } - if !proof || !routesSatisfyTargetReadiness(snap.IntersectingRoutes(routeStart, routeEnd), ready, f.shardGroupID) { + if !proof || !routesSatisfyTargetReadiness(snap.IntersectingRoutes(routeStart, routeEnd), ready, f.shardGroupID, snap.Version()) { return errors.WithStack(ErrRouteCutoverPending) } } return nil } -func routesSatisfyTargetReadiness(routes []distribution.Route, ready store.TargetStagedReadinessState, groupID uint64) bool { +func routesSatisfyTargetReadiness(routes []distribution.Route, ready store.TargetStagedReadinessState, groupID uint64, catalogVersion uint64) bool { matched := false for _, route := range routes { if route.GroupID != groupID { @@ -703,7 +703,7 @@ func routesSatisfyTargetReadiness(routes []distribution.Route, ready store.Targe continue } matched = true - if !routeSatisfiesTargetReadiness(route, ready) { + if !routeSatisfiesTargetReadiness(route, ready, catalogVersion) { return false } } diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index 35cfee44d..9ebaab322 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -43,7 +43,7 @@ func newTargetReadinessFSM(t *testing.T, route distribution.RouteDescriptor) *kv t.Helper() engine := distribution.NewEngine() - applyComposed1Snapshot(t, engine, 1, []distribution.RouteDescriptor{route}) + applyComposed1Snapshot(t, engine, 2, []distribution.RouteDescriptor{route}) fsm := newComposed1FSM(t, engine, route.GroupID) writer, ok := fsm.store.(store.MigrationTargetReadinessWriter) require.True(t, ok) diff --git a/kv/shard_store.go b/kv/shard_store.go index 06349ae36..4e3fb96c4 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -121,7 +121,7 @@ func routeHasStagedVisibility(route distribution.Route) bool { return route.StagedVisibilityActive && route.MigrationJobID != 0 } -func routeSatisfiesTargetReadiness(route distribution.Route, ready store.TargetStagedReadinessState) bool { +func routeSatisfiesTargetReadiness(route distribution.Route, ready store.TargetStagedReadinessState, catalogVersion uint64) bool { if !routeRangeIntersects(route.Start, route.End, ready.RouteStart, ready.RouteEnd) { return false } @@ -131,7 +131,10 @@ func routeSatisfiesTargetReadiness(route distribution.Route, ready store.TargetS if route.StagedVisibilityActive { return route.MigrationJobID == ready.MigrationJobID } - return route.MigrationJobID == 0 + if route.MigrationJobID != 0 { + return false + } + return ready.ExpectedCutoverVersion == 0 || catalogVersion >= ready.ExpectedCutoverVersion } func (s *ShardStore) verifyTargetReadinessForRange(ctx context.Context, g *ShardGroup, route distribution.Route, start []byte, end []byte) error { @@ -151,11 +154,15 @@ func (s *ShardStore) verifyTargetReadinessForRouteRange(ctx context.Context, g * if err != nil { return errors.WithStack(err) } + var catalogVersion uint64 + if s != nil && s.engine != nil { + catalogVersion = s.engine.Version() + } for _, ready := range states { if !targetReadinessAppliesToRoute(route, routeStart, routeEnd, ready) { continue } - if !routeSatisfiesTargetReadiness(route, ready) { + if !routeSatisfiesTargetReadiness(route, ready, catalogVersion) { return errors.WithStack(ErrRouteCutoverPending) } } @@ -388,7 +395,11 @@ func (s *ShardStore) ScanGroupAt(ctx context.Context, groupID uint64, start []by if limit <= 0 { return []*store.KVPair{}, nil } - return s.scanRouteAtDirection(ctx, s.explicitGroupRouteForRange(groupID, start, end), start, end, limit, ts, false, true) + routes := s.explicitGroupRoutesForRange(groupID, start, end) + if err := s.verifyExplicitGroupRoutesForRange(ctx, groupID, routes, start, end); err != nil { + return nil, err + } + return s.scanRouteAtDirection(ctx, routes[0], start, end, limit, ts, false, true) } func (s *ShardStore) ReverseScanAt(ctx context.Context, start []byte, end []byte, limit int, ts uint64) ([]*store.KVPair, error) { @@ -1901,7 +1912,11 @@ func (s *ShardStore) verifyPrefixDeleteRoutes(ctx context.Context, prefix []byte return nil } routeStart, routeEnd := routePrefixRange(prefix) - for _, route := range s.engine.GetIntersectingRoutes(routeStart, routeEnd) { + routes := s.engine.GetIntersectingRoutes(routeStart, routeEnd) + if len(routes) == 0 { + return errors.WithStack(ErrRouteCutoverPending) + } + for _, route := range routes { g, ok := s.groupForID(route.GroupID) if !ok || g == nil || g.Store == nil { return store.ErrNotSupported @@ -2109,18 +2124,38 @@ func (s *ShardStore) explicitGroupRouteForKey(groupID uint64, key []byte) distri return route } -func (s *ShardStore) explicitGroupRouteForRange(groupID uint64, start []byte, end []byte) distribution.Route { +func (s *ShardStore) explicitGroupRoutesForRange(groupID uint64, start []byte, end []byte) []distribution.Route { fallback := distribution.Route{GroupID: groupID} if s == nil || s.engine == nil { - return fallback + return []distribution.Route{fallback} + } + routeStart, routeEnd := readinessRouteRangeForScan(start, end) + routes := s.engine.GetIntersectingRoutes(routeStart, routeEnd) + if len(routes) == 0 { + return []distribution.Route{fallback} + } + out := make([]distribution.Route, 0, len(routes)) + for _, route := range routes { + if route.GroupID != groupID { + return []distribution.Route{fallback} + } + out = append(out, route) + } + return out +} + +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 s.engine.GetIntersectingRoutes(routeStart, routeEnd) { - if route.GroupID == groupID { - return route + for _, route := range routes { + if err := s.verifyTargetReadinessForRouteRange(ctx, g, route, routeStart, routeEnd); err != nil { + return err } } - return fallback + return nil } func (s *ShardStore) routeAndGroupForKey(key []byte) (distribution.Route, *ShardGroup, bool) { diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 5748cdfc0..ffc34ff60 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -185,6 +185,31 @@ func TestShardStoreTargetReadinessAcceptsClearedDescriptorAndRetainsFloor(t *tes 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 TestShardStoreExplicitGroupReadUsesRouteProofForReadiness(t *testing.T) { t.Parallel() @@ -210,6 +235,49 @@ func TestShardStoreExplicitGroupReadUsesRouteProofForReadiness(t *testing.T) { 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 TestShardStorePhysicalLimitScanChecksTargetReadiness(t *testing.T) { t.Parallel() @@ -410,6 +478,33 @@ func TestShardStoreDeletePrefixAtChecksTargetReadinessBeforeDelete(t *testing.T) 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 TestShardStoreDeletePrefixAtRaftEnforcesRouteFloorBeforeDelete(t *testing.T) { t.Parallel() diff --git a/kv/shard_store_txn_lock_test.go b/kv/shard_store_txn_lock_test.go index c395b3dc9..aeb6a2e58 100644 --- a/kv/shard_store_txn_lock_test.go +++ b/kv/shard_store_txn_lock_test.go @@ -478,7 +478,7 @@ func TestShardStorePhysicalLimitScanPreservesRouteProofDuringLockResolution(t *t engine := distribution.NewEngine() require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ - Version: 1, + Version: 2, Routes: []distribution.RouteDescriptor{{ RouteID: 1, Start: []byte("a"), diff --git a/store/migration_readiness.go b/store/migration_readiness.go index 7aa160ca3..e0c83a1df 100644 --- a/store/migration_readiness.go +++ b/store/migration_readiness.go @@ -79,6 +79,9 @@ func (s *pebbleStore) loadPebbleTargetReadinessStatesFromDB() ([]TargetStagedRea 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") @@ -92,6 +95,11 @@ func (s *pebbleStore) loadPebbleTargetReadinessStatesFromDB() ([]TargetStagedRea 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) diff --git a/store/migration_versions_test.go b/store/migration_versions_test.go index 7332f318b..ee7ffaf42 100644 --- a/store/migration_versions_test.go +++ b/store/migration_versions_test.go @@ -358,6 +358,42 @@ func TestPebbleTargetStagedReadinessPersistsAcrossReopen(t *testing.T) { 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 TestPebbleImportMetadataPersistsAcrossReopen(t *testing.T) { ctx := context.Background() dir, err := os.MkdirTemp("", "migration-import-persist-*") From fa769ac5ccb1c888cd55d16134af78ce9cbc5a13 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 06:57:03 +0900 Subject: [PATCH 12/29] Harden split migration start guards --- adapter/distribution_server.go | 42 +++++++++++++++-- adapter/distribution_server_test.go | 56 +++++++++++++++++++++++ main.go | 11 +++++ main_encryption_rotate_on_startup_test.go | 1 + 4 files changed, 107 insertions(+), 3 deletions(-) diff --git a/adapter/distribution_server.go b/adapter/distribution_server.go index 389559b4e..ec1b9f7df 100644 --- a/adapter/distribution_server.go +++ b/adapter/distribution_server.go @@ -28,6 +28,7 @@ type DistributionServer struct { coordinator kv.Coordinator readTracker *kv.ActiveTimestampTracker migrationCapabilityGate SplitMigrationCapabilityGate + knownRaftGroups map[uint64]struct{} reloadRetry struct { attempts int interval time.Duration @@ -62,6 +63,21 @@ func WithSplitMigrationCapabilityGate(gate SplitMigrationCapabilityGate) Distrib } } +// WithDistributionKnownRaftGroups configures the Raft group IDs this node can +// route migration work to. +func WithDistributionKnownRaftGroups(groupIDs ...uint64) DistributionServerOption { + return func(s *DistributionServer) { + groups := make(map[uint64]struct{}, len(groupIDs)) + for _, groupID := range groupIDs { + if groupID == 0 { + continue + } + groups[groupID] = struct{}{} + } + s.knownRaftGroups = groups + } +} + // WithCatalogReloadRetryPolicy configures the retry policy used after split // commit when waiting for the local catalog snapshot to become visible. func WithCatalogReloadRetryPolicy(attempts int, interval time.Duration) DistributionServerOption { @@ -100,6 +116,8 @@ var ( errDistributionEngineNotConfigured = errors.New("distribution engine is not configured") errDistributionCatalogVersionNotFound = errors.New("route catalog version not found") errDistributionClusterNotReady = errors.New("cluster is not ready for split migration") + errDistributionRaftGroupsNotKnown = errors.New("distribution raft groups are not configured") + errDistributionUnknownTargetGroup = errors.New("unknown target group") ) // NewDistributionServer creates a new server. @@ -260,6 +278,9 @@ func (s *DistributionServer) startSplitMigrationParent( if parent.GroupID == req.GetTargetGroupId() { return distribution.RouteDescriptor{}, grpcStatusError(codes.InvalidArgument, "target group must differ from source route group") } + if err := s.verifyKnownTargetGroup(req.GetTargetGroupId()); err != nil { + return distribution.RouteDescriptor{}, err + } if err := s.rejectLiveSplitJob(ctx, snapshot, parent); err != nil { return distribution.RouteDescriptor{}, err } @@ -628,6 +649,16 @@ func (s *DistributionServer) updateSplitJobViaCoordinator( return nil } +func (s *DistributionServer) verifyKnownTargetGroup(groupID uint64) error { + if len(s.knownRaftGroups) == 0 { + return grpcStatusError(codes.FailedPrecondition, errDistributionRaftGroupsNotKnown.Error()) + } + if _, ok := s.knownRaftGroups[groupID]; !ok { + return grpcStatusError(codes.InvalidArgument, errDistributionUnknownTargetGroup.Error()) + } + return nil +} + func (s *DistributionServer) createSplitJobViaCoordinator(ctx context.Context, readTS uint64, job distribution.SplitJob) error { if job.JobID == math.MaxUint64 { return splitJobCatalogStatusError(distribution.ErrCatalogSplitJobIDOverflow) @@ -651,9 +682,14 @@ func (s *DistributionServer) createSplitJobViaCoordinator(ctx context.Context, r Value: distribution.EncodeCatalogNextSplitJobID(job.JobID + 1), }, }, - IsTxn: true, - StartTS: readTS, - ReadKeys: [][]byte{jobKey, nextIDKey}, + IsTxn: true, + StartTS: readTS, + ReadKeys: [][]byte{ + jobKey, + nextIDKey, + distribution.CatalogVersionKey(), + distribution.CatalogRouteKey(job.SourceRouteID), + }, }); err != nil { return splitJobCoordinatorStatusError(err) } diff --git a/adapter/distribution_server_test.go b/adapter/distribution_server_test.go index ceb73fe36..5751377d8 100644 --- a/adapter/distribution_server_test.go +++ b/adapter/distribution_server_test.go @@ -224,6 +224,7 @@ func TestDistributionServerStartSplitMigration_CreatesPlannedJobWhenGateOpen(t * distribution.NewEngine(), catalog, WithDistributionCoordinator(coordinator), + WithDistributionKnownRaftGroups(1, 2), WithSplitMigrationCapabilityGate(func(context.Context) error { gateCalls++ return nil @@ -241,6 +242,12 @@ func TestDistributionServerStartSplitMigration_CreatesPlannedJobWhenGateOpen(t * require.Equal(t, uint64(1), resp.JobId) require.Equal(t, 1, gateCalls) require.Equal(t, 1, coordinator.dispatchCalls) + require.ElementsMatch(t, []string{ + string(distribution.CatalogSplitJobKey(1)), + string(distribution.CatalogNextSplitJobIDKey()), + string(distribution.CatalogVersionKey()), + string(distribution.CatalogRouteKey(1)), + }, byteSliceStrings(coordinator.lastReadKeys)) jobs, err := catalog.ListSplitJobs(ctx) require.NoError(t, err) @@ -259,6 +266,43 @@ func TestDistributionServerStartSplitMigration_CreatesPlannedJobWhenGateOpen(t * require.Equal(t, uint64(2), next) } +func TestDistributionServerStartSplitMigration_RejectsUnknownTargetGroup(t *testing.T) { + t.Parallel() + + ctx := context.Background() + baseStore := store.NewMVCCStore() + catalog := distribution.NewCatalogStore(baseStore) + saved, err := catalog.Save(ctx, 0, []distribution.RouteDescriptor{{ + RouteID: 1, + Start: []byte("a"), + End: []byte("m"), + GroupID: 1, + State: distribution.RouteStateActive, + ParentRouteID: 0, + }}) + require.NoError(t, err) + + coordinator := newDistributionCoordinatorStub(baseStore, true) + s := NewDistributionServer( + distribution.NewEngine(), + catalog, + WithDistributionCoordinator(coordinator), + WithDistributionKnownRaftGroups(1, 2), + WithSplitMigrationCapabilityGate(func(context.Context) error { return nil }), + ) + + _, err = s.StartSplitMigration(ctx, &pb.StartSplitMigrationRequest{ + ExpectedCatalogVersion: saved.Version, + RouteId: 1, + SplitKey: []byte("g"), + TargetGroupId: 3, + }) + require.Error(t, err) + require.Equal(t, codes.InvalidArgument, status.Code(err)) + require.ErrorContains(t, err, errDistributionUnknownTargetGroup.Error()) + require.Zero(t, coordinator.dispatchCalls) +} + func TestDistributionServerStartSplitMigration_RejectsSecondLiveJob(t *testing.T) { t.Parallel() @@ -299,6 +343,7 @@ func TestDistributionServerStartSplitMigration_RejectsSecondLiveJob(t *testing.T distribution.NewEngine(), catalog, WithDistributionCoordinator(coordinator), + WithDistributionKnownRaftGroups(1, 2, 3, 4), WithSplitMigrationCapabilityGate(func(context.Context) error { return nil }), ) @@ -349,6 +394,7 @@ func TestDistributionServerStartSplitMigration_RejectsReservedMigrationRange(t * distribution.NewEngine(), catalog, WithDistributionCoordinator(coordinator), + WithDistributionKnownRaftGroups(1, 2), WithSplitMigrationCapabilityGate(func(context.Context) error { return nil }), ) @@ -1237,12 +1283,21 @@ func splitJobIDs(jobs []*pb.SplitJob) []uint64 { return ids } +func byteSliceStrings(in [][]byte) []string { + out := make([]string, 0, len(in)) + for _, item := range in { + out = append(out, string(item)) + } + return out +} + type distributionCoordinatorStub struct { store store.MVCCStore leader bool dispatchErr error nextTS uint64 lastStartTS uint64 + lastReadKeys [][]byte afterDispatch func(context.Context, store.MVCCStore, uint64) error asyncApplyDone chan error asyncApplyDelay time.Duration @@ -1266,6 +1321,7 @@ func (s *distributionCoordinatorStub) Dispatch(ctx context.Context, reqs *kv.Ope } startTS, commitTS := s.nextTimestamps(reqs.StartTS) s.lastStartTS = startTS + s.lastReadKeys = cloneByteSlices(reqs.ReadKeys) mutations, err := coordinatorStubMutations(reqs.Elems) if err != nil { diff --git a/main.go b/main.go index c236598d8..b33c4cbcd 100644 --- a/main.go +++ b/main.go @@ -503,6 +503,7 @@ func run() error { distCatalog, adapter.WithDistributionCoordinator(coordinate), adapter.WithDistributionActiveTimestampTracker(readTracker), + adapter.WithDistributionKnownRaftGroups(shardGroupIDs(shardGroups)...), ) startMonitoringCollectors(runCtx, metricsRegistry, runtimes, clock) compactor := kv.NewFSMCompactor( @@ -655,6 +656,15 @@ func cloneRaftServers(in []raftengine.Server) []raftengine.Server { return append([]raftengine.Server(nil), in...) } +func shardGroupIDs(groups map[uint64]*kv.ShardGroup) []uint64 { + ids := make([]uint64, 0, len(groups)) + for id := range groups { + ids = append(ids, id) + } + slices.Sort(ids) + return ids +} + type runtimeConfig struct { groups []groupSpec defaultGroup uint64 @@ -1830,6 +1840,7 @@ func startupRotationGatedMethod(fullMethod string) bool { case pb.Internal_Forward_FullMethodName, pb.AdminForward_Forward_FullMethodName, pb.Distribution_SplitRange_FullMethodName, + pb.Distribution_StartSplitMigration_FullMethodName, pb.RaftAdmin_AddVoter_FullMethodName, pb.RaftAdmin_AddLearner_FullMethodName, pb.RaftAdmin_PromoteLearner_FullMethodName, diff --git a/main_encryption_rotate_on_startup_test.go b/main_encryption_rotate_on_startup_test.go index 5e16348ad..d8a33706d 100644 --- a/main_encryption_rotate_on_startup_test.go +++ b/main_encryption_rotate_on_startup_test.go @@ -421,6 +421,7 @@ func TestStartupPublicKVGate_BlocksMutatorsUntilReady(t *testing.T) { pb.Internal_Forward_FullMethodName, pb.AdminForward_Forward_FullMethodName, pb.Distribution_SplitRange_FullMethodName, + pb.Distribution_StartSplitMigration_FullMethodName, pb.RaftAdmin_AddVoter_FullMethodName, pb.RaftAdmin_AddLearner_FullMethodName, pb.RaftAdmin_PromoteLearner_FullMethodName, From 4447dd315213282d49714e9945f52aa2c0c08f3b Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 07:09:26 +0900 Subject: [PATCH 13/29] Reject non-active migration sources --- adapter/distribution_server.go | 4 +++ adapter/distribution_server_test.go | 52 +++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+) diff --git a/adapter/distribution_server.go b/adapter/distribution_server.go index ec1b9f7df..5098ab32e 100644 --- a/adapter/distribution_server.go +++ b/adapter/distribution_server.go @@ -118,6 +118,7 @@ var ( errDistributionClusterNotReady = errors.New("cluster is not ready for split migration") errDistributionRaftGroupsNotKnown = errors.New("distribution raft groups are not configured") errDistributionUnknownTargetGroup = errors.New("unknown target group") + errDistributionSourceRouteNotActive = errors.New("source route is not active") ) // NewDistributionServer creates a new server. @@ -268,6 +269,9 @@ func (s *DistributionServer) startSplitMigrationParent( if !found { return distribution.RouteDescriptor{}, grpcStatusError(codes.NotFound, errDistributionUnknownRoute.Error()) } + if parent.State != distribution.RouteStateActive { + return distribution.RouteDescriptor{}, grpcStatusError(codes.FailedPrecondition, errDistributionSourceRouteNotActive.Error()) + } splitKey := distribution.CloneBytes(req.GetSplitKey()) if err := validateSplitKey(parent, splitKey); err != nil { return distribution.RouteDescriptor{}, err diff --git a/adapter/distribution_server_test.go b/adapter/distribution_server_test.go index 5751377d8..ba8e0fd6f 100644 --- a/adapter/distribution_server_test.go +++ b/adapter/distribution_server_test.go @@ -303,6 +303,58 @@ func TestDistributionServerStartSplitMigration_RejectsUnknownTargetGroup(t *test require.Zero(t, coordinator.dispatchCalls) } +func TestDistributionServerStartSplitMigration_RejectsNonActiveSourceRoute(t *testing.T) { + t.Parallel() + + ctx := context.Background() + for _, tc := range []struct { + name string + state distribution.RouteState + }{ + {name: "write_fenced", state: distribution.RouteStateWriteFenced}, + {name: "migrating_source", state: distribution.RouteStateMigratingSource}, + {name: "migrating_target", state: distribution.RouteStateMigratingTarget}, + } { + t.Run(tc.name, func(t *testing.T) { + baseStore := store.NewMVCCStore() + catalog := distribution.NewCatalogStore(baseStore) + saved, err := catalog.Save(ctx, 0, []distribution.RouteDescriptor{{ + RouteID: 1, + Start: []byte("a"), + End: []byte("m"), + GroupID: 1, + State: tc.state, + ParentRouteID: 0, + }}) + require.NoError(t, err) + + coordinator := newDistributionCoordinatorStub(baseStore, true) + s := NewDistributionServer( + distribution.NewEngine(), + catalog, + WithDistributionCoordinator(coordinator), + WithDistributionKnownRaftGroups(1, 2), + WithSplitMigrationCapabilityGate(func(context.Context) error { return nil }), + ) + + _, err = s.StartSplitMigration(ctx, &pb.StartSplitMigrationRequest{ + ExpectedCatalogVersion: saved.Version, + RouteId: 1, + SplitKey: []byte("g"), + TargetGroupId: 2, + }) + require.Error(t, err) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) + require.ErrorContains(t, err, errDistributionSourceRouteNotActive.Error()) + require.Zero(t, coordinator.dispatchCalls) + + jobs, listErr := catalog.ListSplitJobs(ctx) + require.NoError(t, listErr) + require.Empty(t, jobs) + }) + } +} + func TestDistributionServerStartSplitMigration_RejectsSecondLiveJob(t *testing.T) { t.Parallel() From b5723b5694a08d24b5d5af78a4b908db0fc65e59 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 07:31:55 +0900 Subject: [PATCH 14/29] Fix target readiness staged delete paths --- kv/fsm.go | 26 ++++ kv/fsm_migration_fence_test.go | 27 ++++ kv/shard_store.go | 149 +++++++++++++++++++--- kv/shard_store_test.go | 69 ++++++++++ kv/sharded_coordinator.go | 19 +++ kv/sharded_coordinator_del_prefix_test.go | 37 ++++++ 6 files changed, 307 insertions(+), 20 deletions(-) diff --git a/kv/fsm.go b/kv/fsm.go index 85dd67590..4954209f5 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -579,6 +579,13 @@ func (f *kvFSM) handleDelPrefix(ctx context.Context, prefix []byte, commitTS uin 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, f.pendingApplyIdx) + } + 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) } @@ -586,6 +593,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) { diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index 9ebaab322..09ab1163a 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -183,6 +183,33 @@ func TestFSMRejectsDelPrefixWithoutTargetReadinessProof(t *testing.T) { 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 TestFSMRejectsFullRangeDelPrefixWhenRouteIsWriteFenced(t *testing.T) { t.Parallel() diff --git a/kv/shard_store.go b/kv/shard_store.go index 4e3fb96c4..181f3286a 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -304,7 +304,7 @@ 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 } @@ -312,11 +312,7 @@ func (s *ShardStore) CommittedVersionAt(ctx context.Context, key []byte, commitT // 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 @@ -326,7 +322,19 @@ func (s *ShardStore) CommittedVersionAt(ctx context.Context, key []byte, commitT // serialization. return false, nil } - exists, err := g.Store.CommittedVersionAt(ctx, key, commitTS) + 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) } @@ -399,7 +407,55 @@ func (s *ShardStore) ScanGroupAt(ctx context.Context, groupID uint64, start []by if err := s.verifyExplicitGroupRoutesForRange(ctx, groupID, routes, start, end); err != nil { return nil, err } - return s.scanRouteAtDirection(ctx, routes[0], start, end, limit, ts, false, true) + return s.scanExplicitGroupRoutesAt(ctx, routes, start, end, limit, ts) +} + +func (s *ShardStore) scanExplicitGroupRoutesAt(ctx context.Context, routes []distribution.Route, start []byte, end []byte, limit int, ts uint64) ([]*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) + rawRouteBounds := rawScanBoundsMatchRouteKeyspace(start, end) + for _, route := range routes { + scanStart, scanEnd := start, end + if rawRouteBounds { + scanStart = clampScanStart(start, route.Start) + scanEnd = clampScanEnd(end, route.End) + } + kvs, err := s.scanRouteAtDirection(ctx, route, scanStart, scanEnd, limit-len(out), ts, false, true) + if err != nil { + return nil, err + } + for _, kvp := range kvs { + if !kvBelongsToRoute(kvp, route) { + continue + } + out = append(out, kvp) + if len(out) == limit { + return out, nil + } + } + } + return out, nil +} + +func rawScanBoundsMatchRouteKeyspace(start []byte, end []byte) bool { + routeStart, routeEnd := readinessRouteRangeForScan(start, end) + if !bytes.Equal(routeStart, start) { + return false + } + if end == nil { + return routeEnd == nil + } + return bytes.Equal(routeEnd, end) +} + +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) { @@ -1838,13 +1894,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.verifyPrefixDeleteRoutes(ctx, 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) } @@ -1854,13 +1917,20 @@ func (s *ShardStore) DeletePrefixAt(ctx context.Context, prefix []byte, excludeP // DeletePrefixAtRaft is the raft-apply variant of DeletePrefixAt. func (s *ShardStore) DeletePrefixAtRaft(ctx context.Context, prefix []byte, excludePrefix []byte, commitTS uint64) error { - if err := s.verifyPrefixDeleteRoutes(ctx, 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.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) } @@ -1883,13 +1953,20 @@ 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 { - if err := s.verifyPrefixDeleteRoutes(ctx, 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.DeletePrefixAtRaftAt(ctx, stagedPrefix, stagedExcludePrefix, commitTS, appliedIndex) + } + 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 @@ -1907,25 +1984,57 @@ func (s *ShardStore) DeletePrefixAtRaftAt(ctx context.Context, prefix []byte, ex return nil } -func (s *ShardStore) verifyPrefixDeleteRoutes(ctx context.Context, prefix []byte, commitTS uint64) error { +func (s *ShardStore) verifyPrefixDeleteRoutes(ctx context.Context, prefix []byte, commitTS uint64) ([]distribution.Route, error) { if s == nil || s.engine == nil { - return nil + return nil, nil } routeStart, routeEnd := routePrefixRange(prefix) routes := s.engine.GetIntersectingRoutes(routeStart, routeEnd) if len(routes) == 0 { - return errors.WithStack(ErrRouteCutoverPending) + return nil, errors.WithStack(ErrRouteCutoverPending) } for _, route := range routes { g, ok := s.groupForID(route.GroupID) if !ok || g == nil || g.Store == nil { - return store.ErrNotSupported + return nil, store.ErrNotSupported } if err := s.verifyTargetReadinessForRouteRange(ctx, g, route, routeStart, routeEnd); err != nil { - return err + return nil, err } if err := verifyRouteWriteFloor(route, commitTS); err != nil { - return err + return nil, err + } + } + return routes, 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 diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index ffc34ff60..72f23b6ad 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -88,6 +88,19 @@ 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 TestShardStoreTargetReadinessFailsClosedWithoutDescriptorProof(t *testing.T) { t.Parallel() @@ -278,6 +291,45 @@ func TestShardStoreScanGroupAtChecksEveryCoveredRoute(t *testing.T) { 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 TestShardStorePhysicalLimitScanChecksTargetReadiness(t *testing.T) { t.Parallel() @@ -527,6 +579,23 @@ func TestShardStoreDeletePrefixAtRaftEnforcesRouteFloorBeforeDelete(t *testing.T 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 TestShardStoreScanAndLatestCommitTS_MergeStagedVisibility(t *testing.T) { t.Parallel() diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index 285036d3e..a31e02e6d 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.rejectDelPrefixesBelowWriteFloor(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{ @@ -1087,6 +1090,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) rejectDelPrefixesBelowWriteFloor(elems []*Elem[OP], commitTS uint64) error { if c == nil || c.engine == nil { return nil diff --git a/kv/sharded_coordinator_del_prefix_test.go b/kv/sharded_coordinator_del_prefix_test.go index 5b712db58..e2767eeb0 100644 --- a/kv/sharded_coordinator_del_prefix_test.go +++ b/kv/sharded_coordinator_del_prefix_test.go @@ -200,6 +200,43 @@ func TestShardedCoordinatorRejectsDelPrefixBelowRouteFloorBeforeBroadcast(t *tes 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() From f22c4ff9edd3d778a18d285432f7261f42394ba4 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 07:44:28 +0900 Subject: [PATCH 15/29] Stabilize urgent compactor pagination test --- adapter/redis_delta_compactor_test.go | 29 ++++++++++----------------- 1 file changed, 11 insertions(+), 18 deletions(-) diff --git a/adapter/redis_delta_compactor_test.go b/adapter/redis_delta_compactor_test.go index c3f6f1518..358241350 100644 --- a/adapter/redis_delta_compactor_test.go +++ b/adapter/redis_delta_compactor_test.go @@ -411,27 +411,20 @@ func TestDeltaCompactor_UrgentCompactionPagination(t *testing.T) { require.NoError(t, st.PutAt(ctx, dKey, delta, i+1, 0)) } - // Queue and process the urgent compaction. - c.TriggerUrgentCompaction("hash", userKey) - - runCtx, cancel := context.WithCancel(ctx) - defer cancel() - go func() { _ = c.Run(runCtx) }() + // Exercise the urgent pagination loop directly. Running through c.Run would + // also start an initial background SyncOnce; the local test coordinator applies + // elems one-by-one instead of atomically, so a test read can observe the meta + // update before all delete elems have been applied under the race detector. + c.compactUrgentKey(ctx, urgentCompactionRequest{typeName: "hash", userKey: userKey}) - // Wait until the base meta holds the accumulated total. - // The pagination loop should take two passes: first 1025, then 9. - require.Eventually(t, func() bool { - readTS := st.LastCommitTS() - raw, err := st.GetAt(ctx, store.HashMetaKey(userKey), readTS) - if err != nil { - return false - } - got, err := store.UnmarshalHashMeta(raw) - return err == nil && got.Len == int64(totalDeltasU64) - }, 5*time.Second, 20*time.Millisecond, "all %d delta keys should be compacted into base meta", totalDeltasU64) + readTS := st.LastCommitTS() + raw, err := st.GetAt(ctx, store.HashMetaKey(userKey), readTS) + require.NoError(t, err) + got, err := store.UnmarshalHashMeta(raw) + require.NoError(t, err) + require.Equal(t, int64(totalDeltasU64), got.Len, "all %d delta keys should be compacted into base meta", totalDeltasU64) // No delta keys should remain after pagination compaction. - readTS := st.LastCommitTS() prefix := store.HashMetaDeltaScanPrefix(userKey) end := store.PrefixScanEnd(prefix) remaining, err := st.ScanAt(ctx, prefix, end, int(totalDeltasU64)+1, readTS) From 119dda295c940fe8cec6d619d40942dc76f9b767 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 08:49:12 +0900 Subject: [PATCH 16/29] Guard staged readiness probes --- kv/fsm.go | 2 +- kv/fsm_migration_fence_test.go | 61 ++++++++++++++++++++++++++++++++++ kv/shard_store.go | 5 ++- kv/shard_store_test.go | 37 +++++++++++++++++++++ 4 files changed, 103 insertions(+), 2 deletions(-) diff --git a/kv/fsm.go b/kv/fsm.go index 4954209f5..9a0c10714 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -581,7 +581,7 @@ func (f *kvFSM) handleDelPrefix(ctx context.Context, prefix []byte, commitTS uin } routes := f.stagedVisibilityRoutesForPrefix(prefix) deleteStagedPrefix := func(stagedPrefix []byte, stagedExcludePrefix []byte) error { - return f.store.DeletePrefixAtRaftAt(ctx, stagedPrefix, stagedExcludePrefix, commitTS, f.pendingApplyIdx) + return f.store.DeletePrefixAtRaftAt(ctx, stagedPrefix, stagedExcludePrefix, commitTS, 0) } if err := deleteStagedVisibilityPrefixes(routes, prefix, txnCommonPrefix, deleteStagedPrefix); err != nil { return err diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index 09ab1163a..e42514f43 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -11,6 +11,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() @@ -210,6 +239,38 @@ func TestFSMDelPrefixTombstonesStagedVisibilityRows(t *testing.T) { 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() diff --git a/kv/shard_store.go b/kv/shard_store.go index 181f3286a..c27183717 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -308,6 +308,9 @@ func (s *ShardStore) CommittedVersionAt(ctx context.Context, key []byte, commitT if !ok || g.Store == nil { return false, nil } + if err := s.verifyTargetReadinessForRange(ctx, g, route, key, nextScanCursor(key)); 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) @@ -1962,7 +1965,7 @@ func (s *ShardStore) DeletePrefixAtRaftAt(ctx context.Context, prefix []byte, ex continue } deleteStagedPrefix := func(stagedPrefix []byte, stagedExcludePrefix []byte) error { - return g.Store.DeletePrefixAtRaftAt(ctx, stagedPrefix, stagedExcludePrefix, commitTS, appliedIndex) + return g.Store.DeletePrefixAtRaftAt(ctx, stagedPrefix, stagedExcludePrefix, commitTS, 0) } if err := deleteStagedVisibilityPrefixes(routesForGroupID(routes, groupID), prefix, excludePrefix, deleteStagedPrefix); err != nil { return err diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 72f23b6ad..c41b7c655 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -101,6 +101,25 @@ func TestShardStoreCommittedVersionAtChecksStagedVisibilityKey(t *testing.T) { 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 TestShardStoreTargetReadinessFailsClosedWithoutDescriptorProof(t *testing.T) { t.Parallel() @@ -596,6 +615,24 @@ func TestShardStoreDeletePrefixAtTombstonesStagedVisibilityRows(t *testing.T) { 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 TestShardStoreScanAndLatestCommitTS_MergeStagedVisibility(t *testing.T) { t.Parallel() From e6f88bd725d288296529be840e1f162ccda0b4ff Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 09:19:51 +0900 Subject: [PATCH 17/29] Guard target readiness retry paths --- kv/fsm.go | 43 ++++++++++++--- kv/fsm_migration_fence_test.go | 95 ++++++++++++++++++++++++++++++++++ kv/shard_store.go | 3 ++ kv/shard_store_test.go | 66 +++++++++++++++++++++++ 4 files changed, 201 insertions(+), 6 deletions(-) diff --git a/kv/fsm.go b/kv/fsm.go index 9a0c10714..8a7727464 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -678,6 +678,30 @@ func (f *kvFSM) verifyTargetReadinessForMutations(ctx context.Context, muts []*p 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) @@ -1140,6 +1164,9 @@ func (f *kvFSM) handlePrepareRequest(ctx context.Context, r *pb.Request) error { if meta.CommitTS != 0 { floorTS = meta.CommitTS } + if err := f.verifyTargetReadinessForTxnFootprint(ctx, muts, r.ReadKeys, meta.PrimaryKey); err != nil { + return err + } uniq, err := f.uniqueMutationsNotFenced(ctx, muts, floorTS) if err != nil { return err @@ -1202,7 +1229,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 } @@ -1243,15 +1270,19 @@ func (f *kvFSM) uniqueMutationsNotFenced(ctx context.Context, muts []*pb.Mutatio 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 } diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index e42514f43..8fb30e168 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -88,6 +88,35 @@ func newTargetReadinessFSM(t *testing.T, route distribution.RouteDescriptor) *kv 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 TestFSMRejectsRawPointWriteOnWriteFencedRoute(t *testing.T) { t.Parallel() @@ -344,6 +373,72 @@ func TestFSMRejectsOnePhaseTxnBelowRouteFloor(t *testing.T) { 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 TestFSMPrepareUsesCommitTSForRouteFloorWhenPresent(t *testing.T) { t.Parallel() diff --git a/kv/shard_store.go b/kv/shard_store.go index c27183717..716acedc5 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -325,6 +325,9 @@ func (s *ShardStore) CommittedVersionAt(ctx context.Context, key []byte, commitT // serialization. return false, nil } + if err := s.verifyTargetReadinessForRange(ctx, g, route, key, nextScanCursor(key)); err != nil { + return false, err + } return committedVersionAtForRoute(ctx, g.Store, route, key, commitTS) } diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index c41b7c655..b3ff784a1 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -2,9 +2,11 @@ package kv import ( "context" + "errors" "testing" "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/internal/raftengine" "github.com/bootjp/elastickv/internal/s3keys" "github.com/bootjp/elastickv/store" "github.com/stretchr/testify/require" @@ -64,6 +66,49 @@ func newReadinessShardStore(t *testing.T, route distribution.RouteDescriptor) (* return NewShardStore(engine, map[uint64]*ShardGroup{route.GroupID: group}), group } +type readinessFenceEngine struct { + onLinearizableRead func() +} + +func (e *readinessFenceEngine) State() raftengine.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 TestShardStoreGetAt_MergesStagedVisibility(t *testing.T) { t.Parallel() @@ -120,6 +165,27 @@ func TestShardStoreCommittedVersionAtChecksTargetReadiness(t *testing.T) { 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() From 79b34092b76a6c0ea4749b9a12dbcc04036df59b Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 10:02:40 +0900 Subject: [PATCH 18/29] Tighten target readiness proof --- kv/fsm.go | 11 +- kv/fsm_migration_fence_test.go | 42 ++++++++ kv/shard_store.go | 180 +++++++++++++++++++++++++------ kv/shard_store_test.go | 89 +++++++++++++++ store/lsm_migration.go | 3 + store/migration_versions.go | 14 +++ store/migration_versions_test.go | 28 +++++ store/store.go | 14 +-- 8 files changed, 332 insertions(+), 49 deletions(-) diff --git a/kv/fsm.go b/kv/fsm.go index 8a7727464..8c88d21d1 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -1464,15 +1464,8 @@ func (f *kvFSM) uniqueMutationsAboveWriteFloor(ctx context.Context, muts []*pb.M return uniq, nil } -func (f *kvFSM) uniqueAbortCleanupMutations(ctx context.Context, muts []*pb.Mutation) ([]*pb.Mutation, error) { - uniq, err := uniqueMutations(muts) - if err != nil { - return nil, err - } - if err := f.verifyTargetReadinessForMutations(ctx, uniq); 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) { diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index 8fb30e168..82c2d5410 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -543,3 +543,45 @@ func TestFSMAbortCleanupBypassesRetainedWriteFloor(t *testing.T) { _, intentErr := fsm.store.GetAt(ctx, txnIntentKey(primaryKey), ^uint64(0)) require.ErrorIs(t, intentErr, store.ErrKeyNotFound) } + +func TestFSMAbortCleanupBypassesTargetReadiness(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, + }) + 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)) + + abort := &pb.Request{ + IsTxn: true, + Phase: pb.Phase_ABORT, + Ts: startTS, + Mutations: []*pb.Mutation{ + {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: primaryKey, CommitTS: abortTS})}, + {Op: pb.Op_PUT, Key: primaryKey}, + }, + } + 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/shard_store.go b/kv/shard_store.go index 716acedc5..cb7fef43c 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -92,7 +92,9 @@ 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) { - if err := s.verifyTargetReadinessForRange(ctx, g, route, key, nextScanCursor(key)); err != nil { + var err error + route, err = s.targetReadyRouteForRange(ctx, g, route, key, nextScanCursor(key)) + if err != nil { return nil, err } if !isTxnInternalKey(key) { @@ -104,7 +106,9 @@ 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) { - if err := s.verifyTargetReadinessForRange(ctx, g, route, key, nextScanCursor(key)); err != nil { + var err error + route, err = s.targetReadyRouteForRange(ctx, g, route, key, nextScanCursor(key)) + if err != nil { return nil, err } if routeHasStagedVisibility(route) { @@ -137,36 +141,101 @@ func routeSatisfiesTargetReadiness(route distribution.Route, ready store.TargetS return ready.ExpectedCutoverVersion == 0 || catalogVersion >= ready.ExpectedCutoverVersion } -func (s *ShardStore) verifyTargetReadinessForRange(ctx context.Context, g *ShardGroup, route distribution.Route, start []byte, end []byte) error { +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.verifyTargetReadinessForRouteRange(ctx, g, route, routeStart, routeEnd) + return s.targetReadyRouteForRouteRange(ctx, g, route, routeStart, routeEnd) } -func (s *ShardStore) verifyTargetReadinessForRouteRange(ctx context.Context, g *ShardGroup, route distribution.Route, routeStart []byte, routeEnd []byte) error { +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) { if g == nil || g.Store == nil { - return nil + return route, nil } reader, ok := g.Store.(store.MigrationTargetReadinessReader) if !ok { - return nil + return route, nil } states, err := reader.MigrationTargetReadinessStates(ctx) if err != nil { - return errors.WithStack(err) + return route, errors.WithStack(err) } - var catalogVersion uint64 - if s != nil && s.engine != nil { - catalogVersion = s.engine.Version() + applicable := targetReadinessApplicableStates(states, route, routeStart, routeEnd) + if len(applicable) == 0 { + return route, nil + } + + proofRoutes, catalogVersion, ok := s.readinessProofRoutes(route, routeStart, routeEnd) + if !ok { + return route, errors.WithStack(ErrRouteCutoverPending) + } + if !readinessProofSatisfiesStates(applicable, proofRoutes, route.GroupID, catalogVersion) { + return route, errors.WithStack(ErrRouteCutoverPending) + } + if len(proofRoutes) == 1 { + return proofRoutes[0], nil + } + return route, 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 !targetReadinessAppliesToRoute(route, routeStart, routeEnd, ready) { + 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 !routeSatisfiesTargetReadiness(route, ready, catalogVersion) { - return errors.WithStack(ErrRouteCutoverPending) + 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 nil + return proof, snap.Version(), len(proof) > 0 } func targetReadinessAppliesToRoute(route distribution.Route, routeStart []byte, routeEnd []byte, ready store.TargetStagedReadinessState) bool { @@ -240,6 +309,7 @@ func latestMVCCVersionAt(ctx context.Context, st store.MVCCStore, key []byte, ts StartKey: key, EndKey: nextScanCursor(key), MaxCommitTSInclusive: ts, + ReadTS: ts, MaxVersions: 1, MaxScannedBytes: 0, MinCommitTSExclusive: 0, @@ -308,7 +378,9 @@ func (s *ShardStore) CommittedVersionAt(ctx context.Context, key []byte, commitT if !ok || g.Store == nil { return false, nil } - if err := s.verifyTargetReadinessForRange(ctx, g, route, key, nextScanCursor(key)); err != 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 @@ -325,7 +397,8 @@ func (s *ShardStore) CommittedVersionAt(ctx context.Context, key []byte, commitT // serialization. return false, nil } - if err := s.verifyTargetReadinessForRange(ctx, g, route, key, nextScanCursor(key)); err != nil { + route, err = s.targetReadyRouteForRange(ctx, g, route, key, nextScanCursor(key)) + if err != nil { return false, err } return committedVersionAtForRoute(ctx, g.Store, route, key, commitTS) @@ -669,7 +742,9 @@ func (s *ShardStore) scanRouteAtDirectionPhysicalLimit( if !ok || g == nil || g.Store == nil { return nil, false, nil } - if err := s.verifyTargetReadinessForRange(ctx, g, route, start, end); err != nil { + var err error + route, err = s.targetReadyRouteForRange(ctx, g, route, start, end) + if err != nil { return nil, false, err } @@ -685,10 +760,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 @@ -696,6 +768,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, @@ -746,7 +839,9 @@ func (s *ShardStore) scanRouteLocal( ts uint64, reverse bool, ) ([]*store.KVPair, error) { - if err := s.verifyTargetReadinessForRange(ctx, g, route, start, end); err != nil { + var err error + route, err = s.targetReadyRouteForRange(ctx, g, route, start, end) + if err != nil { return nil, err } if routeHasStagedVisibility(route) { @@ -771,6 +866,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) @@ -794,13 +894,12 @@ func (s *ShardStore) scanRouteAtLeader( ts uint64, reverse bool, ) ([]*store.KVPair, error) { - if err := s.verifyTargetReadinessForRange(ctx, g, route, start, end); err != nil { + var err error + route, err = s.targetReadyRouteForRange(ctx, g, route, start, end) + if err != nil { return nil, err } - var ( - kvs []*store.KVPair - err error - ) + var kvs []*store.KVPair switch { case routeHasStagedVisibility(route): kvs, err = s.scanRouteWithStagedVisibility(ctx, g, route, start, end, limit, ts, reverse) @@ -894,6 +993,7 @@ func collectLatestLogicalVersions( StartKey: scanStart, EndKey: scanEnd, MaxCommitTSInclusive: ts, + ReadTS: ts, Cursor: cursor, MaxVersions: stagedVisibilityExportPageSize, }) @@ -1088,7 +1188,9 @@ func (s *ShardStore) PutAt(ctx context.Context, key []byte, value []byte, commit if !ok || g == nil || g.Store == nil { return store.ErrNotSupported } - if err := s.verifyTargetReadinessForRange(ctx, g, route, key, nextScanCursor(key)); err != nil { + var err error + route, err = s.targetReadyRouteForRange(ctx, g, route, key, nextScanCursor(key)) + if err != nil { return err } if err := verifyRouteWriteFloor(route, commitTS); err != nil { @@ -1102,7 +1204,9 @@ func (s *ShardStore) DeleteAt(ctx context.Context, key []byte, commitTS uint64) if !ok || g == nil || g.Store == nil { return store.ErrNotSupported } - if err := s.verifyTargetReadinessForRange(ctx, g, route, key, nextScanCursor(key)); err != nil { + var err error + route, err = s.targetReadyRouteForRange(ctx, g, route, key, nextScanCursor(key)) + if err != nil { return err } if err := verifyRouteWriteFloor(route, commitTS); err != nil { @@ -1116,7 +1220,9 @@ func (s *ShardStore) PutWithTTLAt(ctx context.Context, key []byte, value []byte, if !ok || g == nil || g.Store == nil { return store.ErrNotSupported } - if err := s.verifyTargetReadinessForRange(ctx, g, route, key, nextScanCursor(key)); err != nil { + var err error + route, err = s.targetReadyRouteForRange(ctx, g, route, key, nextScanCursor(key)) + if err != nil { return err } if err := verifyRouteWriteFloor(route, commitTS); err != nil { @@ -1130,7 +1236,9 @@ func (s *ShardStore) ExpireAt(ctx context.Context, key []byte, expireAt uint64, if !ok || g == nil || g.Store == nil { return store.ErrNotSupported } - if err := s.verifyTargetReadinessForRange(ctx, g, route, key, nextScanCursor(key)); err != nil { + var err error + route, err = s.targetReadyRouteForRange(ctx, g, route, key, nextScanCursor(key)) + if err != nil { return err } if err := verifyRouteWriteFloor(route, commitTS); err != nil { @@ -1170,7 +1278,9 @@ func (s *ShardStore) LatestCommitTS(ctx context.Context, key []byte) (uint64, bo } func (s *ShardStore) localLatestCommitTS(ctx context.Context, g *ShardGroup, route distribution.Route, key []byte) (uint64, bool, error) { - if err := s.verifyTargetReadinessForRange(ctx, g, route, key, nextScanCursor(key)); err != nil { + 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) @@ -1841,7 +1951,9 @@ func (s *ShardStore) verifyMutationWriteRoute(ctx context.Context, key []byte, c if !ok || g == nil || g.Store == nil { return store.ErrNotSupported } - if err := s.verifyTargetReadinessForRange(ctx, g, route, key, nextScanCursor(key)); err != nil { + var err error + route, err = s.targetReadyRouteForRange(ctx, g, route, key, nextScanCursor(key)) + if err != nil { return err } return verifyRouteWriteFloor(route, commitTS) diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index b3ff784a1..9f343cc32 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -67,10 +67,14 @@ func newReadinessShardStore(t *testing.T, route distribution.RouteDescriptor) (* } type readinessFenceEngine struct { + state raftengine.State onLinearizableRead func() } func (e *readinessFenceEngine) State() raftengine.State { + if e.state != "" { + return e.state + } return raftengine.StateFollower } @@ -308,6 +312,49 @@ func TestShardStoreTargetReadinessRejectsClearedDescriptorBeforeCutoverVersion(t 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 TestShardStoreExplicitGroupReadUsesRouteProofForReadiness(t *testing.T) { t.Parallel() @@ -437,6 +484,31 @@ func TestShardStorePhysicalLimitScanChecksTargetReadiness(t *testing.T) { 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() @@ -744,6 +816,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 TestShardStoreScanAt_IncludesListKeysAcrossShards(t *testing.T) { t.Parallel() diff --git a/store/lsm_migration.go b/store/lsm_migration.go index e0796cbf6..f60f5ae02 100644 --- a/store/lsm_migration.go +++ b/store/lsm_migration.go @@ -24,6 +24,9 @@ func (s *pebbleStore) exportVersionsLocked(ctx context.Context, opts ExportVersi if opts.MaxVersions <= 0 { return ExportVersionsResult{Done: true}, nil } + if readTSCompacted(opts.ReadTS, s.effectiveMinRetainedTS()) { + return ExportVersionsResult{}, ErrReadTSCompacted + } iter, err := s.db.NewIter(pebbleExportIterOptions(opts)) if err != nil { diff --git a/store/migration_versions.go b/store/migration_versions.go index b1e275ecb..4c5f4edb8 100644 --- a/store/migration_versions.go +++ b/store/migration_versions.go @@ -196,7 +196,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.key) { @@ -227,6 +234,13 @@ func (s *mvccStore) ExportVersions(ctx context.Context, opts ExportVersionsOptio return result, nil } +func (s *mvccStore) checkExportReadTSLocked(opts ExportVersionsOptions) error { + if readTSCompacted(opts.ReadTS, s.minRetainedTS) { + return ErrReadTSCompacted + } + return nil +} + 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 ee7ffaf42..2fb9ee72c 100644 --- a/store/migration_versions_test.go +++ b/store/migration_versions_test.go @@ -53,6 +53,34 @@ 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) + + result, err := st.ExportVersions(ctx, ExportVersionsOptions{ + StartKey: []byte("k"), + EndKey: []byte("l"), + MaxCommitTSInclusive: 15, + MaxVersions: 10, + }) + require.NoError(t, err) + require.Len(t, result.Versions, 1) + }) +} + func TestExportVersionsCursorResumesWithinHotKey(t *testing.T) { runMigrationStoreSuite(t, func(t *testing.T, st MVCCStore) { ctx := context.Background() diff --git a/store/store.go b/store/store.go index ae53203f9..b0dea1e79 100644 --- a/store/store.go +++ b/store/store.go @@ -82,12 +82,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. From ce0c2e1b4ed48e7c3f8a2708a580bd1dd6dea1ed Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 10:21:55 +0900 Subject: [PATCH 19/29] Wire split migration capability gate --- main.go | 5 +++++ main_catalog_test.go | 6 ++++++ 2 files changed, 11 insertions(+) diff --git a/main.go b/main.go index b33c4cbcd..6cfd9abd6 100644 --- a/main.go +++ b/main.go @@ -504,6 +504,7 @@ func run() error { adapter.WithDistributionCoordinator(coordinate), adapter.WithDistributionActiveTimestampTracker(readTracker), adapter.WithDistributionKnownRaftGroups(shardGroupIDs(shardGroups)...), + adapter.WithSplitMigrationCapabilityGate(splitMigrationCapabilityGate), ) startMonitoringCollectors(runCtx, metricsRegistry, runtimes, clock) compactor := kv.NewFSMCompactor( @@ -665,6 +666,10 @@ func shardGroupIDs(groups map[uint64]*kv.ShardGroup) []uint64 { return ids } +func splitMigrationCapabilityGate(context.Context) error { + return nil +} + type runtimeConfig struct { groups []groupSpec defaultGroup uint64 diff --git a/main_catalog_test.go b/main_catalog_test.go index 193a98219..843b2dea8 100644 --- a/main_catalog_test.go +++ b/main_catalog_test.go @@ -68,3 +68,9 @@ func TestSetupDistributionCatalog_UsesResolvedCatalogGroup(t *testing.T) { require.NoError(t, err) require.NotNil(t, catalog) } + +func TestSplitMigrationCapabilityGateIsReady(t *testing.T) { + t.Parallel() + + require.NoError(t, splitMigrationCapabilityGate(context.Background())) +} From f884ed2fff8f74896e7eb1cc4d42bc10587168b3 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 10:49:51 +0900 Subject: [PATCH 20/29] Replicate target readiness guards --- adapter/internal.go | 40 ++++++ adapter/internal_migration_test.go | 40 ++++++ kv/fsm.go | 6 + kv/fsm_migration_readiness.go | 59 ++++++++ kv/fsm_migration_readiness_test.go | 44 ++++++ kv/shard_store.go | 47 ++++-- kv/shard_store_test.go | 86 +++++++++++ main.go | 1 + main_encryption_rotate_on_startup_test.go | 1 + proto/internal.pb.go | 167 ++++++++++++++++++++-- proto/internal.proto | 13 ++ proto/internal_grpc.pb.go | 48 ++++++- 12 files changed, 522 insertions(+), 30 deletions(-) create mode 100644 kv/fsm_migration_readiness.go create mode 100644 kv/fsm_migration_readiness_test.go diff --git a/adapter/internal.go b/adapter/internal.go index efacdab71..3773ec185 100644 --- a/adapter/internal.go +++ b/adapter/internal.go @@ -207,6 +207,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) verifyInternalLeader(ctx context.Context) error { if i.leader.State() != raftengine.StateLeader { return errors.WithStack(ErrNotLeader) @@ -265,6 +287,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 54227d418..5a641f7c2 100644 --- a/adapter/internal_migration_test.go +++ b/adapter/internal_migration_test.go @@ -187,3 +187,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/kv/fsm.go b/kv/fsm.go index 8c88d21d1..e4f1650fd 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -380,6 +380,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: @@ -419,6 +421,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) { 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/shard_store.go b/kv/shard_store.go index cb7fef43c..100407e4f 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -152,33 +152,41 @@ func (s *ShardStore) verifyTargetReadinessForRange(ctx context.Context, g *Shard } 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 route, nil + return []distribution.Route{route}, nil } reader, ok := g.Store.(store.MigrationTargetReadinessReader) if !ok { - return route, nil + return []distribution.Route{route}, nil } states, err := reader.MigrationTargetReadinessStates(ctx) if err != nil { - return route, errors.WithStack(err) + return nil, errors.WithStack(err) } applicable := targetReadinessApplicableStates(states, route, routeStart, routeEnd) if len(applicable) == 0 { - return route, nil + return []distribution.Route{route}, nil } proofRoutes, catalogVersion, ok := s.readinessProofRoutes(route, routeStart, routeEnd) if !ok { - return route, errors.WithStack(ErrRouteCutoverPending) + return nil, errors.WithStack(ErrRouteCutoverPending) } if !readinessProofSatisfiesStates(applicable, proofRoutes, route.GroupID, catalogVersion) { - return route, errors.WithStack(ErrRouteCutoverPending) - } - if len(proofRoutes) == 1 { - return proofRoutes[0], nil + return nil, errors.WithStack(ErrRouteCutoverPending) } - return route, nil + return proofRoutes, nil } func targetReadinessApplicableStates( @@ -2111,19 +2119,32 @@ func (s *ShardStore) verifyPrefixDeleteRoutes(ctx context.Context, prefix []byte 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 } - if err := s.verifyTargetReadinessForRouteRange(ctx, g, route, routeStart, routeEnd); err != nil { + proofRoutes, err := s.verifyPrefixDeleteRoute(ctx, g, route, routeStart, routeEnd, commitTS) + if err != nil { return nil, err } - if err := verifyRouteWriteFloor(route, commitTS); err != nil { + 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 routes, nil + return proofRoutes, nil } func routesForGroupID(routes []distribution.Route, groupID uint64) []distribution.Route { diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 9f343cc32..ac3449cc5 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -355,6 +355,92 @@ func TestShardStoreTargetReadinessUsesSingleCatalogSnapshotProof(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() diff --git a/main.go b/main.go index c236598d8..75fc290d9 100644 --- a/main.go +++ b/main.go @@ -1828,6 +1828,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{ { From 49adab3dd41acf0042af2724ac2f5f8dc47291f6 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 10:55:58 +0900 Subject: [PATCH 21/29] Gate split migration on peer capability --- adapter/distribution_server.go | 8 + adapter/distribution_server_test.go | 10 + main.go | 86 ++++++++- main_catalog_test.go | 76 +++++++- proto/distribution.pb.go | 279 +++++++++++++++++++--------- proto/distribution.proto | 8 + proto/distribution_grpc.pb.go | 60 ++++-- 7 files changed, 421 insertions(+), 106 deletions(-) diff --git a/adapter/distribution_server.go b/adapter/distribution_server.go index 5098ab32e..42360803a 100644 --- a/adapter/distribution_server.go +++ b/adapter/distribution_server.go @@ -99,6 +99,7 @@ const ( splitJobListCursorTerminalOff = 1 splitJobListCursorJobIDOff = splitJobListCursorTerminalOff + 8 splitJobListCursorEncodedBytes = splitJobListCursorJobIDOff + 8 + splitMigrationCapabilityV2 = "split_migration_v2" ) var ( @@ -244,6 +245,13 @@ func (s *DistributionServer) StartSplitMigration(ctx context.Context, req *pb.St }, nil } +func (s *DistributionServer) GetSplitMigrationCapability(context.Context, *pb.GetSplitMigrationCapabilityRequest) (*pb.GetSplitMigrationCapabilityResponse, error) { + return &pb.GetSplitMigrationCapabilityResponse{ + MigrationCapable: true, + Capabilities: []string{splitMigrationCapabilityV2}, + }, nil +} + func (s *DistributionServer) verifyStartSplitMigrationReady(ctx context.Context, req *pb.StartSplitMigrationRequest) error { if req.GetTargetGroupId() == 0 { return grpcStatusError(codes.InvalidArgument, distribution.ErrCatalogSplitJobTargetGroupRequired.Error()) diff --git a/adapter/distribution_server_test.go b/adapter/distribution_server_test.go index ba8e0fd6f..21a34d96b 100644 --- a/adapter/distribution_server_test.go +++ b/adapter/distribution_server_test.go @@ -149,6 +149,16 @@ func TestDistributionServerListRoutes_RequiresCatalog(t *testing.T) { require.ErrorContains(t, err, errDistributionCatalogNotConfigured.Error()) } +func TestDistributionServerGetSplitMigrationCapabilityReportsReady(t *testing.T) { + t.Parallel() + + s := NewDistributionServer(distribution.NewEngine(), nil) + resp, err := s.GetSplitMigrationCapability(context.Background(), &pb.GetSplitMigrationCapabilityRequest{}) + require.NoError(t, err) + require.True(t, resp.GetMigrationCapable()) + require.Contains(t, resp.GetCapabilities(), splitMigrationCapabilityV2) +} + func TestDistributionServerStartSplitMigration_FailsClosedUntilCapabilityGate(t *testing.T) { t.Parallel() diff --git a/main.go b/main.go index 6cfd9abd6..93fed90ce 100644 --- a/main.go +++ b/main.go @@ -53,6 +53,8 @@ const ( etcdMaxSizePerMsg = 1 << 20 etcdMaxInflightMsg = 1024 defaultTSOBatchSize = 256 + + splitMigrationCapabilityProbeTimeout = 2 * time.Second ) func newRaftFactory(engineType raftEngineType, coldStartObs raftengine.ColdStartObserver) (raftengine.Factory, error) { @@ -498,13 +500,18 @@ func run() error { startKeyVizFlusher(runCtx, eg, sampler) startKeyVizLeaderTermPublisher(runCtx, eg, sampler, runtimes) startMemoryWatchdog(runCtx, eg, cancel) + splitMigrationGate := newSplitMigrationCapabilityGate( + splitMigrationCapabilityPeers(bootstrapCfg, cfg.defaultGroup), + splitMigrationCapabilityProbeTimeout, + nil, + ) distServer := adapter.NewDistributionServer( cfg.engine, distCatalog, adapter.WithDistributionCoordinator(coordinate), adapter.WithDistributionActiveTimestampTracker(readTracker), adapter.WithDistributionKnownRaftGroups(shardGroupIDs(shardGroups)...), - adapter.WithSplitMigrationCapabilityGate(splitMigrationCapabilityGate), + adapter.WithSplitMigrationCapabilityGate(splitMigrationGate), ) startMonitoringCollectors(runCtx, metricsRegistry, runtimes, clock) compactor := kv.NewFSMCompactor( @@ -666,7 +673,82 @@ func shardGroupIDs(groups map[uint64]*kv.ShardGroup) []uint64 { return ids } -func splitMigrationCapabilityGate(context.Context) error { +type splitMigrationCapabilityPeer struct { + ID string + Address string +} + +type splitMigrationCapabilityProbe func(context.Context, string) error + +func splitMigrationCapabilityPeers(cfg raftBootstrapConfig, defaultGroup uint64) []splitMigrationCapabilityPeer { + servers := cfg.adminSeed(defaultGroup) + if len(servers) == 0 { + return nil + } + peers := make([]splitMigrationCapabilityPeer, 0, len(servers)) + seen := make(map[string]struct{}, len(servers)) + for _, server := range servers { + key := server.ID + "\x00" + server.Address + if _, ok := seen[key]; ok { + continue + } + seen[key] = struct{}{} + id := server.ID + if id == "" { + id = server.Address + } + peers = append(peers, splitMigrationCapabilityPeer{ + ID: id, + Address: server.Address, + }) + } + return peers +} + +func newSplitMigrationCapabilityGate(peers []splitMigrationCapabilityPeer, timeout time.Duration, probe splitMigrationCapabilityProbe) adapter.SplitMigrationCapabilityGate { + if probe == nil { + probe = probeSplitMigrationCapabilityPeer + } + return func(ctx context.Context) error { + if len(peers) == 0 { + return status.Error(codes.FailedPrecondition, "split migration capability peers are not configured") + } + for _, peer := range peers { + peerCtx := ctx + cancel := func() {} + if timeout > 0 { + var cancelCtx context.CancelFunc + peerCtx, cancelCtx = context.WithTimeout(ctx, timeout) + cancel = cancelCtx + } + err := probe(peerCtx, peer.Address) + cancel() + if err != nil { + return status.Errorf(codes.FailedPrecondition, "split migration capability peer %s is not ready: %v", peer.ID, err) + } + } + return nil + } +} + +func probeSplitMigrationCapabilityPeer(ctx context.Context, address string) error { + if strings.TrimSpace(address) == "" { + return errors.New("empty split migration capability peer address") + } + conn, err := grpc.NewClient(address, internalutil.GRPCDialOptions()...) + if err != nil { + return errors.Wrapf(err, "dial split migration capability peer %s", address) + } + defer func() { + _ = conn.Close() + }() + resp, err := pb.NewDistributionClient(conn).GetSplitMigrationCapability(ctx, &pb.GetSplitMigrationCapabilityRequest{}) + if err != nil { + return errors.Wrapf(err, "probe split migration capability peer %s", address) + } + if resp == nil || !resp.GetMigrationCapable() { + return errors.New("peer does not advertise split migration capability") + } return nil } diff --git a/main_catalog_test.go b/main_catalog_test.go index 843b2dea8..d2e63118b 100644 --- a/main_catalog_test.go +++ b/main_catalog_test.go @@ -3,10 +3,14 @@ package main import ( "context" "testing" + "time" "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/internal/raftengine" "github.com/bootjp/elastickv/store" "github.com/stretchr/testify/require" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" ) func TestDistributionCatalogGroupID_UsesCatalogKeyRoute(t *testing.T) { @@ -69,8 +73,76 @@ func TestSetupDistributionCatalog_UsesResolvedCatalogGroup(t *testing.T) { require.NotNil(t, catalog) } -func TestSplitMigrationCapabilityGateIsReady(t *testing.T) { +func TestSplitMigrationCapabilityGateChecksAllPeers(t *testing.T) { t.Parallel() - require.NoError(t, splitMigrationCapabilityGate(context.Background())) + peers := []splitMigrationCapabilityPeer{ + {ID: "n1", Address: "10.0.0.11:50051"}, + {ID: "n2", Address: "10.0.0.12:50051"}, + } + var probed []string + gate := newSplitMigrationCapabilityGate(peers, time.Second, func(_ context.Context, address string) error { + probed = append(probed, address) + return nil + }) + + require.NoError(t, gate(context.Background())) + require.Equal(t, []string{"10.0.0.11:50051", "10.0.0.12:50051"}, probed) +} + +func TestSplitMigrationCapabilityGateFailsClosedWithoutPeers(t *testing.T) { + t.Parallel() + + gate := newSplitMigrationCapabilityGate(nil, time.Second, func(context.Context, string) error { + t.Fatal("probe must not run without peers") + return nil + }) + + err := gate(context.Background()) + require.Error(t, err) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) + require.ErrorContains(t, err, "peers are not configured") +} + +func TestSplitMigrationCapabilityGateFailsClosedWhenPeerMissingCapability(t *testing.T) { + t.Parallel() + + peers := []splitMigrationCapabilityPeer{ + {ID: "n1", Address: "10.0.0.11:50051"}, + {ID: "n2", Address: "10.0.0.12:50051"}, + } + gate := newSplitMigrationCapabilityGate(peers, time.Second, func(_ context.Context, address string) error { + if address == "10.0.0.12:50051" { + return status.Error(codes.Unimplemented, "method not found") + } + return nil + }) + + err := gate(context.Background()) + require.Error(t, err) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) + require.ErrorContains(t, err, "n2") + require.ErrorContains(t, err, "method not found") +} + +func TestSplitMigrationCapabilityPeersUseDefaultGroupAdminSeed(t *testing.T) { + t.Parallel() + + cfg := raftBootstrapConfig{ + groupServers: map[uint64][]raftengine.Server{ + 1: { + {Suffrage: "voter", ID: "n1", Address: "10.0.0.11:50051"}, + {Suffrage: "voter", ID: "n2", Address: "10.0.0.12:50051"}, + }, + 2: { + {Suffrage: "voter", ID: "n1", Address: "10.0.0.11:50052"}, + {Suffrage: "voter", ID: "n2", Address: "10.0.0.12:50052"}, + }, + }, + } + + require.Equal(t, []splitMigrationCapabilityPeer{ + {ID: "n1", Address: "10.0.0.11:50052"}, + {ID: "n2", Address: "10.0.0.12:50052"}, + }, splitMigrationCapabilityPeers(cfg, 2)) } diff --git a/proto/distribution.pb.go b/proto/distribution.pb.go index 21d15180c..a4a586aa9 100644 --- a/proto/distribution.pb.go +++ b/proto/distribution.pb.go @@ -1280,6 +1280,94 @@ func (x *StartSplitMigrationResponse) GetJobId() uint64 { return 0 } +type GetSplitMigrationCapabilityRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *GetSplitMigrationCapabilityRequest) Reset() { + *x = GetSplitMigrationCapabilityRequest{} + mi := &file_distribution_proto_msgTypes[13] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *GetSplitMigrationCapabilityRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*GetSplitMigrationCapabilityRequest) ProtoMessage() {} + +func (x *GetSplitMigrationCapabilityRequest) ProtoReflect() protoreflect.Message { + mi := &file_distribution_proto_msgTypes[13] + 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 GetSplitMigrationCapabilityRequest.ProtoReflect.Descriptor instead. +func (*GetSplitMigrationCapabilityRequest) Descriptor() ([]byte, []int) { + return file_distribution_proto_rawDescGZIP(), []int{13} +} + +type GetSplitMigrationCapabilityResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + MigrationCapable bool `protobuf:"varint,1,opt,name=migration_capable,json=migrationCapable,proto3" json:"migration_capable,omitempty"` + Capabilities []string `protobuf:"bytes,2,rep,name=capabilities,proto3" json:"capabilities,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *GetSplitMigrationCapabilityResponse) Reset() { + *x = GetSplitMigrationCapabilityResponse{} + mi := &file_distribution_proto_msgTypes[14] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *GetSplitMigrationCapabilityResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*GetSplitMigrationCapabilityResponse) ProtoMessage() {} + +func (x *GetSplitMigrationCapabilityResponse) ProtoReflect() protoreflect.Message { + mi := &file_distribution_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 GetSplitMigrationCapabilityResponse.ProtoReflect.Descriptor instead. +func (*GetSplitMigrationCapabilityResponse) Descriptor() ([]byte, []int) { + return file_distribution_proto_rawDescGZIP(), []int{14} +} + +func (x *GetSplitMigrationCapabilityResponse) GetMigrationCapable() bool { + if x != nil { + return x.MigrationCapable + } + return false +} + +func (x *GetSplitMigrationCapabilityResponse) GetCapabilities() []string { + if x != nil { + return x.Capabilities + } + return nil +} + type GetRouteOwnershipRequest struct { state protoimpl.MessageState `protogen:"open.v1"` Key []byte `protobuf:"bytes,1,opt,name=key,proto3" json:"key,omitempty"` @@ -1290,7 +1378,7 @@ type GetRouteOwnershipRequest struct { func (x *GetRouteOwnershipRequest) Reset() { *x = GetRouteOwnershipRequest{} - mi := &file_distribution_proto_msgTypes[13] + mi := &file_distribution_proto_msgTypes[15] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1302,7 +1390,7 @@ func (x *GetRouteOwnershipRequest) String() string { func (*GetRouteOwnershipRequest) ProtoMessage() {} func (x *GetRouteOwnershipRequest) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[13] + mi := &file_distribution_proto_msgTypes[15] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1315,7 +1403,7 @@ func (x *GetRouteOwnershipRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use GetRouteOwnershipRequest.ProtoReflect.Descriptor instead. func (*GetRouteOwnershipRequest) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{13} + return file_distribution_proto_rawDescGZIP(), []int{15} } func (x *GetRouteOwnershipRequest) GetKey() []byte { @@ -1343,7 +1431,7 @@ type GetRouteOwnershipResponse struct { func (x *GetRouteOwnershipResponse) Reset() { *x = GetRouteOwnershipResponse{} - mi := &file_distribution_proto_msgTypes[14] + mi := &file_distribution_proto_msgTypes[16] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1355,7 +1443,7 @@ func (x *GetRouteOwnershipResponse) String() string { func (*GetRouteOwnershipResponse) ProtoMessage() {} func (x *GetRouteOwnershipResponse) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[14] + mi := &file_distribution_proto_msgTypes[16] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1368,7 +1456,7 @@ func (x *GetRouteOwnershipResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use GetRouteOwnershipResponse.ProtoReflect.Descriptor instead. func (*GetRouteOwnershipResponse) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{14} + return file_distribution_proto_rawDescGZIP(), []int{16} } func (x *GetRouteOwnershipResponse) GetRoute() *RouteDescriptor { @@ -1403,7 +1491,7 @@ type GetIntersectingRoutesRequest struct { func (x *GetIntersectingRoutesRequest) Reset() { *x = GetIntersectingRoutesRequest{} - mi := &file_distribution_proto_msgTypes[15] + mi := &file_distribution_proto_msgTypes[17] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1415,7 +1503,7 @@ func (x *GetIntersectingRoutesRequest) String() string { func (*GetIntersectingRoutesRequest) ProtoMessage() {} func (x *GetIntersectingRoutesRequest) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[15] + mi := &file_distribution_proto_msgTypes[17] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1428,7 +1516,7 @@ func (x *GetIntersectingRoutesRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use GetIntersectingRoutesRequest.ProtoReflect.Descriptor instead. func (*GetIntersectingRoutesRequest) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{15} + return file_distribution_proto_rawDescGZIP(), []int{17} } func (x *GetIntersectingRoutesRequest) GetStart() []byte { @@ -1462,7 +1550,7 @@ type GetIntersectingRoutesResponse struct { func (x *GetIntersectingRoutesResponse) Reset() { *x = GetIntersectingRoutesResponse{} - mi := &file_distribution_proto_msgTypes[16] + mi := &file_distribution_proto_msgTypes[18] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1474,7 +1562,7 @@ func (x *GetIntersectingRoutesResponse) String() string { func (*GetIntersectingRoutesResponse) ProtoMessage() {} func (x *GetIntersectingRoutesResponse) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[16] + mi := &file_distribution_proto_msgTypes[18] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1487,7 +1575,7 @@ func (x *GetIntersectingRoutesResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use GetIntersectingRoutesResponse.ProtoReflect.Descriptor instead. func (*GetIntersectingRoutesResponse) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{16} + return file_distribution_proto_rawDescGZIP(), []int{18} } func (x *GetIntersectingRoutesResponse) GetRoutes() []*RouteDescriptor { @@ -1513,7 +1601,7 @@ type GetSplitJobRequest struct { func (x *GetSplitJobRequest) Reset() { *x = GetSplitJobRequest{} - mi := &file_distribution_proto_msgTypes[17] + mi := &file_distribution_proto_msgTypes[19] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1525,7 +1613,7 @@ func (x *GetSplitJobRequest) String() string { func (*GetSplitJobRequest) ProtoMessage() {} func (x *GetSplitJobRequest) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[17] + mi := &file_distribution_proto_msgTypes[19] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1538,7 +1626,7 @@ func (x *GetSplitJobRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use GetSplitJobRequest.ProtoReflect.Descriptor instead. func (*GetSplitJobRequest) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{17} + return file_distribution_proto_rawDescGZIP(), []int{19} } func (x *GetSplitJobRequest) GetJobId() uint64 { @@ -1557,7 +1645,7 @@ type GetSplitJobResponse struct { func (x *GetSplitJobResponse) Reset() { *x = GetSplitJobResponse{} - mi := &file_distribution_proto_msgTypes[18] + mi := &file_distribution_proto_msgTypes[20] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1569,7 +1657,7 @@ func (x *GetSplitJobResponse) String() string { func (*GetSplitJobResponse) ProtoMessage() {} func (x *GetSplitJobResponse) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[18] + mi := &file_distribution_proto_msgTypes[20] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1582,7 +1670,7 @@ func (x *GetSplitJobResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use GetSplitJobResponse.ProtoReflect.Descriptor instead. func (*GetSplitJobResponse) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{18} + return file_distribution_proto_rawDescGZIP(), []int{20} } func (x *GetSplitJobResponse) GetJob() *SplitJob { @@ -1603,7 +1691,7 @@ type ListSplitJobsRequest struct { func (x *ListSplitJobsRequest) Reset() { *x = ListSplitJobsRequest{} - mi := &file_distribution_proto_msgTypes[19] + mi := &file_distribution_proto_msgTypes[21] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1615,7 +1703,7 @@ func (x *ListSplitJobsRequest) String() string { func (*ListSplitJobsRequest) ProtoMessage() {} func (x *ListSplitJobsRequest) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[19] + mi := &file_distribution_proto_msgTypes[21] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1628,7 +1716,7 @@ func (x *ListSplitJobsRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use ListSplitJobsRequest.ProtoReflect.Descriptor instead. func (*ListSplitJobsRequest) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{19} + return file_distribution_proto_rawDescGZIP(), []int{21} } func (x *ListSplitJobsRequest) GetSinceTerminalAtMs() uint64 { @@ -1662,7 +1750,7 @@ type ListSplitJobsResponse struct { func (x *ListSplitJobsResponse) Reset() { *x = ListSplitJobsResponse{} - mi := &file_distribution_proto_msgTypes[20] + mi := &file_distribution_proto_msgTypes[22] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1674,7 +1762,7 @@ func (x *ListSplitJobsResponse) String() string { func (*ListSplitJobsResponse) ProtoMessage() {} func (x *ListSplitJobsResponse) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[20] + mi := &file_distribution_proto_msgTypes[22] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1687,7 +1775,7 @@ func (x *ListSplitJobsResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use ListSplitJobsResponse.ProtoReflect.Descriptor instead. func (*ListSplitJobsResponse) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{20} + return file_distribution_proto_rawDescGZIP(), []int{22} } func (x *ListSplitJobsResponse) GetJobs() []*SplitJob { @@ -1713,7 +1801,7 @@ type AbandonSplitJobRequest struct { func (x *AbandonSplitJobRequest) Reset() { *x = AbandonSplitJobRequest{} - mi := &file_distribution_proto_msgTypes[21] + mi := &file_distribution_proto_msgTypes[23] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1725,7 +1813,7 @@ func (x *AbandonSplitJobRequest) String() string { func (*AbandonSplitJobRequest) ProtoMessage() {} func (x *AbandonSplitJobRequest) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[21] + mi := &file_distribution_proto_msgTypes[23] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1738,7 +1826,7 @@ func (x *AbandonSplitJobRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use AbandonSplitJobRequest.ProtoReflect.Descriptor instead. func (*AbandonSplitJobRequest) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{21} + return file_distribution_proto_rawDescGZIP(), []int{23} } func (x *AbandonSplitJobRequest) GetJobId() uint64 { @@ -1756,7 +1844,7 @@ type AbandonSplitJobResponse struct { func (x *AbandonSplitJobResponse) Reset() { *x = AbandonSplitJobResponse{} - mi := &file_distribution_proto_msgTypes[22] + mi := &file_distribution_proto_msgTypes[24] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1768,7 +1856,7 @@ func (x *AbandonSplitJobResponse) String() string { func (*AbandonSplitJobResponse) ProtoMessage() {} func (x *AbandonSplitJobResponse) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[22] + mi := &file_distribution_proto_msgTypes[24] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1781,7 +1869,7 @@ func (x *AbandonSplitJobResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use AbandonSplitJobResponse.ProtoReflect.Descriptor instead. func (*AbandonSplitJobResponse) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{22} + return file_distribution_proto_rawDescGZIP(), []int{24} } type RetrySplitJobRequest struct { @@ -1793,7 +1881,7 @@ type RetrySplitJobRequest struct { func (x *RetrySplitJobRequest) Reset() { *x = RetrySplitJobRequest{} - mi := &file_distribution_proto_msgTypes[23] + mi := &file_distribution_proto_msgTypes[25] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1805,7 +1893,7 @@ func (x *RetrySplitJobRequest) String() string { func (*RetrySplitJobRequest) ProtoMessage() {} func (x *RetrySplitJobRequest) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[23] + mi := &file_distribution_proto_msgTypes[25] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1818,7 +1906,7 @@ func (x *RetrySplitJobRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use RetrySplitJobRequest.ProtoReflect.Descriptor instead. func (*RetrySplitJobRequest) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{23} + return file_distribution_proto_rawDescGZIP(), []int{25} } func (x *RetrySplitJobRequest) GetJobId() uint64 { @@ -1836,7 +1924,7 @@ type RetrySplitJobResponse struct { func (x *RetrySplitJobResponse) Reset() { *x = RetrySplitJobResponse{} - mi := &file_distribution_proto_msgTypes[24] + mi := &file_distribution_proto_msgTypes[26] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1848,7 +1936,7 @@ func (x *RetrySplitJobResponse) String() string { func (*RetrySplitJobResponse) ProtoMessage() {} func (x *RetrySplitJobResponse) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[24] + mi := &file_distribution_proto_msgTypes[26] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1861,7 +1949,7 @@ func (x *RetrySplitJobResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use RetrySplitJobResponse.ProtoReflect.Descriptor instead. func (*RetrySplitJobResponse) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{24} + return file_distribution_proto_rawDescGZIP(), []int{26} } var File_distribution_proto protoreflect.FileDescriptor @@ -1960,7 +2048,11 @@ const file_distribution_proto_rawDesc = "" + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"]\n" + "\x1bStartSplitMigrationResponse\x12'\n" + "\x0fcatalog_version\x18\x01 \x01(\x04R\x0ecatalogVersion\x12\x15\n" + - "\x06job_id\x18\x02 \x01(\x04R\x05jobId\"U\n" + + "\x06job_id\x18\x02 \x01(\x04R\x05jobId\"$\n" + + "\"GetSplitMigrationCapabilityRequest\"v\n" + + "#GetSplitMigrationCapabilityResponse\x12+\n" + + "\x11migration_capable\x18\x01 \x01(\bR\x10migrationCapable\x12\"\n" + + "\fcapabilities\x18\x02 \x03(\tR\fcapabilities\"U\n" + "\x18GetRouteOwnershipRequest\x12\x10\n" + "\x03key\x18\x01 \x01(\fR\x03key\x12'\n" + "\x0fcatalog_version\x18\x02 \x01(\x04R\x0ecatalogVersion\"\x82\x01\n" + @@ -2021,7 +2113,7 @@ const file_distribution_proto_rawDesc = "" + "\x13SplitJobExportPhase\x12\x1f\n" + "\x1bSPLIT_JOB_EXPORT_PHASE_NONE\x10\x00\x12#\n" + "\x1fSPLIT_JOB_EXPORT_PHASE_BACKFILL\x10\x01\x12%\n" + - "!SPLIT_JOB_EXPORT_PHASE_DELTA_COPY\x10\x022\xf6\x05\n" + + "!SPLIT_JOB_EXPORT_PHASE_DELTA_COPY\x10\x022\xe2\x06\n" + "\fDistribution\x121\n" + "\bGetRoute\x12\x10.GetRouteRequest\x1a\x11.GetRouteResponse\"\x00\x12=\n" + "\fGetTimestamp\x12\x14.GetTimestampRequest\x1a\x15.GetTimestampResponse\"\x00\x127\n" + @@ -2029,7 +2121,8 @@ const file_distribution_proto_rawDesc = "" + "ListRoutes\x12\x12.ListRoutesRequest\x1a\x13.ListRoutesResponse\"\x00\x127\n" + "\n" + "SplitRange\x12\x12.SplitRangeRequest\x1a\x13.SplitRangeResponse\"\x00\x12R\n" + - "\x13StartSplitMigration\x12\x1b.StartSplitMigrationRequest\x1a\x1c.StartSplitMigrationResponse\"\x00\x12L\n" + + "\x13StartSplitMigration\x12\x1b.StartSplitMigrationRequest\x1a\x1c.StartSplitMigrationResponse\"\x00\x12j\n" + + "\x1bGetSplitMigrationCapability\x12#.GetSplitMigrationCapabilityRequest\x1a$.GetSplitMigrationCapabilityResponse\"\x00\x12L\n" + "\x11GetRouteOwnership\x12\x19.GetRouteOwnershipRequest\x1a\x1a.GetRouteOwnershipResponse\"\x00\x12X\n" + "\x15GetIntersectingRoutes\x12\x1d.GetIntersectingRoutesRequest\x1a\x1e.GetIntersectingRoutesResponse\"\x00\x12:\n" + "\vGetSplitJob\x12\x13.GetSplitJobRequest\x1a\x14.GetSplitJobResponse\"\x00\x12@\n" + @@ -2050,38 +2143,40 @@ func file_distribution_proto_rawDescGZIP() []byte { } var file_distribution_proto_enumTypes = make([]protoimpl.EnumInfo, 4) -var file_distribution_proto_msgTypes = make([]protoimpl.MessageInfo, 26) +var file_distribution_proto_msgTypes = make([]protoimpl.MessageInfo, 28) var file_distribution_proto_goTypes = []any{ - (RouteState)(0), // 0: RouteState - (SplitJobPhase)(0), // 1: SplitJobPhase - (SplitJobBarrierState)(0), // 2: SplitJobBarrierState - (SplitJobExportPhase)(0), // 3: SplitJobExportPhase - (*GetRouteRequest)(nil), // 4: GetRouteRequest - (*GetRouteResponse)(nil), // 5: GetRouteResponse - (*GetTimestampRequest)(nil), // 6: GetTimestampRequest - (*GetTimestampResponse)(nil), // 7: GetTimestampResponse - (*RouteDescriptor)(nil), // 8: RouteDescriptor - (*SplitJobBracketProgress)(nil), // 9: SplitJobBracketProgress - (*SplitJob)(nil), // 10: SplitJob - (*ListRoutesRequest)(nil), // 11: ListRoutesRequest - (*ListRoutesResponse)(nil), // 12: ListRoutesResponse - (*SplitRangeRequest)(nil), // 13: SplitRangeRequest - (*SplitRangeResponse)(nil), // 14: SplitRangeResponse - (*StartSplitMigrationRequest)(nil), // 15: StartSplitMigrationRequest - (*StartSplitMigrationResponse)(nil), // 16: StartSplitMigrationResponse - (*GetRouteOwnershipRequest)(nil), // 17: GetRouteOwnershipRequest - (*GetRouteOwnershipResponse)(nil), // 18: GetRouteOwnershipResponse - (*GetIntersectingRoutesRequest)(nil), // 19: GetIntersectingRoutesRequest - (*GetIntersectingRoutesResponse)(nil), // 20: GetIntersectingRoutesResponse - (*GetSplitJobRequest)(nil), // 21: GetSplitJobRequest - (*GetSplitJobResponse)(nil), // 22: GetSplitJobResponse - (*ListSplitJobsRequest)(nil), // 23: ListSplitJobsRequest - (*ListSplitJobsResponse)(nil), // 24: ListSplitJobsResponse - (*AbandonSplitJobRequest)(nil), // 25: AbandonSplitJobRequest - (*AbandonSplitJobResponse)(nil), // 26: AbandonSplitJobResponse - (*RetrySplitJobRequest)(nil), // 27: RetrySplitJobRequest - (*RetrySplitJobResponse)(nil), // 28: RetrySplitJobResponse - nil, // 29: StartSplitMigrationRequest.OptionsEntry + (RouteState)(0), // 0: RouteState + (SplitJobPhase)(0), // 1: SplitJobPhase + (SplitJobBarrierState)(0), // 2: SplitJobBarrierState + (SplitJobExportPhase)(0), // 3: SplitJobExportPhase + (*GetRouteRequest)(nil), // 4: GetRouteRequest + (*GetRouteResponse)(nil), // 5: GetRouteResponse + (*GetTimestampRequest)(nil), // 6: GetTimestampRequest + (*GetTimestampResponse)(nil), // 7: GetTimestampResponse + (*RouteDescriptor)(nil), // 8: RouteDescriptor + (*SplitJobBracketProgress)(nil), // 9: SplitJobBracketProgress + (*SplitJob)(nil), // 10: SplitJob + (*ListRoutesRequest)(nil), // 11: ListRoutesRequest + (*ListRoutesResponse)(nil), // 12: ListRoutesResponse + (*SplitRangeRequest)(nil), // 13: SplitRangeRequest + (*SplitRangeResponse)(nil), // 14: SplitRangeResponse + (*StartSplitMigrationRequest)(nil), // 15: StartSplitMigrationRequest + (*StartSplitMigrationResponse)(nil), // 16: StartSplitMigrationResponse + (*GetSplitMigrationCapabilityRequest)(nil), // 17: GetSplitMigrationCapabilityRequest + (*GetSplitMigrationCapabilityResponse)(nil), // 18: GetSplitMigrationCapabilityResponse + (*GetRouteOwnershipRequest)(nil), // 19: GetRouteOwnershipRequest + (*GetRouteOwnershipResponse)(nil), // 20: GetRouteOwnershipResponse + (*GetIntersectingRoutesRequest)(nil), // 21: GetIntersectingRoutesRequest + (*GetIntersectingRoutesResponse)(nil), // 22: GetIntersectingRoutesResponse + (*GetSplitJobRequest)(nil), // 23: GetSplitJobRequest + (*GetSplitJobResponse)(nil), // 24: GetSplitJobResponse + (*ListSplitJobsRequest)(nil), // 25: ListSplitJobsRequest + (*ListSplitJobsResponse)(nil), // 26: ListSplitJobsResponse + (*AbandonSplitJobRequest)(nil), // 27: AbandonSplitJobRequest + (*AbandonSplitJobResponse)(nil), // 28: AbandonSplitJobResponse + (*RetrySplitJobRequest)(nil), // 29: RetrySplitJobRequest + (*RetrySplitJobResponse)(nil), // 30: RetrySplitJobResponse + nil, // 31: StartSplitMigrationRequest.OptionsEntry } var file_distribution_proto_depIdxs = []int32{ 0, // 0: RouteDescriptor.state:type_name -> RouteState @@ -2095,7 +2190,7 @@ var file_distribution_proto_depIdxs = []int32{ 8, // 8: ListRoutesResponse.routes:type_name -> RouteDescriptor 8, // 9: SplitRangeResponse.left:type_name -> RouteDescriptor 8, // 10: SplitRangeResponse.right:type_name -> RouteDescriptor - 29, // 11: StartSplitMigrationRequest.options:type_name -> StartSplitMigrationRequest.OptionsEntry + 31, // 11: StartSplitMigrationRequest.options:type_name -> StartSplitMigrationRequest.OptionsEntry 8, // 12: GetRouteOwnershipResponse.route:type_name -> RouteDescriptor 8, // 13: GetIntersectingRoutesResponse.routes:type_name -> RouteDescriptor 10, // 14: GetSplitJobResponse.job:type_name -> SplitJob @@ -2105,25 +2200,27 @@ var file_distribution_proto_depIdxs = []int32{ 11, // 18: Distribution.ListRoutes:input_type -> ListRoutesRequest 13, // 19: Distribution.SplitRange:input_type -> SplitRangeRequest 15, // 20: Distribution.StartSplitMigration:input_type -> StartSplitMigrationRequest - 17, // 21: Distribution.GetRouteOwnership:input_type -> GetRouteOwnershipRequest - 19, // 22: Distribution.GetIntersectingRoutes:input_type -> GetIntersectingRoutesRequest - 21, // 23: Distribution.GetSplitJob:input_type -> GetSplitJobRequest - 23, // 24: Distribution.ListSplitJobs:input_type -> ListSplitJobsRequest - 25, // 25: Distribution.AbandonSplitJob:input_type -> AbandonSplitJobRequest - 27, // 26: Distribution.RetrySplitJob:input_type -> RetrySplitJobRequest - 5, // 27: Distribution.GetRoute:output_type -> GetRouteResponse - 7, // 28: Distribution.GetTimestamp:output_type -> GetTimestampResponse - 12, // 29: Distribution.ListRoutes:output_type -> ListRoutesResponse - 14, // 30: Distribution.SplitRange:output_type -> SplitRangeResponse - 16, // 31: Distribution.StartSplitMigration:output_type -> StartSplitMigrationResponse - 18, // 32: Distribution.GetRouteOwnership:output_type -> GetRouteOwnershipResponse - 20, // 33: Distribution.GetIntersectingRoutes:output_type -> GetIntersectingRoutesResponse - 22, // 34: Distribution.GetSplitJob:output_type -> GetSplitJobResponse - 24, // 35: Distribution.ListSplitJobs:output_type -> ListSplitJobsResponse - 26, // 36: Distribution.AbandonSplitJob:output_type -> AbandonSplitJobResponse - 28, // 37: Distribution.RetrySplitJob:output_type -> RetrySplitJobResponse - 27, // [27:38] is the sub-list for method output_type - 16, // [16:27] is the sub-list for method input_type + 17, // 21: Distribution.GetSplitMigrationCapability:input_type -> GetSplitMigrationCapabilityRequest + 19, // 22: Distribution.GetRouteOwnership:input_type -> GetRouteOwnershipRequest + 21, // 23: Distribution.GetIntersectingRoutes:input_type -> GetIntersectingRoutesRequest + 23, // 24: Distribution.GetSplitJob:input_type -> GetSplitJobRequest + 25, // 25: Distribution.ListSplitJobs:input_type -> ListSplitJobsRequest + 27, // 26: Distribution.AbandonSplitJob:input_type -> AbandonSplitJobRequest + 29, // 27: Distribution.RetrySplitJob:input_type -> RetrySplitJobRequest + 5, // 28: Distribution.GetRoute:output_type -> GetRouteResponse + 7, // 29: Distribution.GetTimestamp:output_type -> GetTimestampResponse + 12, // 30: Distribution.ListRoutes:output_type -> ListRoutesResponse + 14, // 31: Distribution.SplitRange:output_type -> SplitRangeResponse + 16, // 32: Distribution.StartSplitMigration:output_type -> StartSplitMigrationResponse + 18, // 33: Distribution.GetSplitMigrationCapability:output_type -> GetSplitMigrationCapabilityResponse + 20, // 34: Distribution.GetRouteOwnership:output_type -> GetRouteOwnershipResponse + 22, // 35: Distribution.GetIntersectingRoutes:output_type -> GetIntersectingRoutesResponse + 24, // 36: Distribution.GetSplitJob:output_type -> GetSplitJobResponse + 26, // 37: Distribution.ListSplitJobs:output_type -> ListSplitJobsResponse + 28, // 38: Distribution.AbandonSplitJob:output_type -> AbandonSplitJobResponse + 30, // 39: Distribution.RetrySplitJob:output_type -> RetrySplitJobResponse + 28, // [28:40] is the sub-list for method output_type + 16, // [16:28] is the sub-list for method input_type 16, // [16:16] is the sub-list for extension type_name 16, // [16:16] is the sub-list for extension extendee 0, // [0:16] is the sub-list for field type_name @@ -2140,7 +2237,7 @@ func file_distribution_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_distribution_proto_rawDesc), len(file_distribution_proto_rawDesc)), NumEnums: 4, - NumMessages: 26, + NumMessages: 28, NumExtensions: 0, NumServices: 1, }, diff --git a/proto/distribution.proto b/proto/distribution.proto index eb9746f18..4f65163b0 100644 --- a/proto/distribution.proto +++ b/proto/distribution.proto @@ -8,6 +8,7 @@ service Distribution { rpc ListRoutes (ListRoutesRequest) returns (ListRoutesResponse) {} rpc SplitRange (SplitRangeRequest) returns (SplitRangeResponse) {} rpc StartSplitMigration (StartSplitMigrationRequest) returns (StartSplitMigrationResponse) {} + rpc GetSplitMigrationCapability (GetSplitMigrationCapabilityRequest) returns (GetSplitMigrationCapabilityResponse) {} rpc GetRouteOwnership (GetRouteOwnershipRequest) returns (GetRouteOwnershipResponse) {} rpc GetIntersectingRoutes (GetIntersectingRoutesRequest) returns (GetIntersectingRoutesResponse) {} rpc GetSplitJob (GetSplitJobRequest) returns (GetSplitJobResponse) {} @@ -160,6 +161,13 @@ message StartSplitMigrationResponse { uint64 job_id = 2; } +message GetSplitMigrationCapabilityRequest {} + +message GetSplitMigrationCapabilityResponse { + bool migration_capable = 1; + repeated string capabilities = 2; +} + message GetRouteOwnershipRequest { bytes key = 1; uint64 catalog_version = 2; diff --git a/proto/distribution_grpc.pb.go b/proto/distribution_grpc.pb.go index 66ebf8511..ec02eabee 100644 --- a/proto/distribution_grpc.pb.go +++ b/proto/distribution_grpc.pb.go @@ -19,17 +19,18 @@ import ( const _ = grpc.SupportPackageIsVersion9 const ( - Distribution_GetRoute_FullMethodName = "/Distribution/GetRoute" - Distribution_GetTimestamp_FullMethodName = "/Distribution/GetTimestamp" - Distribution_ListRoutes_FullMethodName = "/Distribution/ListRoutes" - Distribution_SplitRange_FullMethodName = "/Distribution/SplitRange" - Distribution_StartSplitMigration_FullMethodName = "/Distribution/StartSplitMigration" - Distribution_GetRouteOwnership_FullMethodName = "/Distribution/GetRouteOwnership" - Distribution_GetIntersectingRoutes_FullMethodName = "/Distribution/GetIntersectingRoutes" - Distribution_GetSplitJob_FullMethodName = "/Distribution/GetSplitJob" - Distribution_ListSplitJobs_FullMethodName = "/Distribution/ListSplitJobs" - Distribution_AbandonSplitJob_FullMethodName = "/Distribution/AbandonSplitJob" - Distribution_RetrySplitJob_FullMethodName = "/Distribution/RetrySplitJob" + Distribution_GetRoute_FullMethodName = "/Distribution/GetRoute" + Distribution_GetTimestamp_FullMethodName = "/Distribution/GetTimestamp" + Distribution_ListRoutes_FullMethodName = "/Distribution/ListRoutes" + Distribution_SplitRange_FullMethodName = "/Distribution/SplitRange" + Distribution_StartSplitMigration_FullMethodName = "/Distribution/StartSplitMigration" + Distribution_GetSplitMigrationCapability_FullMethodName = "/Distribution/GetSplitMigrationCapability" + Distribution_GetRouteOwnership_FullMethodName = "/Distribution/GetRouteOwnership" + Distribution_GetIntersectingRoutes_FullMethodName = "/Distribution/GetIntersectingRoutes" + Distribution_GetSplitJob_FullMethodName = "/Distribution/GetSplitJob" + Distribution_ListSplitJobs_FullMethodName = "/Distribution/ListSplitJobs" + Distribution_AbandonSplitJob_FullMethodName = "/Distribution/AbandonSplitJob" + Distribution_RetrySplitJob_FullMethodName = "/Distribution/RetrySplitJob" ) // DistributionClient is the client API for Distribution service. @@ -41,6 +42,7 @@ type DistributionClient interface { ListRoutes(ctx context.Context, in *ListRoutesRequest, opts ...grpc.CallOption) (*ListRoutesResponse, error) SplitRange(ctx context.Context, in *SplitRangeRequest, opts ...grpc.CallOption) (*SplitRangeResponse, error) StartSplitMigration(ctx context.Context, in *StartSplitMigrationRequest, opts ...grpc.CallOption) (*StartSplitMigrationResponse, error) + GetSplitMigrationCapability(ctx context.Context, in *GetSplitMigrationCapabilityRequest, opts ...grpc.CallOption) (*GetSplitMigrationCapabilityResponse, error) GetRouteOwnership(ctx context.Context, in *GetRouteOwnershipRequest, opts ...grpc.CallOption) (*GetRouteOwnershipResponse, error) GetIntersectingRoutes(ctx context.Context, in *GetIntersectingRoutesRequest, opts ...grpc.CallOption) (*GetIntersectingRoutesResponse, error) GetSplitJob(ctx context.Context, in *GetSplitJobRequest, opts ...grpc.CallOption) (*GetSplitJobResponse, error) @@ -107,6 +109,16 @@ func (c *distributionClient) StartSplitMigration(ctx context.Context, in *StartS return out, nil } +func (c *distributionClient) GetSplitMigrationCapability(ctx context.Context, in *GetSplitMigrationCapabilityRequest, opts ...grpc.CallOption) (*GetSplitMigrationCapabilityResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(GetSplitMigrationCapabilityResponse) + err := c.cc.Invoke(ctx, Distribution_GetSplitMigrationCapability_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + func (c *distributionClient) GetRouteOwnership(ctx context.Context, in *GetRouteOwnershipRequest, opts ...grpc.CallOption) (*GetRouteOwnershipResponse, error) { cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) out := new(GetRouteOwnershipResponse) @@ -176,6 +188,7 @@ type DistributionServer interface { ListRoutes(context.Context, *ListRoutesRequest) (*ListRoutesResponse, error) SplitRange(context.Context, *SplitRangeRequest) (*SplitRangeResponse, error) StartSplitMigration(context.Context, *StartSplitMigrationRequest) (*StartSplitMigrationResponse, error) + GetSplitMigrationCapability(context.Context, *GetSplitMigrationCapabilityRequest) (*GetSplitMigrationCapabilityResponse, error) GetRouteOwnership(context.Context, *GetRouteOwnershipRequest) (*GetRouteOwnershipResponse, error) GetIntersectingRoutes(context.Context, *GetIntersectingRoutesRequest) (*GetIntersectingRoutesResponse, error) GetSplitJob(context.Context, *GetSplitJobRequest) (*GetSplitJobResponse, error) @@ -207,6 +220,9 @@ func (UnimplementedDistributionServer) SplitRange(context.Context, *SplitRangeRe func (UnimplementedDistributionServer) StartSplitMigration(context.Context, *StartSplitMigrationRequest) (*StartSplitMigrationResponse, error) { return nil, status.Error(codes.Unimplemented, "method StartSplitMigration not implemented") } +func (UnimplementedDistributionServer) GetSplitMigrationCapability(context.Context, *GetSplitMigrationCapabilityRequest) (*GetSplitMigrationCapabilityResponse, error) { + return nil, status.Error(codes.Unimplemented, "method GetSplitMigrationCapability not implemented") +} func (UnimplementedDistributionServer) GetRouteOwnership(context.Context, *GetRouteOwnershipRequest) (*GetRouteOwnershipResponse, error) { return nil, status.Error(codes.Unimplemented, "method GetRouteOwnership not implemented") } @@ -336,6 +352,24 @@ func _Distribution_StartSplitMigration_Handler(srv interface{}, ctx context.Cont return interceptor(ctx, in, info, handler) } +func _Distribution_GetSplitMigrationCapability_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(GetSplitMigrationCapabilityRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(DistributionServer).GetSplitMigrationCapability(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: Distribution_GetSplitMigrationCapability_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(DistributionServer).GetSplitMigrationCapability(ctx, req.(*GetSplitMigrationCapabilityRequest)) + } + return interceptor(ctx, in, info, handler) +} + func _Distribution_GetRouteOwnership_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { in := new(GetRouteOwnershipRequest) if err := dec(in); err != nil { @@ -471,6 +505,10 @@ var Distribution_ServiceDesc = grpc.ServiceDesc{ MethodName: "StartSplitMigration", Handler: _Distribution_StartSplitMigration_Handler, }, + { + MethodName: "GetSplitMigrationCapability", + Handler: _Distribution_GetSplitMigrationCapability_Handler, + }, { MethodName: "GetRouteOwnership", Handler: _Distribution_GetRouteOwnership_Handler, From c64ff5fef248e4c6d94f69b3c1a00965dfac7001 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 11:38:13 +0900 Subject: [PATCH 22/29] migration: harden split capability gate --- adapter/distribution_server.go | 29 +++++++++-- adapter/distribution_server_test.go | 75 ++++++++++++++++++++++++++++- main.go | 36 ++++++++++++-- main_catalog_test.go | 56 ++++++++++++++------- 4 files changed, 169 insertions(+), 27 deletions(-) diff --git a/adapter/distribution_server.go b/adapter/distribution_server.go index 42360803a..40e76ba03 100644 --- a/adapter/distribution_server.go +++ b/adapter/distribution_server.go @@ -28,6 +28,7 @@ type DistributionServer struct { coordinator kv.Coordinator readTracker *kv.ActiveTimestampTracker migrationCapabilityGate SplitMigrationCapabilityGate + splitJobRunnerReady bool knownRaftGroups map[uint64]struct{} reloadRetry struct { attempts int @@ -63,6 +64,14 @@ func WithSplitMigrationCapabilityGate(gate SplitMigrationCapabilityGate) Distrib } } +// WithSplitJobRunnerReady opens the split migration capability probe once this +// node can advance planned split jobs. +func WithSplitJobRunnerReady() DistributionServerOption { + return func(s *DistributionServer) { + s.splitJobRunnerReady = true + } +} + // WithDistributionKnownRaftGroups configures the Raft group IDs this node can // route migration work to. func WithDistributionKnownRaftGroups(groupIDs ...uint64) DistributionServerOption { @@ -215,10 +224,14 @@ func (s *DistributionServer) GetIntersectingRoutes(ctx context.Context, req *pb. } func (s *DistributionServer) StartSplitMigration(ctx context.Context, req *pb.StartSplitMigrationRequest) (*pb.StartSplitMigrationResponse, error) { + if err := s.verifyStartSplitMigrationPreflight(ctx, req); err != nil { + return nil, err + } + s.mu.Lock() defer s.mu.Unlock() - if err := s.verifyStartSplitMigrationReady(ctx, req); err != nil { + if err := s.verifyStartSplitMigrationCatalogReady(ctx); err != nil { return nil, err } snapshot, err := s.loadCatalogSnapshot(ctx) @@ -246,23 +259,33 @@ func (s *DistributionServer) StartSplitMigration(ctx context.Context, req *pb.St } func (s *DistributionServer) GetSplitMigrationCapability(context.Context, *pb.GetSplitMigrationCapabilityRequest) (*pb.GetSplitMigrationCapabilityResponse, error) { + if !s.splitJobRunnerReady { + return &pb.GetSplitMigrationCapabilityResponse{}, nil + } return &pb.GetSplitMigrationCapabilityResponse{ MigrationCapable: true, Capabilities: []string{splitMigrationCapabilityV2}, }, nil } -func (s *DistributionServer) verifyStartSplitMigrationReady(ctx context.Context, req *pb.StartSplitMigrationRequest) error { +func (s *DistributionServer) verifyStartSplitMigrationPreflight(ctx context.Context, req *pb.StartSplitMigrationRequest) error { if req.GetTargetGroupId() == 0 { return grpcStatusError(codes.InvalidArgument, distribution.ErrCatalogSplitJobTargetGroupRequired.Error()) } + if err := s.verifyStartSplitMigrationCatalogReady(ctx); err != nil { + return err + } + return s.verifySplitMigrationCapability(ctx) +} + +func (s *DistributionServer) verifyStartSplitMigrationCatalogReady(ctx context.Context) error { if err := s.verifyCatalogLeader(ctx); err != nil { return err } if s.catalog == nil { return grpcStatusError(codes.FailedPrecondition, errDistributionCatalogNotConfigured.Error()) } - return s.verifySplitMigrationCapability(ctx) + return nil } func (s *DistributionServer) startSplitMigrationParent( diff --git a/adapter/distribution_server_test.go b/adapter/distribution_server_test.go index 21a34d96b..5190498e6 100644 --- a/adapter/distribution_server_test.go +++ b/adapter/distribution_server_test.go @@ -149,12 +149,22 @@ func TestDistributionServerListRoutes_RequiresCatalog(t *testing.T) { require.ErrorContains(t, err, errDistributionCatalogNotConfigured.Error()) } -func TestDistributionServerGetSplitMigrationCapabilityReportsReady(t *testing.T) { +func TestDistributionServerGetSplitMigrationCapabilityReportsNotReadyUntilRunnerReady(t *testing.T) { t.Parallel() s := NewDistributionServer(distribution.NewEngine(), nil) resp, err := s.GetSplitMigrationCapability(context.Background(), &pb.GetSplitMigrationCapabilityRequest{}) require.NoError(t, err) + require.False(t, resp.GetMigrationCapable()) + require.NotContains(t, resp.GetCapabilities(), splitMigrationCapabilityV2) +} + +func TestDistributionServerGetSplitMigrationCapabilityReportsReadyWhenRunnerReady(t *testing.T) { + t.Parallel() + + s := NewDistributionServer(distribution.NewEngine(), nil, WithSplitJobRunnerReady()) + resp, err := s.GetSplitMigrationCapability(context.Background(), &pb.GetSplitMigrationCapabilityRequest{}) + require.NoError(t, err) require.True(t, resp.GetMigrationCapable()) require.Contains(t, resp.GetCapabilities(), splitMigrationCapabilityV2) } @@ -212,6 +222,69 @@ func TestDistributionServerStartSplitMigrationReturnsCapabilityGateError(t *test require.Zero(t, coordinator.dispatchCalls) } +func TestDistributionServerStartSplitMigrationCapabilityGateRunsOutsideCatalogLock(t *testing.T) { + ctx := context.Background() + baseStore := store.NewMVCCStore() + catalog := distribution.NewCatalogStore(baseStore) + saved, err := catalog.Save(ctx, 0, []distribution.RouteDescriptor{{ + RouteID: 1, + Start: []byte("a"), + End: []byte("m"), + GroupID: 1, + State: distribution.RouteStateActive, + ParentRouteID: 0, + }}) + require.NoError(t, err) + + coordinator := newDistributionCoordinatorStub(baseStore, true) + gateEntered := make(chan struct{}) + releaseGate := make(chan struct{}) + s := NewDistributionServer( + distribution.NewEngine(), + catalog, + WithDistributionCoordinator(coordinator), + WithDistributionKnownRaftGroups(1, 2), + WithSplitMigrationCapabilityGate(func(context.Context) error { + close(gateEntered) + <-releaseGate + return status.Error(codes.Unavailable, "split migration capability not ready") + }), + ) + + startDone := make(chan error, 1) + go func() { + _, err := s.StartSplitMigration(ctx, &pb.StartSplitMigrationRequest{ + ExpectedCatalogVersion: saved.Version, + RouteId: 1, + SplitKey: []byte("g"), + TargetGroupId: 2, + }) + startDone <- err + }() + <-gateEntered + + splitDone := make(chan error, 1) + go func() { + _, err := s.SplitRange(ctx, &pb.SplitRangeRequest{ + ExpectedCatalogVersion: saved.Version, + RouteId: 1, + SplitKey: []byte("g"), + }) + splitDone <- err + }() + select { + case err := <-splitDone: + require.NoError(t, err) + case <-time.After(500 * time.Millisecond): + t.Fatal("SplitRange blocked behind StartSplitMigration capability gate") + } + + close(releaseGate) + err = <-startDone + require.Error(t, err) + require.Equal(t, codes.Unavailable, status.Code(err)) +} + func TestDistributionServerStartSplitMigration_CreatesPlannedJobWhenGateOpen(t *testing.T) { t.Parallel() diff --git a/main.go b/main.go index 93fed90ce..628ced4e9 100644 --- a/main.go +++ b/main.go @@ -500,8 +500,9 @@ func run() error { startKeyVizFlusher(runCtx, eg, sampler) startKeyVizLeaderTermPublisher(runCtx, eg, sampler, runtimes) startMemoryWatchdog(runCtx, eg, cancel) + defaultRuntime := findDefaultGroupRuntime(runtimes, cfg.defaultGroup) splitMigrationGate := newSplitMigrationCapabilityGate( - splitMigrationCapabilityPeers(bootstrapCfg, cfg.defaultGroup), + splitMigrationCapabilityPeerSourceForRuntime(defaultRuntime), splitMigrationCapabilityProbeTimeout, nil, ) @@ -529,7 +530,6 @@ func run() error { // group), in which case raftadmin.Server skips the pre-step. encryptionConfChangeInterceptor := newEncryptionPreRegister( coordinate, shardGroups[cfg.defaultGroup], encWiring.cache, *encryptionSidecarPath, etcdraftengine.DeriveNodeID) - defaultRuntime := findDefaultGroupRuntime(runtimes, cfg.defaultGroup) rotateOnStartupDeregister, waitRotateOnStartup := installEncryptionRotateOnStartup( runCtx, *encryptionRotateOnStartup, @@ -678,16 +678,35 @@ type splitMigrationCapabilityPeer struct { Address string } +type splitMigrationCapabilityPeerSource func(context.Context) ([]splitMigrationCapabilityPeer, error) + type splitMigrationCapabilityProbe func(context.Context, string) error -func splitMigrationCapabilityPeers(cfg raftBootstrapConfig, defaultGroup uint64) []splitMigrationCapabilityPeer { - servers := cfg.adminSeed(defaultGroup) +func splitMigrationCapabilityPeerSourceForRuntime(rt *raftGroupRuntime) splitMigrationCapabilityPeerSource { + return func(ctx context.Context) ([]splitMigrationCapabilityPeer, error) { + engine := rt.snapshotEngine() + if engine == nil { + return nil, errors.New("default raft group engine is not configured") + } + cfg, err := engine.Configuration(ctx) + if err != nil { + return nil, errors.Wrap(err, "default raft group configuration") + } + return splitMigrationCapabilityPeersFromConfiguration(cfg), nil + } +} + +func splitMigrationCapabilityPeersFromConfiguration(cfg raftengine.Configuration) []splitMigrationCapabilityPeer { + servers := cfg.Servers if len(servers) == 0 { return nil } peers := make([]splitMigrationCapabilityPeer, 0, len(servers)) seen := make(map[string]struct{}, len(servers)) for _, server := range servers { + if server.Suffrage == etcdraftengine.SuffrageLearner { + continue + } key := server.ID + "\x00" + server.Address if _, ok := seen[key]; ok { continue @@ -705,11 +724,18 @@ func splitMigrationCapabilityPeers(cfg raftBootstrapConfig, defaultGroup uint64) return peers } -func newSplitMigrationCapabilityGate(peers []splitMigrationCapabilityPeer, timeout time.Duration, probe splitMigrationCapabilityProbe) adapter.SplitMigrationCapabilityGate { +func newSplitMigrationCapabilityGate(source splitMigrationCapabilityPeerSource, timeout time.Duration, probe splitMigrationCapabilityProbe) adapter.SplitMigrationCapabilityGate { if probe == nil { probe = probeSplitMigrationCapabilityPeer } return func(ctx context.Context) error { + if source == nil { + return status.Error(codes.FailedPrecondition, "split migration capability peers are not configured") + } + peers, err := source(ctx) + if err != nil { + return status.Errorf(codes.FailedPrecondition, "split migration capability peers are not available: %v", err) + } if len(peers) == 0 { return status.Error(codes.FailedPrecondition, "split migration capability peers are not configured") } diff --git a/main_catalog_test.go b/main_catalog_test.go index d2e63118b..968a0fefb 100644 --- a/main_catalog_test.go +++ b/main_catalog_test.go @@ -2,6 +2,7 @@ package main import ( "context" + "errors" "testing" "time" @@ -81,7 +82,7 @@ func TestSplitMigrationCapabilityGateChecksAllPeers(t *testing.T) { {ID: "n2", Address: "10.0.0.12:50051"}, } var probed []string - gate := newSplitMigrationCapabilityGate(peers, time.Second, func(_ context.Context, address string) error { + gate := newSplitMigrationCapabilityGate(staticSplitMigrationCapabilityPeerSource(peers), time.Second, func(_ context.Context, address string) error { probed = append(probed, address) return nil }) @@ -111,7 +112,7 @@ func TestSplitMigrationCapabilityGateFailsClosedWhenPeerMissingCapability(t *tes {ID: "n1", Address: "10.0.0.11:50051"}, {ID: "n2", Address: "10.0.0.12:50051"}, } - gate := newSplitMigrationCapabilityGate(peers, time.Second, func(_ context.Context, address string) error { + gate := newSplitMigrationCapabilityGate(staticSplitMigrationCapabilityPeerSource(peers), time.Second, func(_ context.Context, address string) error { if address == "10.0.0.12:50051" { return status.Error(codes.Unimplemented, "method not found") } @@ -125,24 +126,43 @@ func TestSplitMigrationCapabilityGateFailsClosedWhenPeerMissingCapability(t *tes require.ErrorContains(t, err, "method not found") } -func TestSplitMigrationCapabilityPeersUseDefaultGroupAdminSeed(t *testing.T) { +func TestSplitMigrationCapabilityGateFailsClosedWhenPeerSourceErrors(t *testing.T) { t.Parallel() - cfg := raftBootstrapConfig{ - groupServers: map[uint64][]raftengine.Server{ - 1: { - {Suffrage: "voter", ID: "n1", Address: "10.0.0.11:50051"}, - {Suffrage: "voter", ID: "n2", Address: "10.0.0.12:50051"}, - }, - 2: { - {Suffrage: "voter", ID: "n1", Address: "10.0.0.11:50052"}, - {Suffrage: "voter", ID: "n2", Address: "10.0.0.12:50052"}, - }, - }, - } + errPeerSourceUnavailable := errors.New("configuration unavailable") + gate := newSplitMigrationCapabilityGate(func(context.Context) ([]splitMigrationCapabilityPeer, error) { + return nil, errPeerSourceUnavailable + }, time.Second, func(context.Context, string) error { + t.Fatal("probe must not run when peer source fails") + return nil + }) + + err := gate(context.Background()) + require.Error(t, err) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) + require.ErrorContains(t, err, "peers are not available") +} + +func TestSplitMigrationCapabilityPeersFromConfigurationUsesCurrentVoters(t *testing.T) { + t.Parallel() + + cfg := raftengine.Configuration{Servers: []raftengine.Server{ + {Suffrage: "voter", ID: "n1", Address: "10.0.0.11:50051"}, + {Suffrage: "learner", ID: "n2", Address: "10.0.0.12:50051"}, + {Suffrage: "", ID: "n3", Address: "10.0.0.13:50051"}, + {Suffrage: "voter", ID: "n1", Address: "10.0.0.11:50051"}, + }} require.Equal(t, []splitMigrationCapabilityPeer{ - {ID: "n1", Address: "10.0.0.11:50052"}, - {ID: "n2", Address: "10.0.0.12:50052"}, - }, splitMigrationCapabilityPeers(cfg, 2)) + {ID: "n1", Address: "10.0.0.11:50051"}, + {ID: "n3", Address: "10.0.0.13:50051"}, + }, splitMigrationCapabilityPeersFromConfiguration(cfg)) +} + +func staticSplitMigrationCapabilityPeerSource(peers []splitMigrationCapabilityPeer) splitMigrationCapabilityPeerSource { + return func(context.Context) ([]splitMigrationCapabilityPeer, error) { + out := make([]splitMigrationCapabilityPeer, len(peers)) + copy(out, peers) + return out, nil + } } From a775fdaaa41f8073a77a725e17ddaababa01a320 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 11:53:08 +0900 Subject: [PATCH 23/29] Trim Lua negative cache bound test --- adapter/redis_lua_context.go | 26 ++++++++------ adapter/redis_lua_negative_type_cache_test.go | 35 ++++++------------- 2 files changed, 27 insertions(+), 34 deletions(-) diff --git a/adapter/redis_lua_context.go b/adapter/redis_lua_context.go index 4a57051b5..447423ee4 100644 --- a/adapter/redis_lua_context.go +++ b/adapter/redis_lua_context.go @@ -478,20 +478,26 @@ func (c *luaScriptContext) keyType(key []byte) (redisValueType, error) { if err != nil { return redisTypeNone, err } - if typ == redisTypeNone && len(c.negativeType) < maxNegativeTypeCacheEntries { - // Pin the absence result for the rest of this Eval so repeated - // BullMQ-style polling of a missing key (e.g. a "delayed" zset) - // does not re-run the ~8-seek rawKeyTypeAt probe on every - // redis.call. - // - // Bounded to keep adversarial scripts from growing the map - // unboundedly; once full, subsequent misses correctly fall - // through to the server probe without caching. - c.negativeType[string(key)] = true + if typ == redisTypeNone { + c.rememberNegativeType(key) } return typ, nil } +func (c *luaScriptContext) rememberNegativeType(key []byte) { + if len(c.negativeType) >= maxNegativeTypeCacheEntries { + return + } + // Pin the absence result for the rest of this Eval so repeated + // BullMQ-style polling of a missing key (e.g. a "delayed" zset) does + // not re-run the ~8-seek rawKeyTypeAt probe on every redis.call. + // + // Bounded to keep adversarial scripts from growing the map + // unboundedly; once full, subsequent misses correctly fall through to + // the server probe without caching. + c.negativeType[string(key)] = true +} + func (c *luaScriptContext) ensureKeyNotExpired(key []byte) error { ttl, err := c.loadTTL(key) if err != nil { diff --git a/adapter/redis_lua_negative_type_cache_test.go b/adapter/redis_lua_negative_type_cache_test.go index 6253f7e30..f44aec35a 100644 --- a/adapter/redis_lua_negative_type_cache_test.go +++ b/adapter/redis_lua_negative_type_cache_test.go @@ -149,44 +149,31 @@ func TestLuaNegativeTypeCache_SingleProbePerKey(t *testing.T) { // itself stays bounded. func TestLuaNegativeTypeCache_BoundedSize(t *testing.T) { t.Parallel() - nodes, _, _ := createNode(t, 3) - defer shutdown(nodes) - ctx := context.Background() - sc, err := newLuaScriptContext(ctx, nodes[0].redisServer) - require.NoError(t, err) - defer sc.Close() + sc := &luaScriptContext{negativeType: map[string]bool{}} // Probe cap+overflow unique missing keys. Each probe is a miss // (redisTypeNone); only the first `cap` should be memoized. const overflow = 50 total := maxNegativeTypeCacheEntries + overflow for i := 0; i < total; i++ { - typ, kerr := sc.keyType([]byte(fmt.Sprintf("lua:neg:cap:%d", i))) - require.NoError(t, kerr) - require.Equal(t, redisTypeNone, typ) + sc.rememberNegativeType([]byte(fmt.Sprintf("lua:neg:cap:%d", i))) } require.Equal(t, maxNegativeTypeCacheEntries, len(sc.negativeType), "negativeType map must be capped at maxNegativeTypeCacheEntries") - // Each unique key above required exactly one probe on first access. - require.Equal(t, total, sc.keyTypeProbeCount, - "each unique key must have triggered exactly one server probe") - - // Re-probing one of the first `cap` keys must hit the cache - // (no additional server probe). Re-probing an overflow key must - // miss the cache and issue another server probe. + // One of the first `cap` keys must be cached. An overflow key must + // remain uncached so keyType falls back to a server probe in real Eval + // execution while the map size stays bounded. cachedKey := []byte("lua:neg:cap:0") - _, kerr := sc.keyType(cachedKey) - require.NoError(t, kerr) - require.Equal(t, total, sc.keyTypeProbeCount, - "a key inserted before the cap must remain cached") + typ, ok := sc.cachedType(cachedKey) + require.True(t, ok, "a key inserted before the cap must remain cached") + require.Equal(t, redisTypeNone, typ) overflowKey := []byte(fmt.Sprintf("lua:neg:cap:%d", maxNegativeTypeCacheEntries+1)) - _, kerr = sc.keyType(overflowKey) - require.NoError(t, kerr) - require.Equal(t, total+1, sc.keyTypeProbeCount, - "a key probed after the cap was reached must fall back to the server probe") + _, ok = sc.cachedType(overflowKey) + require.False(t, ok, "a key probed after the cap must fall back to the server probe") + sc.rememberNegativeType(overflowKey) require.Equal(t, maxNegativeTypeCacheEntries, len(sc.negativeType), "fallback probe must NOT grow the bounded cache") } From c29806025400c0896df536c8ab1eca6b2038284c Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 12:38:44 +0900 Subject: [PATCH 24/29] Enable split migration capability gate --- adapter/distribution_server.go | 2 +- main.go | 42 ++++++++++++++------ main_catalog_test.go | 72 +++++++++++++++++++++++++++++++++- 3 files changed, 101 insertions(+), 15 deletions(-) diff --git a/adapter/distribution_server.go b/adapter/distribution_server.go index 40e76ba03..481159753 100644 --- a/adapter/distribution_server.go +++ b/adapter/distribution_server.go @@ -108,7 +108,7 @@ const ( splitJobListCursorTerminalOff = 1 splitJobListCursorJobIDOff = splitJobListCursorTerminalOff + 8 splitJobListCursorEncodedBytes = splitJobListCursorJobIDOff + 8 - splitMigrationCapabilityV2 = "split_migration_v2" + splitMigrationCapabilityV2 = "cap_migration_v2" ) var ( diff --git a/main.go b/main.go index 628ced4e9..47c44af55 100644 --- a/main.go +++ b/main.go @@ -500,9 +500,8 @@ func run() error { startKeyVizFlusher(runCtx, eg, sampler) startKeyVizLeaderTermPublisher(runCtx, eg, sampler, runtimes) startMemoryWatchdog(runCtx, eg, cancel) - defaultRuntime := findDefaultGroupRuntime(runtimes, cfg.defaultGroup) splitMigrationGate := newSplitMigrationCapabilityGate( - splitMigrationCapabilityPeerSourceForRuntime(defaultRuntime), + splitMigrationCapabilityPeerSourceForRuntimes(runtimes), splitMigrationCapabilityProbeTimeout, nil, ) @@ -513,7 +512,9 @@ func run() error { adapter.WithDistributionActiveTimestampTracker(readTracker), adapter.WithDistributionKnownRaftGroups(shardGroupIDs(shardGroups)...), adapter.WithSplitMigrationCapabilityGate(splitMigrationGate), + adapter.WithSplitJobRunnerReady(), ) + defaultRuntime := findDefaultGroupRuntime(runtimes, cfg.defaultGroup) startMonitoringCollectors(runCtx, metricsRegistry, runtimes, clock) compactor := kv.NewFSMCompactor( fsmCompactionRuntimes(runtimes), @@ -682,17 +683,35 @@ type splitMigrationCapabilityPeerSource func(context.Context) ([]splitMigrationC type splitMigrationCapabilityProbe func(context.Context, string) error -func splitMigrationCapabilityPeerSourceForRuntime(rt *raftGroupRuntime) splitMigrationCapabilityPeerSource { +func splitMigrationCapabilityPeerSourceForRuntimes(runtimes []*raftGroupRuntime) splitMigrationCapabilityPeerSource { return func(ctx context.Context) ([]splitMigrationCapabilityPeer, error) { - engine := rt.snapshotEngine() - if engine == nil { - return nil, errors.New("default raft group engine is not configured") + if len(runtimes) == 0 { + return nil, errors.New("raft group runtimes are not configured") } - cfg, err := engine.Configuration(ctx) - if err != nil { - return nil, errors.Wrap(err, "default raft group configuration") + peers := make([]splitMigrationCapabilityPeer, 0) + seen := make(map[string]struct{}) + for _, rt := range runtimes { + if rt == nil { + return nil, errors.New("raft group runtime is not configured") + } + engine := rt.snapshotEngine() + if engine == nil { + return nil, errors.Errorf("raft group %d engine is not configured", rt.spec.id) + } + cfg, err := engine.Configuration(ctx) + if err != nil { + return nil, errors.Wrapf(err, "raft group %d configuration", rt.spec.id) + } + for _, peer := range splitMigrationCapabilityPeersFromConfiguration(cfg) { + key := peer.ID + "\x00" + peer.Address + if _, ok := seen[key]; ok { + continue + } + seen[key] = struct{}{} + peers = append(peers, peer) + } } - return splitMigrationCapabilityPeersFromConfiguration(cfg), nil + return peers, nil } } @@ -704,9 +723,6 @@ func splitMigrationCapabilityPeersFromConfiguration(cfg raftengine.Configuration peers := make([]splitMigrationCapabilityPeer, 0, len(servers)) seen := make(map[string]struct{}, len(servers)) for _, server := range servers { - if server.Suffrage == etcdraftengine.SuffrageLearner { - continue - } key := server.ID + "\x00" + server.Address if _, ok := seen[key]; ok { continue diff --git a/main_catalog_test.go b/main_catalog_test.go index 968a0fefb..497213d36 100644 --- a/main_catalog_test.go +++ b/main_catalog_test.go @@ -143,7 +143,7 @@ func TestSplitMigrationCapabilityGateFailsClosedWhenPeerSourceErrors(t *testing. require.ErrorContains(t, err, "peers are not available") } -func TestSplitMigrationCapabilityPeersFromConfigurationUsesCurrentVoters(t *testing.T) { +func TestSplitMigrationCapabilityPeersFromConfigurationUsesCurrentMembers(t *testing.T) { t.Parallel() cfg := raftengine.Configuration{Servers: []raftengine.Server{ @@ -155,10 +155,40 @@ func TestSplitMigrationCapabilityPeersFromConfigurationUsesCurrentVoters(t *test require.Equal(t, []splitMigrationCapabilityPeer{ {ID: "n1", Address: "10.0.0.11:50051"}, + {ID: "n2", Address: "10.0.0.12:50051"}, {ID: "n3", Address: "10.0.0.13:50051"}, }, splitMigrationCapabilityPeersFromConfiguration(cfg)) } +func TestSplitMigrationCapabilityPeerSourceForRuntimesChecksAllGroups(t *testing.T) { + t.Parallel() + + runtimes := []*raftGroupRuntime{ + { + spec: groupSpec{id: 1}, + engine: capabilityConfigEngine{cfg: raftengine.Configuration{Servers: []raftengine.Server{ + {Suffrage: "voter", ID: "n1", Address: "10.0.0.11:50051"}, + {Suffrage: "learner", ID: "n2", Address: "10.0.0.12:50051"}, + }}}, + }, + { + spec: groupSpec{id: 2}, + engine: capabilityConfigEngine{cfg: raftengine.Configuration{Servers: []raftengine.Server{ + {Suffrage: "voter", ID: "n1", Address: "10.0.0.11:50051"}, + {Suffrage: "voter", ID: "n3", Address: "10.0.0.13:50051"}, + }}}, + }, + } + + peers, err := splitMigrationCapabilityPeerSourceForRuntimes(runtimes)(context.Background()) + require.NoError(t, err) + require.Equal(t, []splitMigrationCapabilityPeer{ + {ID: "n1", Address: "10.0.0.11:50051"}, + {ID: "n2", Address: "10.0.0.12:50051"}, + {ID: "n3", Address: "10.0.0.13:50051"}, + }, peers) +} + func staticSplitMigrationCapabilityPeerSource(peers []splitMigrationCapabilityPeer) splitMigrationCapabilityPeerSource { return func(context.Context) ([]splitMigrationCapabilityPeer, error) { out := make([]splitMigrationCapabilityPeer, len(peers)) @@ -166,3 +196,43 @@ func staticSplitMigrationCapabilityPeerSource(peers []splitMigrationCapabilityPe return out, nil } } + +type capabilityConfigEngine struct { + cfg raftengine.Configuration +} + +func (e capabilityConfigEngine) Propose(context.Context, []byte) (*raftengine.ProposalResult, error) { + return nil, nil +} + +func (e capabilityConfigEngine) ProposeAdmin(context.Context, []byte) (*raftengine.ProposalResult, error) { + return nil, nil +} + +func (e capabilityConfigEngine) State() raftengine.State { + return raftengine.StateFollower +} + +func (e capabilityConfigEngine) Leader() raftengine.LeaderInfo { + return raftengine.LeaderInfo{} +} + +func (e capabilityConfigEngine) VerifyLeader(context.Context) error { + return nil +} + +func (e capabilityConfigEngine) LinearizableRead(context.Context) (uint64, error) { + return 0, nil +} + +func (e capabilityConfigEngine) Status() raftengine.Status { + return raftengine.Status{} +} + +func (e capabilityConfigEngine) Configuration(context.Context) (raftengine.Configuration, error) { + return e.cfg, nil +} + +func (e capabilityConfigEngine) Close() error { + return nil +} From 70fb630f55ee3e4d3a77ae855c4ab21ec8efa3fc Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 14:05:24 +0900 Subject: [PATCH 25/29] Keep split migration capability fail closed --- adapter/distribution_server.go | 16 ++++--- main.go | 10 ++++- main_catalog_test.go | 80 ++++++++++++++++++++++++++++++++++ 3 files changed, 99 insertions(+), 7 deletions(-) diff --git a/adapter/distribution_server.go b/adapter/distribution_server.go index 481159753..e2683dcbc 100644 --- a/adapter/distribution_server.go +++ b/adapter/distribution_server.go @@ -108,7 +108,8 @@ const ( splitJobListCursorTerminalOff = 1 splitJobListCursorJobIDOff = splitJobListCursorTerminalOff + 8 splitJobListCursorEncodedBytes = splitJobListCursorJobIDOff + 8 - splitMigrationCapabilityV2 = "cap_migration_v2" + SplitMigrationCapabilityV2 = "cap_migration_v2" + splitMigrationCapabilityV2 = SplitMigrationCapabilityV2 ) var ( @@ -238,8 +239,7 @@ func (s *DistributionServer) StartSplitMigration(ctx context.Context, req *pb.St if err != nil { return nil, err } - readPin := s.pinReadTS(snapshot.ReadTS) - defer readPin.Release() + defer releaseReadPin(s.pinReadTS(snapshot.ReadTS)) parent, err := s.startSplitMigrationParent(ctx, snapshot, req) if err != nil { @@ -445,8 +445,7 @@ func (s *DistributionServer) SplitRange(ctx context.Context, req *pb.SplitRangeR if err != nil { return nil, err } - readPin := s.pinReadTS(snapshot.ReadTS) - defer readPin.Release() + defer releaseReadPin(s.pinReadTS(snapshot.ReadTS)) if err := validateExpectedCatalogVersion(snapshot.Version, req.GetExpectedCatalogVersion()); err != nil { return nil, err } @@ -492,6 +491,13 @@ func (s *DistributionServer) pinReadTS(ts uint64) *kv.ActiveTimestampToken { return s.readTracker.Pin(ts) } +func releaseReadPin(token *kv.ActiveTimestampToken) { + if token == nil { + return + } + token.Release() +} + func (s *DistributionServer) verifyCatalogLeader(ctx context.Context) error { if s.coordinator == nil { return grpcStatusError(codes.FailedPrecondition, errDistributionCoordinatorRequired.Error()) diff --git a/main.go b/main.go index 47c44af55..7d2111e8a 100644 --- a/main.go +++ b/main.go @@ -512,7 +512,6 @@ func run() error { adapter.WithDistributionActiveTimestampTracker(readTracker), adapter.WithDistributionKnownRaftGroups(shardGroupIDs(shardGroups)...), adapter.WithSplitMigrationCapabilityGate(splitMigrationGate), - adapter.WithSplitJobRunnerReady(), ) defaultRuntime := findDefaultGroupRuntime(runtimes, cfg.defaultGroup) startMonitoringCollectors(runCtx, metricsRegistry, runtimes, clock) @@ -702,7 +701,11 @@ func splitMigrationCapabilityPeerSourceForRuntimes(runtimes []*raftGroupRuntime) if err != nil { return nil, errors.Wrapf(err, "raft group %d configuration", rt.spec.id) } - for _, peer := range splitMigrationCapabilityPeersFromConfiguration(cfg) { + groupPeers := splitMigrationCapabilityPeersFromConfiguration(cfg) + if len(groupPeers) == 0 { + return nil, errors.Errorf("raft group %d configuration has no split migration capability peers", rt.spec.id) + } + for _, peer := range groupPeers { key := peer.ID + "\x00" + peer.Address if _, ok := seen[key]; ok { continue @@ -791,6 +794,9 @@ func probeSplitMigrationCapabilityPeer(ctx context.Context, address string) erro if resp == nil || !resp.GetMigrationCapable() { return errors.New("peer does not advertise split migration capability") } + if !slices.Contains(resp.GetCapabilities(), adapter.SplitMigrationCapabilityV2) { + return errors.Errorf("peer does not advertise %s", adapter.SplitMigrationCapabilityV2) + } return nil } diff --git a/main_catalog_test.go b/main_catalog_test.go index 497213d36..b4227a9ee 100644 --- a/main_catalog_test.go +++ b/main_catalog_test.go @@ -3,13 +3,17 @@ package main import ( "context" "errors" + "net" "testing" "time" + "github.com/bootjp/elastickv/adapter" "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/internal/raftengine" + pb "github.com/bootjp/elastickv/proto" "github.com/bootjp/elastickv/store" "github.com/stretchr/testify/require" + "google.golang.org/grpc" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" ) @@ -189,6 +193,53 @@ func TestSplitMigrationCapabilityPeerSourceForRuntimesChecksAllGroups(t *testing }, peers) } +func TestSplitMigrationCapabilityPeerSourceForRuntimesFailsClosedOnEmptyGroup(t *testing.T) { + t.Parallel() + + runtimes := []*raftGroupRuntime{ + { + spec: groupSpec{id: 1}, + engine: capabilityConfigEngine{cfg: raftengine.Configuration{Servers: []raftengine.Server{ + {Suffrage: "voter", ID: "n1", Address: "10.0.0.11:50051"}, + }}}, + }, + { + spec: groupSpec{id: 2}, + engine: capabilityConfigEngine{cfg: raftengine.Configuration{}}, + }, + } + + peers, err := splitMigrationCapabilityPeerSourceForRuntimes(runtimes)(context.Background()) + require.Error(t, err) + require.Nil(t, peers) + require.ErrorContains(t, err, "raft group 2") + require.ErrorContains(t, err, "no split migration capability peers") +} + +func TestProbeSplitMigrationCapabilityPeerRequiresDocumentedToken(t *testing.T) { + t.Parallel() + + addr := startSplitMigrationCapabilityTestServer(t, &pb.GetSplitMigrationCapabilityResponse{ + MigrationCapable: true, + Capabilities: []string{"split_migration_v2"}, + }) + + err := probeSplitMigrationCapabilityPeer(context.Background(), addr) + require.Error(t, err) + require.ErrorContains(t, err, adapter.SplitMigrationCapabilityV2) +} + +func TestProbeSplitMigrationCapabilityPeerAcceptsDocumentedToken(t *testing.T) { + t.Parallel() + + addr := startSplitMigrationCapabilityTestServer(t, &pb.GetSplitMigrationCapabilityResponse{ + MigrationCapable: true, + Capabilities: []string{adapter.SplitMigrationCapabilityV2}, + }) + + require.NoError(t, probeSplitMigrationCapabilityPeer(context.Background(), addr)) +} + func staticSplitMigrationCapabilityPeerSource(peers []splitMigrationCapabilityPeer) splitMigrationCapabilityPeerSource { return func(context.Context) ([]splitMigrationCapabilityPeer, error) { out := make([]splitMigrationCapabilityPeer, len(peers)) @@ -197,6 +248,35 @@ func staticSplitMigrationCapabilityPeerSource(peers []splitMigrationCapabilityPe } } +type splitMigrationCapabilityTestServer struct { + pb.UnimplementedDistributionServer + resp *pb.GetSplitMigrationCapabilityResponse +} + +func (s *splitMigrationCapabilityTestServer) GetSplitMigrationCapability(context.Context, *pb.GetSplitMigrationCapabilityRequest) (*pb.GetSplitMigrationCapabilityResponse, error) { + return s.resp, nil +} + +func startSplitMigrationCapabilityTestServer(t *testing.T, resp *pb.GetSplitMigrationCapabilityResponse) string { + t.Helper() + + var listenConfig net.ListenConfig + listener, err := listenConfig.Listen(context.Background(), "tcp", "127.0.0.1:0") + require.NoError(t, err) + grpcServer := grpc.NewServer() + pb.RegisterDistributionServer(grpcServer, &splitMigrationCapabilityTestServer{resp: resp}) + done := make(chan struct{}) + go func() { + defer close(done) + _ = grpcServer.Serve(listener) + }() + t.Cleanup(func() { + grpcServer.Stop() + <-done + }) + return listener.Addr().String() +} + type capabilityConfigEngine struct { cfg raftengine.Configuration } From 9f80fa9d2eae49581353890b149b4ebb6a9ca019 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 16:47:42 +0900 Subject: [PATCH 26/29] migration: tighten readiness guard publication --- kv/sharded_coordinator.go | 2 +- kv/sharded_coordinator_txn_test.go | 10 +++++++--- store/migration_readiness.go | 4 ++-- 3 files changed, 10 insertions(+), 6 deletions(-) diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index a31e02e6d..04083d352 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -1317,7 +1317,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 { diff --git a/kv/sharded_coordinator_txn_test.go b/kv/sharded_coordinator_txn_test.go index 964f2713f..3316fc847 100644 --- a/kv/sharded_coordinator_txn_test.go +++ b/kv/sharded_coordinator_txn_test.go @@ -144,8 +144,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) @@ -164,7 +164,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) } @@ -244,6 +244,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 TestCommitSecondaryWithRetry_RetriesAndSucceeds(t *testing.T) { diff --git a/store/migration_readiness.go b/store/migration_readiness.go index e0c83a1df..d26d25d36 100644 --- a/store/migration_readiness.go +++ b/store/migration_readiness.go @@ -52,11 +52,11 @@ func (s *pebbleStore) ApplyTargetStagedReadiness(_ context.Context, state Target } 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.mtx.Lock() - defer s.mtx.Unlock() s.migrationReadinessCache = upsertTargetStagedReadinessStateCache(s.migrationReadinessCache, state) return nil } From 14aa9e9631ee4b51bc9a88c6f8d2642dd3b65f15 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 18:16:57 +0900 Subject: [PATCH 27/29] migration: guard read-only shard validation --- kv/sharded_coordinator.go | 5 +++ kv/sharded_coordinator_txn_test.go | 57 ++++++++++++++++++++++++++++++ 2 files changed, 62 insertions(+) diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index 04083d352..bcb1d957e 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -1979,7 +1979,12 @@ func (c *ShardedCoordinator) validateReadKeysOnShard(ctx context.Context, gid ui if _, err := linearizableReadEngineCtx(ctx, engineForGroup(g)); err != nil { return errors.WithStack(err) } + readiness := &ShardStore{engine: c.engine, groups: c.groups} for _, key := range keys { + route := readiness.explicitGroupRouteForKey(gid, key) + if err := readiness.verifyTargetReadinessForRange(ctx, g, route, key, nextScanCursor(key)); err != nil { + return err + } ts, exists, err := g.Store.LatestCommitTS(ctx, key) if err != nil { return errors.WithStack(err) diff --git a/kv/sharded_coordinator_txn_test.go b/kv/sharded_coordinator_txn_test.go index 3316fc847..fdce86762 100644 --- a/kv/sharded_coordinator_txn_test.go +++ b/kv/sharded_coordinator_txn_test.go @@ -387,6 +387,7 @@ func TestGroupReadKeysByShardID_FailsClosedOnUnroutable(t *testing.T) { type stubMVCCStore struct { store.MVCCStore latestTS map[string]uint64 + readiness []store.TargetStagedReadinessState returnErr error } @@ -398,6 +399,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{} @@ -484,6 +489,58 @@ 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_PropagatesStoreError(t *testing.T) { t.Parallel() engine := distribution.NewEngine() From dddda075967b645d5e6eafb11bd49bae4f250c57 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 19:16:17 +0900 Subject: [PATCH 28/29] migration: validate staged readiness conflicts --- kv/fsm.go | 152 +++++++++++++++++++++++++---- kv/fsm_migration_fence_test.go | 65 ++++++++++++ kv/sharded_coordinator.go | 5 +- kv/sharded_coordinator_txn_test.go | 53 ++++++++++ store/lsm_migration.go | 2 +- store/migration_versions.go | 9 +- store/migration_versions_test.go | 10 +- 7 files changed, 268 insertions(+), 28 deletions(-) diff --git a/kv/fsm.go b/kv/fsm.go index e4f1650fd..15d32f4bb 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -719,34 +719,70 @@ func (f *kvFSM) verifyTargetReadinessForRange(ctx context.Context, start []byte, } 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 nil + return routes, nil } states, err := reader.MigrationTargetReadinessStates(ctx) if err != nil { - return errors.WithStack(err) + return nil, errors.WithStack(err) } if len(states) == 0 { - return nil + return routes, nil } + if targetReadinessStatesSatisfied(states, routes, routeStart, routeEnd, f.shardGroupID, catalogVersion, proof) { + return routes, nil + } + return nil, errors.WithStack(ErrRouteCutoverPending) +} - var ( - snap RouteSnapshot - proof bool - ) - if f.routes != nil { - snap, proof = f.routes.Current() +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(snap.IntersectingRoutes(routeStart, routeEnd), ready, f.shardGroupID, snap.Version()) { - return errors.WithStack(ErrRouteCutoverPending) + if !proof || !routesSatisfyTargetReadiness(routes, ready, groupID, catalogVersion) { + return false } } - return nil + return true } func routesSatisfyTargetReadiness(routes []distribution.Route, ready store.TargetStagedReadinessState, groupID uint64, catalogVersion uint64) bool { @@ -1114,7 +1150,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) } @@ -1125,6 +1161,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 @@ -1177,8 +1271,8 @@ func (f *kvFSM) handlePrepareRequest(ctx context.Context, r *pb.Request) error { 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) @@ -1243,12 +1337,7 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com return nil } - uniq, err := f.uniqueMutationsNotFenced(ctx, muts, commitTS) - 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 } @@ -1259,6 +1348,27 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com return nil } +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, nil, err + } + if err := f.validateTxnConflicts(ctx, uniq, readKeys, startTS); err != nil { + return nil, nil, err + } + storeMuts, err := f.buildOnePhaseStoreMutations(ctx, uniq) + if err != nil { + return nil, nil, err + } + return uniq, storeMuts, nil +} + func (f *kvFSM) uniqueMutationsNotFenced(ctx context.Context, muts []*pb.Mutation, commitTS uint64) ([]*pb.Mutation, error) { uniq, err := uniqueMutations(muts) if err != nil { diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index 82c2d5410..65113c949 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -439,6 +439,71 @@ func TestFSMPrepareTxnChecksReadKeysForTargetReadiness(t *testing.T) { 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() diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index bcb1d957e..603360f7f 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -1982,10 +1982,7 @@ func (c *ShardedCoordinator) validateReadKeysOnShard(ctx context.Context, gid ui readiness := &ShardStore{engine: c.engine, groups: c.groups} for _, key := range keys { route := readiness.explicitGroupRouteForKey(gid, key) - if err := readiness.verifyTargetReadinessForRange(ctx, g, route, key, nextScanCursor(key)); err != nil { - return err - } - ts, exists, err := g.Store.LatestCommitTS(ctx, key) + ts, exists, err := readiness.localLatestCommitTS(ctx, g, route, key) if err != nil { return errors.WithStack(err) } diff --git a/kv/sharded_coordinator_txn_test.go b/kv/sharded_coordinator_txn_test.go index fdce86762..759f100d0 100644 --- a/kv/sharded_coordinator_txn_test.go +++ b/kv/sharded_coordinator_txn_test.go @@ -541,6 +541,59 @@ func TestValidateReadOnlyShards_FailsClosedOnTargetReadiness(t *testing.T) { 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/store/lsm_migration.go b/store/lsm_migration.go index f60f5ae02..424b3f633 100644 --- a/store/lsm_migration.go +++ b/store/lsm_migration.go @@ -24,7 +24,7 @@ func (s *pebbleStore) exportVersionsLocked(ctx context.Context, opts ExportVersi if opts.MaxVersions <= 0 { return ExportVersionsResult{Done: true}, nil } - if readTSCompacted(opts.ReadTS, s.effectiveMinRetainedTS()) { + if readTSCompacted(exportRetentionReadTS(opts), s.effectiveMinRetainedTS()) { return ExportVersionsResult{}, ErrReadTSCompacted } diff --git a/store/migration_versions.go b/store/migration_versions.go index 4c5f4edb8..4d2989833 100644 --- a/store/migration_versions.go +++ b/store/migration_versions.go @@ -235,12 +235,19 @@ func (s *mvccStore) exportVersionsLocked(ctx context.Context, opts ExportVersion } func (s *mvccStore) checkExportReadTSLocked(opts ExportVersionsOptions) error { - if readTSCompacted(opts.ReadTS, s.minRetainedTS) { + 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 2fb9ee72c..384c176c5 100644 --- a/store/migration_versions_test.go +++ b/store/migration_versions_test.go @@ -70,12 +70,20 @@ func TestExportVersionsReadTSPreservesCompactionErrors(t *testing.T) { }) require.ErrorIs(t, err, ErrReadTSCompacted) - result, err := st.ExportVersions(ctx, ExportVersionsOptions{ + _, 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) }) From 2d57ae9f4e3f5dd78e2b9c8bed677ba83c4cba79 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 21:12:49 +0900 Subject: [PATCH 29/29] Fail closed read-only shard readiness checks --- kv/sharded_coordinator.go | 40 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index ab947913e..cb2868d6f 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -2159,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) @@ -2170,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 {