From eab9622059690e60f87786e872a5da24f446bbd6 Mon Sep 17 00:00:00 2001 From: bootjp Date: Mon, 13 Jul 2026 22:15:41 +0900 Subject: [PATCH 01/26] Add migration fence drain guards --- adapter/distribution_server.go | 66 +++++++++++++++ adapter/distribution_server_test.go | 64 +++++++++++++++ distribution/engine.go | 40 +++++++++- distribution/split_job_catalog.go | 1 + kv/fsm.go | 83 ++++++++++++++++++- kv/fsm_migration_fence_test.go | 97 +++++++++++++++++++++++ kv/migrator_lock_drain.go | 85 ++++++++++++++++++++ kv/migrator_lock_drain_test.go | 45 +++++++++++ kv/route_history.go | 14 ++++ kv/sharded_coordinator.go | 59 ++++++++++++-- kv/sharded_coordinator_del_prefix_test.go | 79 ++++++++++++++++++ 11 files changed, 625 insertions(+), 8 deletions(-) create mode 100644 kv/fsm_migration_fence_test.go create mode 100644 kv/migrator_lock_drain.go create mode 100644 kv/migrator_lock_drain_test.go diff --git a/adapter/distribution_server.go b/adapter/distribution_server.go index 30c111b6b..3001b1329 100644 --- a/adapter/distribution_server.go +++ b/adapter/distribution_server.go @@ -162,6 +162,9 @@ func (s *DistributionServer) SplitRange(ctx context.Context, req *pb.SplitRangeR if err := validateSplitKey(parent, splitKey); err != nil { return nil, err } + if err := s.rejectLiveSplitJobOverlap(ctx, snapshot, parent); err != nil { + return nil, err + } leftID, rightID, err := s.allocateChildRouteIDs(ctx, snapshot.ReadTS, snapshot.Routes) if err != nil { @@ -376,6 +379,69 @@ func validateSplitKey(parent distribution.RouteDescriptor, splitKey []byte) erro return nil } +func (s *DistributionServer) rejectLiveSplitJobOverlap(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) + } + for _, job := range jobs { + if !splitJobIsLive(job) { + continue + } + for _, interval := range liveSplitJobIntervals(job, snapshot.Routes) { + if routeRangeIntersects(parent.Start, parent.End, interval.start, interval.end) { + return grpcStatusError(codes.Aborted, distribution.ErrSplitJobOverlap.Error()) + } + } + } + return nil +} + +func splitJobIsLive(job distribution.SplitJob) bool { + return job.Phase != distribution.SplitJobPhaseDone && job.Phase != distribution.SplitJobPhaseAbandoned +} + +type routeInterval struct { + start []byte + end []byte +} + +const initialLiveSplitJobIntervalCapacity = 2 + +func liveSplitJobIntervals(job distribution.SplitJob, routes []distribution.RouteDescriptor) []routeInterval { + out := make([]routeInterval, 0, initialLiveSplitJobIntervalCapacity) + for _, route := range routes { + switch { + case route.RouteID == job.SourceRouteID: + out = append(out, routeInterval{ + start: distribution.CloneBytes(job.SplitKey), + end: distribution.CloneBytes(route.End), + }) + case route.ParentRouteID == job.SourceRouteID && routeRangeIntersects(route.Start, route.End, job.SplitKey, nil): + out = append(out, routeInterval{ + start: distribution.CloneBytes(route.Start), + end: distribution.CloneBytes(route.End), + }) + case job.JobID != 0 && route.MigrationJobID == job.JobID: + out = append(out, routeInterval{ + start: distribution.CloneBytes(route.Start), + end: distribution.CloneBytes(route.End), + }) + } + } + return out +} + +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 splitCatalogRoutes( parent distribution.RouteDescriptor, splitKey []byte, diff --git a/adapter/distribution_server_test.go b/adapter/distribution_server_test.go index 8112872fc..7604d0e91 100644 --- a/adapter/distribution_server_test.go +++ b/adapter/distribution_server_test.go @@ -269,6 +269,70 @@ func TestDistributionServerSplitRange_VersionConflict(t *testing.T) { require.ErrorContains(t, err, errDistributionCatalogConflict.Error()) } +func TestDistributionServerSplitRange_RejectsLiveSplitJobOverlap(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}, + {RouteID: 2, Start: []byte("m"), End: nil, GroupID: 2, State: distribution.RouteStateActive}, + }) + require.NoError(t, err) + require.NoError(t, catalog.CreateSplitJob(ctx, distribution.SplitJob{ + JobID: 10, + SourceRouteID: 1, + SplitKey: []byte("g"), + TargetGroupID: 8, + Phase: distribution.SplitJobPhaseBackfill, + })) + + coordinator := newDistributionCoordinatorStub(baseStore, true) + s := NewDistributionServer(distribution.NewEngine(), catalog, WithDistributionCoordinator(coordinator)) + _, err = s.SplitRange(ctx, &pb.SplitRangeRequest{ + ExpectedCatalogVersion: saved.Version, + RouteId: 1, + SplitKey: []byte("c"), + }) + require.Error(t, err) + require.Equal(t, codes.Aborted, status.Code(err)) + require.ErrorContains(t, err, distribution.ErrSplitJobOverlap.Error()) + require.Zero(t, coordinator.dispatchCalls) +} + +func TestDistributionServerSplitRange_AllowsDisjointRouteWhileSplitJobLive(t *testing.T) { + t.Parallel() + + ctx := context.Background() + baseStore := store.NewMVCCStore() + catalog := distribution.NewCatalogStore(baseStore) + saved, err := catalog.Save(ctx, 0, []distribution.RouteDescriptor{ + {RouteID: 3, Start: []byte("a"), End: []byte("g"), GroupID: 1, State: distribution.RouteStateActive, ParentRouteID: 1}, + {RouteID: 4, Start: []byte("g"), End: []byte("m"), GroupID: 1, State: distribution.RouteStateWriteFenced, ParentRouteID: 1}, + {RouteID: 2, Start: []byte("m"), End: nil, GroupID: 2, State: distribution.RouteStateActive}, + }) + require.NoError(t, err) + require.NoError(t, catalog.CreateSplitJob(ctx, distribution.SplitJob{ + JobID: 10, + SourceRouteID: 1, + SplitKey: []byte("g"), + TargetGroupID: 8, + Phase: distribution.SplitJobPhaseFence, + })) + + coordinator := newDistributionCoordinatorStub(baseStore, true) + s := NewDistributionServer(distribution.NewEngine(), catalog, WithDistributionCoordinator(coordinator)) + resp, err := s.SplitRange(ctx, &pb.SplitRangeRequest{ + ExpectedCatalogVersion: saved.Version, + RouteId: 3, + SplitKey: []byte("c"), + }) + require.NoError(t, err) + require.Equal(t, uint64(2), resp.CatalogVersion) + require.Equal(t, 1, coordinator.dispatchCalls) +} + func TestDistributionServerSplitRange_UsesCoordinatorForCatalogWrites(t *testing.T) { t.Parallel() diff --git a/distribution/engine.go b/distribution/engine.go index bc6613894..7963d1282 100644 --- a/distribution/engine.go +++ b/distribution/engine.go @@ -158,6 +158,15 @@ func (s RouteHistorySnapshot) Version() uint64 { return s.version } // scan from O(N) to "first non-covering gap" without changing the // resolution semantics). func (s RouteHistorySnapshot) OwnerOf(key []byte) (uint64, bool) { + r, ok := s.RouteOf(key) + if !ok { + return 0, false + } + return r.GroupID, true +} + +// RouteOf returns the route that covered key at this snapshot's version. +func (s RouteHistorySnapshot) RouteOf(key []byte) (Route, bool) { for _, r := range s.routes { if bytes.Compare(key, r.Start) < 0 { break @@ -165,9 +174,25 @@ func (s RouteHistorySnapshot) OwnerOf(key []byte) (uint64, bool) { if r.End != nil && bytes.Compare(key, r.End) >= 0 { continue } - return r.GroupID, true + return cloneRoute(r), true } - return 0, false + return Route{}, false +} + +// IntersectingRoutes returns every route whose range intersects [start, end) +// in this snapshot. A nil end denotes +infinity. +func (s RouteHistorySnapshot) IntersectingRoutes(start, end []byte) []Route { + out := make([]Route, 0) + for _, r := range s.routes { + if r.End != nil && bytes.Compare(r.End, start) <= 0 { + continue + } + if end != nil && bytes.Compare(r.Start, end) >= 0 { + continue + } + out = append(out, cloneRoute(r)) + } + return out } // Current returns the route catalog snapshot at the engine's current @@ -385,6 +410,17 @@ func (e *Engine) GetIntersectingRoutes(start, end []byte) []Route { return result } +func cloneRoute(r Route) Route { + return Route{ + RouteID: r.RouteID, + Start: CloneBytes(r.Start), + End: CloneBytes(r.End), + GroupID: r.GroupID, + State: r.State, + Load: r.Load, + } +} + func (e *Engine) routeIndex(key []byte) int { if len(e.routes) == 0 { return -1 diff --git a/distribution/split_job_catalog.go b/distribution/split_job_catalog.go index d4d322d45..808ad8498 100644 --- a/distribution/split_job_catalog.go +++ b/distribution/split_job_catalog.go @@ -35,6 +35,7 @@ var ( ErrCatalogSplitJobKeyIDMismatch = errors.New("catalog split job key and record job id mismatch") 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") ) // SplitJobPhase is the durable phase of a split migration job. diff --git a/kv/fsm.go b/kv/fsm.go index 421f5fc3b..92ef461a8 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -120,6 +120,12 @@ type RouteSnapshot interface { // OwnerOf returns the Raft group ID that owned key at this // snapshot's version. (0, false) when no route covered key. OwnerOf(key []byte) (uint64, bool) + // WriteFencedForKey reports whether key is currently inside a + // WriteFenced route in this snapshot. + WriteFencedForKey(key []byte) bool + // WriteFencedIntersects reports whether [start, end) intersects + // any WriteFenced route in this snapshot. + WriteFencedIntersects(start, end []byte) bool } // SetApplyIndex implements raftengine.ApplyIndexAware. The engine @@ -261,6 +267,11 @@ var _ raftengine.StateMachine = (*kvFSM)(nil) var ErrUnknownRequestType = errors.New("unknown request type") +// ErrRouteWriteFenced is returned when a mutation targets a route that is in +// WriteFenced state during split migration. Callers should retry after routing +// catches up to the promoted owner. +var ErrRouteWriteFenced = errors.New("route is write-fenced; retry after route migration") + // ErrComposed1Violation is returned by verifyComposed1 when the // transaction's commit cannot proceed on this Raft group because the // txn's read-set or write-set keys are not owned by this group at @@ -485,6 +496,9 @@ func (f *kvFSM) handleRawRequest(ctx context.Context, r *pb.Request, commitTS ui 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 { return err } @@ -517,6 +531,9 @@ func extractDelPrefix(muts []*pb.Mutation) (bool, []byte) { // handleDelPrefix delegates prefix deletion to the store. Transaction-internal // keys are always excluded to preserve transactional integrity. func (f *kvFSM) handleDelPrefix(ctx context.Context, prefix []byte, commitTS uint64) error { + if err := f.verifyRouteNotFencedForPrefix(prefix); err != nil { + return err + } if err := f.store.DeletePrefixAtRaftAt(ctx, prefix, txnCommonPrefix, commitTS, f.pendingApplyIdx); err != nil { return errors.WithStack(err) } @@ -524,6 +541,56 @@ func (f *kvFSM) handleDelPrefix(ctx context.Context, prefix []byte, commitTS uin return nil } +func (f *kvFSM) verifyRouteNotFencedForMutations(muts []*pb.Mutation) error { + for _, mut := range muts { + if mut == nil || len(mut.Key) == 0 || isTxnInternalKey(mut.Key) { + continue + } + if err := f.verifyRouteNotFencedForKey(mut.Key); err != nil { + return err + } + } + return nil +} + +func (f *kvFSM) verifyRouteNotFencedForKey(key []byte) error { + if f.routes == nil { + return nil + } + snap, ok := f.routes.Current() + if !ok { + return nil + } + rkey := routeKey(key) + if !snap.WriteFencedForKey(rkey) { + return nil + } + return errors.Wrapf(ErrRouteWriteFenced, "key %q routeKey %q", key, rkey) +} + +func (f *kvFSM) verifyRouteNotFencedForPrefix(prefix []byte) error { + if f.routes == nil { + return nil + } + snap, ok := f.routes.Current() + if !ok { + return nil + } + start, end := routePrefixRange(prefix) + if !snap.WriteFencedIntersects(start, end) { + return nil + } + return errors.Wrapf(ErrRouteWriteFenced, "prefix %q route range [%q,%q)", prefix, start, end) +} + +func routePrefixRange(prefix []byte) ([]byte, []byte) { + if len(prefix) == 0 { + return []byte(""), nil + } + start := routeKey(prefix) + return start, prefixScanEnd(start) +} + var ErrNotImplemented = errors.New("not implemented") func (f *kvFSM) Snapshot() (raftengine.Snapshot, error) { @@ -836,6 +903,9 @@ func (f *kvFSM) handlePrepareRequest(ctx context.Context, r *pb.Request) error { if err != nil { return err } + if err := f.verifyRouteNotFencedForMutations(uniq); err != nil { + return err + } if err := f.validateConflicts(ctx, uniq, startTS); err != nil { return errors.WithStack(err) } @@ -902,7 +972,7 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com return nil } - uniq, err := uniqueMutations(muts) + uniq, err := f.uniqueMutationsNotFenced(muts) if err != nil { return err } @@ -918,6 +988,17 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com return nil } +func (f *kvFSM) uniqueMutationsNotFenced(muts []*pb.Mutation) ([]*pb.Mutation, error) { + uniq, err := uniqueMutations(muts) + if err != nil { + return nil, err + } + if err := f.verifyRouteNotFencedForMutations(uniq); err != nil { + return nil, err + } + return uniq, nil +} + // dedupProbeOnePhase decides whether handleOnePhaseTxnRequest should no-op // because the entry is a retry whose prior attempt already landed. Extracted // to keep handleOnePhaseTxnRequest under the cyclop budget; the determinism diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go new file mode 100644 index 000000000..a9f4b998f --- /dev/null +++ b/kv/fsm_migration_fence_test.go @@ -0,0 +1,97 @@ +package kv + +import ( + "context" + "testing" + + "github.com/bootjp/elastickv/distribution" + pb "github.com/bootjp/elastickv/proto" + "github.com/bootjp/elastickv/store" + "github.com/cockroachdb/errors" + "github.com/stretchr/testify/require" +) + +func newWriteFencedFSM(t *testing.T) *kvFSM { + t.Helper() + + engine := distribution.NewEngine() + applyComposed1Snapshot(t, engine, 1, []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: []byte("m"), GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: []byte("m"), End: nil, GroupID: 1, State: distribution.RouteStateWriteFenced}, + }) + return newComposed1FSM(t, engine, 1) +} + +func TestFSMRejectsRawPointWriteOnWriteFencedRoute(t *testing.T) { + t.Parallel() + + fsm := newWriteFencedFSM(t) + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("v")}}, + }, 10) + require.ErrorIs(t, err, ErrRouteWriteFenced) + + _, getErr := fsm.store.GetAt(context.Background(), []byte("z"), ^uint64(0)) + require.ErrorIs(t, getErr, store.ErrKeyNotFound) +} + +func TestFSMRejectsDelPrefixIntersectingWriteFencedRoute(t *testing.T) { + t.Parallel() + + fsm := newWriteFencedFSM(t) + require.NoError(t, fsm.store.PutAt(context.Background(), []byte("z"), []byte("v"), 1, 0)) + + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_DEL_PREFIX, Key: []byte("z")}}, + }, 10) + require.ErrorIs(t, err, ErrRouteWriteFenced) + + got, getErr := fsm.store.GetAt(context.Background(), []byte("z"), ^uint64(0)) + require.NoError(t, getErr) + require.Equal(t, []byte("v"), got) +} + +func TestFSMRejectsFullRangeDelPrefixWhenRouteIsWriteFenced(t *testing.T) { + t.Parallel() + + fsm := newWriteFencedFSM(t) + require.NoError(t, fsm.store.PutAt(context.Background(), []byte("z"), []byte("v"), 1, 0)) + + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_DEL_PREFIX, Key: nil}}, + }, 10) + require.ErrorIs(t, err, ErrRouteWriteFenced) + + got, getErr := fsm.store.GetAt(context.Background(), []byte("z"), ^uint64(0)) + require.NoError(t, getErr) + require.Equal(t, []byte("v"), got) +} + +func TestFSMRejectsPrepareOnWriteFencedRouteButAllowsAbort(t *testing.T) { + t.Parallel() + + ctx := context.Background() + fsm := newWriteFencedFSM(t) + prepare := &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("z"), LockTTLms: defaultTxnLockTTLms})}, + {Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("v")}, + }, + } + require.ErrorIs(t, fsm.handleTxnRequest(ctx, prepare, 10), ErrRouteWriteFenced) + + abort := &pb.Request{ + IsTxn: true, + Phase: pb.Phase_ABORT, + Ts: 11, + Mutations: []*pb.Mutation{ + {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: []byte("z"), CommitTS: 11})}, + {Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("v")}, + }, + } + err := fsm.handleTxnRequest(ctx, abort, 11) + require.False(t, errors.Is(err, ErrRouteWriteFenced), "ABORT must keep the narrow cleanup lane open") +} diff --git a/kv/migrator_lock_drain.go b/kv/migrator_lock_drain.go new file mode 100644 index 000000000..67396fc14 --- /dev/null +++ b/kv/migrator_lock_drain.go @@ -0,0 +1,85 @@ +package kv + +import ( + "bytes" + "context" + + "github.com/bootjp/elastickv/store" + "github.com/cockroachdb/errors" +) + +// TxnLockDrainEntry describes one prepared transaction lock that still belongs +// to a migration route. The lock key is cloned so callers can safely retain it +// across drain ticks. +type TxnLockDrainEntry struct { + LockKey []byte + UserKey []byte + StartTS uint64 + TTLExpireAt uint64 + PrimaryKey []byte + IsPrimaryKey bool +} + +// PendingTxnLocksInRoute scans the txn-lock namespace and filters each lock by +// routeKey(userKey). It intentionally does not bracket the scan by +// txnLockKey(routeStart)/txnLockKey(routeEnd): txn locks are sorted by raw user +// key while migration routes are defined in route-key space. +func PendingTxnLocksInRoute(ctx context.Context, st store.MVCCStore, routeStart, routeEnd []byte, ts uint64, limit int) ([]TxnLockDrainEntry, error) { + if st == nil { + return nil, nil + } + if limit <= 0 { + limit = maxTxnLockScanResults + } + start := txnLockKey(nil) + end := prefixScanEnd(start) + filter := RouteKeyFilter(routeStart, routeEnd) + cursor := start + out := make([]TxnLockDrainEntry, 0, min(limit, lockPageLimit)) + + for { + lockKVs, nextCursor, done, err := scanTxnLockPageAt(ctx, st, cursor, end, ts) + if err != nil { + return nil, err + } + for _, kvp := range lockKVs { + entry, ok, err := txnLockDrainEntry(kvp, filter) + if err != nil { + return nil, err + } + if !ok { + continue + } + out = append(out, entry) + if len(out) >= limit { + return out, nil + } + } + if done { + return out, nil + } + cursor = nextCursor + } +} + +func txnLockDrainEntry(kvp *store.KVPair, filter func([]byte) bool) (TxnLockDrainEntry, bool, error) { + if kvp == nil || !bytes.HasPrefix(kvp.Key, txnLockPrefixBytes) { + return TxnLockDrainEntry{}, false, nil + } + userKey := kvp.Key[len(txnLockPrefixBytes):] + if !filter(userKey) { + return TxnLockDrainEntry{}, false, nil + } + lock, err := decodeTxnLock(kvp.Value) + if err != nil { + return TxnLockDrainEntry{}, false, errors.Wrap(err, "decode txn lock during migration drain") + } + return TxnLockDrainEntry{ + LockKey: bytes.Clone(kvp.Key), + UserKey: bytes.Clone(userKey), + StartTS: lock.StartTS, + TTLExpireAt: lock.TTLExpireAt, + PrimaryKey: bytes.Clone(lock.PrimaryKey), + IsPrimaryKey: lock.IsPrimaryKey, + }, true, nil +} diff --git a/kv/migrator_lock_drain_test.go b/kv/migrator_lock_drain_test.go new file mode 100644 index 000000000..916586010 --- /dev/null +++ b/kv/migrator_lock_drain_test.go @@ -0,0 +1,45 @@ +package kv + +import ( + "bytes" + "context" + "testing" + + "github.com/bootjp/elastickv/store" + "github.com/stretchr/testify/require" +) + +func TestPendingTxnLocksInRouteFiltersLocksByRouteKey(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + tableSegment := "tenant-a" + routeStart := []byte(dynamoRoutePrefix + tableSegment) + routeEnd := append(bytes.Clone(routeStart), 0xff) + itemKey := append([]byte(DynamoItemPrefix+tableSegment+"|7|"), []byte("pk\x00\x01")...) + outsideKey := []byte("outside") + + require.Less(t, bytes.Compare(itemKey, routeStart), 0, + "the red-control key must sort outside a txnLockKey(routeStart)/txnLockKey(routeEnd) bracket") + lock := encodeTxnLock(txnLock{ + StartTS: 11, + TTLExpireAt: 99, + PrimaryKey: itemKey, + IsPrimaryKey: true, + }) + require.NoError(t, st.PutAt(ctx, txnLockKey(itemKey), lock, 1, 0)) + require.NoError(t, st.PutAt(ctx, txnLockKey(outsideKey), encodeTxnLock(txnLock{ + StartTS: 12, + PrimaryKey: outsideKey, + }), 2, 0)) + + pending, err := PendingTxnLocksInRoute(ctx, st, routeStart, routeEnd, ^uint64(0), 10) + require.NoError(t, err) + require.Len(t, pending, 1) + require.Equal(t, itemKey, pending[0].UserKey) + require.Equal(t, uint64(11), pending[0].StartTS) + require.Equal(t, uint64(99), pending[0].TTLExpireAt) + require.True(t, pending[0].IsPrimaryKey) + require.Equal(t, txnLockKey(itemKey), pending[0].LockKey) +} diff --git a/kv/route_history.go b/kv/route_history.go index cb768b30b..c6ac85608 100644 --- a/kv/route_history.go +++ b/kv/route_history.go @@ -60,3 +60,17 @@ func (s distributionRouteSnapshot) Version() uint64 { func (s distributionRouteSnapshot) OwnerOf(key []byte) (uint64, bool) { return s.snap.OwnerOf(key) } + +func (s distributionRouteSnapshot) WriteFencedForKey(key []byte) bool { + route, ok := s.snap.RouteOf(key) + return ok && route.State == distribution.RouteStateWriteFenced +} + +func (s distributionRouteSnapshot) WriteFencedIntersects(start, end []byte) bool { + for _, route := range s.snap.IntersectingRoutes(start, end) { + if route.State == distribution.RouteStateWriteFenced { + return true + } + } + return false +} diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index b8b67dba5..7ed2ccb7c 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -700,11 +700,8 @@ func (c *ShardedCoordinator) Dispatch(ctx context.Context, reqs *OperationGroup[ return nil, err } - // DEL_PREFIX cannot be routed to a single shard because the prefix may - // span multiple shards (or be nil, meaning "all keys"). Broadcast the - // operation to every shard group so each FSM scans locally. - if hasDelPrefixElem(reqs.Elems) { - return c.dispatchDelPrefixBroadcast(ctx, reqs.IsTxn, reqs.Elems) + if resp, handled, err := c.dispatchBeforeShardRouting(ctx, reqs); handled { + return resp, err } // Capture whether the caller supplied a non-zero StartTS BEFORE @@ -744,6 +741,20 @@ func (c *ShardedCoordinator) Dispatch(ctx context.Context, reqs *OperationGroup[ return c.dispatchNonTxn(ctx, reqs) } +func (c *ShardedCoordinator) dispatchBeforeShardRouting(ctx context.Context, reqs *OperationGroup[OP]) (*CoordinateResponse, bool, error) { + // DEL_PREFIX cannot be routed to a single shard because the prefix may + // span multiple shards (or be nil, meaning "all keys"). Broadcast the + // operation to every shard group so each FSM scans locally. + if hasDelPrefixElem(reqs.Elems) { + resp, err := c.dispatchDelPrefixBroadcast(ctx, reqs.IsTxn, reqs.Elems) + return resp, true, err + } + if err := c.rejectWriteFencedPointElems(reqs.Elems); err != nil { + return nil, true, err + } + return nil, false, nil +} + // dispatchTxnWithComposed1Retry runs the M4 Composed-1 retry loop // (design doc // docs/design/2026_05_29_implemented_composed1_cross_group_commit_guard.md @@ -1017,6 +1028,9 @@ func (c *ShardedCoordinator) dispatchDelPrefixBroadcast(ctx context.Context, isT if err := validateDelPrefixOnly(elems); err != nil { return nil, err } + if err := c.rejectWriteFencedDelPrefixes(elems); err != nil { + return nil, err + } ts, err := c.allocateTimestamp(ctx, "allocate DEL_PREFIX broadcast ts") if err != nil { @@ -1035,6 +1049,41 @@ func (c *ShardedCoordinator) dispatchDelPrefixBroadcast(ctx context.Context, isT return c.broadcastToAllGroups(ctx, requests) } +func (c *ShardedCoordinator) rejectWriteFencedPointElems(elems []*Elem[OP]) error { + if c == nil || c.engine == nil { + return nil + } + for _, elem := range elems { + if elem == nil || len(elem.Key) == 0 { + continue + } + route, ok := c.engine.GetRoute(routeKey(elem.Key)) + if !ok || route.State != distribution.RouteStateWriteFenced { + continue + } + return errors.Wrapf(ErrRouteWriteFenced, "key %q routeKey %q", elem.Key, routeKey(elem.Key)) + } + return nil +} + +func (c *ShardedCoordinator) rejectWriteFencedDelPrefixes(elems []*Elem[OP]) 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 route.State == distribution.RouteStateWriteFenced { + return errors.Wrapf(ErrRouteWriteFenced, "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 11de9de78..01e63961a 100644 --- a/kv/sharded_coordinator_del_prefix_test.go +++ b/kv/sharded_coordinator_del_prefix_test.go @@ -121,6 +121,85 @@ func TestShardedCoordinator_DelPrefixBroadcastsToAllGroups(t *testing.T) { "same DEL_PREFIX element must use the same timestamp across shards") } +func TestShardedCoordinatorRejectsPointWriteOnWriteFencedRoute(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.RouteStateWriteFenced}, + }, + })) + + g2Txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: &recordingTransactional{}}, + 2: {Txn: g2Txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: Put, Key: []byte("z"), Value: []byte("v")}}, + }) + require.ErrorIs(t, err, ErrRouteWriteFenced) + require.Empty(t, g2Txn.requests, "coordinator must reject before proposing to the fenced shard") +} + +func TestShardedCoordinatorRejectsDelPrefixIntersectingWriteFencedRoute(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.RouteStateWriteFenced}, + }, + })) + + 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, ErrRouteWriteFenced) + require.Empty(t, g1Txn.requests) + require.Empty(t, g2Txn.requests) +} + +func TestShardedCoordinatorRejectsFullRangeDelPrefixWhenRouteIsWriteFenced(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.RouteStateWriteFenced}, + }, + })) + + 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: nil}}, + }) + require.ErrorIs(t, err, ErrRouteWriteFenced) + require.Empty(t, g1Txn.requests) + require.Empty(t, g2Txn.requests) +} + // TestShardedCoordinator_DelPrefixRejectsTxn verifies that DEL_PREFIX inside // a transactional group is rejected. func TestShardedCoordinator_DelPrefixRejectsTxn(t *testing.T) { From 07a428ee880738fe5aaa6c7306690e4e2d26dc21 Mon Sep 17 00:00:00 2001 From: bootjp Date: Mon, 13 Jul 2026 22:23:16 +0900 Subject: [PATCH 02/26] Address migration guard review notes --- distribution/engine.go | 4 ++-- kv/migrator_lock_drain.go | 9 ++++++++- kv/migrator_lock_drain_test.go | 11 +++++++++++ 3 files changed, 21 insertions(+), 3 deletions(-) diff --git a/distribution/engine.go b/distribution/engine.go index 7963d1282..dc55ed74e 100644 --- a/distribution/engine.go +++ b/distribution/engine.go @@ -188,7 +188,7 @@ func (s RouteHistorySnapshot) IntersectingRoutes(start, end []byte) []Route { continue } if end != nil && bytes.Compare(r.Start, end) >= 0 { - continue + break } out = append(out, cloneRoute(r)) } @@ -395,7 +395,7 @@ func (e *Engine) GetIntersectingRoutes(start, end []byte) []Route { } // Route starts at or after scan ends: end != nil && rStart >= end if end != nil && bytes.Compare(r.Start, end) >= 0 { - continue + break } // Route intersects with scan range result = append(result, Route{ diff --git a/kv/migrator_lock_drain.go b/kv/migrator_lock_drain.go index 67396fc14..ef7a7e531 100644 --- a/kv/migrator_lock_drain.go +++ b/kv/migrator_lock_drain.go @@ -38,7 +38,7 @@ func PendingTxnLocksInRoute(ctx context.Context, st store.MVCCStore, routeStart, out := make([]TxnLockDrainEntry, 0, min(limit, lockPageLimit)) for { - lockKVs, nextCursor, done, err := scanTxnLockPageAt(ctx, st, cursor, end, ts) + lockKVs, nextCursor, done, err := scanTxnLockDrainPage(ctx, st, cursor, end, ts) if err != nil { return nil, err } @@ -62,6 +62,13 @@ func PendingTxnLocksInRoute(ctx context.Context, st store.MVCCStore, routeStart, } } +func scanTxnLockDrainPage(ctx context.Context, st store.MVCCStore, cursor, end []byte, ts uint64) ([]*store.KVPair, []byte, bool, error) { + if err := ctx.Err(); err != nil { + return nil, nil, false, errors.WithStack(err) + } + return scanTxnLockPageAt(ctx, st, cursor, end, ts) +} + func txnLockDrainEntry(kvp *store.KVPair, filter func([]byte) bool) (TxnLockDrainEntry, bool, error) { if kvp == nil || !bytes.HasPrefix(kvp.Key, txnLockPrefixBytes) { return TxnLockDrainEntry{}, false, nil diff --git a/kv/migrator_lock_drain_test.go b/kv/migrator_lock_drain_test.go index 916586010..d790e85c0 100644 --- a/kv/migrator_lock_drain_test.go +++ b/kv/migrator_lock_drain_test.go @@ -43,3 +43,14 @@ func TestPendingTxnLocksInRouteFiltersLocksByRouteKey(t *testing.T) { require.True(t, pending[0].IsPrimaryKey) require.Equal(t, txnLockKey(itemKey), pending[0].LockKey) } + +func TestPendingTxnLocksInRouteHonorsCanceledContext(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + pending, err := PendingTxnLocksInRoute(ctx, store.NewMVCCStore(), nil, nil, ^uint64(0), 10) + require.ErrorIs(t, err, context.Canceled) + require.Nil(t, pending) +} From 7c1780323eedc9d489ba55c5794b7c65c1d7e635 Mon Sep 17 00:00:00 2001 From: bootjp Date: Mon, 13 Jul 2026 23:29:50 +0900 Subject: [PATCH 03/26] kv: fence broad mapped prefix deletes --- kv/fsm.go | 53 ++++++++++++++++++++ kv/fsm_migration_fence_test.go | 17 +++++++ kv/shard_key_test.go | 59 +++++++++++++++++++++++ kv/sharded_coordinator_del_prefix_test.go | 36 ++++++++++++++ 4 files changed, 165 insertions(+) diff --git a/kv/fsm.go b/kv/fsm.go index 92ef461a8..5cd95d05d 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -11,6 +11,7 @@ import ( "github.com/bootjp/elastickv/internal/encryption/fsmwire" "github.com/bootjp/elastickv/internal/raftengine" + "github.com/bootjp/elastickv/internal/s3keys" pb "github.com/bootjp/elastickv/proto" "github.com/bootjp/elastickv/store" "github.com/cockroachdb/errors" @@ -587,10 +588,62 @@ func routePrefixRange(prefix []byte) ([]byte, []byte) { if len(prefix) == 0 { return []byte(""), nil } + if routeKeyspaceWideRawPrefix(prefix) { + return []byte(""), nil + } start := routeKey(prefix) return start, prefixScanEnd(start) } +func routeKeyspaceWideRawPrefix(prefix []byte) bool { + if !rawPrefixMayContainRouteMappedKeys(prefix) { + return false + } + return bytes.Equal(routeKey(prefix), prefix) +} + +func rawPrefixMayContainRouteMappedKeys(prefix []byte) bool { + for _, mappedPrefix := range routeMappedRawPrefixes { + if bytes.HasPrefix(prefix, mappedPrefix) || bytes.HasPrefix(mappedPrefix, prefix) { + return true + } + } + return false +} + +var routeMappedRawPrefixes = [][]byte{ + []byte(redisInternalRoutePrefix), + []byte(DynamoTableMetaPrefix), + []byte(DynamoTableGenerationPrefix), + []byte(DynamoItemPrefix), + []byte(DynamoGSIPrefix), + []byte(sqsInternalPrefix), + []byte(store.ListMetaPrefix), + []byte(store.ListItemPrefix), + []byte(store.ListMetaDeltaPrefix), + []byte(store.ListClaimPrefix), + []byte(store.HashMetaPrefix), + []byte(store.HashFieldPrefix), + []byte(store.HashMetaDeltaPrefix), + []byte(store.SetMetaPrefix), + []byte(store.SetMemberPrefix), + []byte(store.SetMetaDeltaPrefix), + []byte(store.ZSetMetaPrefix), + []byte(store.ZSetMemberPrefix), + []byte(store.ZSetScorePrefix), + []byte(store.ZSetMetaDeltaPrefix), + []byte(store.StreamMetaPrefix), + []byte(store.StreamEntryPrefix), + []byte(s3keys.BucketMetaPrefix), + []byte(s3keys.BucketGenerationPrefix), + []byte(s3keys.ObjectManifestPrefix), + []byte(s3keys.UploadMetaPrefix), + []byte(s3keys.UploadPartPrefix), + []byte(s3keys.BlobPrefix), + []byte(s3keys.GCUploadPrefix), + []byte(s3keys.RoutePrefix), +} + var ErrNotImplemented = errors.New("not implemented") func (f *kvFSM) Snapshot() (raftengine.Snapshot, error) { diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index a9f4b998f..e47e71f95 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -67,6 +67,23 @@ func TestFSMRejectsFullRangeDelPrefixWhenRouteIsWriteFenced(t *testing.T) { require.Equal(t, []byte("v"), got) } +func TestFSMRejectsBroadInternalDelPrefixWhenRouteIsWriteFenced(t *testing.T) { + t.Parallel() + + fsm := newWriteFencedFSM(t) + key := []byte("!redis|string|z") + require.NoError(t, fsm.store.PutAt(context.Background(), key, []byte("v"), 1, 0)) + + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_DEL_PREFIX, Key: []byte("!redis|")}}, + }, 10) + require.ErrorIs(t, err, ErrRouteWriteFenced) + + got, getErr := fsm.store.GetAt(context.Background(), key, ^uint64(0)) + require.NoError(t, getErr) + require.Equal(t, []byte("v"), got) +} + func TestFSMRejectsPrepareOnWriteFencedRouteButAllowsAbort(t *testing.T) { t.Parallel() diff --git a/kv/shard_key_test.go b/kv/shard_key_test.go index e7d6c55d7..fb282236a 100644 --- a/kv/shard_key_test.go +++ b/kv/shard_key_test.go @@ -174,6 +174,65 @@ func TestRouteKey_S3DecoderIsConcreteOnly(t *testing.T) { require.Equal(t, rawUser, routeKey(rawUser), "adapter-looking raw user key must stay on its raw route") } +func TestRoutePrefixRangeTreatsBroadMappedPrefixesAsFullKeyspace(t *testing.T) { + t.Parallel() + + tableSegment := base64.RawURLEncoding.EncodeToString([]byte("users")) + dynamoTableRoute := dynamoRouteTableKey([]byte(tableSegment)) + + for _, tc := range []struct { + name string + prefix []byte + wantStart []byte + wantEnd []byte + }{ + { + name: "raw user prefix", + prefix: []byte("ab"), + wantStart: []byte("ab"), + wantEnd: prefixScanEnd([]byte("ab")), + }, + { + name: "concrete redis prefix", + prefix: []byte("!redis|string|ab"), + wantStart: []byte("ab"), + wantEnd: prefixScanEnd([]byte("ab")), + }, + { + name: "dynamo table cleanup prefix", + prefix: []byte(DynamoItemPrefix + tableSegment + "|7|"), + wantStart: dynamoTableRoute, + wantEnd: prefixScanEnd(dynamoTableRoute), + }, + { + name: "broad redis namespace", + prefix: []byte("!redis|"), + wantStart: []byte(""), + wantEnd: nil, + }, + { + name: "broad wide-column namespace", + prefix: []byte("!lst|"), + wantStart: []byte(""), + wantEnd: nil, + }, + { + name: "s3 bucket cleanup prefix", + prefix: s3keys.ObjectManifestPrefixForBucket("bucket", 2), + wantStart: []byte(""), + wantEnd: nil, + }, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + start, end := routePrefixRange(tc.prefix) + require.Equal(t, tc.wantStart, start) + require.Equal(t, tc.wantEnd, end) + }) + } +} + func TestRouteKeyFilterTreatsNilAndEmptyEndAsInfinity(t *testing.T) { t.Parallel() diff --git a/kv/sharded_coordinator_del_prefix_test.go b/kv/sharded_coordinator_del_prefix_test.go index 01e63961a..4b83dc131 100644 --- a/kv/sharded_coordinator_del_prefix_test.go +++ b/kv/sharded_coordinator_del_prefix_test.go @@ -200,6 +200,42 @@ func TestShardedCoordinatorRejectsFullRangeDelPrefixWhenRouteIsWriteFenced(t *te require.Empty(t, g2Txn.requests) } +func TestShardedCoordinatorRejectsBroadInternalDelPrefixWhenRouteIsWriteFenced(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.RouteStateWriteFenced}, + }, + })) + + for _, prefix := range [][]byte{ + []byte("!redis|"), + []byte("!lst|"), + } { + t.Run(string(prefix), func(t *testing.T) { + t.Parallel() + + 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: prefix}}, + }) + require.ErrorIs(t, err, ErrRouteWriteFenced) + require.Empty(t, g1Txn.requests) + require.Empty(t, g2Txn.requests) + }) + } +} + // TestShardedCoordinator_DelPrefixRejectsTxn verifies that DEL_PREFIX inside // a transactional group is rejected. func TestShardedCoordinator_DelPrefixRejectsTxn(t *testing.T) { From 5acb33f4261e3445b275eb221d4039fbada82252 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 18:01:29 +0900 Subject: [PATCH 04/26] migration: fence S3 bucket metadata writes --- kv/fsm.go | 9 +++-- kv/fsm_migration_fence_test.go | 40 +++++++++++++++++++++++ kv/migrator_filter.go | 18 +++++++--- kv/sharded_coordinator.go | 23 ++++++++++--- kv/sharded_coordinator_del_prefix_test.go | 31 ++++++++++++++++++ 5 files changed, 109 insertions(+), 12 deletions(-) diff --git a/kv/fsm.go b/kv/fsm.go index 5cd95d05d..a9ff4cea5 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -563,10 +563,13 @@ func (f *kvFSM) verifyRouteNotFencedForKey(key []byte) error { return nil } rkey := routeKey(key) - if !snap.WriteFencedForKey(rkey) { - return nil + if snap.WriteFencedForKey(rkey) { + return errors.Wrapf(ErrRouteWriteFenced, "key %q routeKey %q", key, rkey) + } + if start, end, ok := s3BucketAuxiliaryRouteRange(key); ok && snap.WriteFencedIntersects(start, end) { + return errors.Wrapf(ErrRouteWriteFenced, "key %q route range [%q,%q)", key, start, end) } - return errors.Wrapf(ErrRouteWriteFenced, "key %q routeKey %q", key, rkey) + return nil } func (f *kvFSM) verifyRouteNotFencedForPrefix(prefix []byte) error { diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index e47e71f95..7aff2b4f4 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -5,6 +5,7 @@ import ( "testing" "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/internal/s3keys" pb "github.com/bootjp/elastickv/proto" "github.com/bootjp/elastickv/store" "github.com/cockroachdb/errors" @@ -22,6 +23,24 @@ func newWriteFencedFSM(t *testing.T) *kvFSM { return newComposed1FSM(t, engine, 1) } +func s3BucketAuxiliaryFenceRoutes(bucket string, rawGroupID, fencedGroupID uint64) []distribution.RouteDescriptor { + start := s3keys.RoutePrefixForBucketAnyGeneration(bucket) + end := prefixScanEnd(start) + return []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: start, GroupID: rawGroupID, State: distribution.RouteStateActive}, + {RouteID: 2, Start: start, End: end, GroupID: fencedGroupID, State: distribution.RouteStateWriteFenced}, + {RouteID: 3, Start: end, End: nil, GroupID: rawGroupID, State: distribution.RouteStateActive}, + } +} + +func newS3BucketAuxiliaryWriteFencedFSM(t *testing.T, bucket string) *kvFSM { + t.Helper() + + engine := distribution.NewEngine() + applyComposed1Snapshot(t, engine, 1, s3BucketAuxiliaryFenceRoutes(bucket, 1, 1)) + return newComposed1FSM(t, engine, 1) +} + func TestFSMRejectsRawPointWriteOnWriteFencedRoute(t *testing.T) { t.Parallel() @@ -35,6 +54,27 @@ func TestFSMRejectsRawPointWriteOnWriteFencedRoute(t *testing.T) { require.ErrorIs(t, getErr, store.ErrKeyNotFound) } +func TestFSMRejectsS3BucketAuxiliaryPointWriteOnWriteFencedRoute(t *testing.T) { + t.Parallel() + + ctx := context.Background() + const bucket = "bucket-a" + fsm := newS3BucketAuxiliaryWriteFencedFSM(t, bucket) + + for _, key := range [][]byte{ + s3keys.BucketMetaKey(bucket), + s3keys.BucketGenerationKey(bucket), + } { + err := fsm.handleRawRequest(ctx, &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: key, Value: []byte("v")}}, + }, 10) + require.ErrorIs(t, err, ErrRouteWriteFenced) + + _, getErr := fsm.store.GetAt(ctx, key, ^uint64(0)) + require.ErrorIs(t, getErr, store.ErrKeyNotFound) + } +} + func TestFSMRejectsDelPrefixIntersectingWriteFencedRoute(t *testing.T) { t.Parallel() diff --git a/kv/migrator_filter.go b/kv/migrator_filter.go index 8a1da5526..acbf1a796 100644 --- a/kv/migrator_filter.go +++ b/kv/migrator_filter.go @@ -37,18 +37,26 @@ func RouteKeyFilterForGroup(rangeStart, rangeEnd []byte, sourceGroupID uint64, r } func s3BucketAuxiliaryRouteInRange(rawKey, routeStart, routeEnd []byte) bool { - bucket, ok := s3keys.ParseBucketMetaKey(rawKey) - if !ok { - bucket, ok = s3keys.ParseBucketGenerationKey(rawKey) - } + bucketRouteStart, bucketRouteEnd, ok := s3BucketAuxiliaryRouteRange(rawKey) if !ok { return false } if keyInMigrationRouteRange(rawKey, routeStart, routeEnd) { return true } + return migrationRouteRangesIntersect(routeStart, routeEnd, bucketRouteStart, bucketRouteEnd) +} + +func s3BucketAuxiliaryRouteRange(rawKey []byte) ([]byte, []byte, bool) { + bucket, ok := s3keys.ParseBucketMetaKey(rawKey) + if !ok { + bucket, ok = s3keys.ParseBucketGenerationKey(rawKey) + } + if !ok { + return nil, nil, false + } bucketRouteStart := s3keys.RoutePrefixForBucketAnyGeneration(bucket) - return migrationRouteRangesIntersect(routeStart, routeEnd, bucketRouteStart, prefixScanEnd(bucketRouteStart)) + return bucketRouteStart, prefixScanEnd(bucketRouteStart), true } func keyInMigrationRouteRange(key, routeStart, routeEnd []byte) bool { diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index 7ed2ccb7c..f345f2e63 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -1057,11 +1057,26 @@ func (c *ShardedCoordinator) rejectWriteFencedPointElems(elems []*Elem[OP]) erro if elem == nil || len(elem.Key) == 0 { continue } - route, ok := c.engine.GetRoute(routeKey(elem.Key)) - if !ok || route.State != distribution.RouteStateWriteFenced { - continue + if err := c.rejectWriteFencedPointKey(elem.Key); err != nil { + return err + } + } + return nil +} + +func (c *ShardedCoordinator) rejectWriteFencedPointKey(key []byte) error { + rkey := routeKey(key) + if route, ok := c.engine.GetRoute(rkey); ok && route.State == distribution.RouteStateWriteFenced { + return errors.Wrapf(ErrRouteWriteFenced, "key %q routeKey %q", key, rkey) + } + start, end, ok := s3BucketAuxiliaryRouteRange(key) + if !ok { + return nil + } + for _, route := range c.engine.GetIntersectingRoutes(start, end) { + if route.State == distribution.RouteStateWriteFenced { + return errors.Wrapf(ErrRouteWriteFenced, "key %q route range [%q,%q)", key, start, end) } - return errors.Wrapf(ErrRouteWriteFenced, "key %q routeKey %q", elem.Key, routeKey(elem.Key)) } return nil } diff --git a/kv/sharded_coordinator_del_prefix_test.go b/kv/sharded_coordinator_del_prefix_test.go index 4b83dc131..5cc0f845c 100644 --- a/kv/sharded_coordinator_del_prefix_test.go +++ b/kv/sharded_coordinator_del_prefix_test.go @@ -6,6 +6,7 @@ import ( "testing" "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/internal/s3keys" pb "github.com/bootjp/elastickv/proto" "github.com/bootjp/elastickv/store" "github.com/stretchr/testify/require" @@ -146,6 +147,36 @@ func TestShardedCoordinatorRejectsPointWriteOnWriteFencedRoute(t *testing.T) { require.Empty(t, g2Txn.requests, "coordinator must reject before proposing to the fenced shard") } +func TestShardedCoordinatorRejectsS3BucketAuxiliaryPointWriteOnWriteFencedRoute(t *testing.T) { + t.Parallel() + + const bucket = "bucket-a" + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: s3BucketAuxiliaryFenceRoutes(bucket, 1, 2), + })) + + g1Txn := &recordingTransactional{} + g2Txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + }, 1, NewHLC(), nil) + + for _, key := range [][]byte{ + s3keys.BucketMetaKey(bucket), + s3keys.BucketGenerationKey(bucket), + } { + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: Put, Key: key, Value: []byte("v")}}, + }) + require.ErrorIs(t, err, ErrRouteWriteFenced) + require.Empty(t, g1Txn.requests, "coordinator must reject before proposing to the raw-key shard") + require.Empty(t, g2Txn.requests, "coordinator must reject before proposing to the fenced shard") + } +} + func TestShardedCoordinatorRejectsDelPrefixIntersectingWriteFencedRoute(t *testing.T) { t.Parallel() From f1aeeeb396903a904e62fdfee9203cfc11aaafe3 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 20:07:14 +0900 Subject: [PATCH 05/26] Guard split fence retry paths --- adapter/distribution_server.go | 39 ++++++++++--- adapter/distribution_server_test.go | 79 ++++++++++++++++++++++++++- adapter/dynamodb_transact.go | 2 +- adapter/redis_retry.go | 2 +- adapter/redis_retry_test.go | 1 + adapter/retryable_write_fence_test.go | 16 ++++++ adapter/s3.go | 2 +- distribution/engine.go | 15 +++-- distribution/engine_test.go | 15 +++++ 9 files changed, 150 insertions(+), 21 deletions(-) create mode 100644 adapter/retryable_write_fence_test.go diff --git a/adapter/distribution_server.go b/adapter/distribution_server.go index fa2b73a7d..0fcacba58 100644 --- a/adapter/distribution_server.go +++ b/adapter/distribution_server.go @@ -10,6 +10,7 @@ import ( "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" @@ -162,7 +163,8 @@ func (s *DistributionServer) SplitRange(ctx context.Context, req *pb.SplitRangeR if err := validateSplitKey(parent, splitKey); err != nil { return nil, err } - if err := s.rejectLiveSplitJobOverlap(ctx, snapshot, parent); err != nil { + splitJobReadKeys, err := s.splitJobOverlapReadKeys(ctx, snapshot, parent) + if err != nil { return nil, err } @@ -172,7 +174,7 @@ func (s *DistributionServer) SplitRange(ctx context.Context, req *pb.SplitRangeR } left, right := splitCatalogRoutes(parent, splitKey, leftID, rightID) - saved, err := s.saveSplitResultViaCoordinator(ctx, snapshot.ReadTS, req.GetExpectedCatalogVersion(), parent.RouteID, left, right) + saved, err := s.saveSplitResultViaCoordinator(ctx, snapshot.ReadTS, req.GetExpectedCatalogVersion(), parent.RouteID, splitJobReadKeys, left, right) if err != nil { return nil, err } @@ -213,6 +215,7 @@ func (s *DistributionServer) saveSplitResultViaCoordinator( readTS uint64, expectedVersion uint64, parentID uint64, + readKeys [][]byte, left distribution.RouteDescriptor, right distribution.RouteDescriptor, ) (distribution.CatalogSnapshot, error) { @@ -230,10 +233,14 @@ func (s *DistributionServer) saveSplitResultViaCoordinator( return distribution.CatalogSnapshot{}, grpcStatusErrorf(codes.Internal, "build split mutations: %v", err) } if _, err := s.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ - Elems: ops, - IsTxn: true, - StartTS: readTS, + Elems: ops, + IsTxn: true, + StartTS: readTS, + ReadKeys: readKeys, }); err != nil { + if errors.Is(err, store.ErrWriteConflict) { + return distribution.CatalogSnapshot{}, grpcStatusError(codes.Aborted, errDistributionCatalogConflict.Error()) + } return distribution.CatalogSnapshot{}, grpcStatusErrorf(codes.Internal, "commit split mutations: %v", err) } return s.loadCatalogSnapshotAtLeastVersion(ctx, nextVersion) @@ -380,22 +387,36 @@ func validateSplitKey(parent distribution.RouteDescriptor, splitKey []byte) erro return nil } -func (s *DistributionServer) rejectLiveSplitJobOverlap(ctx context.Context, snapshot distribution.CatalogSnapshot, parent distribution.RouteDescriptor) error { +func (s *DistributionServer) splitJobOverlapReadKeys(ctx context.Context, snapshot distribution.CatalogSnapshot, parent distribution.RouteDescriptor) ([][]byte, error) { jobs, err := s.catalog.ListSplitJobsAt(ctx, snapshot.ReadTS) if err != nil { - return grpcStatusErrorf(codes.Internal, "load split jobs: %v", err) + return nil, grpcStatusErrorf(codes.Internal, "load split jobs: %v", err) } + readKeys := splitJobReadFenceKeys(jobs) for _, job := range jobs { if !splitJobIsLive(job) { continue } for _, interval := range liveSplitJobIntervals(job, snapshot.Routes) { if routeRangeIntersects(parent.Start, parent.End, interval.start, interval.end) { - return grpcStatusError(codes.Aborted, distribution.ErrSplitJobOverlap.Error()) + return nil, grpcStatusError(codes.Aborted, distribution.ErrSplitJobOverlap.Error()) } } } - return nil + return readKeys, nil +} + +func splitJobReadFenceKeys(jobs []distribution.SplitJob) [][]byte { + readKeys := make([][]byte, 0, len(jobs)+1) + readKeys = append(readKeys, distribution.CatalogNextSplitJobIDKey()) + for _, job := range jobs { + if job.TerminalAtMs > 0 { + readKeys = append(readKeys, distribution.CatalogSplitJobHistoryKey(job.TerminalAtMs, job.JobID)) + continue + } + readKeys = append(readKeys, distribution.CatalogSplitJobKey(job.JobID)) + } + return readKeys } func splitJobIsLive(job distribution.SplitJob) bool { diff --git a/adapter/distribution_server_test.go b/adapter/distribution_server_test.go index 392cb581b..e69171029 100644 --- a/adapter/distribution_server_test.go +++ b/adapter/distribution_server_test.go @@ -1,6 +1,7 @@ package adapter import ( + "bytes" "context" "testing" "time" @@ -338,6 +339,47 @@ func TestDistributionServerSplitRange_AllowsDisjointRouteWhileSplitJobLive(t *te require.NoError(t, err) require.Equal(t, uint64(2), resp.CatalogVersion) require.Equal(t, 1, coordinator.dispatchCalls) + requireReadKeysContain(t, coordinator.lastReadKeys, distribution.CatalogNextSplitJobIDKey()) + requireReadKeysContain(t, coordinator.lastReadKeys, distribution.CatalogSplitJobKey(10)) +} + +func TestDistributionServerSplitRange_ConflictsWhenSplitJobCreatedAfterOverlapScan(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}, + {RouteID: 2, Start: []byte("m"), End: nil, GroupID: 2, State: distribution.RouteStateActive}, + }) + require.NoError(t, err) + + coordinator := newDistributionCoordinatorStub(baseStore, true) + coordinator.beforeApply = func(ctx context.Context, _ store.MVCCStore) error { + return catalog.CreateSplitJob(ctx, distribution.SplitJob{ + JobID: 10, + SourceRouteID: 1, + SplitKey: []byte("g"), + TargetGroupID: 8, + Phase: distribution.SplitJobPhaseBackfill, + }) + } + s := NewDistributionServer(distribution.NewEngine(), catalog, WithDistributionCoordinator(coordinator)) + + _, err = s.SplitRange(ctx, &pb.SplitRangeRequest{ + ExpectedCatalogVersion: saved.Version, + RouteId: 1, + SplitKey: []byte("c"), + }) + require.Error(t, err) + require.Equal(t, codes.Aborted, status.Code(err)) + require.ErrorContains(t, err, errDistributionCatalogConflict.Error()) + require.Equal(t, 1, coordinator.dispatchCalls) + + snapshot, err := catalog.Snapshot(ctx) + require.NoError(t, err) + require.Equal(t, saved.Version, snapshot.Version) } func TestDistributionServerSplitRange_UsesCoordinatorForCatalogWrites(t *testing.T) { @@ -669,6 +711,8 @@ type distributionCoordinatorStub struct { leader bool nextTS uint64 lastStartTS uint64 + lastReadKeys [][]byte + beforeApply func(context.Context, store.MVCCStore) error afterDispatch func(context.Context, store.MVCCStore, uint64) error asyncApplyDone chan error asyncApplyDelay time.Duration @@ -689,6 +733,8 @@ func (s *distributionCoordinatorStub) Dispatch(ctx context.Context, reqs *kv.Ope s.dispatchCalls++ startTS, commitTS := s.nextTimestamps(reqs.StartTS) s.lastStartTS = startTS + readKeys := cloneDistributionReadKeys(reqs.ReadKeys) + s.lastReadKeys = readKeys mutations, err := coordinatorStubMutations(reqs.Elems) if err != nil { @@ -699,14 +745,14 @@ func (s *distributionCoordinatorStub) Dispatch(ctx context.Context, reqs *kv.Ope delay := s.asyncApplyDelay 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, readKeys, startTS, commitTS); err != nil { return nil, err } return &kv.CoordinateResponse{CommitIndex: commitTS}, nil @@ -737,10 +783,16 @@ 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 s.beforeApply != nil { + if err := s.beforeApply(ctx, s.store); err != nil { + return err + } + } + if err := s.store.ApplyMutations(ctx, mutations, readKeys, startTS, commitTS); err != nil { return err } if s.afterDispatch != nil { @@ -751,6 +803,27 @@ func (s *distributionCoordinatorStub) applyDispatch( return nil } +func cloneDistributionReadKeys(in [][]byte) [][]byte { + if len(in) == 0 { + return nil + } + out := make([][]byte, len(in)) + for i := range in { + out[i] = distribution.CloneBytes(in[i]) + } + return out +} + +func requireReadKeysContain(t *testing.T, readKeys [][]byte, want []byte) { + t.Helper() + for _, key := range readKeys { + if bytes.Equal(key, want) { + return + } + } + t.Fatalf("expected read keys to contain %q, got %q", want, readKeys) +} + func coordinatorStubMutations(elems []*kv.Elem[kv.OP]) ([]*store.KVPairMutation, error) { mutations := make([]*store.KVPairMutation, 0, len(elems)) for _, elem := range elems { diff --git a/adapter/dynamodb_transact.go b/adapter/dynamodb_transact.go index 3f58a688a..7ccaaac59 100644 --- a/adapter/dynamodb_transact.go +++ b/adapter/dynamodb_transact.go @@ -1105,7 +1105,7 @@ func (d *DynamoDBServer) resolveTransactTableSchema(ctx context.Context, cache m } func isRetryableTransactWriteError(err error) bool { - return errors.Is(err, store.ErrWriteConflict) || errors.Is(err, kv.ErrTxnLocked) + return errors.Is(err, store.ErrWriteConflict) || errors.Is(err, kv.ErrTxnLocked) || errors.Is(err, kv.ErrRouteWriteFenced) } func waitTransactRetryBackoff(ctx context.Context, delay time.Duration) error { diff --git a/adapter/redis_retry.go b/adapter/redis_retry.go index 206bc0930..0c38ddf41 100644 --- a/adapter/redis_retry.go +++ b/adapter/redis_retry.go @@ -41,7 +41,7 @@ var ( ) func isRetryableRedisTxnErr(err error) bool { - return errors.Is(err, store.ErrWriteConflict) || errors.Is(err, kv.ErrTxnLocked) + return errors.Is(err, store.ErrWriteConflict) || errors.Is(err, kv.ErrTxnLocked) || errors.Is(err, kv.ErrRouteWriteFenced) } func retryPolicyForRedisTxnErr(err error) redisTxnRetryPolicy { diff --git a/adapter/redis_retry_test.go b/adapter/redis_retry_test.go index 7c1ce8491..b74a73277 100644 --- a/adapter/redis_retry_test.go +++ b/adapter/redis_retry_test.go @@ -353,6 +353,7 @@ func TestRetryPolicyForRedisTxnErr(t *testing.T) { require.Equal(t, redisWriteConflictRetryPolicy, retryPolicyForRedisTxnErr(store.ErrWriteConflict)) require.Equal(t, redisTxnLockedRetryPolicy, retryPolicyForRedisTxnErr(kv.ErrTxnLocked)) + require.Equal(t, redisWriteConflictRetryPolicy, retryPolicyForRedisTxnErr(kv.ErrRouteWriteFenced)) } // TestZCard_LegacyBlobZSet verifies that ZCARD inside a Lua script returns the diff --git a/adapter/retryable_write_fence_test.go b/adapter/retryable_write_fence_test.go new file mode 100644 index 000000000..fb9f9a5e0 --- /dev/null +++ b/adapter/retryable_write_fence_test.go @@ -0,0 +1,16 @@ +package adapter + +import ( + "testing" + + "github.com/bootjp/elastickv/kv" + "github.com/stretchr/testify/require" +) + +func TestWriteFenceErrorsAreAdapterRetryable(t *testing.T) { + t.Parallel() + + require.True(t, isRetryableRedisTxnErr(kv.ErrRouteWriteFenced)) + require.True(t, isRetryableS3MutationErr(kv.ErrRouteWriteFenced)) + require.True(t, isRetryableTransactWriteError(kv.ErrRouteWriteFenced)) +} diff --git a/adapter/s3.go b/adapter/s3.go index e267d7ccf..c01165c31 100644 --- a/adapter/s3.go +++ b/adapter/s3.go @@ -2459,7 +2459,7 @@ func (s *S3Server) nextTxnCommitTS(ctx context.Context, startTS uint64) (uint64, } func isRetryableS3MutationErr(err error) bool { - return errors.Is(err, store.ErrWriteConflict) || errors.Is(err, kv.ErrTxnLocked) + return errors.Is(err, store.ErrWriteConflict) || errors.Is(err, kv.ErrTxnLocked) || errors.Is(err, kv.ErrRouteWriteFenced) } func waitS3RetryBackoff(ctx context.Context, delay time.Duration) bool { diff --git a/distribution/engine.go b/distribution/engine.go index 972bc2c4d..96d9a4c8e 100644 --- a/distribution/engine.go +++ b/distribution/engine.go @@ -424,12 +424,15 @@ func (e *Engine) GetIntersectingRoutes(start, end []byte) []Route { func cloneRoute(r Route) Route { return Route{ - RouteID: r.RouteID, - Start: CloneBytes(r.Start), - End: CloneBytes(r.End), - GroupID: r.GroupID, - State: r.State, - Load: r.Load, + RouteID: r.RouteID, + Start: CloneBytes(r.Start), + End: CloneBytes(r.End), + GroupID: r.GroupID, + State: r.State, + StagedVisibilityActive: r.StagedVisibilityActive, + MigrationJobID: r.MigrationJobID, + MinWriteTSExclusive: r.MinWriteTSExclusive, + Load: r.Load, } } diff --git a/distribution/engine_test.go b/distribution/engine_test.go index c464ba3fd..b777cb0f3 100644 --- a/distribution/engine_test.go +++ b/distribution/engine_test.go @@ -126,6 +126,21 @@ func TestEngineApplySnapshot_PreservesMigrationRouteFields(t *testing.T) { t.Fatalf("expected 1 intersecting route, got %d", len(intersections)) } requireMigrationRouteFields(t, "GetIntersectingRoutes", intersections[0]) + + snapshot, ok := e.Current() + if !ok { + t.Fatal("expected current history snapshot") + } + historyRoute, ok := snapshot.RouteOf([]byte("m")) + if !ok { + t.Fatal("expected history route") + } + requireMigrationRouteFields(t, "RouteHistorySnapshot.RouteOf", historyRoute) + historyIntersections := snapshot.IntersectingRoutes([]byte("b"), []byte("c")) + if len(historyIntersections) != 1 { + t.Fatalf("expected 1 history intersecting route, got %d", len(historyIntersections)) + } + requireMigrationRouteFields(t, "RouteHistorySnapshot.IntersectingRoutes", historyIntersections[0]) } func requireMigrationRouteFields(t *testing.T, label string, route Route) { From 8d111618e5b0caa8d451f48676b160c45b6b83d9 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 23:19:21 +0900 Subject: [PATCH 06/26] Guard snapshot spooling disk headroom --- internal/raftengine/etcd/snapshot_spool.go | 70 +++++++++++++++++-- .../etcd/snapshot_spool_space_other.go | 9 +++ .../etcd/snapshot_spool_space_unix.go | 30 ++++++++ .../raftengine/etcd/snapshot_spool_test.go | 23 ++++++ 4 files changed, 125 insertions(+), 7 deletions(-) create mode 100644 internal/raftengine/etcd/snapshot_spool_space_other.go create mode 100644 internal/raftengine/etcd/snapshot_spool_space_unix.go diff --git a/internal/raftengine/etcd/snapshot_spool.go b/internal/raftengine/etcd/snapshot_spool.go index 4065c4526..40afbead3 100644 --- a/internal/raftengine/etcd/snapshot_spool.go +++ b/internal/raftengine/etcd/snapshot_spool.go @@ -4,6 +4,7 @@ import ( "encoding/binary" "io" "log/slog" + "math" "os" "path/filepath" "strconv" @@ -30,6 +31,10 @@ const defaultMaxSnapshotPayloadBytes int64 = 16 << 30 // 16 GiB const maxSnapshotPayloadBytesEnvVar = "ELASTICKV_RAFT_MAX_SNAPSHOT_PAYLOAD_BYTES" +const defaultSnapshotSpoolMinFreeBytes int64 = 2 << 30 // 2 GiB + +const snapshotSpoolMinFreeBytesEnvVar = "ELASTICKV_RAFT_SNAPSHOT_SPOOL_MIN_FREE_BYTES" + // resolveMaxSnapshotPayloadBytes evaluates the env override once per spool // creation. Snapshots are infrequent enough that one Getenv + ParseInt per // spool is invisible in profiles, and resolving at construction means tests @@ -48,8 +53,24 @@ func resolveMaxSnapshotPayloadBytes() int64 { return n } +func resolveSnapshotSpoolMinFreeBytes() int64 { + v := strings.TrimSpace(os.Getenv(snapshotSpoolMinFreeBytesEnvVar)) + if v == "" { + return defaultSnapshotSpoolMinFreeBytes + } + n, err := strconv.ParseInt(v, 10, 64) + if err != nil || n < 0 { + slog.Warn("invalid ELASTICKV_RAFT_SNAPSHOT_SPOOL_MIN_FREE_BYTES; using default", + "value", v, "default_bytes", defaultSnapshotSpoolMinFreeBytes) + return defaultSnapshotSpoolMinFreeBytes + } + return n +} + var errSnapshotPayloadTooLarge = errors.New("etcd raft snapshot payload exceeds limit") +var errSnapshotSpoolDiskHeadroom = errors.New("etcd raft snapshot spool insufficient disk headroom") + // snapshotSyncDir indirects fsync-on-directory through a package var so a // fault-injection test can simulate "rename succeeded but the fsync that // would persist the directory entry failed" — that's the partial-failure @@ -57,13 +78,16 @@ var errSnapshotPayloadTooLarge = errors.New("etcd raft snapshot payload exceeds // and without an injection seam there's no portable way to reproduce it. var snapshotSyncDir = syncDir +var snapshotSpoolAvailableBytes = snapshotSpoolAvailableBytesFS + const snapshotSpoolPattern = "elastickv-etcd-snapshot-*" type snapshotSpool struct { - file *os.File - path string - size int64 - maxSize int64 + file *os.File + path string + size int64 + maxSize int64 + minFreeBytes int64 } func newSnapshotSpool(dir string) (*snapshotSpool, error) { @@ -72,9 +96,10 @@ func newSnapshotSpool(dir string) (*snapshotSpool, error) { return nil, errors.WithStack(err) } return &snapshotSpool{ - file: file, - path: file.Name(), - maxSize: resolveMaxSnapshotPayloadBytes(), + file: file, + path: file.Name(), + maxSize: resolveMaxSnapshotPayloadBytes(), + minFreeBytes: resolveSnapshotSpoolMinFreeBytes(), }, nil } @@ -87,6 +112,9 @@ func (s *snapshotSpool) Write(p []byte) (int, error) { if int64(len(p)) > s.maxSize-s.size { return 0, errors.Wrapf(errSnapshotPayloadTooLarge, "adding %d bytes to current %d would exceed limit %d", len(p), s.size, s.maxSize) } + if err := s.checkDiskHeadroom(len(p)); err != nil { + return 0, err + } n, err := s.file.Write(p) s.size += int64(n) if err != nil { @@ -95,6 +123,31 @@ func (s *snapshotSpool) Write(p []byte) (int, error) { return n, nil } +func (s *snapshotSpool) checkDiskHeadroom(writeBytes int) error { + if writeBytes <= 0 { + return nil + } + available, err := snapshotSpoolAvailableBytes(filepath.Dir(s.path)) + if err != nil { + return errors.WithStack(err) + } + if available < 0 { + available = 0 + } + if s.minFreeBytes > math.MaxInt64-int64(writeBytes) { + return errors.Wrapf(errSnapshotSpoolDiskHeadroom, + "write_bytes=%d min_free_bytes=%d available_bytes=%d", + writeBytes, s.minFreeBytes, available) + } + required := s.minFreeBytes + int64(writeBytes) + if available < required { + return errors.Wrapf(errSnapshotSpoolDiskHeadroom, + "write_bytes=%d min_free_bytes=%d available_bytes=%d", + writeBytes, s.minFreeBytes, available) + } + return nil +} + func (s *snapshotSpool) Bytes() ([]byte, error) { if _, err := s.file.Seek(0, io.SeekStart); err != nil { return nil, errors.WithStack(err) @@ -160,6 +213,9 @@ func (s *snapshotSpool) FinalizeAsFSMFile(fsmSnapDir string, index uint64, crc32 // so Close() is a no-op. Without the per-step clear, Close() // would attempt os.Remove(s.path) and surface a misleading // ErrNotExist that buries the real syncDir error. + if err := s.checkDiskHeadroom(fsmFooterSize); err != nil { + return err + } if err := binary.Write(s.file, binary.BigEndian, crc32c); err != nil { return errors.WithStack(err) } diff --git a/internal/raftengine/etcd/snapshot_spool_space_other.go b/internal/raftengine/etcd/snapshot_spool_space_other.go new file mode 100644 index 000000000..0e5a1ba33 --- /dev/null +++ b/internal/raftengine/etcd/snapshot_spool_space_other.go @@ -0,0 +1,9 @@ +//go:build !(aix || darwin || dragonfly || freebsd || linux || netbsd || openbsd || solaris) + +package etcd + +import "math" + +func snapshotSpoolAvailableBytesFS(string) (int64, error) { + return math.MaxInt64, nil +} diff --git a/internal/raftengine/etcd/snapshot_spool_space_unix.go b/internal/raftengine/etcd/snapshot_spool_space_unix.go new file mode 100644 index 000000000..194d46e86 --- /dev/null +++ b/internal/raftengine/etcd/snapshot_spool_space_unix.go @@ -0,0 +1,30 @@ +//go:build aix || darwin || dragonfly || freebsd || linux || netbsd || openbsd || solaris + +package etcd + +import ( + "math" + + "github.com/cockroachdb/errors" + "golang.org/x/sys/unix" +) + +func snapshotSpoolAvailableBytesFS(path string) (int64, error) { + var st unix.Statfs_t + if err := unix.Statfs(path, &st); err != nil { + return 0, errors.WithStack(err) + } + if st.Bsize <= 0 { + return 0, nil + } + blockSize := uint64(st.Bsize) + availableBlocks := st.Bavail + if availableBlocks > uint64(math.MaxInt64)/blockSize { + return math.MaxInt64, nil + } + availableBytes := availableBlocks * blockSize + if availableBytes > uint64(math.MaxInt64) { + return math.MaxInt64, nil + } + return int64(availableBytes), nil +} diff --git a/internal/raftengine/etcd/snapshot_spool_test.go b/internal/raftengine/etcd/snapshot_spool_test.go index c8ed62474..01f9edf70 100644 --- a/internal/raftengine/etcd/snapshot_spool_test.go +++ b/internal/raftengine/etcd/snapshot_spool_test.go @@ -95,6 +95,29 @@ func TestSnapshotSpool_OverrideInvalidFallsBack(t *testing.T) { require.Equal(t, defaultMaxSnapshotPayloadBytes, spool.maxSize) } +func TestSnapshotSpoolRejectsWhenReserveWouldBeConsumed(t *testing.T) { + t.Setenv(snapshotSpoolMinFreeBytesEnvVar, "1024") + originalAvailable := snapshotSpoolAvailableBytes + snapshotSpoolAvailableBytes = func(string) (int64, error) { + return 1024, nil + } + t.Cleanup(func() { snapshotSpoolAvailableBytes = originalAvailable }) + + spool, err := newSnapshotSpool(t.TempDir()) + require.NoError(t, err) + t.Cleanup(func() { _ = spool.Close() }) + + n, err := spool.Write([]byte("x")) + require.Zero(t, n) + require.Error(t, err) + require.True(t, errors.Is(err, errSnapshotSpoolDiskHeadroom), "got %v", err) + require.Zero(t, spool.size) + + info, statErr := spool.file.Stat() + require.NoError(t, statErr) + require.Zero(t, info.Size(), "headroom rejection must happen before writing bytes") +} + // TestFinalizeAsFSMFile_PostFinalizeCloseIsNoop pins the gemini-medium // review on PR #747: after a successful FinalizeAsFSMFile, the deferred // caller-side spool.Close() must NOT attempt to remove the renamed file From 0215c389fcf1d43f027b83418dd76e36df9870be Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 23:24:18 +0900 Subject: [PATCH 07/26] Raise snapshot spool headroom reserve --- internal/raftengine/etcd/snapshot_spool.go | 8 +++++++- internal/raftengine/etcd/snapshot_spool_test.go | 10 ++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/internal/raftengine/etcd/snapshot_spool.go b/internal/raftengine/etcd/snapshot_spool.go index 40afbead3..c6e569a63 100644 --- a/internal/raftengine/etcd/snapshot_spool.go +++ b/internal/raftengine/etcd/snapshot_spool.go @@ -31,7 +31,13 @@ const defaultMaxSnapshotPayloadBytes int64 = 16 << 30 // 16 GiB const maxSnapshotPayloadBytesEnvVar = "ELASTICKV_RAFT_MAX_SNAPSHOT_PAYLOAD_BYTES" -const defaultSnapshotSpoolMinFreeBytes int64 = 2 << 30 // 2 GiB +// Keep one full max-size snapshot worth of headroom after receive-side spooling. +// Applying a token snapshot restores the FSM from the completed .fsm into a new +// local store directory, so a node can transiently need old snapshot + incoming +// .fsm + restored store bytes. A small reserve lets the spool complete and only +// fails later in restore, which can leave the host near-full until startup +// orphan cleanup runs. +const defaultSnapshotSpoolMinFreeBytes = defaultMaxSnapshotPayloadBytes const snapshotSpoolMinFreeBytesEnvVar = "ELASTICKV_RAFT_SNAPSHOT_SPOOL_MIN_FREE_BYTES" diff --git a/internal/raftengine/etcd/snapshot_spool_test.go b/internal/raftengine/etcd/snapshot_spool_test.go index 01f9edf70..21e8c2425 100644 --- a/internal/raftengine/etcd/snapshot_spool_test.go +++ b/internal/raftengine/etcd/snapshot_spool_test.go @@ -22,6 +22,7 @@ func TestSnapshotSpool_DefaultCapAcceptsRealisticFSM(t *testing.T) { if testing.Short() { t.Skip("skipping: writes 1.5 GiB to a temp file") } + t.Setenv(snapshotSpoolMinFreeBytesEnvVar, "0") dir := t.TempDir() spool, err := newSnapshotSpool(dir) require.NoError(t, err) @@ -66,6 +67,7 @@ func TestSnapshotSpool_DefaultCapAcceptsRealisticFSM(t *testing.T) { func TestSnapshotSpool_OverrideViaEnv(t *testing.T) { const spoolCap = int64(4096) t.Setenv(maxSnapshotPayloadBytesEnvVar, strconv.FormatInt(spoolCap, 10)) + t.Setenv(snapshotSpoolMinFreeBytesEnvVar, "0") spool, err := newSnapshotSpool(t.TempDir()) require.NoError(t, err) @@ -95,6 +97,14 @@ func TestSnapshotSpool_OverrideInvalidFallsBack(t *testing.T) { require.Equal(t, defaultMaxSnapshotPayloadBytes, spool.maxSize) } +func TestSnapshotSpoolDefaultReserveTracksPayloadCap(t *testing.T) { + spool, err := newSnapshotSpool(t.TempDir()) + require.NoError(t, err) + t.Cleanup(func() { _ = spool.Close() }) + + require.Equal(t, defaultMaxSnapshotPayloadBytes, spool.minFreeBytes) +} + func TestSnapshotSpoolRejectsWhenReserveWouldBeConsumed(t *testing.T) { t.Setenv(snapshotSpoolMinFreeBytesEnvVar, "1024") originalAvailable := snapshotSpoolAvailableBytes From 774779b109682ecffeb16c8fe4108d92b3153917 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 23:53:42 +0900 Subject: [PATCH 08/26] Fix raft startup and snapshot spool guards --- adapter/distribution_server.go | 6 +- adapter/distribution_server_test.go | 36 ++++++++ internal/raftengine/etcd/engine.go | 19 +++- internal/raftengine/etcd/engine_test.go | 90 +++++++++++++++++++ internal/raftengine/etcd/grpc_transport.go | 2 +- internal/raftengine/etcd/snapshot_spool.go | 35 ++++---- .../etcd/snapshot_spool_space_other.go | 2 +- .../etcd/snapshot_spool_space_unix.go | 2 +- .../raftengine/etcd/snapshot_spool_test.go | 26 +++++- 9 files changed, 193 insertions(+), 25 deletions(-) diff --git a/adapter/distribution_server.go b/adapter/distribution_server.go index 0fcacba58..3d3276f5c 100644 --- a/adapter/distribution_server.go +++ b/adapter/distribution_server.go @@ -410,11 +410,9 @@ func splitJobReadFenceKeys(jobs []distribution.SplitJob) [][]byte { readKeys := make([][]byte, 0, len(jobs)+1) readKeys = append(readKeys, distribution.CatalogNextSplitJobIDKey()) for _, job := range jobs { - if job.TerminalAtMs > 0 { - readKeys = append(readKeys, distribution.CatalogSplitJobHistoryKey(job.TerminalAtMs, job.JobID)) - continue + if splitJobIsLive(job) { + readKeys = append(readKeys, distribution.CatalogSplitJobKey(job.JobID)) } - readKeys = append(readKeys, distribution.CatalogSplitJobKey(job.JobID)) } return readKeys } diff --git a/adapter/distribution_server_test.go b/adapter/distribution_server_test.go index e69171029..f7cc59f17 100644 --- a/adapter/distribution_server_test.go +++ b/adapter/distribution_server_test.go @@ -343,6 +343,33 @@ func TestDistributionServerSplitRange_AllowsDisjointRouteWhileSplitJobLive(t *te requireReadKeysContain(t, coordinator.lastReadKeys, distribution.CatalogSplitJobKey(10)) } +func TestSplitJobReadFenceKeysExcludesTerminalHistory(t *testing.T) { + t.Parallel() + + readKeys := splitJobReadFenceKeys([]distribution.SplitJob{ + { + JobID: 10, + Phase: distribution.SplitJobPhaseBackfill, + }, + { + JobID: 11, + Phase: distribution.SplitJobPhaseDone, + TerminalAtMs: 1000, + }, + { + JobID: 12, + Phase: distribution.SplitJobPhaseAbandoned, + TerminalAtMs: 1001, + }, + }) + + require.Len(t, readKeys, 2) + requireReadKeysContain(t, readKeys, distribution.CatalogNextSplitJobIDKey()) + requireReadKeysContain(t, readKeys, distribution.CatalogSplitJobKey(10)) + requireReadKeysNotContain(t, readKeys, distribution.CatalogSplitJobHistoryKey(1000, 11)) + requireReadKeysNotContain(t, readKeys, distribution.CatalogSplitJobHistoryKey(1001, 12)) +} + func TestDistributionServerSplitRange_ConflictsWhenSplitJobCreatedAfterOverlapScan(t *testing.T) { t.Parallel() @@ -824,6 +851,15 @@ func requireReadKeysContain(t *testing.T, readKeys [][]byte, want []byte) { t.Fatalf("expected read keys to contain %q, got %q", want, readKeys) } +func requireReadKeysNotContain(t *testing.T, readKeys [][]byte, want []byte) { + t.Helper() + for _, key := range readKeys { + if bytes.Equal(key, want) { + t.Fatalf("expected read keys not to contain %q, got %q", want, readKeys) + } + } +} + func coordinatorStubMutations(elems []*kv.Elem[kv.OP]) ([]*store.KVPairMutation, error) { mutations := make([]*store.KVPairMutation, 0, len(elems)) for _, elem := range elems { diff --git a/internal/raftengine/etcd/engine.go b/internal/raftengine/etcd/engine.go index df65236b7..7b88b06a8 100644 --- a/internal/raftengine/etcd/engine.go +++ b/internal/raftengine/etcd/engine.go @@ -351,10 +351,15 @@ type Engine struct { snapshotStopCh chan struct{} closeCh chan struct{} doneCh chan struct{} - startedCh chan struct{} + // openCh is closed when run starts and callers may bind public/raft + // listeners. startedCh remains the stronger gate for processing inbound + // raft messages after the startup Ready drain has completed. + openCh chan struct{} + startedCh chan struct{} leaderReady chan struct{} leaderOnce sync.Once + openOnce sync.Once startOnce sync.Once closeOnce sync.Once dispatchOnce sync.Once @@ -656,6 +661,7 @@ func Open(ctx context.Context, cfg OpenConfig) (*Engine, error) { dispatchReportCh: make(chan dispatchReport, inboundQueueCap), closeCh: make(chan struct{}), doneCh: make(chan struct{}), + openCh: make(chan struct{}), startedCh: make(chan struct{}), leaderReady: make(chan struct{}), config: configurationFromConfState(peerMap, confStateValue(prepared.disk.LocalSnap.GetMetadata().GetConfState())), @@ -1685,6 +1691,7 @@ func (e *Engine) run() { ticker := time.NewTicker(e.tickInterval) defer ticker.Stop() + e.markOpen() if err := e.startup(); err != nil { e.fail(err) return @@ -3582,10 +3589,20 @@ func (e *Engine) markStarted() { e.startOnce.Do(func() { close(e.startedCh) }) } +func (e *Engine) markOpen() { + if e.openCh == nil { + return + } + e.openOnce.Do(func() { close(e.openCh) }) +} + func (e *Engine) openReady(waitForLeader bool) <-chan struct{} { if waitForLeader { return e.leaderReady } + if e.openCh != nil { + return e.openCh + } return e.startedCh } diff --git a/internal/raftengine/etcd/engine_test.go b/internal/raftengine/etcd/engine_test.go index 00aa807c2..4139728ff 100644 --- a/internal/raftengine/etcd/engine_test.go +++ b/internal/raftengine/etcd/engine_test.go @@ -108,6 +108,12 @@ type blockingSnapshotStateMachine struct { release chan struct{} } +type blockingApplyStateMachine struct { + started chan struct{} + release chan struct{} + startOnce sync.Once +} + type blockingSnapshot struct { started chan struct{} release chan struct{} @@ -141,6 +147,21 @@ func (s *blockingSnapshotStateMachine) Apply(data []byte) any { return string(data) } +func (s *blockingApplyStateMachine) Apply(data []byte) any { + s.startOnce.Do(func() { close(s.started) }) + <-s.release + return string(data) +} + +func (s *blockingApplyStateMachine) Snapshot() (Snapshot, error) { + return &testSnapshot{}, nil +} + +func (s *blockingApplyStateMachine) Restore(r io.Reader) error { + _, err := io.Copy(io.Discard, r) + return err +} + func (s *blockingSnapshotStateMachine) Snapshot() (Snapshot, error) { return &blockingSnapshot{ started: s.started, @@ -1409,6 +1430,75 @@ func TestOpenRestoresLegacySnapshotState(t *testing.T) { require.Equal(t, [][]byte{[]byte("snap"), []byte("tail")}, fsm.Applied()) } +func TestOpenMultiNodeReturnsBeforeCommittedTailDrain(t *testing.T) { + dir := t.TempDir() + peers := []Peer{ + {NodeID: 1, ID: "n1", Address: "127.0.0.1:7001"}, + {NodeID: 2, ID: "n2", Address: "127.0.0.1:7002"}, + } + require.NoError(t, saveStateFile(stateFilePath(dir), persistedState{ + HardState: testHardState(2, 2), + Snapshot: raftTestSnapshot(1, 1, []uint64{1, 2}, mustEncodeSnapshotData(t, nil)), + Entries: []raftpb.Entry{{ + Type: entryTypePtr(raftpb.EntryNormal), + Term: uint64Ptr(2), + Index: uint64Ptr(2), + Data: encodeProposalEnvelope(1, []byte("tail")), + }}, + })) + + fsm := &blockingApplyStateMachine{ + started: make(chan struct{}), + release: make(chan struct{}), + } + type openResult struct { + engine *Engine + err error + } + done := make(chan openResult, 1) + go func() { + engine, err := Open(context.Background(), OpenConfig{ + NodeID: 1, + LocalID: "n1", + LocalAddress: "127.0.0.1:7001", + DataDir: dir, + Peers: peers, + StateMachine: fsm, + }) + done <- openResult{engine: engine, err: err} + }() + + var result openResult + select { + case result = <-done: + case <-time.After(time.Second): + t.Fatal("multi-node Open blocked on committed tail drain before returning") + } + require.NoError(t, result.err) + require.NotNil(t, result.engine) + defer func() { + require.NoError(t, result.engine.Close()) + }() + + select { + case <-fsm.started: + case <-time.After(time.Second): + t.Fatal("startup committed tail was not being applied") + } + select { + case <-result.engine.startedCh: + t.Fatal("engine marked started before startup committed tail drain completed") + default: + } + + close(fsm.release) + select { + case <-result.engine.startedCh: + case <-time.After(time.Second): + t.Fatal("engine did not mark started after committed tail drain completed") + } +} + func TestOpenMultiNodeReplicatesOverGRPCTransport(t *testing.T) { nodes, peers := newTransportTestNodes(t, 3) startTransportTestServers(nodes, peers) diff --git a/internal/raftengine/etcd/grpc_transport.go b/internal/raftengine/etcd/grpc_transport.go index d3ef7d572..32d82e83e 100644 --- a/internal/raftengine/etcd/grpc_transport.go +++ b/internal/raftengine/etcd/grpc_transport.go @@ -1231,7 +1231,7 @@ func (t *GRPCTransport) receiveSnapshotStream(stream pb.EtcdRaft_SendSnapshotSer if err != nil { return raftpb.Message{}, err } - spool, err := newSnapshotSpool(spoolPlacement) + spool, err := newReceiveSnapshotSpool(spoolPlacement) if err != nil { return raftpb.Message{}, err } diff --git a/internal/raftengine/etcd/snapshot_spool.go b/internal/raftengine/etcd/snapshot_spool.go index c6e569a63..cb1b3bc8e 100644 --- a/internal/raftengine/etcd/snapshot_spool.go +++ b/internal/raftengine/etcd/snapshot_spool.go @@ -31,14 +31,6 @@ const defaultMaxSnapshotPayloadBytes int64 = 16 << 30 // 16 GiB const maxSnapshotPayloadBytesEnvVar = "ELASTICKV_RAFT_MAX_SNAPSHOT_PAYLOAD_BYTES" -// Keep one full max-size snapshot worth of headroom after receive-side spooling. -// Applying a token snapshot restores the FSM from the completed .fsm into a new -// local store directory, so a node can transiently need old snapshot + incoming -// .fsm + restored store bytes. A small reserve lets the spool complete and only -// fails later in restore, which can leave the host near-full until startup -// orphan cleanup runs. -const defaultSnapshotSpoolMinFreeBytes = defaultMaxSnapshotPayloadBytes - const snapshotSpoolMinFreeBytesEnvVar = "ELASTICKV_RAFT_SNAPSHOT_SPOOL_MIN_FREE_BYTES" // resolveMaxSnapshotPayloadBytes evaluates the env override once per spool @@ -59,16 +51,16 @@ func resolveMaxSnapshotPayloadBytes() int64 { return n } -func resolveSnapshotSpoolMinFreeBytes() int64 { +func resolveSnapshotSpoolMinFreeBytes(defaultBytes int64) int64 { v := strings.TrimSpace(os.Getenv(snapshotSpoolMinFreeBytesEnvVar)) if v == "" { - return defaultSnapshotSpoolMinFreeBytes + return defaultBytes } n, err := strconv.ParseInt(v, 10, 64) if err != nil || n < 0 { slog.Warn("invalid ELASTICKV_RAFT_SNAPSHOT_SPOOL_MIN_FREE_BYTES; using default", - "value", v, "default_bytes", defaultSnapshotSpoolMinFreeBytes) - return defaultSnapshotSpoolMinFreeBytes + "value", v, "default_bytes", defaultBytes) + return defaultBytes } return n } @@ -97,6 +89,19 @@ type snapshotSpool struct { } func newSnapshotSpool(dir string) (*snapshotSpool, error) { + return newSnapshotSpoolWithLimits(dir, resolveMaxSnapshotPayloadBytes(), 0) +} + +func newReceiveSnapshotSpool(dir string) (*snapshotSpool, error) { + maxSize := resolveMaxSnapshotPayloadBytes() + // Keep one full max-size snapshot worth of headroom after receive-side + // spooling. Applying a token snapshot restores the FSM from the completed + // .fsm into a new local store directory, so a node can transiently need old + // snapshot + incoming .fsm + restored store bytes. + return newSnapshotSpoolWithLimits(dir, maxSize, resolveSnapshotSpoolMinFreeBytes(maxSize)) +} + +func newSnapshotSpoolWithLimits(dir string, maxSize, minFreeBytes int64) (*snapshotSpool, error) { file, err := os.CreateTemp(dir, snapshotSpoolPattern) if err != nil { return nil, errors.WithStack(err) @@ -104,8 +109,8 @@ func newSnapshotSpool(dir string) (*snapshotSpool, error) { return &snapshotSpool{ file: file, path: file.Name(), - maxSize: resolveMaxSnapshotPayloadBytes(), - minFreeBytes: resolveSnapshotSpoolMinFreeBytes(), + maxSize: maxSize, + minFreeBytes: minFreeBytes, }, nil } @@ -130,7 +135,7 @@ func (s *snapshotSpool) Write(p []byte) (int, error) { } func (s *snapshotSpool) checkDiskHeadroom(writeBytes int) error { - if writeBytes <= 0 { + if writeBytes <= 0 || s.minFreeBytes <= 0 { return nil } available, err := snapshotSpoolAvailableBytes(filepath.Dir(s.path)) diff --git a/internal/raftengine/etcd/snapshot_spool_space_other.go b/internal/raftengine/etcd/snapshot_spool_space_other.go index 0e5a1ba33..f9d733c4a 100644 --- a/internal/raftengine/etcd/snapshot_spool_space_other.go +++ b/internal/raftengine/etcd/snapshot_spool_space_other.go @@ -1,4 +1,4 @@ -//go:build !(aix || darwin || dragonfly || freebsd || linux || netbsd || openbsd || solaris) +//go:build !(aix || darwin || linux) package etcd diff --git a/internal/raftengine/etcd/snapshot_spool_space_unix.go b/internal/raftengine/etcd/snapshot_spool_space_unix.go index 194d46e86..a2c5db370 100644 --- a/internal/raftengine/etcd/snapshot_spool_space_unix.go +++ b/internal/raftengine/etcd/snapshot_spool_space_unix.go @@ -1,4 +1,4 @@ -//go:build aix || darwin || dragonfly || freebsd || linux || netbsd || openbsd || solaris +//go:build aix || darwin || linux package etcd diff --git a/internal/raftengine/etcd/snapshot_spool_test.go b/internal/raftengine/etcd/snapshot_spool_test.go index 21e8c2425..bed5fd223 100644 --- a/internal/raftengine/etcd/snapshot_spool_test.go +++ b/internal/raftengine/etcd/snapshot_spool_test.go @@ -98,11 +98,33 @@ func TestSnapshotSpool_OverrideInvalidFallsBack(t *testing.T) { } func TestSnapshotSpoolDefaultReserveTracksPayloadCap(t *testing.T) { + const spoolCap = int64(4096) + t.Setenv(maxSnapshotPayloadBytesEnvVar, strconv.FormatInt(spoolCap, 10)) + + spool, err := newReceiveSnapshotSpool(t.TempDir()) + require.NoError(t, err) + t.Cleanup(func() { _ = spool.Close() }) + + require.Equal(t, spoolCap, spool.maxSize) + require.Equal(t, spoolCap, spool.minFreeBytes) +} + +func TestSnapshotSpoolMaterializeDoesNotReserveDiskHeadroom(t *testing.T) { + t.Setenv(snapshotSpoolMinFreeBytesEnvVar, "1024") + originalAvailable := snapshotSpoolAvailableBytes + snapshotSpoolAvailableBytes = func(string) (int64, error) { + return 0, nil + } + t.Cleanup(func() { snapshotSpoolAvailableBytes = originalAvailable }) + spool, err := newSnapshotSpool(t.TempDir()) require.NoError(t, err) t.Cleanup(func() { _ = spool.Close() }) - require.Equal(t, defaultMaxSnapshotPayloadBytes, spool.minFreeBytes) + n, err := spool.Write([]byte("x")) + require.NoError(t, err) + require.Equal(t, 1, n) + require.Equal(t, int64(1), spool.size) } func TestSnapshotSpoolRejectsWhenReserveWouldBeConsumed(t *testing.T) { @@ -113,7 +135,7 @@ func TestSnapshotSpoolRejectsWhenReserveWouldBeConsumed(t *testing.T) { } t.Cleanup(func() { snapshotSpoolAvailableBytes = originalAvailable }) - spool, err := newSnapshotSpool(t.TempDir()) + spool, err := newReceiveSnapshotSpool(t.TempDir()) require.NoError(t, err) t.Cleanup(func() { _ = spool.Close() }) From 46722f6a436c72cb84e9ba8e1b1db921784a6fde Mon Sep 17 00:00:00 2001 From: bootjp Date: Wed, 15 Jul 2026 00:04:17 +0900 Subject: [PATCH 09/26] Handle route fences in adapter retries --- adapter/redis_exec_dedup_test.go | 25 +++++++++++ adapter/redis_list_dedup_test.go | 56 +++++++++++++++++++---- adapter/redis_lists.go | 17 ++++--- adapter/redis_retry.go | 4 ++ adapter/redis_txn.go | 20 ++++----- adapter/s3.go | 9 ++-- adapter/s3_admin.go | 29 ++++++------ adapter/s3_admin_test.go | 76 ++++++++++++++++++++++++++++++++ 8 files changed, 188 insertions(+), 48 deletions(-) diff --git a/adapter/redis_exec_dedup_test.go b/adapter/redis_exec_dedup_test.go index 40bc6956b..ea59634b4 100644 --- a/adapter/redis_exec_dedup_test.go +++ b/adapter/redis_exec_dedup_test.go @@ -47,6 +47,31 @@ func TestExecDedup_LandedPriorAttempt_ReturnsCachedResults(t *testing.T) { require.Equal(t, []byte("v1"), val) } +func TestExecDedup_RouteFenceRetryPreservesPriorProbe(t *testing.T) { + t.Parallel() + ctx := context.Background() + st := store.NewMVCCStore() + coord := newDedupTestCoordinator(st, 1, true) + coord.routeFenceAtDispatch = 2 + srv := &RedisServer{store: st, coordinator: coord, scriptCache: map[string]string{}, onePhaseTxnDedup: true} + + queue := []redcon.Command{ + {Args: [][]byte{[]byte(cmdSet), []byte("k"), []byte("v1")}}, + } + results, err := srv.runTransaction(queue) + require.NoError(t, err) + require.Len(t, results, 1) + require.Equal(t, "OK", results[0].str) + require.Equal(t, 3, coord.dispatches) + require.Equal(t, 1, coord.probeNoOps, "route-fenced reuse must not replace the prior landed probe") + + rawVal, err := st.GetAt(ctx, redisStrKey([]byte("k")), snapshotTS(coord.Clock(), st)) + require.NoError(t, err) + val, _, err := decodeRedisStr(rawVal) + require.NoError(t, err) + require.Equal(t, []byte("v1"), val) +} + // TestExecDedup_PriorAttemptDidNotLand_Applies covers the truncated case for // MULTI/EXEC: attempt 1 errored without committing (OCC-style pre-reject), // so the probe misses and the reuse applies the same write set at a fresh diff --git a/adapter/redis_list_dedup_test.go b/adapter/redis_list_dedup_test.go index 64575e93f..5e5c46303 100644 --- a/adapter/redis_list_dedup_test.go +++ b/adapter/redis_list_dedup_test.go @@ -39,8 +39,12 @@ type dedupTestCoordinator struct { // WITHOUT applying — an ambiguous retryable error distinct from // WriteConflict, exercising the "advance pending.commitTS and retry" branch. txnLockedAtDispatch int - dispatches int - probeNoOps int + // routeFenceAtDispatch makes the named dispatch return ErrRouteWriteFenced + // before the FSM dedup probe or apply. It is retryable but cannot be + // treated as an ambiguous landing. + routeFenceAtDispatch int + dispatches int + probeNoOps int // beforeDispatch, if set, runs at the start of each Dispatch with the // 1-based dispatch number — lets a test inject a concurrent commit // between the adapter's attempts. @@ -258,16 +262,14 @@ func (c *dedupTestCoordinator) Dispatch(ctx context.Context, req *kv.OperationGr if c.beforeDispatch != nil { c.beforeDispatch(n) } + if c.shouldRouteFence(n) { + return nil, kv.ErrRouteWriteFenced + } if handled, resp, err := c.maybeProbe(ctx, req); handled { return resp, err } - if n == c.ambiguousDispatch && !c.ambiguousLands { - // OCC-style pre-reject: nothing is written, definitely did not land. - return nil, store.ErrWriteConflict - } - if n == c.txnLockedAtDispatch { - // Ambiguous lock error, nothing written: definitely did not land. - return nil, kv.ErrTxnLocked + if err := c.preApplyError(n); err != nil { + return nil, err } resp, err := c.occAdapterCoordinator.Dispatch(ctx, req) if err != nil { @@ -287,6 +289,21 @@ func (c *dedupTestCoordinator) Dispatch(ctx context.Context, req *kv.OperationGr return resp, nil } +func (c *dedupTestCoordinator) shouldRouteFence(dispatch int) bool { + return dispatch == c.routeFenceAtDispatch +} + +func (c *dedupTestCoordinator) preApplyError(dispatch int) error { + switch { + case dispatch == c.ambiguousDispatch && !c.ambiguousLands: + return store.ErrWriteConflict + case dispatch == c.txnLockedAtDispatch: + return kv.ErrTxnLocked + default: + return nil + } +} + // maybeProbe mimics handleOnePhaseTxnRequest's exact-ts dedup check. It // returns handled=true when the probe owns the response (hit → no-op success // or probe error), and handled=false to fall through to the normal apply. @@ -365,6 +382,27 @@ func TestListPushDedup_LandedPriorAttempt_NoDuplicate(t *testing.T) { require.Equal(t, []byte("v"), val) } +func TestListPushDedup_RouteFenceRetryPreservesPriorProbe(t *testing.T) { + t.Parallel() + ctx := context.Background() + st := store.NewMVCCStore() + coord := newDedupTestCoordinator(st, 1, true) + coord.routeFenceAtDispatch = 2 + srv := &RedisServer{store: st, coordinator: coord, scriptCache: map[string]string{}, onePhaseTxnDedup: true} + + key := []byte("mylist") + n, err := srv.listRPush(ctx, key, [][]byte{[]byte("v")}) + require.NoError(t, err) + require.Equal(t, int64(1), n) + require.Equal(t, 3, coord.dispatches, "attempt 1 landed, route-fenced reuse, then dedup probe retry") + require.Equal(t, 1, coord.probeNoOps, "route-fenced reuse must not replace the prior landed probe") + + readTS := snapshotTS(coord.Clock(), st) + meta, _, err := srv.resolveListMeta(ctx, key, readTS) + require.NoError(t, err) + require.Equal(t, int64(1), meta.Len, "route-fence retry must not append a duplicate") +} + // TestListPushDedup_PriorAttemptDidNotLand_Applies covers the truncated case: // attempt 1 errored without committing, so the probe misses and the reuse // applies the same write set at a fresh commit_ts. The element lands exactly diff --git a/adapter/redis_lists.go b/adapter/redis_lists.go index 4380073cd..e36cde008 100644 --- a/adapter/redis_lists.go +++ b/adapter/redis_lists.go @@ -243,13 +243,10 @@ func (r *RedisServer) dispatchListPushReuse(ctx context.Context, key []byte, pen // iteration recomputes from a fresh meta read. return 0, true, errors.WithStack(dispErr) } - // Still ambiguous (lock / other retryable): this reuse may itself - // have landed, so the next retry must probe THIS commit_ts. Only - // advance pending.commitTS if retryRedisWrite will actually loop - // (non-retryable errors escape to the client; pending is then - // discarded with the goroutine, so the update is wasted and the - // stale value would be misleading if some future caller reads it). - if isRetryableRedisTxnErr(dispErr) { + // Still ambiguous (lock / other retryable): this reuse may itself have + // landed, so the next retry must probe THIS commit_ts. Route-fence + // rejections are retryable but pre-apply, so keep the older witness. + if shouldPreserveRedisTxnAttempt(dispErr) { pending.commitTS = commitTS } return 0, false, errors.WithStack(dispErr) @@ -411,7 +408,9 @@ func (r *RedisServer) listPushCoreWithDedup(ctx context.Context, key []byte, val return nil } // Only remember the attempt for reuse if retryRedisWrite will actually - // loop — i.e. the error is one of WriteConflict / TxnLocked. For + // loop and the attempt may have landed. Route-fence rejections are + // retryable but happen before this write set can apply, so preserving + // that commitTS would overwrite an older ambiguous witness. For // errors that escape the loop (transient-leader, context deadline, // FSM apply error, etc.), `pending` would be discarded with the // goroutine, and recording it would mislead a future reader about @@ -419,7 +418,7 @@ func (r *RedisServer) listPushCoreWithDedup(ctx context.Context, key []byte, val // retryRedisWrite's retry predicate; ambiguous errors that escape // to the client are a separate problem space (cross-request // idempotency cache) and out of scope for this design. - if isRetryableRedisTxnErr(dispErr) { + if shouldPreserveRedisTxnAttempt(dispErr) { pending = &reusableListPush{ ops: ops, startTS: startTS, diff --git a/adapter/redis_retry.go b/adapter/redis_retry.go index 0c38ddf41..d6f13f0e5 100644 --- a/adapter/redis_retry.go +++ b/adapter/redis_retry.go @@ -44,6 +44,10 @@ func isRetryableRedisTxnErr(err error) bool { return errors.Is(err, store.ErrWriteConflict) || errors.Is(err, kv.ErrTxnLocked) || errors.Is(err, kv.ErrRouteWriteFenced) } +func shouldPreserveRedisTxnAttempt(err error) bool { + return isRetryableRedisTxnErr(err) && !errors.Is(err, kv.ErrRouteWriteFenced) +} + func retryPolicyForRedisTxnErr(err error) redisTxnRetryPolicy { if errors.Is(err, kv.ErrTxnLocked) { return redisTxnLockedRetryPolicy diff --git a/adapter/redis_txn.go b/adapter/redis_txn.go index 750011ca7..53b937ed2 100644 --- a/adapter/redis_txn.go +++ b/adapter/redis_txn.go @@ -2216,12 +2216,10 @@ func (r *RedisServer) dispatchExecReuse(ctx context.Context, pending *reusableEx // iteration rebuilds from a fresh snapshot. return nil, true, errors.WithStack(dispErr) } - // Still ambiguous (lock / other retryable): the reuse may itself - // have landed, so the next retry must probe THIS commit_ts. Only - // advance pending.commitTS if retryRedisWrite will actually loop - // (non-retryable errors escape to the client; pending is then - // discarded with the goroutine). - if isRetryableRedisTxnErr(dispErr) { + // Still ambiguous (lock / other retryable): the reuse may itself have + // landed, so the next retry must probe THIS commit_ts. Route-fence + // rejections are retryable but pre-apply, so keep the older witness. + if shouldPreserveRedisTxnAttempt(dispErr) { pending.commitTS = commitTS } return nil, false, errors.WithStack(dispErr) @@ -2347,12 +2345,10 @@ func (r *RedisServer) firstExecAttempt(dispatchCtx context.Context, queue []redc } if _, dispErr := r.coordinator.Dispatch(prepared.ctx, group); dispErr != nil { // Only remember the attempt for reuse if retryRedisWrite will - // actually loop. Mirrors listPushCoreWithDedup's gating - // rationale — errors that escape the loop (transient-leader, - // context deadline, FSM apply error) leave pending pointing at - // state wasted with the goroutine; ambiguous errors that - // escape to the client are out of scope for this loop. - if isRetryableRedisTxnErr(dispErr) { + // actually loop and the attempt may have landed. Mirrors + // listPushCoreWithDedup's gating rationale; route-fence + // rejections are retryable but pre-apply. + if shouldPreserveRedisTxnAttempt(dispErr) { return nil, &reusableExecTxn{ elems: prepared.elems, startTS: txn.startTS, diff --git a/adapter/s3.go b/adapter/s3.go index c01165c31..d08f0e4d6 100644 --- a/adapter/s3.go +++ b/adapter/s3.go @@ -728,9 +728,12 @@ func (s *S3Server) deleteBucket(w http.ResponseWriter, r *http.Request, bucket s writeS3MutationError(w, err, bucket, "") return } - // Phase 2: best-effort DEL_PREFIX safety net. See - // AdminDeleteBucket / runBucketDeleteSafetyNet for the contract. - s.runBucketDeleteSafetyNet(r.Context(), bucket, deletedGeneration) + // Phase 2: DEL_PREFIX safety net. See AdminDeleteBucket / + // runBucketDeleteSafetyNet for the contract. + if err := s.runBucketDeleteSafetyNet(r.Context(), bucket, deletedGeneration); err != nil { + writeS3MutationError(w, err, bucket, "") + return + } w.WriteHeader(http.StatusNoContent) } diff --git a/adapter/s3_admin.go b/adapter/s3_admin.go index a9eb8915a..ccc65bd1f 100644 --- a/adapter/s3_admin.go +++ b/adapter/s3_admin.go @@ -423,12 +423,7 @@ func (s *S3Server) AdminDeleteBucket(ctx context.Context, principal AdminPrincip if err != nil { return err //nolint:wrapcheck // sentinel errors propagate as-is. } - // Phase 2: best-effort safety-net DEL_PREFIX. Outside the - // retryS3Mutation closure because retrying after Phase 1 - // committed would 404 at loadBucketMetaAt; we want the error - // (if any) logged but not propagated to the operator. - s.runBucketDeleteSafetyNet(ctx, name, deletedGeneration) - return nil + return s.runBucketDeleteSafetyNet(ctx, name, deletedGeneration) } // adminDeleteBucketTxnBody is the per-attempt body retryS3Mutation @@ -501,22 +496,26 @@ func bucketDeleteSafetyNetElems(bucket string, generation uint64) []*kv.Elem[kv. } } -// runBucketDeleteSafetyNet runs the Phase-2 DEL_PREFIX dispatch -// and swallows transport / cluster errors after logging — the -// caller has already deleted the bucket meta and the operator- -// visible state is consistent with that. Shared between admin and -// SigV4 paths. -func (s *S3Server) runBucketDeleteSafetyNet(ctx context.Context, bucket string, generation uint64) { - if _, err := s.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ - Elems: bucketDeleteSafetyNetElems(bucket, generation), - }); err != nil { +// runBucketDeleteSafetyNet runs the Phase-2 DEL_PREFIX dispatch. Phase 1 has +// already deleted the bucket meta, so this helper retries transient fencing in +// place instead of re-entering the full delete transaction. +func (s *S3Server) runBucketDeleteSafetyNet(ctx context.Context, bucket string, generation uint64) error { + err := s.retryS3Mutation(ctx, func() error { + _, err := s.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + Elems: bucketDeleteSafetyNetElems(bucket, generation), + }) + return errors.WithStack(err) + }) + if err != nil { slog.WarnContext(ctx, "bucket delete safety-net DEL_PREFIX failed; bucket meta is gone but orphan sweep incomplete", slog.String("bucket", bucket), slog.Uint64("generation", generation), slog.String("error", err.Error()), ) + return err } + return nil } // adminCanonicalACL normalises an empty input to the canned diff --git a/adapter/s3_admin_test.go b/adapter/s3_admin_test.go index 3a83677f2..23d4549fd 100644 --- a/adapter/s3_admin_test.go +++ b/adapter/s3_admin_test.go @@ -264,6 +264,56 @@ func TestS3Server_AdminDeleteBucket_HappyPath(t *testing.T) { require.False(t, exists) } +func TestS3Server_AdminDeleteBucket_RetriesSafetyNetRouteFence(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + coord := &bucketDeleteSafetyNetFenceCoordinator{ + localAdapterCoordinator: newLocalAdapterCoordinator(st), + failuresRemaining: 1, + } + server := NewS3Server(nil, "", st, coord, nil) + ctx := context.Background() + + summary, err := server.AdminCreateBucket(ctx, + fullAdminBucketsPrincipal(), "to-delete", s3AclPrivate) + require.NoError(t, err) + orphan := s3keys.RouteKey("to-delete", summary.Generation, "orphan") + _, err = coord.localAdapterCoordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + Elems: []*kv.Elem[kv.OP]{{Op: kv.Put, Key: orphan, Value: []byte("orphan")}}, + }) + require.NoError(t, err) + + err = server.AdminDeleteBucket(ctx, + fullAdminBucketsPrincipal(), "to-delete") + require.NoError(t, err) + require.Equal(t, 2, coord.safetyNetCalls) + + _, err = st.GetAt(ctx, orphan, snapshotTS(coord.Clock(), st)) + require.ErrorIs(t, err, store.ErrKeyNotFound, "safety-net retry must sweep the orphan before acknowledging delete") +} + +func TestS3Server_AdminDeleteBucket_PropagatesPersistentSafetyNetRouteFence(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + coord := &bucketDeleteSafetyNetFenceCoordinator{ + localAdapterCoordinator: newLocalAdapterCoordinator(st), + failuresRemaining: s3TxnRetryMaxAttempts, + } + server := NewS3Server(nil, "", st, coord, nil) + ctx := context.Background() + + _, err := server.AdminCreateBucket(ctx, + fullAdminBucketsPrincipal(), "to-delete", s3AclPrivate) + require.NoError(t, err) + + err = server.AdminDeleteBucket(ctx, + fullAdminBucketsPrincipal(), "to-delete") + require.ErrorIs(t, err, kv.ErrRouteWriteFenced) + require.Equal(t, s3TxnRetryMaxAttempts, coord.safetyNetCalls) +} + func TestS3Server_AdminDeleteBucket_MissingBucket(t *testing.T) { t.Parallel() @@ -275,6 +325,32 @@ func TestS3Server_AdminDeleteBucket_MissingBucket(t *testing.T) { require.ErrorIs(t, err, ErrAdminBucketNotFound) } +type bucketDeleteSafetyNetFenceCoordinator struct { + *localAdapterCoordinator + failuresRemaining int + safetyNetCalls int +} + +func (c *bucketDeleteSafetyNetFenceCoordinator) Dispatch(ctx context.Context, req *kv.OperationGroup[kv.OP]) (*kv.CoordinateResponse, error) { + if req != nil && operationGroupHasDelPrefix(req.Elems) { + c.safetyNetCalls++ + if c.failuresRemaining > 0 { + c.failuresRemaining-- + return nil, kv.ErrRouteWriteFenced + } + } + return c.localAdapterCoordinator.Dispatch(ctx, req) +} + +func operationGroupHasDelPrefix(elems []*kv.Elem[kv.OP]) bool { + for _, elem := range elems { + if elem != nil && elem.Op == kv.DelPrefix { + return true + } + } + return false +} + func TestS3Server_AdminDeleteBucket_RejectsReadOnly(t *testing.T) { t.Parallel() From 70cf6f86bed772630fe224e4eda41e56fd1daf12 Mon Sep 17 00:00:00 2001 From: bootjp Date: Wed, 15 Jul 2026 02:21:02 +0900 Subject: [PATCH 10/26] Stabilize raft snapshot dispatch and maintenance gates --- ...4_27_implemented_keyviz_cluster_fanout.md} | 11 +++- ...04_28_implemented_keyviz_adapter_labels.md | 2 +- .../2026_04_29_proposed_logical_backup.md | 2 +- ...ial_hotspot_split_milestone3_automation.md | 2 +- .../2026_06_23_proposed_scaling_roadmap.md | 2 +- internal/admin/keyviz_fanout.go | 2 +- internal/admin/keyviz_handler.go | 2 +- .../raftengine/etcd/dispatch_report_test.go | 33 ++++++++++++ internal/raftengine/etcd/engine.go | 50 +++++++++++++------ main.go | 49 ++++++++++++++---- main_maintenance_env_test.go | 33 ++++++++++++ scripts/rolling-update.env.example | 2 +- 12 files changed, 157 insertions(+), 33 deletions(-) rename docs/design/{2026_04_27_proposed_keyviz_cluster_fanout.md => 2026_04_27_implemented_keyviz_cluster_fanout.md} (98%) create mode 100644 main_maintenance_env_test.go diff --git a/docs/design/2026_04_27_proposed_keyviz_cluster_fanout.md b/docs/design/2026_04_27_implemented_keyviz_cluster_fanout.md similarity index 98% rename from docs/design/2026_04_27_proposed_keyviz_cluster_fanout.md rename to docs/design/2026_04_27_implemented_keyviz_cluster_fanout.md index 6e1b8f762..0ceafb300 100644 --- a/docs/design/2026_04_27_proposed_keyviz_cluster_fanout.md +++ b/docs/design/2026_04_27_implemented_keyviz_cluster_fanout.md @@ -1,5 +1,5 @@ --- -status: proposed +status: implemented phase: 2-C parent_design: docs/admin_ui_key_visualizer_design.md date: 2026-04-27 @@ -47,6 +47,15 @@ operator-visible value (cluster-wide heatmap, degraded-node banner) without requiring the full proto extension §9.1 calls for. +## 1.1 Implementation status + +Implemented in `internal/admin/keyviz_fanout.go`, +`internal/admin/keyviz_handler.go`, `web/admin/src/pages/KeyViz.tsx`, +and `main.go`'s `--keyvizFanoutNodes` / `--keyvizFanoutTimeout` +flags. Tests cover fan-out merge and degraded-node handling in +`internal/admin/keyviz_fanout_test.go` and the admin KeyViz UI +suite. + ## 2. Scope ### 2.1 In scope diff --git a/docs/design/2026_04_28_implemented_keyviz_adapter_labels.md b/docs/design/2026_04_28_implemented_keyviz_adapter_labels.md index 831f8cd18..b68dca189 100644 --- a/docs/design/2026_04_28_implemented_keyviz_adapter_labels.md +++ b/docs/design/2026_04_28_implemented_keyviz_adapter_labels.md @@ -587,7 +587,7 @@ from PR #694: Claude bot critical, Gemini high.) The fan-out aggregator's per-cell merge key gains the label: - Phase 2-C (current): `(bucketID, raftGroupID, leaderTerm, - windowStart)` per design `2026_04_27_proposed_keyviz_cluster_fanout.md` + windowStart)` per design `2026_04_27_implemented_keyviz_cluster_fanout.md` §4. - With labels: same tuple — but `bucketID` itself now carries the label via the §5 composite (`route:1:dynamo`). The merge key diff --git a/docs/design/2026_04_29_proposed_logical_backup.md b/docs/design/2026_04_29_proposed_logical_backup.md index 24c033cd8..39350d188 100644 --- a/docs/design/2026_04_29_proposed_logical_backup.md +++ b/docs/design/2026_04_29_proposed_logical_backup.md @@ -1570,7 +1570,7 @@ Scope: out of this proposal; mentioned only to draw the boundary. - Concurrent multi-cluster fan-out (one logical backup spanning shards on different physical clusters) — depends on the `Distribution` control plane being fan-out-aware (see - `2026_04_27_proposed_keyviz_cluster_fanout.md`). + `2026_04_27_implemented_keyviz_cluster_fanout.md`). ## Required Tests diff --git a/docs/design/2026_06_11_partial_hotspot_split_milestone3_automation.md b/docs/design/2026_06_11_partial_hotspot_split_milestone3_automation.md index 2f8a1c799..89c2e85a4 100644 --- a/docs/design/2026_06_11_partial_hotspot_split_milestone3_automation.md +++ b/docs/design/2026_06_11_partial_hotspot_split_milestone3_automation.md @@ -161,7 +161,7 @@ Options considered: **Follower-forwarded write caveat (codex P2).** `ShardedCoordinator.observeMutation` runs on the node that received the client request before `ShardRouter.Commit` forwards (`kv/sharded_coordinator.go:1795-1844`), and the follower's `Internal.Forward` handler calls `transactionManager.Commit` directly (`adapter/internal.go:52`) rather than re-entering the leader's `ShardedCoordinator` observation hook. Therefore even a route whose shard group is led by the default-group leader can be **under-observed** if most clients connect through followers. M3 does not treat "missing local rows" as evidence of low load: local evidence can promote a split, but absence of local evidence only means "unknown" and is fail-closed for target selection (§6.2) and conservative for candidacy (fewer splits). Full coverage for follower-forwarded writes needs follower reports / fan-out into the same detector interface and is tracked under OQ-5. -**Why not the keyviz cluster fan-out** (`docs/design/2026_04_27_proposed_keyviz_cluster_fanout.md`): that aggregates across nodes for the admin heatmap UI and dedupes write samples by `(groupID, leaderTerm)`. Pulling cross-node data in would add latency and a dependency on the fan-out being configured, so M3 stays leader-local and fail-closed on unknown load. If a future milestone wants complete cluster-wide scoring, including follower-forwarded writes, it can read the same fan-out aggregate behind the same detector interface (§5) — recorded as OQ-5. +**Why not the keyviz cluster fan-out** (`docs/design/2026_04_27_implemented_keyviz_cluster_fanout.md`): that aggregates across nodes for the admin heatmap UI and dedupes write samples by `(groupID, leaderTerm)`. Pulling cross-node data in would add latency and a dependency on the fan-out being configured, so M3 stays leader-local and fail-closed on unknown load. If a future milestone wants complete cluster-wide scoring, including follower-forwarded writes, it can read the same fan-out aggregate behind the same detector interface (§5) — recorded as OQ-5. **Known limitation — multi-leader (multi-group) deployments.** Leader-local consumption is only authoritative for routes whose shard group is led by the **same node** that runs the detector — i.e. the default-group leader (where the scheduler must run, §6.1). In a multi-group cluster (`--raftGroups` with G1, G2, …), the default-group leader is **not necessarily** the leader of G1/G2/…: each node runs a single `ShardedCoordinator` and `MemSampler` and only `Observe`s traffic it *receives* (`kv/sharded_coordinator.go:1795-1824`). Concretely, in a 3-node cluster where node A leads the default group and node B leads G1, **node A's sampler has no data for routes in G1 served through node B**, so the detector on A cannot score or split them. This is **broader than "heavy follower-forwarded reads"**: it is a structural gap for any route whose group leader differs from the default-group leader, not just an edge case of read forwarding. diff --git a/docs/design/2026_06_23_proposed_scaling_roadmap.md b/docs/design/2026_06_23_proposed_scaling_roadmap.md index b51827adb..d30405c20 100644 --- a/docs/design/2026_06_23_proposed_scaling_roadmap.md +++ b/docs/design/2026_06_23_proposed_scaling_roadmap.md @@ -260,7 +260,7 @@ memory each group's private cache/memtable pins. - **keyviz** — the per-route load sampler is wired and allocation-free on the write hot path (`observeMutation`, `kv/sharded_coordinator.go:1841-1846`), with proposed extensions for cluster fan-out - (`2026_04_27_proposed_keyviz_cluster_fanout.md`), + (`2026_04_27_implemented_keyviz_cluster_fanout.md`), subrange sampling (`2026_05_25_implemented_keyviz_subrange_sampling.md`), hot-key top-K (`2026_05_28_implemented_keyviz_hot_key_topk.md`), and per-cell conflict (implemented). It is the detection signal M3 reuses. Current diff --git a/internal/admin/keyviz_fanout.go b/internal/admin/keyviz_fanout.go index 472c2fc3f..6e3dfa7af 100644 --- a/internal/admin/keyviz_fanout.go +++ b/internal/admin/keyviz_fanout.go @@ -71,7 +71,7 @@ const keyVizMergeBucketHint = 64 // a stable row order; Responded counts ok=true entries; Expected is // the configured peer count plus self. // -// See docs/design/2026_04_27_proposed_keyviz_cluster_fanout.md 5. +// See docs/design/2026_04_27_implemented_keyviz_cluster_fanout.md 5. type FanoutResult struct { Nodes []FanoutNodeStatus `json:"nodes"` Responded int `json:"responded"` diff --git a/internal/admin/keyviz_handler.go b/internal/admin/keyviz_handler.go index d4192d693..db54f812b 100644 --- a/internal/admin/keyviz_handler.go +++ b/internal/admin/keyviz_handler.go @@ -56,7 +56,7 @@ const keyVizRowBudgetCap = 1024 // // Fanout is non-nil when the handler is configured for cluster-wide // fan-out (Phase 2-C): it carries per-node status so the SPA can -// surface degraded responses inline (see design 2026_04_27_proposed_keyviz_cluster_fanout.md). +// surface degraded responses inline (see design 2026_04_27_implemented_keyviz_cluster_fanout.md). // The field is omitted from the wire form when fan-out is disabled // so old clients keep working unchanged. type KeyVizMatrix struct { diff --git a/internal/raftengine/etcd/dispatch_report_test.go b/internal/raftengine/etcd/dispatch_report_test.go index d0ba97049..dfd1dfd67 100644 --- a/internal/raftengine/etcd/dispatch_report_test.go +++ b/internal/raftengine/etcd/dispatch_report_test.go @@ -33,6 +33,39 @@ func TestPostDispatchReport_DeliversWhenChannelHasSpace(t *testing.T) { } } +func TestReportSuccessfulDispatchReportsSnapshotFinish(t *testing.T) { + t.Parallel() + e := &Engine{ + dispatchReportCh: make(chan dispatchReport, 1), + closeCh: make(chan struct{}), + } + + e.reportSuccessfulDispatch(raftpb.Message{Type: messageTypePtr(raftpb.MsgSnap), To: uint64Ptr(2)}) + + select { + case got := <-e.dispatchReportCh: + require.Equal(t, dispatchReport{to: 2, msgType: raftpb.MsgSnap, snapshotFinish: true}, got) + default: + t.Fatal("expected successful MsgSnap dispatch to report SnapshotFinish input") + } +} + +func TestReportSuccessfulDispatchIgnoresRegularMessage(t *testing.T) { + t.Parallel() + e := &Engine{ + dispatchReportCh: make(chan dispatchReport, 1), + closeCh: make(chan struct{}), + } + + e.reportSuccessfulDispatch(raftpb.Message{Type: messageTypePtr(raftpb.MsgApp), To: uint64Ptr(2)}) + + select { + case got := <-e.dispatchReportCh: + t.Fatalf("unexpected dispatch report for regular message: %+v", got) + default: + } +} + // TestPostDispatchReport_DropsWhenChannelFull asserts the non-blocking // contract: dispatch workers must not stall because the event loop is busy. // The worst case is an eventually-consistent gap that raft will fix on the diff --git a/internal/raftengine/etcd/engine.go b/internal/raftengine/etcd/engine.go index 7b88b06a8..57fbee77b 100644 --- a/internal/raftengine/etcd/engine.go +++ b/internal/raftengine/etcd/engine.go @@ -1808,13 +1808,14 @@ func (e *Engine) tryReceivePriorityStep() (raftpb.Message, bool) { } } -// dispatchReport is posted by the dispatch workers when a transport send -// to a peer fails; the engine goroutine drains these and informs etcd/raft -// via rawNode so follower Progress leaves StateReplicate / StateSnapshot on -// unreachable peers and does not silently stall. +// dispatchReport is posted by the dispatch workers after a transport send +// completes or fails; the engine goroutine drains these and informs etcd/raft +// via rawNode so follower Progress leaves StateReplicate / StateSnapshot and +// does not silently stall. type dispatchReport struct { - to uint64 - msgType raftpb.MessageType + to uint64 + msgType raftpb.MessageType + snapshotFinish bool } func (e *Engine) handleDispatchReport(report dispatchReport) { @@ -1827,19 +1828,23 @@ func (e *Engine) handleDispatchReport(report dispatchReport) { // peer from StateReplicate to StateProbe so the next heartbeat response // drives a fresh sendAppend attempt. if report.msgType == raftpb.MsgSnap { - e.rawNode.ReportSnapshot(report.to, etcdraft.SnapshotFailure) + status := etcdraft.SnapshotFailure + if report.snapshotFinish { + status = etcdraft.SnapshotFinish + } + e.rawNode.ReportSnapshot(report.to, status) return } e.rawNode.ReportUnreachable(report.to) } -// postDispatchReport delivers a dispatch failure to the event loop without -// blocking the caller. Dispatch workers use it for transport failures, and the -// event loop uses it for local queue drops before transport. If the channel is -// full (unlikely — the buffer is sized to MaxInflightMsg), the report is -// dropped and logged; this is acceptable because raft will retry on the next -// tick and we only need eventual consistency between transport state and -// Progress state. +// postDispatchReport delivers a dispatch outcome to the event loop without +// blocking the caller. Dispatch workers use it for transport completion or +// failures, and the event loop uses it for local queue drops before transport. +// If the channel is full (unlikely — the buffer is sized to MaxInflightMsg), +// the report is dropped and logged; this is acceptable because raft will retry +// on the next tick and we only need eventual consistency between transport +// state and Progress state. func (e *Engine) postDispatchReport(report dispatchReport) { select { case e.dispatchReportCh <- report: @@ -4423,7 +4428,11 @@ func (e *Engine) handleDispatchRequest(ctx context.Context, req dispatchRequest) if err := req.Close(); err != nil { slog.Error("etcd raft dispatch: failed to close request", "err", err) } - if dispatchErr == nil || errors.Is(dispatchErr, ctx.Err()) { + if dispatchErr == nil { + e.reportSuccessfulDispatch(req.msg) + return + } + if errors.Is(dispatchErr, ctx.Err()) { return } code := dispatchErrorCodeOf(dispatchErr) @@ -4445,6 +4454,17 @@ func (e *Engine) handleDispatchRequest(ctx context.Context, req dispatchRequest) e.postDispatchReport(dispatchReport{to: req.msg.GetTo(), msgType: req.msg.GetType()}) } +func (e *Engine) reportSuccessfulDispatch(msg raftpb.Message) { + if msg.GetType() != raftpb.MsgSnap { + return + } + e.postDispatchReport(dispatchReport{ + to: msg.GetTo(), + msgType: msg.GetType(), + snapshotFinish: true, + }) +} + func (e *Engine) stopDispatchWorkers() { e.dispatchOnce.Do(func() { if e.dispatchCancel != nil { diff --git a/main.go b/main.go index 530dbf4b7..236fcb0c7 100644 --- a/main.go +++ b/main.go @@ -53,6 +53,9 @@ const ( etcdMaxSizePerMsg = 1 << 20 etcdMaxInflightMsg = 1024 defaultTSOBatchSize = 256 + + lockResolverEnabledEnv = "ELASTICKV_LOCK_RESOLVER_ENABLED" + fsmCompactorEnabledEnv = "ELASTICKV_FSM_COMPACTOR_ENABLED" ) func newRaftFactory(engineType raftEngineType, coldStartObs raftengine.ColdStartObserver) (raftengine.Factory, error) { @@ -85,6 +88,18 @@ func durationToTicks(timeout time.Duration, tick time.Duration, min int) int { return ticks } +func optionalBoolEnv(name string, def bool) bool { + raw := strings.TrimSpace(os.Getenv(name)) + if raw == "" { + return def + } + v, err := strconv.ParseBool(raw) + if err != nil { + return def + } + return v +} + var ( myAddr = flag.String("address", "localhost:50051", "TCP host+port for this node") redisAddr = flag.String("redisAddress", "localhost:6379", "TCP host+port for redis") @@ -236,7 +251,7 @@ var ( // HTTP endpoints (host:port or scheme://host:port). When set, // the admin keyviz handler aggregates the local matrix with // peer responses; when empty, behaviour is unchanged - // (single-node view). See docs/design/2026_04_27_proposed_keyviz_cluster_fanout.md. + // (single-node view). See docs/design/2026_04_27_implemented_keyviz_cluster_fanout.md. keyvizFanoutNodes = flag.String("keyvizFanoutNodes", "", "Comma-separated peer admin endpoints (host:port) for keyviz cluster-wide fan-out; empty disables") keyvizFanoutTimeout = flag.Duration("keyvizFanoutTimeout", keyvizFanoutDefaultTimeout, "Per-peer timeout for keyviz fan-out HTTP calls") ) @@ -445,8 +460,7 @@ func run() error { } }) cleanup.Add(cancel) - lockResolver := kv.NewLockResolver(shardStore, shardGroups, nil) - cleanup.Add(func() { lockResolver.Close() }) + startLockResolverIfEnabled(shardStore, shardGroups, &cleanup) sampler := buildKeyVizSampler() coordinate := kv.NewShardedCoordinator(cfg.engine, shardGroups, cfg.defaultGroup, clock, shardStore). WithLeaseReadObserver(metricsRegistry.LeaseReadObserver()). @@ -505,13 +519,7 @@ func run() error { adapter.WithDistributionActiveTimestampTracker(readTracker), ) startMonitoringCollectors(runCtx, metricsRegistry, runtimes, clock) - compactor := kv.NewFSMCompactor( - fsmCompactionRuntimes(runtimes), - kv.WithFSMCompactorActiveTimestampTracker(readTracker), - ) - eg.Go(func() error { - return compactor.Run(runCtx) - }) + startFSMCompactorIfEnabled(runCtx, eg, runtimes, readTracker) // Stage 7c §3.1: build the encryption-aware // MembershipChangeInterceptor here where the concrete @@ -2029,6 +2037,27 @@ func writeConflictMonitorSources(runtimes []*raftGroupRuntime) []monitoring.Writ return out } +func startLockResolverIfEnabled(ss *kv.ShardStore, groups map[uint64]*kv.ShardGroup, cleanup *internalutil.CleanupStack) { + if !optionalBoolEnv(lockResolverEnabledEnv, true) { + return + } + lockResolver := kv.NewLockResolver(ss, groups, nil) + cleanup.Add(func() { lockResolver.Close() }) +} + +func startFSMCompactorIfEnabled(ctx context.Context, eg *errgroup.Group, runtimes []*raftGroupRuntime, readTracker *kv.ActiveTimestampTracker) { + if !optionalBoolEnv(fsmCompactorEnabledEnv, true) { + return + } + compactor := kv.NewFSMCompactor( + fsmCompactionRuntimes(runtimes), + kv.WithFSMCompactorActiveTimestampTracker(readTracker), + ) + eg.Go(func() error { + return compactor.Run(ctx) + }) +} + func fsmCompactionRuntimes(runtimes []*raftGroupRuntime) []kv.FSMCompactRuntime { out := make([]kv.FSMCompactRuntime, 0, len(runtimes)) for _, runtime := range runtimes { diff --git a/main_maintenance_env_test.go b/main_maintenance_env_test.go new file mode 100644 index 000000000..bee78cdeb --- /dev/null +++ b/main_maintenance_env_test.go @@ -0,0 +1,33 @@ +package main + +import "testing" + +func TestOptionalBoolEnv(t *testing.T) { + t.Setenv("ELASTICKV_TEST_BOOL", "") + if got := optionalBoolEnv("ELASTICKV_TEST_BOOL", true); !got { + t.Fatal("empty env should use default true") + } + if got := optionalBoolEnv("ELASTICKV_TEST_BOOL", false); got { + t.Fatal("empty env should use default false") + } + + for _, tc := range []struct { + name string + raw string + want bool + }{ + {name: "true", raw: "true", want: true}, + {name: "one", raw: "1", want: true}, + {name: "false", raw: "false", want: false}, + {name: "zero", raw: "0", want: false}, + {name: "trimmed", raw: " false ", want: false}, + {name: "invalid uses default", raw: "disabled", want: true}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Setenv("ELASTICKV_TEST_BOOL", tc.raw) + if got := optionalBoolEnv("ELASTICKV_TEST_BOOL", true); got != tc.want { + t.Fatalf("optionalBoolEnv()=%v, want %v", got, tc.want) + } + }) + } +} diff --git a/scripts/rolling-update.env.example b/scripts/rolling-update.env.example index 6ee72dd60..f71cfee29 100644 --- a/scripts/rolling-update.env.example +++ b/scripts/rolling-update.env.example @@ -122,7 +122,7 @@ ADMIN_ENABLED="false" # role allow-lists must be configured cluster-wide. Peers without # --adminEnabled expose an unauthenticated keyviz endpoint and # respond unconditionally. -# See docs/design/2026_04_27_proposed_keyviz_cluster_fanout.md for the +# See docs/design/2026_04_27_implemented_keyviz_cluster_fanout.md for the # full design. KEYVIZ_ENABLED="false" # KEYVIZ_FANOUT_NODES="10.0.0.1:8080,10.0.0.2:8080,10.0.0.3:8080" From a2418094fe690c3e4b2d86ae1d497b56c741050f Mon Sep 17 00:00:00 2001 From: bootjp Date: Wed, 15 Jul 2026 02:46:03 +0900 Subject: [PATCH 11/26] Honor partition resolver in fence prechecks --- kv/sharded_coordinator.go | 13 +++++ kv/sharded_coordinator_partition_test.go | 70 ++++++++++++++++++++++++ 2 files changed, 83 insertions(+) diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index f345f2e63..f70d5ecd8 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -1065,6 +1065,9 @@ func (c *ShardedCoordinator) rejectWriteFencedPointElems(elems []*Elem[OP]) erro } func (c *ShardedCoordinator) rejectWriteFencedPointKey(key []byte) error { + if c.partitionResolverRecognisesPointKey(key) { + return nil + } rkey := routeKey(key) if route, ok := c.engine.GetRoute(rkey); ok && route.State == distribution.RouteStateWriteFenced { return errors.Wrapf(ErrRouteWriteFenced, "key %q routeKey %q", key, rkey) @@ -1081,6 +1084,16 @@ func (c *ShardedCoordinator) rejectWriteFencedPointKey(key []byte) error { return nil } +func (c *ShardedCoordinator) partitionResolverRecognisesPointKey(key []byte) bool { + if c == nil || c.router == nil || c.router.partitionResolver == nil || len(key) == 0 { + return false + } + if _, ok := c.router.partitionResolver.ResolveGroup(key); ok { + return true + } + return c.router.partitionResolver.RecognisesPartitionedKey(key) +} + func (c *ShardedCoordinator) rejectWriteFencedDelPrefixes(elems []*Elem[OP]) error { if c == nil || c.engine == nil { return nil diff --git a/kv/sharded_coordinator_partition_test.go b/kv/sharded_coordinator_partition_test.go index ae46301a1..e9b5c16bb 100644 --- a/kv/sharded_coordinator_partition_test.go +++ b/kv/sharded_coordinator_partition_test.go @@ -2,6 +2,7 @@ package kv import ( "context" + "errors" "sync" "testing" @@ -115,6 +116,75 @@ func TestShardedCoordinator_DispatchHonoursPartitionResolver(t *testing.T) { require.Equal(t, []byte("!sqs|msg|data|p|partitioned-key"), calls[0]) } +func TestShardedCoordinatorWriteFencePrecheckHonoursPartitionResolver(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: nil, GroupID: 1, State: distribution.RouteStateWriteFenced}, + }, + })) + + g1 := &recordingTransactional{ + responses: []*TransactionResponse{{CommitIndex: 1}}, + } + g42 := &recordingTransactional{ + responses: []*TransactionResponse{{CommitIndex: 42}}, + } + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1, Store: store.NewMVCCStore()}, + 42: {Txn: g42, Store: store.NewMVCCStore()}, + }, 1, NewHLC(), nil) + + key := []byte("!sqs|msg|data|p|partitioned-key") + coord.WithPartitionResolver(&stubResolver{claim: map[string]uint64{ + string(key): 42, + }}) + + resp, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: Put, Key: key, Value: []byte("v")}}, + }) + require.NoError(t, err) + require.NotNil(t, resp) + require.Equal(t, uint64(42), resp.CommitIndex) + require.Empty(t, g1.requests, "engine route fence must not preempt resolver-owned keys") + require.Len(t, g42.requests, 1) +} + +func TestShardedCoordinatorWriteFencePrecheckSkipsUnresolvedPartitionKey(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: nil, GroupID: 1, State: distribution.RouteStateWriteFenced}, + }, + })) + + g1 := &recordingTransactional{ + responses: []*TransactionResponse{{CommitIndex: 1}}, + } + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1, Store: store.NewMVCCStore()}, + }, 1, NewHLC(), nil) + + key := []byte("!sqs|msg|data|p|unknown-partition-key") + coord.WithPartitionResolver(&stubResolver{ + recognisedPrefix: []byte("!sqs|msg|data|p|"), + }) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: Put, Key: key, Value: []byte("v")}}, + }) + require.ErrorIs(t, err, ErrInvalidRequest) + require.False(t, errors.Is(err, ErrRouteWriteFenced), + "resolver-recognised keys must fail through resolver routing, not engine route fences") + require.Empty(t, g1.requests) +} + // TestShardedCoordinator_DispatchSplitsMutationsByResolverGroup is // the genuine regression for the Gemini-HIGH groupMutations // bypass: a Dispatch with mutations belonging to TWO different From ee363c4b5fa8034c7da23e17d4902208486e0177 Mon Sep 17 00:00:00 2001 From: bootjp Date: Wed, 15 Jul 2026 03:39:44 +0900 Subject: [PATCH 12/26] Make raft startup and fence checks deterministic --- internal/raftengine/etcd/engine.go | 50 ++++++++---------- internal/raftengine/etcd/engine_test.go | 35 +++++++------ kv/fsm.go | 67 +------------------------ kv/fsm_migration_fence_test.go | 53 +++++++++---------- 4 files changed, 67 insertions(+), 138 deletions(-) diff --git a/internal/raftengine/etcd/engine.go b/internal/raftengine/etcd/engine.go index 57fbee77b..437d56a32 100644 --- a/internal/raftengine/etcd/engine.go +++ b/internal/raftengine/etcd/engine.go @@ -351,15 +351,13 @@ type Engine struct { snapshotStopCh chan struct{} closeCh chan struct{} doneCh chan struct{} - // openCh is closed when run starts and callers may bind public/raft - // listeners. startedCh remains the stronger gate for processing inbound - // raft messages after the startup Ready drain has completed. - openCh chan struct{} + // startedCh is closed after startup has drained committed Ready entries. + // Open waits on this gate so callers never observe a store-ready engine + // before replay has caught the local state machine up to the raft log. startedCh chan struct{} leaderReady chan struct{} leaderOnce sync.Once - openOnce sync.Once startOnce sync.Once closeOnce sync.Once dispatchOnce sync.Once @@ -661,7 +659,6 @@ func Open(ctx context.Context, cfg OpenConfig) (*Engine, error) { dispatchReportCh: make(chan dispatchReport, inboundQueueCap), closeCh: make(chan struct{}), doneCh: make(chan struct{}), - openCh: make(chan struct{}), startedCh: make(chan struct{}), leaderReady: make(chan struct{}), config: configurationFromConfState(peerMap, confStateValue(prepared.disk.LocalSnap.GetMetadata().GetConfState())), @@ -919,17 +916,30 @@ func newRawNode(cfg OpenConfig, storage *etcdraft.MemoryStorage, applied uint64) } func waitForOpen(ctx context.Context, engine *Engine, waitForLeader bool) (*Engine, error) { + if err := waitForEngineSignal(ctx, engine, engine.startedCh); err != nil { + return nil, err + } + if !waitForLeader { + return engine, nil + } + if err := waitForEngineSignal(ctx, engine, engine.leaderReady); err != nil { + return nil, err + } + return engine, nil +} + +func waitForEngineSignal(ctx context.Context, engine *Engine, ready <-chan struct{}) error { select { case <-ctx.Done(): _ = engine.Close() - return nil, errors.WithStack(ctx.Err()) - case <-engine.openReady(waitForLeader): - return engine, nil + return errors.WithStack(ctx.Err()) + case <-ready: + return nil case <-engine.doneCh: if err := engine.currentError(); err != nil { - return nil, err + return err } - return nil, errors.WithStack(errClosed) + return errors.WithStack(errClosed) } } @@ -1691,7 +1701,6 @@ func (e *Engine) run() { ticker := time.NewTicker(e.tickInterval) defer ticker.Stop() - e.markOpen() if err := e.startup(); err != nil { e.fail(err) return @@ -3594,23 +3603,6 @@ func (e *Engine) markStarted() { e.startOnce.Do(func() { close(e.startedCh) }) } -func (e *Engine) markOpen() { - if e.openCh == nil { - return - } - e.openOnce.Do(func() { close(e.openCh) }) -} - -func (e *Engine) openReady(waitForLeader bool) <-chan struct{} { - if waitForLeader { - return e.leaderReady - } - if e.openCh != nil { - return e.openCh - } - return e.startedCh -} - func (e *Engine) requestShutdown() { e.closeOnce.Do(func() { close(e.closeCh) diff --git a/internal/raftengine/etcd/engine_test.go b/internal/raftengine/etcd/engine_test.go index 4139728ff..3ac45cd11 100644 --- a/internal/raftengine/etcd/engine_test.go +++ b/internal/raftengine/etcd/engine_test.go @@ -1430,7 +1430,7 @@ func TestOpenRestoresLegacySnapshotState(t *testing.T) { require.Equal(t, [][]byte{[]byte("snap"), []byte("tail")}, fsm.Applied()) } -func TestOpenMultiNodeReturnsBeforeCommittedTailDrain(t *testing.T) { +func TestOpenMultiNodeWaitsForCommittedTailDrain(t *testing.T) { dir := t.TempDir() peers := []Peer{ {NodeID: 1, ID: "n1", Address: "127.0.0.1:7001"}, @@ -1468,30 +1468,35 @@ func TestOpenMultiNodeReturnsBeforeCommittedTailDrain(t *testing.T) { done <- openResult{engine: engine, err: err} }() + select { + case <-fsm.started: + case <-time.After(time.Second): + t.Fatal("startup committed tail was not being applied") + } + + select { + case result := <-done: + if result.engine != nil { + require.NoError(t, result.engine.Close()) + } + require.NoError(t, result.err) + t.Fatal("multi-node Open returned before committed tail drain completed") + case <-time.After(20 * time.Millisecond): + } + + close(fsm.release) + var result openResult select { case result = <-done: case <-time.After(time.Second): - t.Fatal("multi-node Open blocked on committed tail drain before returning") + t.Fatal("multi-node Open did not return after committed tail drain completed") } require.NoError(t, result.err) require.NotNil(t, result.engine) defer func() { require.NoError(t, result.engine.Close()) }() - - select { - case <-fsm.started: - case <-time.After(time.Second): - t.Fatal("startup committed tail was not being applied") - } - select { - case <-result.engine.startedCh: - t.Fatal("engine marked started before startup committed tail drain completed") - default: - } - - close(fsm.release) select { case <-result.engine.startedCh: case <-time.After(time.Second): diff --git a/kv/fsm.go b/kv/fsm.go index a9ff4cea5..bdb7fc025 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -497,9 +497,6 @@ func (f *kvFSM) handleRawRequest(ctx context.Context, r *pb.Request, commitTS ui 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 { return err } @@ -532,9 +529,6 @@ func extractDelPrefix(muts []*pb.Mutation) (bool, []byte) { // handleDelPrefix delegates prefix deletion to the store. Transaction-internal // keys are always excluded to preserve transactional integrity. func (f *kvFSM) handleDelPrefix(ctx context.Context, prefix []byte, commitTS uint64) error { - if err := f.verifyRouteNotFencedForPrefix(prefix); err != nil { - return err - } if err := f.store.DeletePrefixAtRaftAt(ctx, prefix, txnCommonPrefix, commitTS, f.pendingApplyIdx); err != nil { return errors.WithStack(err) } @@ -542,51 +536,6 @@ func (f *kvFSM) handleDelPrefix(ctx context.Context, prefix []byte, commitTS uin return nil } -func (f *kvFSM) verifyRouteNotFencedForMutations(muts []*pb.Mutation) error { - for _, mut := range muts { - if mut == nil || len(mut.Key) == 0 || isTxnInternalKey(mut.Key) { - continue - } - if err := f.verifyRouteNotFencedForKey(mut.Key); err != nil { - return err - } - } - return nil -} - -func (f *kvFSM) verifyRouteNotFencedForKey(key []byte) error { - if f.routes == nil { - return nil - } - snap, ok := f.routes.Current() - if !ok { - return nil - } - rkey := routeKey(key) - if snap.WriteFencedForKey(rkey) { - return errors.Wrapf(ErrRouteWriteFenced, "key %q routeKey %q", key, rkey) - } - if start, end, ok := s3BucketAuxiliaryRouteRange(key); ok && snap.WriteFencedIntersects(start, end) { - return errors.Wrapf(ErrRouteWriteFenced, "key %q route range [%q,%q)", key, start, end) - } - return nil -} - -func (f *kvFSM) verifyRouteNotFencedForPrefix(prefix []byte) error { - if f.routes == nil { - return nil - } - snap, ok := f.routes.Current() - if !ok { - return nil - } - start, end := routePrefixRange(prefix) - if !snap.WriteFencedIntersects(start, end) { - return nil - } - return errors.Wrapf(ErrRouteWriteFenced, "prefix %q route range [%q,%q)", prefix, start, end) -} - func routePrefixRange(prefix []byte) ([]byte, []byte) { if len(prefix) == 0 { return []byte(""), nil @@ -959,9 +908,6 @@ func (f *kvFSM) handlePrepareRequest(ctx context.Context, r *pb.Request) error { if err != nil { return err } - if err := f.verifyRouteNotFencedForMutations(uniq); err != nil { - return err - } if err := f.validateConflicts(ctx, uniq, startTS); err != nil { return errors.WithStack(err) } @@ -1028,7 +974,7 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com return nil } - uniq, err := f.uniqueMutationsNotFenced(muts) + uniq, err := uniqueMutations(muts) if err != nil { return err } @@ -1044,17 +990,6 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com return nil } -func (f *kvFSM) uniqueMutationsNotFenced(muts []*pb.Mutation) ([]*pb.Mutation, error) { - uniq, err := uniqueMutations(muts) - if err != nil { - return nil, err - } - if err := f.verifyRouteNotFencedForMutations(uniq); err != nil { - return nil, err - } - return uniq, nil -} - // dedupProbeOnePhase decides whether handleOnePhaseTxnRequest should no-op // because the entry is a retry whose prior attempt already landed. Extracted // to keep handleOnePhaseTxnRequest under the cyclop budget; the determinism diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index 7aff2b4f4..7279a1e83 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -7,8 +7,6 @@ import ( "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/internal/s3keys" pb "github.com/bootjp/elastickv/proto" - "github.com/bootjp/elastickv/store" - "github.com/cockroachdb/errors" "github.com/stretchr/testify/require" ) @@ -41,20 +39,21 @@ func newS3BucketAuxiliaryWriteFencedFSM(t *testing.T, bucket string) *kvFSM { return newComposed1FSM(t, engine, 1) } -func TestFSMRejectsRawPointWriteOnWriteFencedRoute(t *testing.T) { +func TestFSMAppliesRawPointWriteDespiteLocalWriteFencedRoute(t *testing.T) { t.Parallel() fsm := newWriteFencedFSM(t) err := fsm.handleRawRequest(context.Background(), &pb.Request{ Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("v")}}, }, 10) - require.ErrorIs(t, err, ErrRouteWriteFenced) + require.NoError(t, err) - _, getErr := fsm.store.GetAt(context.Background(), []byte("z"), ^uint64(0)) - require.ErrorIs(t, getErr, store.ErrKeyNotFound) + got, getErr := fsm.store.GetAt(context.Background(), []byte("z"), ^uint64(0)) + require.NoError(t, getErr) + require.Equal(t, []byte("v"), got) } -func TestFSMRejectsS3BucketAuxiliaryPointWriteOnWriteFencedRoute(t *testing.T) { +func TestFSMAppliesS3BucketAuxiliaryPointWriteDespiteLocalWriteFencedRoute(t *testing.T) { t.Parallel() ctx := context.Background() @@ -68,14 +67,15 @@ func TestFSMRejectsS3BucketAuxiliaryPointWriteOnWriteFencedRoute(t *testing.T) { err := fsm.handleRawRequest(ctx, &pb.Request{ Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: key, Value: []byte("v")}}, }, 10) - require.ErrorIs(t, err, ErrRouteWriteFenced) + require.NoError(t, err) - _, getErr := fsm.store.GetAt(ctx, key, ^uint64(0)) - require.ErrorIs(t, getErr, store.ErrKeyNotFound) + got, getErr := fsm.store.GetAt(ctx, key, ^uint64(0)) + require.NoError(t, getErr) + require.Equal(t, []byte("v"), got) } } -func TestFSMRejectsDelPrefixIntersectingWriteFencedRoute(t *testing.T) { +func TestFSMAppliesDelPrefixDespiteLocalWriteFencedRoute(t *testing.T) { t.Parallel() fsm := newWriteFencedFSM(t) @@ -84,14 +84,13 @@ func TestFSMRejectsDelPrefixIntersectingWriteFencedRoute(t *testing.T) { err := fsm.handleRawRequest(context.Background(), &pb.Request{ Mutations: []*pb.Mutation{{Op: pb.Op_DEL_PREFIX, Key: []byte("z")}}, }, 10) - require.ErrorIs(t, err, ErrRouteWriteFenced) + require.NoError(t, err) - got, getErr := fsm.store.GetAt(context.Background(), []byte("z"), ^uint64(0)) - require.NoError(t, getErr) - require.Equal(t, []byte("v"), got) + _, getErr := fsm.store.GetAt(context.Background(), []byte("z"), ^uint64(0)) + require.Error(t, getErr) } -func TestFSMRejectsFullRangeDelPrefixWhenRouteIsWriteFenced(t *testing.T) { +func TestFSMAppliesFullRangeDelPrefixDespiteLocalWriteFencedRoute(t *testing.T) { t.Parallel() fsm := newWriteFencedFSM(t) @@ -100,14 +99,13 @@ func TestFSMRejectsFullRangeDelPrefixWhenRouteIsWriteFenced(t *testing.T) { err := fsm.handleRawRequest(context.Background(), &pb.Request{ Mutations: []*pb.Mutation{{Op: pb.Op_DEL_PREFIX, Key: nil}}, }, 10) - require.ErrorIs(t, err, ErrRouteWriteFenced) + require.NoError(t, err) - got, getErr := fsm.store.GetAt(context.Background(), []byte("z"), ^uint64(0)) - require.NoError(t, getErr) - require.Equal(t, []byte("v"), got) + _, getErr := fsm.store.GetAt(context.Background(), []byte("z"), ^uint64(0)) + require.Error(t, getErr) } -func TestFSMRejectsBroadInternalDelPrefixWhenRouteIsWriteFenced(t *testing.T) { +func TestFSMAppliesBroadInternalDelPrefixDespiteLocalWriteFencedRoute(t *testing.T) { t.Parallel() fsm := newWriteFencedFSM(t) @@ -117,14 +115,13 @@ func TestFSMRejectsBroadInternalDelPrefixWhenRouteIsWriteFenced(t *testing.T) { err := fsm.handleRawRequest(context.Background(), &pb.Request{ Mutations: []*pb.Mutation{{Op: pb.Op_DEL_PREFIX, Key: []byte("!redis|")}}, }, 10) - require.ErrorIs(t, err, ErrRouteWriteFenced) + require.NoError(t, err) - got, getErr := fsm.store.GetAt(context.Background(), key, ^uint64(0)) - require.NoError(t, getErr) - require.Equal(t, []byte("v"), got) + _, getErr := fsm.store.GetAt(context.Background(), key, ^uint64(0)) + require.Error(t, getErr) } -func TestFSMRejectsPrepareOnWriteFencedRouteButAllowsAbort(t *testing.T) { +func TestFSMAppliesPrepareAndAbortDespiteLocalWriteFencedRoute(t *testing.T) { t.Parallel() ctx := context.Background() @@ -138,7 +135,7 @@ func TestFSMRejectsPrepareOnWriteFencedRouteButAllowsAbort(t *testing.T) { {Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("v")}, }, } - require.ErrorIs(t, fsm.handleTxnRequest(ctx, prepare, 10), ErrRouteWriteFenced) + require.NoError(t, fsm.handleTxnRequest(ctx, prepare, 10)) abort := &pb.Request{ IsTxn: true, @@ -150,5 +147,5 @@ 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") + require.NotErrorIs(t, err, ErrRouteWriteFenced, "ABORT must keep the narrow cleanup lane open") } From 54ad99937cf43df1e4769559393bbc28cf12d566 Mon Sep 17 00:00:00 2001 From: bootjp Date: Wed, 15 Jul 2026 04:02:00 +0900 Subject: [PATCH 13/26] Gate public startup after raft replay --- internal/raftengine/engine.go | 6 ++ internal/raftengine/etcd/engine.go | 17 +++-- internal/raftengine/etcd/engine_test.go | 85 ++++++++++++++++--------- main.go | 20 ++++++ 4 files changed, 94 insertions(+), 34 deletions(-) diff --git a/internal/raftengine/engine.go b/internal/raftengine/engine.go index 75ceacc8c..02895d9f5 100644 --- a/internal/raftengine/engine.go +++ b/internal/raftengine/engine.go @@ -244,6 +244,12 @@ type Lifecycle interface { Err() error } +// StartupBarrier is an optional capability for engines that can report when +// their local startup replay has drained far enough for user traffic. +type StartupBarrier interface { + WaitStarted(ctx context.Context) error +} + type Admin interface { LeaderView StatusReader diff --git a/internal/raftengine/etcd/engine.go b/internal/raftengine/etcd/engine.go index 437d56a32..296bb00f1 100644 --- a/internal/raftengine/etcd/engine.go +++ b/internal/raftengine/etcd/engine.go @@ -352,8 +352,8 @@ type Engine struct { closeCh chan struct{} doneCh chan struct{} // startedCh is closed after startup has drained committed Ready entries. - // Open waits on this gate so callers never observe a store-ready engine - // before replay has caught the local state machine up to the raft log. + // Multi-node Open must return before this so callers can register the + // transport listener; service startup waits through WaitStarted instead. startedCh chan struct{} leaderReady chan struct{} @@ -916,18 +916,25 @@ func newRawNode(cfg OpenConfig, storage *etcdraft.MemoryStorage, applied uint64) } func waitForOpen(ctx context.Context, engine *Engine, waitForLeader bool) (*Engine, error) { - if err := waitForEngineSignal(ctx, engine, engine.startedCh); err != nil { - return nil, err - } if !waitForLeader { return engine, nil } + if err := waitForEngineSignal(ctx, engine, engine.startedCh); err != nil { + return nil, err + } if err := waitForEngineSignal(ctx, engine, engine.leaderReady); err != nil { return nil, err } return engine, nil } +func (e *Engine) WaitStarted(ctx context.Context) error { + if e == nil { + return errors.WithStack(errNilEngine) + } + return waitForEngineSignal(ctx, e, e.startedCh) +} + func waitForEngineSignal(ctx context.Context, engine *Engine, ready <-chan struct{}) error { select { case <-ctx.Done(): diff --git a/internal/raftengine/etcd/engine_test.go b/internal/raftengine/etcd/engine_test.go index 3ac45cd11..e379282f7 100644 --- a/internal/raftengine/etcd/engine_test.go +++ b/internal/raftengine/etcd/engine_test.go @@ -1430,7 +1430,7 @@ func TestOpenRestoresLegacySnapshotState(t *testing.T) { require.Equal(t, [][]byte{[]byte("snap"), []byte("tail")}, fsm.Applied()) } -func TestOpenMultiNodeWaitsForCommittedTailDrain(t *testing.T) { +func TestOpenMultiNodeWaitStartedWaitsForCommittedTailDrain(t *testing.T) { dir := t.TempDir() peers := []Peer{ {NodeID: 1, ID: "n1", Address: "127.0.0.1:7001"}, @@ -1451,10 +1451,6 @@ func TestOpenMultiNodeWaitsForCommittedTailDrain(t *testing.T) { started: make(chan struct{}), release: make(chan struct{}), } - type openResult struct { - engine *Engine - err error - } done := make(chan openResult, 1) go func() { engine, err := Open(context.Background(), OpenConfig{ @@ -1468,39 +1464,70 @@ func TestOpenMultiNodeWaitsForCommittedTailDrain(t *testing.T) { done <- openResult{engine: engine, err: err} }() + result := requireOpenResult(t, done, time.Second, "multi-node Open did not return before committed tail drain completed") + require.NoError(t, result.err) + require.NotNil(t, result.engine) + defer func() { + require.NoError(t, result.engine.Close()) + }() + + requireSignal(t, fsm.started, time.Second, "startup committed tail was not being applied") + + waitDone := make(chan error, 1) + go func() { + waitDone <- result.engine.WaitStarted(context.Background()) + }() + + requireNoStartupWaitResult(t, waitDone, 20*time.Millisecond, "WaitStarted returned before committed tail drain completed") + + close(fsm.release) + + requireStartupWaitResult(t, waitDone, time.Second, "WaitStarted did not return after committed tail drain completed") + requireSignal(t, result.engine.startedCh, time.Second, "engine did not mark started after committed tail drain completed") +} + +type openResult struct { + engine *Engine + err error +} + +func requireOpenResult(t *testing.T, ch <-chan openResult, timeout time.Duration, msg string) openResult { + t.Helper() select { - case <-fsm.started: - case <-time.After(time.Second): - t.Fatal("startup committed tail was not being applied") + case result := <-ch: + return result + case <-time.After(timeout): + t.Fatal(msg) + return openResult{} } +} +func requireSignal(t *testing.T, ch <-chan struct{}, timeout time.Duration, msg string) { + t.Helper() select { - case result := <-done: - if result.engine != nil { - require.NoError(t, result.engine.Close()) - } - require.NoError(t, result.err) - t.Fatal("multi-node Open returned before committed tail drain completed") - case <-time.After(20 * time.Millisecond): + case <-ch: + case <-time.After(timeout): + t.Fatal(msg) } +} - close(fsm.release) - - var result openResult +func requireNoStartupWaitResult(t *testing.T, ch <-chan error, timeout time.Duration, msg string) { + t.Helper() select { - case result = <-done: - case <-time.After(time.Second): - t.Fatal("multi-node Open did not return after committed tail drain completed") + case err := <-ch: + require.NoError(t, err) + t.Fatal(msg) + case <-time.After(timeout): } - require.NoError(t, result.err) - require.NotNil(t, result.engine) - defer func() { - require.NoError(t, result.engine.Close()) - }() +} + +func requireStartupWaitResult(t *testing.T, ch <-chan error, timeout time.Duration, msg string) { + t.Helper() select { - case <-result.engine.startedCh: - case <-time.After(time.Second): - t.Fatal("engine did not mark started after committed tail drain completed") + case err := <-ch: + require.NoError(t, err) + case <-time.After(timeout): + t.Fatal(msg) } } diff --git a/main.go b/main.go index 236fcb0c7..06f2b1d93 100644 --- a/main.go +++ b/main.go @@ -1543,6 +1543,9 @@ func startServersAfterStartupRotation(waitRotateOnStartup startupRotationWaiter, if err := runner.startRaftTransport(); err != nil { return err } + if err := waitForRaftStartupAfterTransport(in.ctx, in.runtimes); err != nil { + return runner.startupFailure(err) + } if err := runner.preparePublicServices(); err != nil { return runner.startupFailure(err) } @@ -1580,6 +1583,23 @@ func startServersAfterStartupRotation(waitRotateOnStartup startupRotationWaiter, return nil } +func waitForRaftStartupAfterTransport(ctx context.Context, runtimes []*raftGroupRuntime) error { + for _, rt := range runtimes { + if rt == nil { + continue + } + engine := rt.snapshotEngine() + barrier, ok := engine.(raftengine.StartupBarrier) + if !ok { + continue + } + if err := barrier.WaitStarted(ctx); err != nil { + return errors.Wrapf(err, "wait for raft group %d startup", rt.spec.id) + } + } + return nil +} + func configureCoordinatorTSO(coordinate *kv.ShardedCoordinator) error { if !*tsoEnabled { return nil From e20475f700172d112d0a91aaeca3f67891340bc3 Mon Sep 17 00:00:00 2001 From: bootjp Date: Wed, 15 Jul 2026 04:39:54 +0900 Subject: [PATCH 14/26] Stabilize raft dispatch and S3 cleanup fences --- internal/raftengine/etcd/engine.go | 152 ++++++++++++++++------ internal/raftengine/etcd/engine_test.go | 119 +++++++++++------ internal/s3keys/keys.go | 35 +++++ kv/fsm.go | 3 + kv/shard_key_test.go | 10 +- kv/sharded_coordinator_del_prefix_test.go | 40 ++++++ 6 files changed, 275 insertions(+), 84 deletions(-) diff --git a/internal/raftengine/etcd/engine.go b/internal/raftengine/etcd/engine.go index 296bb00f1..b3f704396 100644 --- a/internal/raftengine/etcd/engine.go +++ b/internal/raftengine/etcd/engine.go @@ -71,8 +71,8 @@ const ( priorityStepBurstLimit = 64 // defaultHeartbeatBufPerPeer is the capacity of the priority dispatch channel. // It carries low-frequency control traffic: heartbeats, votes, read-index, - // leader-transfer, and their corresponding response messages - // (MsgHeartbeatResp, MsgReadIndexResp, MsgVoteResp, MsgPreVoteResp). + // leader-transfer, and their corresponding response messages except for + // MsgHeartbeatResp, which uses its own coalescing response lane. // MsgAppResp is intentionally kept in the normal channel: followers — the // only senders of MsgAppResp — do not send MsgApp, so there is no // head-of-line blocking risk there. @@ -84,8 +84,15 @@ const ( // upside is that a ~5 s transient pause (election-timeout scale) // no longer drops heartbeats and forces the peers' lease to expire. defaultHeartbeatBufPerPeer = 512 + // defaultHeartbeatRespBufPerPeer sizes the dedicated follower-to-leader + // heartbeat response lane. When a follower is receiving a large snapshot, + // the leader may continue to send heartbeats while the follower's outbound + // transport is slow. Heartbeat responses are superseded by newer heartbeat + // responses for the same peer, so enqueueDispatchMessage coalesces this lane + // instead of reporting a dropped raft message and starving the leader lease. + defaultHeartbeatRespBufPerPeer = 128 // defaultSnapshotLaneBufPerPeer sizes the per-peer MsgSnap lane when the - // 4-lane dispatcher mode is enabled (see ELASTICKV_RAFT_DISPATCHER_LANES). + // opt-in multi-lane dispatcher is enabled (see ELASTICKV_RAFT_DISPATCHER_LANES). // MsgSnap is rare and bulky; 4 is enough to absorb a retry or two without // holding up MsgApp replication behind a multi-MiB payload. defaultSnapshotLaneBufPerPeer = 4 @@ -93,11 +100,11 @@ const ( // types not classified as heartbeat/replication/snapshot (e.g. surprise // locally-addressed control types). Small buffer: traffic volume is tiny. defaultOtherLaneBufPerPeer = 16 - // dispatcherLanesEnvVar toggles the 4-lane dispatcher (heartbeat / - // replication / snapshot / other). When unset or "0", the legacy - // 2-lane layout (heartbeat + normal) is used. Opt-in by design: the - // raft hot path is high blast radius and a regression here can cause - // cluster-wide elections. + // dispatcherLanesEnvVar toggles the multi-lane dispatcher (heartbeat / + // heartbeatResp / replication / snapshot / other). When unset or "0", the + // legacy 3-lane layout (heartbeat + heartbeatResp + normal) is used. + // Opt-in by design: the raft hot path is high blast radius and a regression + // here can cause cluster-wide elections. dispatcherLanesEnvVar = "ELASTICKV_RAFT_DISPATCHER_LANES" // defaultSnapshotEvery is the fallback trigger threshold: take an FSM // snapshot once the applied index has advanced this many entries past @@ -328,7 +335,7 @@ type Engine struct { dispatchReportCh chan dispatchReport peerDispatchers map[uint64]*peerQueues perPeerQueueSize int - // dispatcherLanesEnabled toggles the 4-lane dispatcher layout. Captured + // dispatcherLanesEnabled toggles the opt-in multi-lane dispatcher layout. Captured // once at Open from ELASTICKV_RAFT_DISPATCHER_LANES so the run-time code // path is branch-free per message and does not need to re-read env vars. dispatcherLanesEnabled bool @@ -568,22 +575,23 @@ type dispatchRequest struct { // peerQueues holds separate dispatch channels per peer so that heartbeats // are never blocked behind large log-entry RPCs. // -// Legacy 2-lane layout (default): heartbeat + normal. +// Legacy 3-lane layout (default): heartbeat + heartbeatResp + normal. // -// 4-lane layout (opt-in via ELASTICKV_RAFT_DISPATCHER_LANES=1): heartbeat + -// replication (MsgApp/MsgAppResp) + snapshot (MsgSnap) + other. Each lane -// gets its own goroutine so a bulky MsgSnap transfer cannot stall MsgApp -// replication and vice versa. Per-peer ordering within a given message type -// is preserved because a single peer's MsgApp stream all share one lane and -// one worker. +// 5-lane layout (opt-in via ELASTICKV_RAFT_DISPATCHER_LANES=1): heartbeat + +// heartbeatResp + replication (MsgApp/MsgAppResp) + snapshot (MsgSnap) + +// other. Each lane gets its own goroutine so a bulky MsgSnap transfer cannot +// stall MsgApp replication and vice versa. Per-peer ordering within a given +// message type is preserved because a single peer's MsgApp stream all share +// one lane and one worker. type peerQueues struct { - normal chan dispatchRequest - heartbeat chan dispatchRequest - replication chan dispatchRequest // 4-lane mode only; nil otherwise - snapshot chan dispatchRequest // 4-lane mode only; nil otherwise - other chan dispatchRequest // 4-lane mode only; nil otherwise - ctx context.Context - cancel context.CancelFunc + normal chan dispatchRequest + heartbeat chan dispatchRequest + heartbeatResp chan dispatchRequest + replication chan dispatchRequest // 5-lane mode only; nil otherwise + snapshot chan dispatchRequest // 5-lane mode only; nil otherwise + other chan dispatchRequest // 5-lane mode only; nil otherwise + ctx context.Context + cancel context.CancelFunc } type preparedOpenState struct { @@ -2370,6 +2378,9 @@ func (e *Engine) enqueueDispatchMessage(msg raftpb.Message) error { // is already full. The len/cap check is safe here because this function is // only ever called from the single engine event-loop goroutine. if len(ch) >= cap(ch) { + if msg.GetType() == raftpb.MsgHeartbeatResp && coalesceHeartbeatResp(ch, msg) { + return nil + } e.recordDroppedDispatch(msg) return nil } @@ -2384,12 +2395,61 @@ func (e *Engine) enqueueDispatchMessage(msg raftpb.Message) error { } } +func coalesceHeartbeatResp(ch chan dispatchRequest, msg raftpb.Message) bool { + if ch == nil || cap(ch) == 0 { + return false + } + stale, ok := tryReceiveDispatchRequest(ch) + if !ok { + return false + } + if stale.msg.GetType() != raftpb.MsgHeartbeatResp { + restoreDispatchRequest(ch, stale) + return false + } + closeDispatchRequest(stale) + return tryEnqueueDispatchRequest(ch, prepareDispatchRequest(msg)) +} + +func tryReceiveDispatchRequest(ch chan dispatchRequest) (dispatchRequest, bool) { + select { + case req := <-ch: + return req, true + default: + return dispatchRequest{}, false + } +} + +func restoreDispatchRequest(ch chan dispatchRequest, req dispatchRequest) { + select { + case ch <- req: + default: + closeDispatchRequest(req) + } +} + +func tryEnqueueDispatchRequest(ch chan dispatchRequest, req dispatchRequest) bool { + select { + case ch <- req: + return true + default: + closeDispatchRequest(req) + return false + } +} + +func closeDispatchRequest(req dispatchRequest) { + if err := req.Close(); err != nil { + slog.Error("etcd raft dispatch: failed to close request", "err", err) + } +} + // isPriorityMsg returns true for small, low-frequency control messages that // must not be queued behind large MsgApp payloads in the normal channel. // MsgAppResp is intentionally excluded: it is sent by followers, which never // send MsgApp, so it faces no head-of-line blocking in the normal channel. -// Keeping it out of the priority queue preserves the low-frequency invariant -// that justifies defaultHeartbeatBufPerPeer = 64. +// Keeping MsgHeartbeatResp on its own coalescing lane preserves the +// low-frequency invariant that justifies defaultHeartbeatBufPerPeer. func isPriorityMsg(t raftpb.MessageType) bool { return t == raftpb.MsgHeartbeat || t == raftpb.MsgHeartbeatResp || t == raftpb.MsgReadIndex || t == raftpb.MsgReadIndexResp || @@ -2399,11 +2459,15 @@ func isPriorityMsg(t raftpb.MessageType) bool { } // selectDispatchLane picks the per-peer channel for msgType. In the legacy -// 2-lane layout it returns pd.heartbeat for priority control traffic and -// pd.normal for everything else. In the 4-lane layout it additionally -// partitions the non-heartbeat traffic so that MsgApp/MsgAppResp and MsgSnap -// do not share a goroutine and cannot block each other. +// layout it returns pd.heartbeatResp for MsgHeartbeatResp, pd.heartbeat for +// other priority control traffic, and pd.normal for everything else. In the +// opt-in multi-lane layout it additionally partitions the non-heartbeat traffic +// so that MsgApp/MsgAppResp and MsgSnap do not share a goroutine and cannot +// block each other. func (e *Engine) selectDispatchLane(pd *peerQueues, msgType raftpb.MessageType) chan dispatchRequest { + if msgType == raftpb.MsgHeartbeatResp && pd.heartbeatResp != nil { + return pd.heartbeatResp + } // Priority control traffic (heartbeats, votes, read-index, timeout-now) // always rides the heartbeat lane in both layouts so it keeps its // low-latency treatment and is never stuck behind MsgApp payloads. @@ -4296,24 +4360,26 @@ func (e *Engine) startPeerDispatcher(nodeID uint64) { } ctx, cancel := context.WithCancel(baseCtx) pd := &peerQueues{ - heartbeat: make(chan dispatchRequest, defaultHeartbeatBufPerPeer), - ctx: ctx, - cancel: cancel, + heartbeat: make(chan dispatchRequest, defaultHeartbeatBufPerPeer), + heartbeatResp: make(chan dispatchRequest, defaultHeartbeatRespBufPerPeer), + ctx: ctx, + cancel: cancel, } var workers []chan dispatchRequest if e.dispatcherLanesEnabled { - // 4-lane layout: split MsgApp/MsgAppResp (replication), MsgSnap - // (snapshot), and misc (other) onto independent goroutines so a - // bulky snapshot transfer cannot stall replication. Each channel - // still serves a single peer, so within-type ordering (the raft - // invariant we care about for MsgApp) is preserved. + // 5-lane layout: split MsgHeartbeatResp, MsgApp/MsgAppResp + // (replication), MsgSnap (snapshot), and misc (other) onto independent + // goroutines so a bulky snapshot transfer cannot stall replication or + // follower heartbeat responses. Each channel still serves a single + // peer, so within-type ordering (the raft invariant we care about for + // MsgApp) is preserved. pd.replication = make(chan dispatchRequest, size) pd.snapshot = make(chan dispatchRequest, defaultSnapshotLaneBufPerPeer) pd.other = make(chan dispatchRequest, defaultOtherLaneBufPerPeer) - workers = []chan dispatchRequest{pd.heartbeat, pd.replication, pd.snapshot, pd.other} + workers = []chan dispatchRequest{pd.heartbeat, pd.heartbeatResp, pd.replication, pd.snapshot, pd.other} } else { pd.normal = make(chan dispatchRequest, size) - workers = []chan dispatchRequest{pd.normal, pd.heartbeat} + workers = []chan dispatchRequest{pd.normal, pd.heartbeat, pd.heartbeatResp} } e.peerDispatchers[nodeID] = pd e.dispatchWG.Add(len(workers)) @@ -4370,7 +4436,7 @@ func snapshotEveryFromEnv() uint64 { return n } -// dispatcherLanesEnabledFromEnv returns true when the 4-lane dispatcher has +// dispatcherLanesEnabledFromEnv returns true when the multi-lane dispatcher has // been explicitly opted into via ELASTICKV_RAFT_DISPATCHER_LANES. The value // is parsed with strconv.ParseBool, which accepts the standard tokens // (1, t, T, TRUE, true, True enable; 0, f, F, FALSE, false, False disable). @@ -4385,10 +4451,10 @@ func dispatcherLanesEnabledFromEnv() bool { } // closePeerLanes closes every non-nil dispatch channel on pd so that the -// drain loops in runDispatchWorker exit. It is safe to call with either the -// 2-lane or 4-lane layout because unused lanes are nil. +// drain loops in runDispatchWorker exit. It is safe to call with any dispatch +// layout because unused lanes are nil. func closePeerLanes(pd *peerQueues) { - for _, ch := range []chan dispatchRequest{pd.heartbeat, pd.normal, pd.replication, pd.snapshot, pd.other} { + for _, ch := range []chan dispatchRequest{pd.heartbeat, pd.heartbeatResp, pd.normal, pd.replication, pd.snapshot, pd.other} { if ch != nil { close(ch) } diff --git a/internal/raftengine/etcd/engine_test.go b/internal/raftengine/etcd/engine_test.go index e379282f7..8a1f82b4a 100644 --- a/internal/raftengine/etcd/engine_test.go +++ b/internal/raftengine/etcd/engine_test.go @@ -853,6 +853,7 @@ func TestUpsertPeerStartsDispatcherAndAcceptsMessages(t *testing.T) { pd, ok := engine.peerDispatchers[2] require.True(t, ok, "dispatcher must be created on upsert") require.Equal(t, defaultHeartbeatBufPerPeer, cap(pd.heartbeat)) + require.Equal(t, defaultHeartbeatRespBufPerPeer, cap(pd.heartbeatResp)) require.Equal(t, 4, cap(pd.normal)) require.NoError(t, engine.enqueueDispatchMessage(raftpb.Message{Type: messageTypePtr(raftpb.MsgHeartbeat), To: uint64Ptr(2)})) @@ -872,10 +873,11 @@ func TestRemovePeerClosesDispatcherAndDropsSubsequentMessages(t *testing.T) { stopCh := make(chan struct{}) ctx, cancel := context.WithCancel(context.Background()) pd := &peerQueues{ - normal: make(chan dispatchRequest, 4), - heartbeat: make(chan dispatchRequest, 4), - ctx: ctx, - cancel: cancel, + normal: make(chan dispatchRequest, 4), + heartbeat: make(chan dispatchRequest, 4), + heartbeatResp: make(chan dispatchRequest, 4), + ctx: ctx, + cancel: cancel, } engine := &Engine{ nodeID: 1, @@ -883,9 +885,10 @@ func TestRemovePeerClosesDispatcherAndDropsSubsequentMessages(t *testing.T) { peerDispatchers: map[uint64]*peerQueues{2: pd}, dispatchStopCh: stopCh, } - engine.dispatchWG.Add(2) + engine.dispatchWG.Add(3) go engine.runDispatchWorker(ctx, pd.normal) go engine.runDispatchWorker(ctx, pd.heartbeat) + go engine.runDispatchWorker(ctx, pd.heartbeatResp) engine.removePeer(2) @@ -1390,6 +1393,37 @@ func TestPrepareDispatchRequestClonesSnapshotPayload(t *testing.T) { require.Equal(t, []uint64{1, 2}, req.msg.Snapshot.GetMetadata().GetConfState().GetVoters()) } +func TestEnqueueDispatchMessageCoalescesHeartbeatResponses(t *testing.T) { + t.Parallel() + pd := &peerQueues{ + heartbeat: make(chan dispatchRequest, 1), + heartbeatResp: make(chan dispatchRequest, 1), + } + engine := &Engine{ + nodeID: 1, + peerDispatchers: map[uint64]*peerQueues{ + 2: pd, + }, + } + pd.heartbeatResp <- prepareDispatchRequest(raftpb.Message{ + Type: messageTypePtr(raftpb.MsgHeartbeatResp), + To: uint64Ptr(2), + Context: []byte("old"), + }) + + require.NoError(t, engine.enqueueDispatchMessage(raftpb.Message{ + Type: messageTypePtr(raftpb.MsgHeartbeatResp), + To: uint64Ptr(2), + Context: []byte("new"), + })) + + require.Zero(t, engine.DispatchDropCount()) + require.Len(t, pd.heartbeatResp, 1) + req := <-pd.heartbeatResp + require.Equal(t, raftpb.MsgHeartbeatResp, req.msg.GetType()) + require.Equal(t, []byte("new"), req.msg.Context) +} + func TestMaxAppliedIndexStartsFromSnapshotIndex(t *testing.T) { storage := etcdraft.NewMemoryStorage() snap := raftTestSnapshot(5, 2, []uint64{1}, nil) @@ -2073,20 +2107,22 @@ func TestErrNotLeaderMatchesRaftEngineSentinel(t *testing.T) { require.True(t, errors.Is(errors.WithStack(errLeadershipTransferConfChangePending), raftengine.ErrLeadershipTransferConfChangePending)) } -// TestSelectDispatchLane_LegacyTwoLane verifies that, when the 4-lane -// dispatcher is disabled (default), messages are routed exactly as before: -// priority control traffic → heartbeat lane, everything else → normal lane. -func TestSelectDispatchLane_LegacyTwoLane(t *testing.T) { +// TestSelectDispatchLane_LegacyThreeLane verifies that, when the opt-in +// multi-lane dispatcher is disabled (default), priority control traffic uses +// the heartbeat lane, heartbeat responses use their coalescing response lane, +// and everything else uses the normal lane. +func TestSelectDispatchLane_LegacyThreeLane(t *testing.T) { t.Parallel() engine := &Engine{dispatcherLanesEnabled: false} pd := &peerQueues{ - normal: make(chan dispatchRequest, 1), - heartbeat: make(chan dispatchRequest, 1), + normal: make(chan dispatchRequest, 1), + heartbeat: make(chan dispatchRequest, 1), + heartbeatResp: make(chan dispatchRequest, 1), } cases := map[raftpb.MessageType]chan dispatchRequest{ raftpb.MsgHeartbeat: pd.heartbeat, - raftpb.MsgHeartbeatResp: pd.heartbeat, + raftpb.MsgHeartbeatResp: pd.heartbeatResp, raftpb.MsgReadIndex: pd.heartbeat, raftpb.MsgReadIndexResp: pd.heartbeat, raftpb.MsgVote: pd.heartbeat, @@ -2104,22 +2140,24 @@ func TestSelectDispatchLane_LegacyTwoLane(t *testing.T) { } } -// TestSelectDispatchLane_FourLane verifies that, when ELASTICKV_RAFT_DISPATCHER_LANES +// TestSelectDispatchLane_FiveLane verifies that, when ELASTICKV_RAFT_DISPATCHER_LANES // is enabled, MsgApp/MsgAppResp goes to the replication lane, MsgSnap goes to -// the snapshot lane, and heartbeats/votes/read-index share the priority lane. -func TestSelectDispatchLane_FourLane(t *testing.T) { +// the snapshot lane, heartbeat responses get their own lane, and +// heartbeats/votes/read-index share the priority lane. +func TestSelectDispatchLane_FiveLane(t *testing.T) { t.Parallel() engine := &Engine{dispatcherLanesEnabled: true} pd := &peerQueues{ - heartbeat: make(chan dispatchRequest, 1), - replication: make(chan dispatchRequest, 1), - snapshot: make(chan dispatchRequest, 1), - other: make(chan dispatchRequest, 1), + heartbeat: make(chan dispatchRequest, 1), + heartbeatResp: make(chan dispatchRequest, 1), + replication: make(chan dispatchRequest, 1), + snapshot: make(chan dispatchRequest, 1), + other: make(chan dispatchRequest, 1), } cases := map[raftpb.MessageType]chan dispatchRequest{ raftpb.MsgHeartbeat: pd.heartbeat, - raftpb.MsgHeartbeatResp: pd.heartbeat, + raftpb.MsgHeartbeatResp: pd.heartbeatResp, raftpb.MsgVote: pd.heartbeat, raftpb.MsgVoteResp: pd.heartbeat, raftpb.MsgPreVote: pd.heartbeat, @@ -2133,7 +2171,7 @@ func TestSelectDispatchLane_FourLane(t *testing.T) { } for mt, want := range cases { got := engine.selectDispatchLane(pd, mt) - require.Equalf(t, want, got, "4-lane mode routing for %s", mt) + require.Equalf(t, want, got, "multi-lane mode routing for %s", mt) } } @@ -2146,10 +2184,11 @@ func TestSelectDispatchLane_MsgPropReachesDefaultFallback(t *testing.T) { t.Parallel() engine := &Engine{nodeID: 1, dispatcherLanesEnabled: true} pd := &peerQueues{ - heartbeat: make(chan dispatchRequest, 1), - replication: make(chan dispatchRequest, 1), - snapshot: make(chan dispatchRequest, 1), - other: make(chan dispatchRequest, 1), + heartbeat: make(chan dispatchRequest, 1), + heartbeatResp: make(chan dispatchRequest, 1), + replication: make(chan dispatchRequest, 1), + snapshot: make(chan dispatchRequest, 1), + other: make(chan dispatchRequest, 1), } require.NotPanics(t, func() { got := engine.selectDispatchLane(pd, raftpb.MsgProp) @@ -2157,11 +2196,11 @@ func TestSelectDispatchLane_MsgPropReachesDefaultFallback(t *testing.T) { }) } -// TestFourLaneDispatcher_SnapshotDoesNotBlockReplication exercises the key -// correctness invariant for the 4-lane layout: a stuck MsgSnap transfer must +// TestMultiLaneDispatcher_SnapshotDoesNotBlockReplication exercises the key +// correctness invariant for the multi-lane layout: a stuck MsgSnap transfer must // not prevent MsgApp from being dispatched, because they now run on // independent goroutines. -func TestFourLaneDispatcher_SnapshotDoesNotBlockReplication(t *testing.T) { +func TestMultiLaneDispatcher_SnapshotDoesNotBlockReplication(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) t.Cleanup(cancel) @@ -2214,19 +2253,20 @@ func TestFourLaneDispatcher_SnapshotDoesNotBlockReplication(t *testing.T) { engine.dispatchWG.Wait() } -// TestFourLaneDispatcher_RemovePeerClosesAllLanes confirms removePeer closes +// TestMultiLaneDispatcher_RemovePeerClosesAllLanes confirms removePeer closes // every lane (not just normal/heartbeat) so no worker goroutine leaks under -// the opt-in 4-lane layout. -func TestFourLaneDispatcher_RemovePeerClosesAllLanes(t *testing.T) { +// the opt-in multi-lane layout. +func TestMultiLaneDispatcher_RemovePeerClosesAllLanes(t *testing.T) { stopCh := make(chan struct{}) ctx, cancel := context.WithCancel(context.Background()) pd := &peerQueues{ - heartbeat: make(chan dispatchRequest, 4), - replication: make(chan dispatchRequest, 4), - snapshot: make(chan dispatchRequest, 4), - other: make(chan dispatchRequest, 4), - ctx: ctx, - cancel: cancel, + heartbeat: make(chan dispatchRequest, 4), + heartbeatResp: make(chan dispatchRequest, 4), + replication: make(chan dispatchRequest, 4), + snapshot: make(chan dispatchRequest, 4), + other: make(chan dispatchRequest, 4), + ctx: ctx, + cancel: cancel, } engine := &Engine{ nodeID: 1, @@ -2235,8 +2275,9 @@ func TestFourLaneDispatcher_RemovePeerClosesAllLanes(t *testing.T) { dispatchStopCh: stopCh, dispatcherLanesEnabled: true, } - engine.dispatchWG.Add(4) + engine.dispatchWG.Add(5) go engine.runDispatchWorker(ctx, pd.heartbeat) + go engine.runDispatchWorker(ctx, pd.heartbeatResp) go engine.runDispatchWorker(ctx, pd.replication) go engine.runDispatchWorker(ctx, pd.snapshot) go engine.runDispatchWorker(ctx, pd.other) @@ -2251,7 +2292,7 @@ func TestFourLaneDispatcher_RemovePeerClosesAllLanes(t *testing.T) { select { case <-done: case <-time.After(time.Second): - t.Fatal("4-lane dispatch workers did not exit after peer removal") + t.Fatal("multi-lane dispatch workers did not exit after peer removal") } // Subsequent sends to the removed peer must be dropped without panic. diff --git a/internal/s3keys/keys.go b/internal/s3keys/keys.go index c90324219..d3ef66178 100644 --- a/internal/s3keys/keys.go +++ b/internal/s3keys/keys.go @@ -245,6 +245,41 @@ func RoutePrefixForBucketAnyGeneration(bucket string) []byte { return out } +func BucketGenerationRoutePrefixForCleanupPrefix(prefix []byte) ([]byte, bool) { + familyPrefix := bucketGenerationFamilyPrefix(prefix) + if familyPrefix == nil { + return nil, false + } + bucketRaw, next, ok := decodeSegment(prefix, len(familyPrefix)) + if !ok { + return nil, false + } + generation, next, ok := readU64(prefix, next) + if !ok || next != len(prefix) { + return nil, false + } + return RoutePrefixForBucket(string(bucketRaw), generation), true +} + +func bucketGenerationFamilyPrefix(key []byte) []byte { + switch { + case bytes.HasPrefix(key, objectManifestPrefixBytes): + return objectManifestPrefixBytes + case bytes.HasPrefix(key, uploadMetaPrefixBytes): + return uploadMetaPrefixBytes + case bytes.HasPrefix(key, uploadPartPrefixBytes): + return uploadPartPrefixBytes + case bytes.HasPrefix(key, blobPrefixBytes): + return blobPrefixBytes + case bytes.HasPrefix(key, gcUploadPrefixBytes): + return gcUploadPrefixBytes + case bytes.HasPrefix(key, routePrefixBytes): + return routePrefixBytes + default: + return nil + } +} + func bucketScopedPrefix(prefix []byte, bucket string, generation uint64) []byte { out := make([]byte, 0, len(prefix)+len(bucket)+u64Bytes+segmentEscapeOverhead) out = append(out, prefix...) diff --git a/kv/fsm.go b/kv/fsm.go index bdb7fc025..62d29647b 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -540,6 +540,9 @@ func routePrefixRange(prefix []byte) ([]byte, []byte) { if len(prefix) == 0 { return []byte(""), nil } + if start, ok := s3keys.BucketGenerationRoutePrefixForCleanupPrefix(prefix); ok { + return start, prefixScanEnd(start) + } if routeKeyspaceWideRawPrefix(prefix) { return []byte(""), nil } diff --git a/kv/shard_key_test.go b/kv/shard_key_test.go index 67d18ad7b..817a1d809 100644 --- a/kv/shard_key_test.go +++ b/kv/shard_key_test.go @@ -323,8 +323,14 @@ func TestRoutePrefixRangeTreatsBroadMappedPrefixesAsFullKeyspace(t *testing.T) { { name: "s3 bucket cleanup prefix", prefix: s3keys.ObjectManifestPrefixForBucket("bucket", 2), - wantStart: []byte(""), - wantEnd: nil, + wantStart: s3keys.RoutePrefixForBucket("bucket", 2), + wantEnd: prefixScanEnd(s3keys.RoutePrefixForBucket("bucket", 2)), + }, + { + name: "s3 bucket cleanup route prefix", + prefix: s3keys.RoutePrefixForBucket("bucket", 2), + wantStart: s3keys.RoutePrefixForBucket("bucket", 2), + wantEnd: prefixScanEnd(s3keys.RoutePrefixForBucket("bucket", 2)), }, } { t.Run(tc.name, func(t *testing.T) { diff --git a/kv/sharded_coordinator_del_prefix_test.go b/kv/sharded_coordinator_del_prefix_test.go index 5cc0f845c..1b8a31cb2 100644 --- a/kv/sharded_coordinator_del_prefix_test.go +++ b/kv/sharded_coordinator_del_prefix_test.go @@ -267,6 +267,46 @@ func TestShardedCoordinatorRejectsBroadInternalDelPrefixWhenRouteIsWriteFenced(t } } +func TestShardedCoordinatorAllowsS3BucketDelPrefixWhenUnrelatedRouteIsWriteFenced(t *testing.T) { + t.Parallel() + + const ( + activeBucket = "bucket-a" + fencedBucket = "bucket-b" + generation = uint64(7) + ) + activeStart := s3keys.RoutePrefixForBucketAnyGeneration(activeBucket) + activeEnd := prefixScanEnd(activeStart) + fencedStart := s3keys.RoutePrefixForBucketAnyGeneration(fencedBucket) + fencedEnd := prefixScanEnd(fencedStart) + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: activeStart, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: activeStart, End: activeEnd, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 3, Start: activeEnd, End: fencedStart, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 4, Start: fencedStart, End: fencedEnd, GroupID: 2, State: distribution.RouteStateWriteFenced}, + {RouteID: 5, Start: fencedEnd, End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }, + })) + + 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: s3keys.ObjectManifestPrefixForBucket(activeBucket, generation)}}, + }) + require.NoError(t, err) + require.NotEmpty(t, g1Txn.requests) + require.NotEmpty(t, g2Txn.requests, "DEL_PREFIX still broadcasts to every group after the narrow fence precheck passes") +} + // TestShardedCoordinator_DelPrefixRejectsTxn verifies that DEL_PREFIX inside // a transactional group is rejected. func TestShardedCoordinator_DelPrefixRejectsTxn(t *testing.T) { From c3a8ecff66c409963a80a244fc600a581c176a7e Mon Sep 17 00:00:00 2001 From: bootjp Date: Wed, 15 Jul 2026 05:28:28 +0900 Subject: [PATCH 15/26] Stabilize raft startup and write fences --- adapter/grpc_test.go | 6 +- .../raftengine/etcd/dispatch_report_test.go | 57 +++++ internal/raftengine/etcd/engine.go | 19 +- kv/coordinator.go | 19 +- kv/coordinator_dispatch_test.go | 17 +- kv/fsm.go | 82 +++++++ kv/fsm_migration_fence_test.go | 89 +++++++ kv/sharded_coordinator.go | 33 ++- kv/sharded_coordinator_del_prefix_test.go | 46 ++++ main.go | 218 ++++++++++++------ main_encryption_registration.go | 43 +++- 11 files changed, 528 insertions(+), 101 deletions(-) diff --git a/adapter/grpc_test.go b/adapter/grpc_test.go index 781aa9f47..c70818adb 100644 --- a/adapter/grpc_test.go +++ b/adapter/grpc_test.go @@ -19,6 +19,8 @@ import ( goproto "google.golang.org/protobuf/proto" ) +const consistencySequenceIterations = 512 + func Test_value_can_be_deleted(t *testing.T) { t.Parallel() nodes, adders, _ := createNode(t, 3) @@ -466,7 +468,7 @@ func Test_consistency_satisfy_write_after_read_sequence(t *testing.T) { // not abort the test. The post-RPC assert.Equal still pins the // consistency invariant: once Put eventually succeeds, the // subsequent Get must return the same value, otherwise we fail. - for i := range 9999 { + for i := range consistencySequenceIterations { want := []byte("sequence" + strconv.Itoa(i)) err := retryNotLeader(ctx, func() error { _, perr := c.RawPut(ctx, &pb.RawPutRequest{Key: key, Value: want}) @@ -521,7 +523,7 @@ func Test_grpc_transaction(t *testing.T) { // _sequence: tolerate transient leader churn (purely availability, // not consistency) while keeping the Put → Get → Delete → Get // invariants strict. - for i := range 9999 { + for i := range consistencySequenceIterations { want := []byte("sequence" + strconv.Itoa(i)) err := retryNotLeader(ctx, func() error { _, perr := c.Put(ctx, &pb.PutRequest{Key: key, Value: want}) diff --git a/internal/raftengine/etcd/dispatch_report_test.go b/internal/raftengine/etcd/dispatch_report_test.go index dfd1dfd67..6b122952b 100644 --- a/internal/raftengine/etcd/dispatch_report_test.go +++ b/internal/raftengine/etcd/dispatch_report_test.go @@ -50,6 +50,63 @@ func TestReportSuccessfulDispatchReportsSnapshotFinish(t *testing.T) { } } +func TestReportSuccessfulDispatchWaitsForSnapshotFinishSlot(t *testing.T) { + t.Parallel() + e := &Engine{ + dispatchReportCh: make(chan dispatchReport, 1), + closeCh: make(chan struct{}), + } + blockingReport := dispatchReport{to: 1, msgType: raftpb.MsgApp} + e.dispatchReportCh <- blockingReport + + done := make(chan struct{}) + go func() { + e.reportSuccessfulDispatch(raftpb.Message{Type: messageTypePtr(raftpb.MsgSnap), To: uint64Ptr(2)}) + close(done) + }() + + select { + case <-done: + t.Fatal("snapshot finish report returned while dispatchReportCh was full") + case <-time.After(50 * time.Millisecond): + } + + require.Equal(t, blockingReport, <-e.dispatchReportCh) + select { + case <-done: + case <-time.After(testDispatchReportTimeout): + t.Fatal("snapshot finish report did not complete after dispatchReportCh had space") + } + select { + case got := <-e.dispatchReportCh: + require.Equal(t, dispatchReport{to: 2, msgType: raftpb.MsgSnap, snapshotFinish: true}, got) + default: + t.Fatal("expected reliable successful MsgSnap dispatch report") + } +} + +func TestReportSuccessfulDispatchAbortsOnCloseWhenReportFull(t *testing.T) { + t.Parallel() + e := &Engine{ + dispatchReportCh: make(chan dispatchReport, 1), + closeCh: make(chan struct{}), + } + e.dispatchReportCh <- dispatchReport{to: 1, msgType: raftpb.MsgApp} + close(e.closeCh) + + done := make(chan struct{}) + go func() { + e.reportSuccessfulDispatch(raftpb.Message{Type: messageTypePtr(raftpb.MsgSnap), To: uint64Ptr(2)}) + close(done) + }() + + select { + case <-done: + case <-time.After(testDispatchReportTimeout): + t.Fatal("snapshot finish report did not abort when closeCh was signalled") + } +} + func TestReportSuccessfulDispatchIgnoresRegularMessage(t *testing.T) { t.Parallel() e := &Engine{ diff --git a/internal/raftengine/etcd/engine.go b/internal/raftengine/etcd/engine.go index b3f704396..0326b66d1 100644 --- a/internal/raftengine/etcd/engine.go +++ b/internal/raftengine/etcd/engine.go @@ -1868,7 +1868,8 @@ func (e *Engine) handleDispatchReport(report dispatchReport) { // If the channel is full (unlikely — the buffer is sized to MaxInflightMsg), // the report is dropped and logged; this is acceptable because raft will retry // on the next tick and we only need eventual consistency between transport -// state and Progress state. +// state and Progress state. Successful MsgSnap reports are the exception; see +// postReliableDispatchReport. func (e *Engine) postDispatchReport(report dispatchReport) { select { case e.dispatchReportCh <- report: @@ -1881,6 +1882,20 @@ func (e *Engine) postDispatchReport(report dispatchReport) { } } +// postReliableDispatchReport delivers a dispatch outcome that must not be +// dropped. SnapshotFinish is not eventually consistent with ordinary +// unreachable reports: if it is lost, raft can keep the follower's Progress in +// StateSnapshot after the follower accepted the snapshot. +func (e *Engine) postReliableDispatchReport(report dispatchReport) { + if e.dispatchReportCh == nil { + return + } + select { + case e.dispatchReportCh <- report: + case <-e.closeCh: + } +} + func (e *Engine) handleProposal(req proposalRequest) { if err := contextErr(req.ctx); err != nil { req.done <- proposalResult{err: err} @@ -4523,7 +4538,7 @@ func (e *Engine) reportSuccessfulDispatch(msg raftpb.Message) { if msg.GetType() != raftpb.MsgSnap { return } - e.postDispatchReport(dispatchReport{ + e.postReliableDispatchReport(dispatchReport{ to: msg.GetTo(), msgType: msg.GetType(), snapshotFinish: true, diff --git a/kv/coordinator.go b/kv/coordinator.go index f1b8a5627..174c3b636 100644 --- a/kv/coordinator.go +++ b/kv/coordinator.go @@ -1145,12 +1145,13 @@ func (c *Coordinate) dispatchRaw(ctx context.Context, req []*Elem[OP]) (*Coordin // The returned Request is structurally identical to the pre-stamping // shape the leader's stampRawTimestamps already handles for Ts == 0 // (see adapter/internal.go). -func (c *Coordinate) toRawRequest(req *Elem[OP]) *pb.Request { +func (c *Coordinate) toRawRequest(req *Elem[OP], observedRouteVersion uint64) *pb.Request { switch req.Op { case Put: return &pb.Request{ - IsTxn: false, - Phase: pb.Phase_NONE, + IsTxn: false, + Phase: pb.Phase_NONE, + ObservedRouteVersion: observedRouteVersion, Mutations: []*pb.Mutation{ { Op: pb.Op_PUT, @@ -1162,8 +1163,9 @@ func (c *Coordinate) toRawRequest(req *Elem[OP]) *pb.Request { case Del: return &pb.Request{ - IsTxn: false, - Phase: pb.Phase_NONE, + IsTxn: false, + Phase: pb.Phase_NONE, + ObservedRouteVersion: observedRouteVersion, Mutations: []*pb.Mutation{ { Op: pb.Op_DEL, @@ -1174,8 +1176,9 @@ func (c *Coordinate) toRawRequest(req *Elem[OP]) *pb.Request { case DelPrefix: return &pb.Request{ - IsTxn: false, - Phase: pb.Phase_NONE, + IsTxn: false, + Phase: pb.Phase_NONE, + ObservedRouteVersion: observedRouteVersion, Mutations: []*pb.Mutation{ { Op: pb.Op_DEL_PREFIX, @@ -1238,7 +1241,7 @@ func (c *Coordinate) buildRedirectRequests(reqs *OperationGroup[OP]) ([]*pb.Requ if !reqs.IsTxn { requests := make([]*pb.Request, 0, len(reqs.Elems)) for _, req := range reqs.Elems { - requests = append(requests, c.toRawRequest(req)) + requests = append(requests, c.toRawRequest(req, reqs.ObservedRouteVersion)) } return requests, nil } diff --git a/kv/coordinator_dispatch_test.go b/kv/coordinator_dispatch_test.go index a559c6405..ae1ffdf30 100644 --- a/kv/coordinator_dispatch_test.go +++ b/kv/coordinator_dispatch_test.go @@ -213,7 +213,7 @@ func TestToRawRequestLeavesTsForLeaderStamping(t *testing.T) { } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { - r := c.toRawRequest(tc.req) + r := c.toRawRequest(tc.req, 0) require.NotNil(t, r) require.Equal(t, uint64(0), r.Ts, "forwarded raw requests must arrive with Ts==0 so the leader's stampRawTimestamps assigns the canonical ts (HLC leader-only invariant + HLC-4 (iii) fence)") @@ -221,6 +221,21 @@ func TestToRawRequestLeavesTsForLeaderStamping(t *testing.T) { } } +func TestBuildRedirectRequests_PreservesRawObservedRouteVersion(t *testing.T) { + t.Parallel() + + c := &Coordinate{} + got, err := c.buildRedirectRequests(&OperationGroup[OP]{ + ObservedRouteVersion: 17, + Elems: []*Elem[OP]{ + {Op: Put, Key: []byte("k"), Value: []byte("v")}, + }, + }) + require.NoError(t, err) + require.Len(t, got, 1) + require.Equal(t, uint64(17), got[0].GetObservedRouteVersion()) +} + // TestBuildRedirectRequestsSurvivesStaleFollowerCeiling exercises the // follower's redirect path end-to-end with an expired ceiling: it // confirms the follower hands off a Ts==0 request to the leader diff --git a/kv/fsm.go b/kv/fsm.go index 62d29647b..bc5da0b36 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -482,6 +482,9 @@ func (f *kvFSM) handleRequest(ctx context.Context, r *pb.Request, commitTS uint6 } func (f *kvFSM) handleRawRequest(ctx context.Context, r *pb.Request, commitTS uint64) error { + if err := f.verifyWriteFence(r); err != nil { + return err + } // DEL_PREFIX mutations are handled by the store's DeletePrefixAt which // scans and writes tombstones locally. A DEL_PREFIX request must be the // sole mutation in a request (enforced by the coordinator's toRawRequest). @@ -720,6 +723,9 @@ func (f *kvFSM) IsVolatileOnlyPayload(payload []byte) bool { } func (f *kvFSM) handleTxnRequest(ctx context.Context, r *pb.Request, commitTS uint64) error { + if err := f.verifyWriteFence(r); err != nil { + return err + } if err := f.verifyComposed1(r); err != nil { return err } @@ -817,6 +823,82 @@ func (f *kvFSM) verifyComposed1(r *pb.Request) error { return f.verifyOwnerFromSnapshot(r.GetMutations(), currentSnap, currentSnap.Version(), "current") } +func (f *kvFSM) verifyWriteFence(r *pb.Request) error { + if requestBypassesWriteFence(r) { + return nil + } + observedVer := r.GetObservedRouteVersion() + if !f.writeFenceHistoryReady(observedVer) { + return nil + } + observedSnap, ok := f.routes.SnapshotAt(observedVer) + if !ok { + return errors.WithStack(ErrComposed1VersionGCd) + } + if err := verifyWriteFenceFromSnapshot(r.GetMutations(), observedSnap, observedVer, "observed"); err != nil { + return err + } + + currentSnap, ok := f.routes.Current() + if !ok || currentSnap.Version() == observedSnap.Version() { + return nil + } + return verifyWriteFenceFromSnapshot(r.GetMutations(), currentSnap, currentSnap.Version(), "current") +} + +func requestBypassesWriteFence(r *pb.Request) bool { + if !r.GetIsTxn() { + return false + } + switch r.GetPhase() { + case pb.Phase_COMMIT, pb.Phase_ABORT: + return true + case pb.Phase_NONE, pb.Phase_PREPARE: + return false + } + return false +} + +func (f *kvFSM) writeFenceHistoryReady(observedVer uint64) bool { + return f.routes != nil && f.shardGroupID != 0 && observedVer != 0 +} + +func verifyWriteFenceFromSnapshot(mutations []*pb.Mutation, snap RouteSnapshot, snapVer uint64, phase string) error { + for _, mut := range mutations { + if mut == nil { + continue + } + if isTxnInternalKey(mut.Key) { + continue + } + if mut.GetOp() == pb.Op_DEL_PREFIX { + start, end := routePrefixRange(mut.Key) + if snap.WriteFencedIntersects(start, end) { + return errors.Wrapf(ErrRouteWriteFenced, + "%s-version v=%d: prefix %q route range [%q,%q)", + phase, snapVer, mut.Key, start, end) + } + continue + } + if len(mut.Key) == 0 { + continue + } + rKey := routeKey(mut.Key) + if snap.WriteFencedForKey(rKey) { + return errors.Wrapf(ErrRouteWriteFenced, + "%s-version v=%d: key %q routeKey %q", + phase, snapVer, mut.Key, rKey) + } + start, end, ok := s3BucketAuxiliaryRouteRange(mut.Key) + if ok && snap.WriteFencedIntersects(start, end) { + return errors.Wrapf(ErrRouteWriteFenced, + "%s-version v=%d: key %q route range [%q,%q)", + phase, snapVer, mut.Key, start, end) + } + } + return nil +} + // verifyOwnerFromSnapshot is the shared per-mutation owner-check // loop used by verifyComposed1's observed-version and current- // version passes. `phase` is the diagnostic label ("observed" / diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index 7279a1e83..959358513 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -53,6 +53,37 @@ func TestFSMAppliesRawPointWriteDespiteLocalWriteFencedRoute(t *testing.T) { require.Equal(t, []byte("v"), got) } +func TestFSMRejectsObservedWriteFencedRawPointWrite(t *testing.T) { + t.Parallel() + + fsm := newWriteFencedFSM(t) + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + ObservedRouteVersion: 1, + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("v")}}, + }, 10) + require.ErrorIs(t, err, ErrRouteWriteFenced) +} + +func TestFSMRejectsCurrentWriteFenceAfterObservedActiveRawPointWrite(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + applyComposed1Snapshot(t, engine, 1, []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }) + fsm := newComposed1FSM(t, engine, 1) + applyComposed1Snapshot(t, engine, 2, []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: []byte("m"), GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: []byte("m"), End: nil, GroupID: 1, State: distribution.RouteStateWriteFenced}, + }) + + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + ObservedRouteVersion: 1, + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("v")}}, + }, 10) + require.ErrorIs(t, err, ErrRouteWriteFenced) +} + func TestFSMAppliesS3BucketAuxiliaryPointWriteDespiteLocalWriteFencedRoute(t *testing.T) { t.Parallel() @@ -75,6 +106,21 @@ func TestFSMAppliesS3BucketAuxiliaryPointWriteDespiteLocalWriteFencedRoute(t *te } } +func TestFSMRejectsObservedWriteFencedS3BucketAuxiliaryPointWrite(t *testing.T) { + t.Parallel() + + const bucket = "bucket-a" + fsm := newS3BucketAuxiliaryWriteFencedFSM(t, bucket) + + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + ObservedRouteVersion: 1, + Mutations: []*pb.Mutation{ + {Op: pb.Op_PUT, Key: s3keys.BucketGenerationKey(bucket), Value: []byte("v")}, + }, + }, 10) + require.ErrorIs(t, err, ErrRouteWriteFenced) +} + func TestFSMAppliesDelPrefixDespiteLocalWriteFencedRoute(t *testing.T) { t.Parallel() @@ -90,6 +136,19 @@ func TestFSMAppliesDelPrefixDespiteLocalWriteFencedRoute(t *testing.T) { require.Error(t, getErr) } +func TestFSMRejectsObservedWriteFencedDelPrefix(t *testing.T) { + t.Parallel() + + fsm := newWriteFencedFSM(t) + require.NoError(t, fsm.store.PutAt(context.Background(), []byte("z"), []byte("v"), 1, 0)) + + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + ObservedRouteVersion: 1, + Mutations: []*pb.Mutation{{Op: pb.Op_DEL_PREFIX, Key: []byte("z")}}, + }, 10) + require.ErrorIs(t, err, ErrRouteWriteFenced) +} + func TestFSMAppliesFullRangeDelPrefixDespiteLocalWriteFencedRoute(t *testing.T) { t.Parallel() @@ -149,3 +208,33 @@ func TestFSMAppliesPrepareAndAbortDespiteLocalWriteFencedRoute(t *testing.T) { err := fsm.handleTxnRequest(ctx, abort, 11) require.NotErrorIs(t, err, ErrRouteWriteFenced, "ABORT must keep the narrow cleanup lane open") } + +func TestFSMRejectsObservedWriteFencedPrepareButAllowsAbort(t *testing.T) { + t.Parallel() + + ctx := context.Background() + fsm := newWriteFencedFSM(t) + prepare := &pb.Request{ + IsTxn: true, + Phase: pb.Phase_PREPARE, + Ts: 10, + ObservedRouteVersion: 1, + Mutations: []*pb.Mutation{ + {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: []byte("z"), LockTTLms: defaultTxnLockTTLms})}, + {Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("v")}, + }, + } + require.ErrorIs(t, fsm.handleTxnRequest(ctx, prepare, 10), ErrRouteWriteFenced) + + abort := &pb.Request{ + IsTxn: true, + Phase: pb.Phase_ABORT, + Ts: 11, + ObservedRouteVersion: 1, + Mutations: []*pb.Mutation{ + {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: []byte("z"), CommitTS: 11})}, + {Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("v")}, + }, + } + require.NotErrorIs(t, fsm.handleTxnRequest(ctx, abort, 11), ErrRouteWriteFenced) +} diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index f70d5ecd8..e8bfb5e80 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -700,6 +700,7 @@ func (c *ShardedCoordinator) Dispatch(ctx context.Context, reqs *OperationGroup[ return nil, err } + c.maybeAutoPinRawObservedRouteVersion(reqs) if resp, handled, err := c.dispatchBeforeShardRouting(ctx, reqs); handled { return resp, err } @@ -746,7 +747,7 @@ func (c *ShardedCoordinator) dispatchBeforeShardRouting(ctx context.Context, req // span multiple shards (or be nil, meaning "all keys"). Broadcast the // operation to every shard group so each FSM scans locally. if hasDelPrefixElem(reqs.Elems) { - resp, err := c.dispatchDelPrefixBroadcast(ctx, reqs.IsTxn, reqs.Elems) + resp, err := c.dispatchDelPrefixBroadcast(ctx, reqs.IsTxn, reqs.Elems, reqs.ObservedRouteVersion) return resp, true, err } if err := c.rejectWriteFencedPointElems(reqs.Elems); err != nil { @@ -837,6 +838,16 @@ func (c *ShardedCoordinator) maybeAutoPinObservedRouteVersion(reqs *OperationGro reqs.ObservedRouteVersion = c.engine.Version() } +func (c *ShardedCoordinator) maybeAutoPinRawObservedRouteVersion(reqs *OperationGroup[OP]) { + if c.engine == nil || reqs == nil || reqs.IsTxn || reqs.ObservedRouteVersion != 0 { + return + } + if c.anyResolverClaimedKey(reqs.Elems) { + return + } + reqs.ObservedRouteVersion = c.engine.Version() +} + // anyResolverClaimedKey reports whether any element's key is // claimed by the partition resolver. Returns false when the // router has no resolver installed (the most common case) so the @@ -1021,7 +1032,7 @@ func validateDelPrefixOnly(elems []*Elem[OP]) error { // to every shard group. Each element becomes a separate pb.Request (the FSM's // extractDelPrefix processes only the first DEL_PREFIX mutation per request). // All requests are batched into a single Commit call per shard group. -func (c *ShardedCoordinator) dispatchDelPrefixBroadcast(ctx context.Context, isTxn bool, elems []*Elem[OP]) (*CoordinateResponse, error) { +func (c *ShardedCoordinator) dispatchDelPrefixBroadcast(ctx context.Context, isTxn bool, elems []*Elem[OP], observedRouteVersion uint64) (*CoordinateResponse, error) { if isTxn { return nil, errors.Wrap(ErrInvalidRequest, "DEL_PREFIX not supported in transactions") } @@ -1039,10 +1050,11 @@ func (c *ShardedCoordinator) dispatchDelPrefixBroadcast(ctx context.Context, isT requests := make([]*pb.Request, 0, len(elems)) for _, elem := range elems { requests = append(requests, &pb.Request{ - IsTxn: false, - Phase: pb.Phase_NONE, - Ts: ts, - Mutations: []*pb.Mutation{elemToMutation(elem)}, + IsTxn: false, + Phase: pb.Phase_NONE, + Ts: ts, + Mutations: []*pb.Mutation{elemToMutation(elem)}, + ObservedRouteVersion: observedRouteVersion, }) } @@ -2013,10 +2025,11 @@ func (c *ShardedCoordinator) rawLogs(ctx context.Context, reqs *OperationGroup[O return nil, err } logs = append(logs, &pb.Request{ - IsTxn: false, - Phase: pb.Phase_NONE, - Ts: ts, - Mutations: grouped[gid], + IsTxn: false, + Phase: pb.Phase_NONE, + Ts: ts, + Mutations: grouped[gid], + ObservedRouteVersion: reqs.ObservedRouteVersion, }) } return logs, nil diff --git a/kv/sharded_coordinator_del_prefix_test.go b/kv/sharded_coordinator_del_prefix_test.go index 1b8a31cb2..7b52fe105 100644 --- a/kv/sharded_coordinator_del_prefix_test.go +++ b/kv/sharded_coordinator_del_prefix_test.go @@ -122,6 +122,52 @@ func TestShardedCoordinator_DelPrefixBroadcastsToAllGroups(t *testing.T) { "same DEL_PREFIX element must use the same timestamp across shards") } +func TestShardedCoordinator_DelPrefixCarriesObservedRouteVersion(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 7, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: nil, End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }, + })) + txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: DelPrefix, Key: []byte("user:")}}, + }) + require.NoError(t, err) + require.Len(t, txn.requests, 1) + require.Equal(t, uint64(7), txn.requests[0].GetObservedRouteVersion()) +} + +func TestShardedCoordinator_RawWriteCarriesObservedRouteVersion(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 9, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: nil, End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }, + })) + txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: Put, Key: []byte("k"), Value: []byte("v")}}, + }) + require.NoError(t, err) + require.Len(t, txn.requests, 1) + require.Equal(t, uint64(9), txn.requests[0].GetObservedRouteVersion()) +} + func TestShardedCoordinatorRejectsPointWriteOnWriteFencedRoute(t *testing.T) { t.Parallel() diff --git a/main.go b/main.go index 06f2b1d93..a630690d0 100644 --- a/main.go +++ b/main.go @@ -486,40 +486,6 @@ func run() error { cleanup.Add(leadershipRefusalDeregister) eg, runCtx := errgroup.WithContext(ctx) startRaftEngineLifecycleWatchers(runCtx, eg, runtimes) - // setupDistributionCatalog + the Stage 7a process-start registration - // gate are bundled so run() has a single startup-fault path: a - // registry-read / behind-epoch failure fails the process - // synchronously here, BEFORE the gRPC servers serve, so writes never - // run with no registration gate installed. - distCatalog, err := setupDistributionAndRegistration( - runCtx, eg, runtimes, cfg.engine, - coordinate, shardGroups[cfg.defaultGroup], encWiring, *raftId, *encryptionSidecarPath) - if err != nil { - cancel() - return err - } - // Seed AFTER setupDistributionCatalog so the sampler picks up the - // catalog-assigned RouteIDs. EnsureCatalogSnapshot inside - // setupDistributionCatalog applies a snapshot back into the engine - // with durable non-zero RouteIDs; seeding earlier would register - // the placeholder zero IDs from buildEngine and Observe would miss - // every dispatched mutation. - seedKeyVizRoutes(sampler, cfg.engine) - - eg.Go(func() error { - return runDistributionCatalogWatcher(runCtx, distCatalog, cfg.engine) - }) - startKeyVizFlusher(runCtx, eg, sampler) - startKeyVizLeaderTermPublisher(runCtx, eg, sampler, runtimes) - startMemoryWatchdog(runCtx, eg, cancel) - distServer := adapter.NewDistributionServer( - cfg.engine, - distCatalog, - adapter.WithDistributionCoordinator(coordinate), - adapter.WithDistributionActiveTimestampTracker(readTracker), - ) - startMonitoringCollectors(runCtx, metricsRegistry, runtimes, clock) - startFSMCompactorIfEnabled(runCtx, eg, runtimes, readTracker) // Stage 7c §3.1: build the encryption-aware // MembershipChangeInterceptor here where the concrete @@ -543,16 +509,41 @@ func run() error { slog.Default(), ) cleanup.Add(rotateOnStartupDeregister) - if err := startServersAfterStartupRotation(waitRotateOnStartup, serversInput{ + + serverInput := serversInput{ ctx: runCtx, eg: eg, cancel: cancel, lc: &lc, runtimes: runtimes, shardGroups: shardGroups, bootstrapServers: bootstrapCfg.adminSeed(cfg.defaultGroup), shardStore: shardStore, coordinate: coordinate, - distServer: distServer, readTracker: readTracker, + readTracker: readTracker, metricsRegistry: metricsRegistry, cfg: cfg, redisApplyObserver: redisApplyObserver, encWiring: encWiring, keyvizSampler: sampler, encryptionConfChangeInterceptor: encryptionConfChangeInterceptor, + } + runtimeStartup, err := prepareDistributionRuntimeServer( + waitRotateOnStartup, serverInput, cfg.engine, runtimes, coordinate, readTracker) + if err != nil { + cancel() + return err + } + if err := startDistributionRuntimeAfterTransport(distributionRuntimeStartupInput{ + runCtx: runCtx, + eg: eg, + cancel: cancel, + runtimes: runtimes, + engine: cfg.engine, + coordinate: coordinate, + defaultGroup: shardGroups[cfg.defaultGroup], + encWiring: encWiring, + raftID: *raftId, + sidecarPath: *encryptionSidecarPath, + sampler: sampler, + clock: clock, + metricsRegistry: metricsRegistry, + readTracker: readTracker, + waitRotateOnStartup: waitRotateOnStartup, + prepared: runtimeStartup, }); err != nil { return err } @@ -1441,14 +1432,103 @@ type serversInput struct { encryptionConfChangeInterceptor internalraftadmin.MembershipChangeInterceptor } -// startServersAfterStartupRotation wires up the AdminServer, starts the -// per-group Raft listeners needed for quorum traffic, waits for any requested -// startup rotation, then opens the public service listeners. -func startServersAfterStartupRotation(waitRotateOnStartup startupRotationWaiter, in serversInput) error { - adminServer, adminGRPCOpts, err := setupAdminService(*raftId, *myAddr, in.runtimes, in.bootstrapServers, in.keyvizSampler) +type preparedRuntimeServer struct { + runner *runtimeServerRunner + connCache *kv.GRPCConnCache +} + +type preparedDistributionRuntime struct { + catalog *distribution.CatalogStore + serverInput serversInput + preparedServer *preparedRuntimeServer +} + +type distributionRuntimeStartupInput struct { + runCtx context.Context + eg *errgroup.Group + cancel context.CancelFunc + runtimes []*raftGroupRuntime + engine *distribution.Engine + coordinate *kv.ShardedCoordinator + defaultGroup *kv.ShardGroup + encWiring encryptionWriteWiring + raftID string + sidecarPath string + sampler *keyviz.MemSampler + clock *kv.HLC + metricsRegistry *monitoring.Registry + readTracker *kv.ActiveTimestampTracker + waitRotateOnStartup startupRotationWaiter + prepared *preparedDistributionRuntime +} + +func prepareDistributionRuntimeServer( + waitRotateOnStartup startupRotationWaiter, + in serversInput, + engine *distribution.Engine, + runtimes []*raftGroupRuntime, + coordinate *kv.ShardedCoordinator, + readTracker *kv.ActiveTimestampTracker, +) (*preparedDistributionRuntime, error) { + distCatalog, err := distributionCatalogStoreForEngine(runtimes, engine) + if err != nil { + return nil, err + } + in.distServer = adapter.NewDistributionServer( + engine, + distCatalog, + adapter.WithDistributionCoordinator(coordinate), + adapter.WithDistributionActiveTimestampTracker(readTracker), + ) + preparedServer, err := prepareRuntimeServerRunner(waitRotateOnStartup, in) if err != nil { + return nil, err + } + return &preparedDistributionRuntime{ + catalog: distCatalog, + serverInput: in, + preparedServer: preparedServer, + }, nil +} + +func startDistributionRuntimeAfterTransport(in distributionRuntimeStartupInput) error { + prepared := in.prepared + if prepared == nil || prepared.preparedServer == nil || prepared.preparedServer.runner == nil { + return errors.New("distribution runtime server is not prepared") + } + if err := prepared.preparedServer.runner.startRaftTransport(); err != nil { return err } + if err := waitForRaftStartupAfterTransport(in.runCtx, in.runtimes); err != nil { + return prepared.preparedServer.runner.startupFailure(err) + } + // setupDistributionCatalog + the Stage 7a process-start registration + // gate are bundled so startup faults flow through one path. This now + // runs after raft transport is bound, but before public listeners serve. + if err := setupDistributionAndRegistration( + in.runCtx, in.eg, prepared.catalog, in.engine, + in.coordinate, in.defaultGroup, in.encWiring, in.raftID, in.sidecarPath); err != nil { + in.cancel() + return err + } + seedKeyVizRoutes(in.sampler, in.engine) + in.eg.Go(func() error { + return runDistributionCatalogWatcher(in.runCtx, prepared.catalog, in.engine) + }) + startKeyVizFlusher(in.runCtx, in.eg, in.sampler) + startKeyVizLeaderTermPublisher(in.runCtx, in.eg, in.sampler, in.runtimes) + startMemoryWatchdog(in.runCtx, in.eg, in.cancel) + startMonitoringCollectors(in.runCtx, in.metricsRegistry, in.runtimes, in.clock) + startFSMCompactorIfEnabled(in.runCtx, in.eg, in.runtimes, in.readTracker) + return startPreparedServerAfterStartupRotation( + in.waitRotateOnStartup, prepared.serverInput, prepared.preparedServer) +} + +func prepareRuntimeServerRunner(waitRotateOnStartup startupRotationWaiter, in serversInput) (*preparedRuntimeServer, error) { + adminServer, adminGRPCOpts, err := setupAdminService(*raftId, *myAddr, in.runtimes, in.bootstrapServers, in.keyvizSampler) + if err != nil { + return nil, err + } // roleStore + connCache are gated on *adminEnabled. With admin // disabled, building either is wasted work AND a security // regression risk: a non-empty -adminFullAccessKeys flag would @@ -1540,12 +1620,14 @@ func startServersAfterStartupRotation(waitRotateOnStartup startupRotationWaiter, encryptionConfChangeInterceptor: in.encryptionConfChangeInterceptor, publicKVGate: publicKVGate, } - if err := runner.startRaftTransport(); err != nil { - return err - } - if err := waitForRaftStartupAfterTransport(in.ctx, in.runtimes); err != nil { - return runner.startupFailure(err) + return &preparedRuntimeServer{runner: &runner, connCache: connCache}, nil +} + +func startPreparedServerAfterStartupRotation(waitRotateOnStartup startupRotationWaiter, in serversInput, prepared *preparedRuntimeServer) error { + if prepared == nil || prepared.runner == nil { + return errors.New("runtime server runner is not prepared") } + runner := prepared.runner if err := runner.preparePublicServices(); err != nil { return runner.startupFailure(err) } @@ -1565,17 +1647,19 @@ func startServersAfterStartupRotation(waitRotateOnStartup startupRotationWaiter, Nodes: parseCSV(*keyvizFanoutNodes), Timeout: *keyvizFanoutTimeout, } - adminHTTP, err := prepareAdminFromFlags(in.ctx, in.lc, in.runtimes, runner.dynamoServer, runner.s3Server, runner.sqsServer, runner.coordinate, connCache, in.keyvizSampler, fanoutCfg) + adminHTTP, err := prepareAdminFromFlags(in.ctx, in.lc, in.runtimes, runner.dynamoServer, runner.s3Server, runner.sqsServer, runner.coordinate, prepared.connCache, in.keyvizSampler, fanoutCfg) if err != nil { return runner.startupFailure(err) } runner.adminHTTP = adminHTTP - publicKVGate.blockMutator = waitRotateOnStartup.BlockMutators + if runner.publicKVGate != nil { + runner.publicKVGate.blockMutator = waitRotateOnStartup.BlockMutators + } if err := waitRotateOnStartup.Wait(in.ctx); err != nil { return runner.startupFailure(errors.Wrap(err, "encryption rotate-on-startup: wait before serving")) } startHLCLeaseRenewal(in.ctx, in.eg, in.coordinate) - publicKVGate.markReady() + runner.publicKVGate.markReady() if err := runner.startPublicServices(); err != nil { return err } @@ -2405,8 +2489,7 @@ func distributionCatalogStoreForGroup(runtimes []*raftGroupRuntime, groupID uint return nil } -func setupDistributionCatalog( - ctx context.Context, +func distributionCatalogStoreForEngine( runtimes []*raftGroupRuntime, engine *distribution.Engine, ) (*distribution.CatalogStore, error) { @@ -2418,25 +2501,20 @@ func setupDistributionCatalog( if distCatalog == nil { return nil, errors.WithStack(errors.Newf("distribution catalog store is not available for group %d", catalogGroupID)) } - // EnsureCatalogSnapshot may Save through the direct (non-raft) write - // path. When the §7.1 storage envelope is active and this load's - // writer registration has not yet committed, that Save fails closed - // with store.ErrWriterNotRegistered (Stage 7a-2). retryUntilRegistered - // retries the bootstrap until the registration goroutine — armed - // before this call in setupDistributionAndRegistration — commits and - // the gate clears. The common cases (populated catalog → no-op Save, - // or pre-cutover → cleartext Save) never hit the gate and return on - // the first attempt. - // - // Idempotency requirement: the retry re-invokes EnsureCatalogSnapshot - // from scratch on each ErrWriterNotRegistered, so it MUST be - // re-entrant — on the populated-catalog path it is a version-unchanged - // no-op Save (no mutation, no nonce), so re-running it is safe. - if err := retryUntilRegistered(ctx, "distribution catalog bootstrap", func() error { - _, e := distribution.EnsureCatalogSnapshot(ctx, distCatalog, engine) - return errors.Wrap(e, "ensure catalog snapshot") - }); err != nil { - return nil, errors.Wrapf(err, "initialize distribution catalog") + return distCatalog, nil +} + +func setupDistributionCatalog( + ctx context.Context, + runtimes []*raftGroupRuntime, + engine *distribution.Engine, +) (*distribution.CatalogStore, error) { + distCatalog, err := distributionCatalogStoreForEngine(runtimes, engine) + if err != nil { + return nil, err + } + if err := ensureDistributionCatalogSnapshot(ctx, distCatalog, engine); err != nil { + return nil, err } return distCatalog, nil } diff --git a/main_encryption_registration.go b/main_encryption_registration.go index 7cdb53f9e..eb5aa68de 100644 --- a/main_encryption_registration.go +++ b/main_encryption_registration.go @@ -126,19 +126,19 @@ func retryUntilRegistered(ctx context.Context, what string, fn func() error) err func setupDistributionAndRegistration( runCtx context.Context, eg *errgroup.Group, - runtimes []*raftGroupRuntime, + distCatalog *distribution.CatalogStore, engine *distribution.Engine, coordinate *kv.ShardedCoordinator, defaultGroup *kv.ShardGroup, w encryptionWriteWiring, raftID string, sidecarPath string, -) (*distribution.CatalogStore, error) { +) error { if err := validateRaftRegistrationStartupEpoch(defaultGroup, w, raftID, sidecarPath); err != nil { - return nil, err + return err } if err := installProcessStartRegistrationGate(runCtx, eg, coordinate, defaultGroup, w, raftID); err != nil { - return nil, err + return err } installRaftRegistrationVerifier(defaultGroup, w, raftID) installRuntimeRaftRegistrationWatcher(runCtx, eg, coordinate, defaultGroup, w, raftID, sidecarPath) @@ -153,11 +153,38 @@ func setupDistributionAndRegistration( installRuntimeRegistrationWatcher(runCtx, eg, coordinate, defaultGroup, w, raftID) // Bootstrap + registration both run under runCtx so a shutdown // cancels the bounded retry rather than hanging. - distCatalog, err := setupDistributionCatalog(runCtx, runtimes, engine) - if err != nil { - return nil, err + return ensureDistributionCatalogSnapshot(runCtx, distCatalog, engine) +} + +func ensureDistributionCatalogSnapshot( + ctx context.Context, + distCatalog *distribution.CatalogStore, + engine *distribution.Engine, +) error { + if distCatalog == nil { + return errors.New("distribution catalog store is not available") + } + // EnsureCatalogSnapshot may Save through the direct (non-raft) write + // path. When the §7.1 storage envelope is active and this load's + // writer registration has not yet committed, that Save fails closed + // with store.ErrWriterNotRegistered (Stage 7a-2). retryUntilRegistered + // retries the bootstrap until the registration goroutine — armed + // before this call in setupDistributionAndRegistration — commits and + // the gate clears. The common cases (populated catalog → no-op Save, + // or pre-cutover → cleartext Save) never hit the gate and return on + // the first attempt. + // + // Idempotency requirement: the retry re-invokes EnsureCatalogSnapshot + // from scratch on each ErrWriterNotRegistered, so it MUST be + // re-entrant — on the populated-catalog path it is a version-unchanged + // no-op Save (no mutation, no nonce), so re-running it is safe. + if err := retryUntilRegistered(ctx, "distribution catalog bootstrap", func() error { + _, e := distribution.EnsureCatalogSnapshot(ctx, distCatalog, engine) + return errors.Wrap(e, "ensure catalog snapshot") + }); err != nil { + return errors.Wrapf(err, "initialize distribution catalog") } - return distCatalog, nil + return nil } func installRaftRegistrationVerifier(defaultGroup *kv.ShardGroup, w encryptionWriteWiring, raftID string) { From 61feed4ce354950d0eee65e3758906a271016ff9 Mon Sep 17 00:00:00 2001 From: bootjp Date: Wed, 15 Jul 2026 05:44:57 +0900 Subject: [PATCH 16/26] Stabilize route-fence cleanup retries --- adapter/dynamodb_cleanup_retry_test.go | 42 ++++++++++++++++++ adapter/dynamodb_item_write.go | 13 +++--- adapter/dynamodb_onephase_dedup_test.go | 18 ++++++++ adapter/dynamodb_schema.go | 24 ++++++++++- adapter/dynamodb_transact.go | 4 ++ adapter/retryable_write_fence_test.go | 1 + adapter/s3.go | 11 ++++- adapter/s3_cleanup_retry_test.go | 57 +++++++++++++++++++++++++ internal/raftengine/etcd/engine.go | 4 ++ internal/raftengine/etcd/engine_test.go | 37 ++++++++++++++-- 10 files changed, 198 insertions(+), 13 deletions(-) create mode 100644 adapter/dynamodb_cleanup_retry_test.go create mode 100644 adapter/s3_cleanup_retry_test.go diff --git a/adapter/dynamodb_cleanup_retry_test.go b/adapter/dynamodb_cleanup_retry_test.go new file mode 100644 index 000000000..71b535e85 --- /dev/null +++ b/adapter/dynamodb_cleanup_retry_test.go @@ -0,0 +1,42 @@ +package adapter + +import ( + "context" + "testing" + + "github.com/bootjp/elastickv/kv" + "github.com/bootjp/elastickv/store" + "github.com/stretchr/testify/require" +) + +func TestDynamoDBCleanupDeletedTableGeneration_RetriesRouteFence(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + coord := &cleanupRouteFenceCoordinator{ + localAdapterCoordinator: newLocalAdapterCoordinator(st), + failuresRemaining: 1, + } + server := NewDynamoDBServer(nil, st, coord) + + err := server.cleanupDeletedTableGeneration(context.Background(), "table-a", 7) + require.NoError(t, err) + require.Equal(t, 3, coord.calls, "first prefix is retried once, second prefix succeeds once") +} + +type cleanupRouteFenceCoordinator struct { + *localAdapterCoordinator + failuresRemaining int + calls int +} + +func (c *cleanupRouteFenceCoordinator) Dispatch(ctx context.Context, req *kv.OperationGroup[kv.OP]) (*kv.CoordinateResponse, error) { + if req != nil && operationGroupHasDelPrefix(req.Elems) { + c.calls++ + if c.failuresRemaining > 0 { + c.failuresRemaining-- + return nil, kv.ErrRouteWriteFenced + } + } + return c.localAdapterCoordinator.Dispatch(ctx, req) +} diff --git a/adapter/dynamodb_item_write.go b/adapter/dynamodb_item_write.go index a5aecda19..b3ad5678f 100644 --- a/adapter/dynamodb_item_write.go +++ b/adapter/dynamodb_item_write.go @@ -285,7 +285,7 @@ func (d *DynamoDBServer) itemWriteFirstAttempt( plan.req.CommitTS = commitTS if dispErr := d.commitItemWrite(ctx, plan.req); dispErr != nil { // dispErr is already wrapped by commitItemWrite; return it raw. - if isRetryableTransactWriteError(dispErr) { + if shouldPreserveTransactWriteAttempt(dispErr) { return nil, &reusableItemWrite{ plan: plan, commitTS: commitTS, @@ -319,10 +319,13 @@ func (d *DynamoDBServer) itemWriteReuseAttempt( return d.resolveReuseWriteConflict(ctx, tableName, pending, commitTS, dispErr) } if isRetryableTransactWriteError(dispErr) { - // Still ambiguous (e.g. TxnLocked): this reuse may itself have landed, - // so the next retry must probe THIS commit_ts. dispErr is already - // wrapped by commitItemWrite; return it raw. - pending.commitTS = commitTS + // Still ambiguous (e.g. TxnLocked): this reuse may itself have + // landed, so the next retry must probe THIS commit_ts. Route-fence + // rejections are retryable but happen before this write set can apply, + // so keep the older witness. + if shouldPreserveTransactWriteAttempt(dispErr) { + pending.commitTS = commitTS + } return nil, pending, dispErr } return nil, nil, dispErr diff --git a/adapter/dynamodb_onephase_dedup_test.go b/adapter/dynamodb_onephase_dedup_test.go index ce7b9c7d4..b28afe3bd 100644 --- a/adapter/dynamodb_onephase_dedup_test.go +++ b/adapter/dynamodb_onephase_dedup_test.go @@ -113,6 +113,24 @@ func TestItemWriteDedup_LandedPriorAttempt_NoDuplicate(t *testing.T) { require.Equal(t, 1, coord.probeNoOps, "the reuse must dedup via the FSM exact-ts probe") } +func TestItemWriteDedup_RouteFenceRetryPreservesPriorProbe(t *testing.T) { + t.Parallel() + ctx := context.Background() + st := store.NewMVCCStore() + coord := newDedupTestCoordinator(st, 1, true) // dispatch 1 lands then errors + coord.routeFenceAtDispatch = 2 + schema, server := newDedupItemWriteServer(st, coord, true) + seedDedupItem(t, st, schema, "1", "2") + + plan, err := server.updateItemWithRetry(ctx, appendListInput()) + require.NoError(t, err) + require.NotNil(t, plan) + + require.Equal(t, []string{"1", "2", "3"}, readListValues(t, server, schema)) + require.Equal(t, 3, coord.dispatches, "attempt 1 landed, route-fenced reuse, then dedup probe retry") + require.Equal(t, 1, coord.probeNoOps, "route-fenced reuse must not replace the prior landed probe") +} + // TestItemWriteDedup_PriorAttemptDidNotLand_Applies: attempt 1 pre-rejects // (definitely did not commit); the reuse's probe misses, so it applies the // reused write set at a fresh commit_ts. One element, no duplicate. diff --git a/adapter/dynamodb_schema.go b/adapter/dynamodb_schema.go index 0ffa0f917..b3843b8d5 100644 --- a/adapter/dynamodb_schema.go +++ b/adapter/dynamodb_schema.go @@ -349,16 +349,36 @@ func (d *DynamoDBServer) cleanupDeletedTableGeneration(ctx context.Context, tabl // scans and writes tombstones locally, avoiding the enumerate-then-batch- // delete loop that previously required many Raft proposals. for _, prefix := range prefixes { + if err := d.dispatchDeletedTableCleanupPrefix(ctx, prefix); err != nil { + return err + } + } + return nil +} + +func (d *DynamoDBServer) dispatchDeletedTableCleanupPrefix(ctx context.Context, prefix []byte) error { + backoff := transactRetryInitialBackoff + deadline := time.Now().Add(transactRetryMaxDuration) + var lastErr error + for range transactRetryMaxAttempts { _, err := d.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ Elems: []*kv.Elem[kv.OP]{ {Op: kv.DelPrefix, Key: prefix}, }, }) - if err != nil { + if err == nil { + return nil + } + if !isRetryableTransactWriteError(err) { return errors.WithStack(err) } + lastErr = err + if waitErr := waitRetryWithDeadline(ctx, deadline, backoff); waitErr != nil { + return errors.Wrap(errors.Join(waitErr, lastErr), "dynamodb delete table cleanup retry canceled") + } + backoff = nextTransactRetryBackoff(backoff) } - return nil + return errors.WithStack(lastErr) } func (d *DynamoDBServer) dispatchDeleteBatch(ctx context.Context, keys [][]byte) error { diff --git a/adapter/dynamodb_transact.go b/adapter/dynamodb_transact.go index 7ccaaac59..9569b6053 100644 --- a/adapter/dynamodb_transact.go +++ b/adapter/dynamodb_transact.go @@ -1108,6 +1108,10 @@ func isRetryableTransactWriteError(err error) bool { return errors.Is(err, store.ErrWriteConflict) || errors.Is(err, kv.ErrTxnLocked) || errors.Is(err, kv.ErrRouteWriteFenced) } +func shouldPreserveTransactWriteAttempt(err error) bool { + return isRetryableTransactWriteError(err) && !errors.Is(err, kv.ErrRouteWriteFenced) +} + func waitTransactRetryBackoff(ctx context.Context, delay time.Duration) error { timer := time.NewTimer(delay) defer timer.Stop() diff --git a/adapter/retryable_write_fence_test.go b/adapter/retryable_write_fence_test.go index fb9f9a5e0..e5b436c06 100644 --- a/adapter/retryable_write_fence_test.go +++ b/adapter/retryable_write_fence_test.go @@ -13,4 +13,5 @@ func TestWriteFenceErrorsAreAdapterRetryable(t *testing.T) { require.True(t, isRetryableRedisTxnErr(kv.ErrRouteWriteFenced)) require.True(t, isRetryableS3MutationErr(kv.ErrRouteWriteFenced)) require.True(t, isRetryableTransactWriteError(kv.ErrRouteWriteFenced)) + require.False(t, shouldPreserveTransactWriteAttempt(kv.ErrRouteWriteFenced)) } diff --git a/adapter/s3.go b/adapter/s3.go index d08f0e4d6..75cc2c074 100644 --- a/adapter/s3.go +++ b/adapter/s3.go @@ -1455,7 +1455,7 @@ func (s *S3Server) cleanupPartBlobsAsync(bucket string, generation uint64, objec if len(pending) == 0 { return } - if _, err := s.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{Elems: pending}); err != nil { + if err := s.dispatchS3CleanupBatch(ctx, pending); err != nil { slog.ErrorContext(ctx, "cleanupPartBlobsAsync: coordinator dispatch failed", "bucket", bucket, "object_key", objectKey, @@ -1519,7 +1519,7 @@ func (s *S3Server) deleteByPrefix(ctx context.Context, prefix []byte, bucket str for _, kvp := range kvs { pending = append(pending, &kv.Elem[kv.OP]{Op: kv.Del, Key: kvp.Key}) } - if _, err := s.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{Elems: pending}); err != nil { + if err := s.dispatchS3CleanupBatch(ctx, pending); err != nil { slog.ErrorContext(ctx, "deleteByPrefix: dispatch failed", "bucket", bucket, "generation", generation, "object_key", objectKey, "upload_id", uploadID, "err", err) @@ -1529,6 +1529,13 @@ func (s *S3Server) deleteByPrefix(ctx context.Context, prefix []byte, bucket str } } +func (s *S3Server) dispatchS3CleanupBatch(ctx context.Context, elems []*kv.Elem[kv.OP]) error { + return s.retryS3Mutation(ctx, func() error { + _, err := s.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{Elems: elems}) + return errors.WithStack(err) + }) +} + func parseS3MaxParts(raw string) int { if strings.TrimSpace(raw) == "" { return s3ListPartsMaxParts diff --git a/adapter/s3_cleanup_retry_test.go b/adapter/s3_cleanup_retry_test.go new file mode 100644 index 000000000..c7591a90b --- /dev/null +++ b/adapter/s3_cleanup_retry_test.go @@ -0,0 +1,57 @@ +package adapter + +import ( + "context" + "testing" + + "github.com/bootjp/elastickv/internal/s3keys" + "github.com/bootjp/elastickv/kv" + "github.com/bootjp/elastickv/store" + "github.com/stretchr/testify/require" +) + +func TestS3DeleteByPrefix_RetriesRouteFence(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + coord := &s3CleanupRouteFenceCoordinator{ + localAdapterCoordinator: newLocalAdapterCoordinator(st), + failuresRemaining: 1, + } + server := NewS3Server(nil, "", st, coord, nil) + key := s3keys.BlobKey("bucket-a", 3, "object-a", "upload-a", 1, 0) + require.NoError(t, st.PutAt(ctx, key, []byte("blob"), 1, 0)) + + server.deleteByPrefix(ctx, s3keys.BlobPrefixForUpload("bucket-a", 3, "object-a", "upload-a"), "bucket-a", 3, "object-a", "upload-a") + + require.Equal(t, 2, coord.calls, "route-fenced cleanup batch should retry") + _, err := st.GetAt(ctx, key, snapshotTS(coord.Clock(), st)) + require.ErrorIs(t, err, store.ErrKeyNotFound) +} + +type s3CleanupRouteFenceCoordinator struct { + *localAdapterCoordinator + failuresRemaining int + calls int +} + +func (c *s3CleanupRouteFenceCoordinator) Dispatch(ctx context.Context, req *kv.OperationGroup[kv.OP]) (*kv.CoordinateResponse, error) { + if req != nil && operationGroupHasDel(req.Elems) { + c.calls++ + if c.failuresRemaining > 0 { + c.failuresRemaining-- + return nil, kv.ErrRouteWriteFenced + } + } + return c.localAdapterCoordinator.Dispatch(ctx, req) +} + +func operationGroupHasDel(elems []*kv.Elem[kv.OP]) bool { + for _, elem := range elems { + if elem != nil && elem.Op == kv.Del { + return true + } + } + return false +} diff --git a/internal/raftengine/etcd/engine.go b/internal/raftengine/etcd/engine.go index 0326b66d1..4df7ed243 100644 --- a/internal/raftengine/etcd/engine.go +++ b/internal/raftengine/etcd/engine.go @@ -2422,6 +2422,10 @@ func coalesceHeartbeatResp(ch chan dispatchRequest, msg raftpb.Message) bool { restoreDispatchRequest(ch, stale) return false } + if len(stale.msg.GetContext()) != 0 { + restoreDispatchRequest(ch, stale) + return false + } closeDispatchRequest(stale) return tryEnqueueDispatchRequest(ch, prepareDispatchRequest(msg)) } diff --git a/internal/raftengine/etcd/engine_test.go b/internal/raftengine/etcd/engine_test.go index 8a1f82b4a..c4ecb1b73 100644 --- a/internal/raftengine/etcd/engine_test.go +++ b/internal/raftengine/etcd/engine_test.go @@ -1393,7 +1393,7 @@ func TestPrepareDispatchRequestClonesSnapshotPayload(t *testing.T) { require.Equal(t, []uint64{1, 2}, req.msg.Snapshot.GetMetadata().GetConfState().GetVoters()) } -func TestEnqueueDispatchMessageCoalescesHeartbeatResponses(t *testing.T) { +func TestEnqueueDispatchMessageCoalescesPlainHeartbeatResponses(t *testing.T) { t.Parallel() pd := &peerQueues{ heartbeat: make(chan dispatchRequest, 1), @@ -1406,9 +1406,8 @@ func TestEnqueueDispatchMessageCoalescesHeartbeatResponses(t *testing.T) { }, } pd.heartbeatResp <- prepareDispatchRequest(raftpb.Message{ - Type: messageTypePtr(raftpb.MsgHeartbeatResp), - To: uint64Ptr(2), - Context: []byte("old"), + Type: messageTypePtr(raftpb.MsgHeartbeatResp), + To: uint64Ptr(2), }) require.NoError(t, engine.enqueueDispatchMessage(raftpb.Message{ @@ -1424,6 +1423,36 @@ func TestEnqueueDispatchMessageCoalescesHeartbeatResponses(t *testing.T) { require.Equal(t, []byte("new"), req.msg.Context) } +func TestEnqueueDispatchMessagePreservesReadIndexHeartbeatResponses(t *testing.T) { + t.Parallel() + pd := &peerQueues{ + heartbeat: make(chan dispatchRequest, 1), + heartbeatResp: make(chan dispatchRequest, 1), + } + engine := &Engine{ + nodeID: 1, + peerDispatchers: map[uint64]*peerQueues{ + 2: pd, + }, + } + pd.heartbeatResp <- prepareDispatchRequest(raftpb.Message{ + Type: messageTypePtr(raftpb.MsgHeartbeatResp), + To: uint64Ptr(2), + Context: []byte("read-index"), + }) + + require.NoError(t, engine.enqueueDispatchMessage(raftpb.Message{ + Type: messageTypePtr(raftpb.MsgHeartbeatResp), + To: uint64Ptr(2), + })) + + require.Equal(t, uint64(1), engine.DispatchDropCount()) + require.Len(t, pd.heartbeatResp, 1) + req := <-pd.heartbeatResp + require.Equal(t, raftpb.MsgHeartbeatResp, req.msg.GetType()) + require.Equal(t, []byte("read-index"), req.msg.Context) +} + func TestMaxAppliedIndexStartsFromSnapshotIndex(t *testing.T) { storage := etcdraft.NewMemoryStorage() snap := raftTestSnapshot(5, 2, []uint64{1}, nil) From f373cd356880f53a1d4622e840ec05eab88efdd9 Mon Sep 17 00:00:00 2001 From: bootjp Date: Wed, 15 Jul 2026 06:17:06 +0900 Subject: [PATCH 17/26] Keep heartbeat responses coalescing under read-index load --- internal/raftengine/etcd/engine.go | 40 +++++++++++----------- internal/raftengine/etcd/engine_test.go | 44 +++++++++++++++++++++++++ 2 files changed, 64 insertions(+), 20 deletions(-) diff --git a/internal/raftengine/etcd/engine.go b/internal/raftengine/etcd/engine.go index 4df7ed243..f33257500 100644 --- a/internal/raftengine/etcd/engine.go +++ b/internal/raftengine/etcd/engine.go @@ -2414,20 +2414,30 @@ func coalesceHeartbeatResp(ch chan dispatchRequest, msg raftpb.Message) bool { if ch == nil || cap(ch) == 0 { return false } - stale, ok := tryReceiveDispatchRequest(ch) - if !ok { + n := len(ch) + if n == 0 { return false } - if stale.msg.GetType() != raftpb.MsgHeartbeatResp { - restoreDispatchRequest(ch, stale) - return false + + drained := make([]dispatchRequest, 0, n) + replaced := false + for i := 0; i < n; i++ { + req, ok := tryReceiveDispatchRequest(ch) + if !ok { + break + } + if !replaced && req.msg.GetType() == raftpb.MsgHeartbeatResp && len(req.msg.GetContext()) == 0 { + closeDispatchRequest(req) + drained = append(drained, prepareDispatchRequest(msg)) + replaced = true + continue + } + drained = append(drained, req) } - if len(stale.msg.GetContext()) != 0 { - restoreDispatchRequest(ch, stale) - return false + for _, req := range drained { + restoreDispatchRequest(ch, req) } - closeDispatchRequest(stale) - return tryEnqueueDispatchRequest(ch, prepareDispatchRequest(msg)) + return replaced } func tryReceiveDispatchRequest(ch chan dispatchRequest) (dispatchRequest, bool) { @@ -2447,16 +2457,6 @@ func restoreDispatchRequest(ch chan dispatchRequest, req dispatchRequest) { } } -func tryEnqueueDispatchRequest(ch chan dispatchRequest, req dispatchRequest) bool { - select { - case ch <- req: - return true - default: - closeDispatchRequest(req) - return false - } -} - func closeDispatchRequest(req dispatchRequest) { if err := req.Close(); err != nil { slog.Error("etcd raft dispatch: failed to close request", "err", err) diff --git a/internal/raftengine/etcd/engine_test.go b/internal/raftengine/etcd/engine_test.go index c4ecb1b73..87e8447df 100644 --- a/internal/raftengine/etcd/engine_test.go +++ b/internal/raftengine/etcd/engine_test.go @@ -1453,6 +1453,50 @@ func TestEnqueueDispatchMessagePreservesReadIndexHeartbeatResponses(t *testing.T require.Equal(t, []byte("read-index"), req.msg.Context) } +func TestEnqueueDispatchMessageCoalescesPlainHeartbeatBehindReadIndex(t *testing.T) { + t.Parallel() + pd := &peerQueues{ + heartbeat: make(chan dispatchRequest, 1), + heartbeatResp: make(chan dispatchRequest, 3), + } + engine := &Engine{ + nodeID: 1, + peerDispatchers: map[uint64]*peerQueues{ + 2: pd, + }, + } + pd.heartbeatResp <- prepareDispatchRequest(raftpb.Message{ + Type: messageTypePtr(raftpb.MsgHeartbeatResp), + To: uint64Ptr(2), + Context: []byte("read-index"), + }) + pd.heartbeatResp <- prepareDispatchRequest(raftpb.Message{ + Type: messageTypePtr(raftpb.MsgHeartbeatResp), + To: uint64Ptr(2), + Index: uint64Ptr(10), + }) + pd.heartbeatResp <- prepareDispatchRequest(raftpb.Message{ + Type: messageTypePtr(raftpb.MsgHeartbeatResp), + To: uint64Ptr(2), + Index: uint64Ptr(11), + }) + + require.NoError(t, engine.enqueueDispatchMessage(raftpb.Message{ + Type: messageTypePtr(raftpb.MsgHeartbeatResp), + To: uint64Ptr(2), + Index: uint64Ptr(99), + })) + + require.Zero(t, engine.DispatchDropCount()) + require.Len(t, pd.heartbeatResp, 3) + req := <-pd.heartbeatResp + require.Equal(t, []byte("read-index"), req.msg.Context) + req = <-pd.heartbeatResp + require.Equal(t, uint64(99), req.msg.GetIndex()) + req = <-pd.heartbeatResp + require.Equal(t, uint64(11), req.msg.GetIndex()) +} + func TestMaxAppliedIndexStartsFromSnapshotIndex(t *testing.T) { storage := etcdraft.NewMemoryStorage() snap := raftTestSnapshot(5, 2, []uint64{1}, nil) From 8f821fe891294eae2c6b136ea21e6e652c2d5929 Mon Sep 17 00:00:00 2001 From: bootjp Date: Wed, 15 Jul 2026 06:58:15 +0900 Subject: [PATCH 18/26] Stabilize raft snapshot recovery checkpoints --- .../raftengine/etcd/dispatch_report_test.go | 6 +- internal/raftengine/etcd/engine.go | 146 +++++++++++++----- .../etcd/engine_applied_index_test.go | 52 ++++++- internal/raftengine/etcd/engine_test.go | 94 ++++++++++- internal/raftengine/statemachine.go | 25 +-- 5 files changed, 264 insertions(+), 59 deletions(-) diff --git a/internal/raftengine/etcd/dispatch_report_test.go b/internal/raftengine/etcd/dispatch_report_test.go index 6b122952b..63c7fb631 100644 --- a/internal/raftengine/etcd/dispatch_report_test.go +++ b/internal/raftengine/etcd/dispatch_report_test.go @@ -176,7 +176,7 @@ func TestPostDispatchReport_AbortsOnClose(t *testing.T) { } } -func TestEnqueueDispatchReportsDroppedSnapshotWhenLaneFull(t *testing.T) { +func TestEnqueueDispatchDefersDroppedSnapshotReportWhenLaneFull(t *testing.T) { t.Parallel() snapshotCh := make(chan dispatchRequest, 1) snapshotCh <- dispatchRequest{msg: raftpb.Message{Type: messageTypePtr(raftpb.MsgSnap), To: uint64Ptr(2)}} @@ -195,10 +195,10 @@ func TestEnqueueDispatchReportsDroppedSnapshotWhenLaneFull(t *testing.T) { require.Equal(t, uint64(1), e.DispatchDropCount()) select { case got := <-e.dispatchReportCh: - require.Equal(t, dispatchReport{to: 2, msgType: raftpb.MsgSnap}, got) + t.Fatalf("dropped MsgSnap should defer without queueing: %+v", got) default: - t.Fatal("expected dropped MsgSnap to report SnapshotFailure input") } + require.Equal(t, []dispatchReport{{to: 2, msgType: raftpb.MsgSnap}}, e.deferredReadyDispatchReports) } func TestEnqueueDispatchReportsDroppedRegularMessageWhenLaneFull(t *testing.T) { diff --git a/internal/raftengine/etcd/engine.go b/internal/raftengine/etcd/engine.go index f33257500..b87a00d13 100644 --- a/internal/raftengine/etcd/engine.go +++ b/internal/raftengine/etcd/engine.go @@ -326,15 +326,16 @@ type Engine struct { nextRequestID atomic.Uint64 - proposeCh chan proposalRequest - readCh chan readRequest - adminCh chan adminRequest - stepCh chan raftpb.Message - priorityStepCh chan raftpb.Message - priorityStepBurst int - dispatchReportCh chan dispatchReport - peerDispatchers map[uint64]*peerQueues - perPeerQueueSize int + proposeCh chan proposalRequest + readCh chan readRequest + adminCh chan adminRequest + stepCh chan raftpb.Message + priorityStepCh chan raftpb.Message + priorityStepBurst int + dispatchReportCh chan dispatchReport + deferredReadyDispatchReports []dispatchReport + peerDispatchers map[uint64]*peerQueues + perPeerQueueSize int // dispatcherLanesEnabled toggles the opt-in multi-lane dispatcher layout. Captured // once at Open from ELASTICKV_RAFT_DISPATCHER_LANES so the run-time code // path is branch-free per message and does not need to re-read env vars. @@ -927,10 +928,10 @@ func waitForOpen(ctx context.Context, engine *Engine, waitForLeader bool) (*Engi if !waitForLeader { return engine, nil } - if err := waitForEngineSignal(ctx, engine, engine.startedCh); err != nil { + if err := waitForOpenSignal(ctx, engine, engine.startedCh); err != nil { return nil, err } - if err := waitForEngineSignal(ctx, engine, engine.leaderReady); err != nil { + if err := waitForOpenSignal(ctx, engine, engine.leaderReady); err != nil { return nil, err } return engine, nil @@ -940,13 +941,19 @@ func (e *Engine) WaitStarted(ctx context.Context) error { if e == nil { return errors.WithStack(errNilEngine) } - return waitForEngineSignal(ctx, e, e.startedCh) + return waitForEngineSignal(ctx, e, e.startedCh, false) } -func waitForEngineSignal(ctx context.Context, engine *Engine, ready <-chan struct{}) error { +func waitForOpenSignal(ctx context.Context, engine *Engine, ready <-chan struct{}) error { + return waitForEngineSignal(ctx, engine, ready, true) +} + +func waitForEngineSignal(ctx context.Context, engine *Engine, ready <-chan struct{}, closeOnCancel bool) error { select { case <-ctx.Done(): - _ = engine.Close() + if closeOnCancel { + _ = engine.Close() + } return errors.WithStack(ctx.Err()) case <-ready: return nil @@ -1720,6 +1727,10 @@ func (e *Engine) run() { e.fail(err) return } + if err := e.persistStartupAppliedIndex(); err != nil { + e.fail(err) + return + } e.markStarted() for { @@ -1868,7 +1879,7 @@ func (e *Engine) handleDispatchReport(report dispatchReport) { // If the channel is full (unlikely — the buffer is sized to MaxInflightMsg), // the report is dropped and logged; this is acceptable because raft will retry // on the next tick and we only need eventual consistency between transport -// state and Progress state. Successful MsgSnap reports are the exception; see +// state and Progress state. MsgSnap reports are the exception; see // postReliableDispatchReport. func (e *Engine) postDispatchReport(report dispatchReport) { select { @@ -1883,9 +1894,10 @@ func (e *Engine) postDispatchReport(report dispatchReport) { } // postReliableDispatchReport delivers a dispatch outcome that must not be -// dropped. SnapshotFinish is not eventually consistent with ordinary -// unreachable reports: if it is lost, raft can keep the follower's Progress in -// StateSnapshot after the follower accepted the snapshot. +// dropped. SnapshotFinish and SnapshotFailure are not eventually consistent +// with ordinary unreachable reports: if either is lost, raft can keep the +// follower's Progress in StateSnapshot after the follower accepted or missed +// the snapshot. func (e *Engine) postReliableDispatchReport(report dispatchReport) { if e.dispatchReportCh == nil { return @@ -2189,6 +2201,7 @@ func (e *Engine) drainReady() error { e.releaseProtectedReceivedFSMSnapshotsUpTo(e.appliedIndex.Load()) e.handleReadStates(rd.ReadStates) e.rawNode.Advance(rd) + e.flushDeferredDispatchReports() if err := e.maybePersistLocalSnapshot(); err != nil { return err } @@ -2229,6 +2242,9 @@ func (e *Engine) persistReadyWithSnapshotLocked(rd etcdraft.Ready) error { if err := persistReadyToWAL(e.persist, rd); err != nil { return err } + if err := e.persistReceivedSnapshotAppliedIndex(rd.Snapshot); err != nil { + return err + } e.releaseProtectedReceivedFSMSnapshotsUpToLocked(snapshotIndex(rd.Snapshot)) return nil } @@ -2547,15 +2563,6 @@ func (e *Engine) applyReadySnapshotLocked(snapshot *raftpb.Snapshot) error { if err != nil { return errors.Wrapf(err, "decode snapshot token index=%d", snapshot.GetMetadata().GetIndex()) } - // B3/follow-up: also call SetDurableAppliedIndex(tok.Index) here - // after Restore so peer-after-InstallSnapshot populates the meta - // key. The local-snapshot persist path already bumps the live - // store (engine.persistLocalSnapshotPayload), but the receiving - // node's restored store inherits the pre-bump value embedded in - // the snapshot artifact. Design Non-Goals § - // docs/design/2026_06_02_implemented_idempotent_snapshot_restore.md#non-goals - // scopes this out of Branch 2; see PR #915 round-4/5 codex P2 on - // engine.go:4077 for the rationale. if err := openAndRestoreFSMSnapshot(e.fsm, fsmSnapPath(e.fsmSnapDir, tok.Index), tok.CRC32C); err != nil { return errors.Wrapf(err, "restore fsm snapshot file index=%d crc=%08x", tok.Index, tok.CRC32C) } @@ -2565,7 +2572,6 @@ func (e *Engine) applyReadySnapshotLocked(snapshot *raftpb.Snapshot) error { return errors.Wrapf(err, "restore fsm from legacy snapshot payload index=%d", snapshot.GetMetadata().GetIndex()) } } - if err := e.storage.ApplySnapshot(snapshot); err != nil { return errors.Wrapf(err, "apply snapshot to raft storage index=%d term=%d", snapshot.GetMetadata().GetIndex(), snapshot.GetMetadata().GetTerm()) @@ -3535,29 +3541,60 @@ func (e *Engine) createConfigSnapshot(index uint64, confState raftpb.ConfState, } } +// setDurableAppliedIndex pins the FSM's durable applied index to a +// known raft log index. FSMs that do not expose +// raftengine.AppliedIndexWriter silently no-op; the skip optimisation +// falls back to full restore for them (legacy test fakes, in-memory +// backends). +func (e *Engine) setDurableAppliedIndex(index uint64) error { + w, ok := e.fsm.(raftengine.AppliedIndexWriter) + if !ok { + return nil + } + return errors.WithStack(w.SetDurableAppliedIndex(index)) +} + // bumpDurableAppliedIndexBeforeSave pins the FSM's durable applied // index to `index` BEFORE the engine calls persist.SaveSnap, so a // successful snapshot persist always implies LastAppliedIndex >= // snap.Metadata.Index — closes the HLC-lease-only / encryption-only // fallback (PR #910 design §6). // -// FSMs that do not expose raftengine.AppliedIndexWriter silently -// no-op; the skip optimisation falls back to full restore for them -// (legacy test fakes, in-memory backends). pebble.Sync is forced on -// the writer side regardless of ELASTICKV_FSM_SYNC_MODE — once -// persist.SaveSnap returns, WAL compaction discards every log entry -// at or before snap.Metadata.Index, so there is no source to replay -// the meta key bump from. +// pebble.Sync is forced on the writer side regardless of +// ELASTICKV_FSM_SYNC_MODE — once persist.SaveSnap returns, WAL +// compaction discards every log entry at or before snap.Metadata.Index, +// so there is no source to replay the meta key bump from. // // Used by BOTH snapshot persist sites: persistCreatedSnapshot (this // file) and e.persistLocalSnapshotPayload (the steady-state // SnapshotCount-triggered hot path). func (e *Engine) bumpDurableAppliedIndexBeforeSave(index uint64) error { - w, ok := e.fsm.(raftengine.AppliedIndexWriter) - if !ok { + return e.setDurableAppliedIndex(index) +} + +func (e *Engine) persistStartupAppliedIndex() error { + if e.applied == 0 { return nil } - return errors.WithStack(w.SetDurableAppliedIndex(index)) + return e.setDurableAppliedIndex(e.applied) +} + +// persistReceivedSnapshotAppliedIndex runs only after persistReadyToWAL has +// saved the incoming raft snapshot. A crash before this point must fall back to +// a full FSM restore instead of letting the FSM meta index get ahead of the +// local durable raft snapshot. +func (e *Engine) persistReceivedSnapshotAppliedIndex(snapshot *raftpb.Snapshot) error { + if etcdraft.IsEmptySnap(snapshot) { + return nil + } + index := snapshot.GetMetadata().GetIndex() + if index == 0 { + return nil + } + if err := e.setDurableAppliedIndex(index); err != nil { + return errors.Wrapf(err, "persist durable applied index from snapshot index=%d", index) + } + return nil } func (e *Engine) persistCreatedSnapshot(snap raftpb.Snapshot) error { @@ -4535,7 +4572,7 @@ func (e *Engine) handleDispatchRequest(ctx context.Context, req dispatchRequest) // out of StateReplicate / StateSnapshot. Without this the leader keeps // Progress stuck and never retries sendAppend/sendSnap for the peer, // leaving the follower indefinitely stale even after heartbeats resume. - e.postDispatchReport(dispatchReport{to: req.msg.GetTo(), msgType: req.msg.GetType()}) + e.reportFailedDispatch(req.msg) } func (e *Engine) reportSuccessfulDispatch(msg raftpb.Message) { @@ -4942,10 +4979,39 @@ func (e *Engine) recordDroppedDispatch(msg raftpb.Message) { } func (e *Engine) reportDroppedDispatch(msg raftpb.Message) { + report := dispatchReport{to: msg.GetTo(), msgType: msg.GetType()} + if msg.GetType() == raftpb.MsgSnap { + e.deferReadyDispatchReport(report) + return + } if e.dispatchReportCh == nil { return } - e.postDispatchReport(dispatchReport{to: msg.GetTo(), msgType: msg.GetType()}) + e.postDispatchReport(report) +} + +func (e *Engine) reportFailedDispatch(msg raftpb.Message) { + report := dispatchReport{to: msg.GetTo(), msgType: msg.GetType()} + if msg.GetType() == raftpb.MsgSnap { + e.postReliableDispatchReport(report) + return + } + e.postDispatchReport(report) +} + +func (e *Engine) deferReadyDispatchReport(report dispatchReport) { + e.deferredReadyDispatchReports = append(e.deferredReadyDispatchReports, report) +} + +func (e *Engine) flushDeferredDispatchReports() { + if len(e.deferredReadyDispatchReports) == 0 { + return + } + reports := e.deferredReadyDispatchReports + e.deferredReadyDispatchReports = nil + for _, report := range reports { + e.handleDispatchReport(report) + } } // dispatchErrorCodeOf extracts the grpc status code name from err, or diff --git a/internal/raftengine/etcd/engine_applied_index_test.go b/internal/raftengine/etcd/engine_applied_index_test.go index d91ef433b..8ebb2ecfe 100644 --- a/internal/raftengine/etcd/engine_applied_index_test.go +++ b/internal/raftengine/etcd/engine_applied_index_test.go @@ -60,7 +60,9 @@ func (f *recordingAppliedIndexFSM) SetDurableAppliedIndex(idx uint64) error { f.failNext = false return f.failErr } - f.rec.record("bump", idx) + if f.rec != nil { + f.rec.record("bump", idx) + } return nil } @@ -338,6 +340,54 @@ func TestPersistCreatedSnapshot_BumpErrorAborts(t *testing.T) { "failed bump MUST NOT have recorded; SaveSnap MUST NOT have run") } +func TestPersistReadyWithSnapshot_BumpsAppliedIndexAfterWALSnapshot(t *testing.T) { + rec := &applyIndexOrderRecorder{} + fsm := &recordingAppliedIndexFSM{rec: rec} + persist := &recordingPersistStorage{rec: rec} + fsmSnapDir := t.TempDir() + const index uint64 = 77 + crc, _ := writeFSMFileForTest(t, fsmSnapDir, index, []byte("snapshot-payload")) + e := &Engine{ + storage: etcdraft.NewMemoryStorage(), + fsm: fsm, + persist: persist, + dataDir: t.TempDir(), + fsmSnapDir: fsmSnapDir, + } + + snap := appliedIndexTestSnapshot(index, encodeSnapshotToken(index, crc)) + require.NoError(t, e.persistReadyWithSnapshotLocked(etcdraft.Ready{Snapshot: &snap})) + + require.Equal(t, []orderEvent{ + {kind: "save", index: index}, + {kind: "bump", index: index}, + }, rec.snapshot(), + "received snapshot restore MUST publish its applied index only after the raft snapshot is durable") + require.Equal(t, index, e.applied) +} + +func TestPersistReadyWithSnapshot_BumpErrorAfterWALSnapshotSurfaces(t *testing.T) { + rec := &applyIndexOrderRecorder{} + fsm := &recordingAppliedIndexFSM{rec: rec, failNext: true, failErr: io.ErrShortBuffer} + persist := &recordingPersistStorage{rec: rec} + fsmSnapDir := t.TempDir() + const index uint64 = 78 + crc, _ := writeFSMFileForTest(t, fsmSnapDir, index, []byte("snapshot-payload")) + e := &Engine{ + storage: etcdraft.NewMemoryStorage(), + fsm: fsm, + persist: persist, + dataDir: t.TempDir(), + fsmSnapDir: fsmSnapDir, + } + + snap := appliedIndexTestSnapshot(index, encodeSnapshotToken(index, crc)) + err := e.persistReadyWithSnapshotLocked(etcdraft.Ready{Snapshot: &snap}) + require.Error(t, err, "received snapshot bump failure MUST be surfaced") + require.Equal(t, []orderEvent{{kind: "save", index: index}}, rec.snapshot(), + "received snapshot path must not bump before the WAL snapshot is durable") +} + // --- Site 2: persistLocalSnapshotPayload (steady-state hot path) --- // // These mirror the Site 1 tests above but exercise the engine's diff --git a/internal/raftengine/etcd/engine_test.go b/internal/raftengine/etcd/engine_test.go index 87e8447df..ca92e1ff6 100644 --- a/internal/raftengine/etcd/engine_test.go +++ b/internal/raftengine/etcd/engine_test.go @@ -114,6 +114,18 @@ type blockingApplyStateMachine struct { startOnce sync.Once } +type recordingStartupApplyStateMachine struct { + *blockingApplyStateMachine + rec *applyIndexOrderRecorder +} + +func (s *recordingStartupApplyStateMachine) SetDurableAppliedIndex(idx uint64) error { + if s.rec != nil { + s.rec.record("bump", idx) + } + return nil +} + type blockingSnapshot struct { started chan struct{} release chan struct{} @@ -789,6 +801,58 @@ func TestSendMessagesDoesNotBlockWhenDispatchQueueIsFull(t *testing.T) { } } +func TestReportFailedDispatchReliablyQueuesMsgSnapFailure(t *testing.T) { + engine := &Engine{ + dispatchReportCh: make(chan dispatchReport, 1), + closeCh: make(chan struct{}), + } + engine.dispatchReportCh <- dispatchReport{to: 2, msgType: raftpb.MsgHeartbeat} + + done := make(chan struct{}) + go func() { + engine.reportFailedDispatch(raftpb.Message{ + Type: messageTypePtr(raftpb.MsgSnap), + To: uint64Ptr(3), + }) + close(done) + }() + + select { + case <-done: + t.Fatal("MsgSnap failure report must wait instead of dropping when report channel is full") + case <-time.After(20 * time.Millisecond): + } + + <-engine.dispatchReportCh + requireSignal(t, done, time.Second, "MsgSnap failure report was not delivered after report channel drained") + report := <-engine.dispatchReportCh + require.Equal(t, uint64(3), report.to) + require.Equal(t, raftpb.MsgSnap, report.msgType) + require.False(t, report.snapshotFinish) +} + +func TestReportDroppedDispatchDefersMsgSnapUntilReadyAdvanced(t *testing.T) { + engine := &Engine{ + dispatchReportCh: make(chan dispatchReport, 1), + closeCh: make(chan struct{}), + } + queued := dispatchReport{to: 2, msgType: raftpb.MsgHeartbeat} + engine.dispatchReportCh <- queued + + engine.reportDroppedDispatch(raftpb.Message{ + Type: messageTypePtr(raftpb.MsgSnap), + To: uint64Ptr(3), + }) + + require.Equal(t, queued, <-engine.dispatchReportCh) + select { + case report := <-engine.dispatchReportCh: + t.Fatalf("dropped MsgSnap report should not enqueue from the event loop: %+v", report) + default: + } + require.Equal(t, []dispatchReport{{to: 3, msgType: raftpb.MsgSnap}}, engine.deferredReadyDispatchReports) +} + func TestStopDispatchWorkersCancelsInflightDispatch(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) t.Cleanup(cancel) @@ -1554,9 +1618,13 @@ func TestOpenMultiNodeWaitStartedWaitsForCommittedTailDrain(t *testing.T) { }}, })) - fsm := &blockingApplyStateMachine{ - started: make(chan struct{}), - release: make(chan struct{}), + rec := &applyIndexOrderRecorder{} + fsm := &recordingStartupApplyStateMachine{ + blockingApplyStateMachine: &blockingApplyStateMachine{ + started: make(chan struct{}), + release: make(chan struct{}), + }, + rec: rec, } done := make(chan openResult, 1) go func() { @@ -1591,6 +1659,8 @@ func TestOpenMultiNodeWaitStartedWaitsForCommittedTailDrain(t *testing.T) { requireStartupWaitResult(t, waitDone, time.Second, "WaitStarted did not return after committed tail drain completed") requireSignal(t, result.engine.startedCh, time.Second, "engine did not mark started after committed tail drain completed") + require.Equal(t, []orderEvent{{kind: "bump", index: 2}}, rec.snapshot(), + "startup must persist the committed tail applied index before marking the engine started") } type openResult struct { @@ -1638,6 +1708,24 @@ func requireStartupWaitResult(t *testing.T, ch <-chan error, timeout time.Durati } } +func TestWaitStartedTimeoutDoesNotCloseEngine(t *testing.T) { + engine := &Engine{ + startedCh: make(chan struct{}), + doneCh: make(chan struct{}), + closeCh: make(chan struct{}), + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + err := engine.WaitStarted(ctx) + require.ErrorIs(t, err, context.Canceled) + select { + case <-engine.closeCh: + t.Fatal("WaitStarted timeout must not close an already-returned engine") + default: + } +} + func TestOpenMultiNodeReplicatesOverGRPCTransport(t *testing.T) { nodes, peers := newTransportTestNodes(t, 3) startTransportTestServers(nodes, peers) diff --git a/internal/raftengine/statemachine.go b/internal/raftengine/statemachine.go index d834ad318..9c796f4f8 100644 --- a/internal/raftengine/statemachine.go +++ b/internal/raftengine/statemachine.go @@ -69,24 +69,25 @@ type AppliedIndexReader interface { } // AppliedIndexWriter is an OPTIONAL extension that lets the engine -// pin the FSM's durable applied-index to a known value at snapshot -// persist time. See docs/design/2026_06_02_implemented_idempotent_snapshot_restore.md +// pin the FSM's durable applied-index to a known value at raft +// durability boundaries. See docs/design/2026_06_02_implemented_idempotent_snapshot_restore.md // §6 "HLC lease entries — checkpoint at snapshot persist". // -// The engine calls SetDurableAppliedIndex(snap.Metadata.Index) -// before it calls persist.SaveSnap, so that on every successful -// snapshot persist the invariant `LastAppliedIndex >= -// snapshot.Metadata.Index` holds unconditionally — closing the -// HLC-lease-only / encryption-only fallback that would otherwise -// leave LastAppliedIndex stuck at the last data-Apply index. +// The engine calls SetDurableAppliedIndex at local snapshot persist, +// after received-snapshot WAL persistence, and at startup +// committed-tail drain boundaries, so the invariant +// `LastAppliedIndex >= the locally durable raft state` holds once the +// engine is store-ready. This closes the HLC-lease-only / +// encryption-only fallback that would otherwise leave LastAppliedIndex +// stuck at the last data-Apply index. // // Implementations MUST persist the value with pebble.Sync (or the // equivalent strong-durability flag for the backing store) // regardless of ELASTICKV_FSM_SYNC_MODE. The checkpoint is the only -// durable carrier of metaAppliedIndex at this point — once -// persist.SaveSnap returns, WAL compaction discards every log entry -// at or before snap.Metadata.Index, so there is no source to replay -// the meta key bump from. +// durable carrier of metaAppliedIndex at local snapshot persist time +// — once persist.SaveSnap returns, WAL compaction discards every log +// entry at or before snap.Metadata.Index, so there is no source to +// replay the meta key bump from. type AppliedIndexWriter interface { SetDurableAppliedIndex(idx uint64) error } From 3e988efca07dd77de287b5241b7a7457e3089e94 Mon Sep 17 00:00:00 2001 From: bootjp Date: Wed, 15 Jul 2026 07:14:15 +0900 Subject: [PATCH 19/26] Enforce current route fences at apply time --- adapter/dynamodb_transact.go | 4 ++ adapter/retryable_write_fence_test.go | 4 ++ adapter/sqs_messages.go | 2 +- adapter/sqs_reaper.go | 4 +- adapter/sqs_redrive.go | 2 +- distribution/engine.go | 13 +++-- kv/fsm.go | 28 ++++++---- kv/fsm_migration_fence_test.go | 41 +++++---------- kv/sharded_coordinator.go | 54 +++++++++++--------- kv/sharded_coordinator_del_prefix_test.go | 62 +++++++++++++++++++++-- 10 files changed, 138 insertions(+), 76 deletions(-) diff --git a/adapter/dynamodb_transact.go b/adapter/dynamodb_transact.go index 9569b6053..e1f6ab221 100644 --- a/adapter/dynamodb_transact.go +++ b/adapter/dynamodb_transact.go @@ -1108,6 +1108,10 @@ func isRetryableTransactWriteError(err error) bool { return errors.Is(err, store.ErrWriteConflict) || errors.Is(err, kv.ErrTxnLocked) || errors.Is(err, kv.ErrRouteWriteFenced) } +func isIgnorableTransactRaceError(err error) bool { + return errors.Is(err, store.ErrWriteConflict) || errors.Is(err, kv.ErrTxnLocked) +} + func shouldPreserveTransactWriteAttempt(err error) bool { return isRetryableTransactWriteError(err) && !errors.Is(err, kv.ErrRouteWriteFenced) } diff --git a/adapter/retryable_write_fence_test.go b/adapter/retryable_write_fence_test.go index e5b436c06..176ac68ca 100644 --- a/adapter/retryable_write_fence_test.go +++ b/adapter/retryable_write_fence_test.go @@ -4,6 +4,7 @@ import ( "testing" "github.com/bootjp/elastickv/kv" + "github.com/bootjp/elastickv/store" "github.com/stretchr/testify/require" ) @@ -14,4 +15,7 @@ func TestWriteFenceErrorsAreAdapterRetryable(t *testing.T) { require.True(t, isRetryableS3MutationErr(kv.ErrRouteWriteFenced)) require.True(t, isRetryableTransactWriteError(kv.ErrRouteWriteFenced)) require.False(t, shouldPreserveTransactWriteAttempt(kv.ErrRouteWriteFenced)) + require.False(t, isIgnorableTransactRaceError(kv.ErrRouteWriteFenced)) + require.True(t, isIgnorableTransactRaceError(store.ErrWriteConflict)) + require.True(t, isIgnorableTransactRaceError(kv.ErrTxnLocked)) } diff --git a/adapter/sqs_messages.go b/adapter/sqs_messages.go index 516d53d87..83bfe6009 100644 --- a/adapter/sqs_messages.go +++ b/adapter/sqs_messages.go @@ -1292,7 +1292,7 @@ func (s *SQSServer) expireMessage(ctx context.Context, queueName string, meta *s Elems: elems, } if _, err := s.coordinator.Dispatch(ctx, req); err != nil { - if isRetryableTransactWriteError(err) { + if isIgnorableTransactRaceError(err) { return nil } return errors.WithStack(err) diff --git a/adapter/sqs_reaper.go b/adapter/sqs_reaper.go index ae7f6c20d..a25b83ca6 100644 --- a/adapter/sqs_reaper.go +++ b/adapter/sqs_reaper.go @@ -664,7 +664,7 @@ func (s *SQSServer) reapOneRecord(ctx context.Context, queueName string, meta *s return err } if _, err := s.coordinator.Dispatch(ctx, req); err != nil { - if isRetryableTransactWriteError(err) { + if isIgnorableTransactRaceError(err) { return nil } return errors.WithStack(err) @@ -873,7 +873,7 @@ func (s *SQSServer) dispatchDedupDelete(ctx context.Context, key []byte, readTS }, } if _, err := s.coordinator.Dispatch(ctx, req); err != nil { - if isRetryableTransactWriteError(err) { + if isIgnorableTransactRaceError(err) { return nil } return errors.WithStack(err) diff --git a/adapter/sqs_redrive.go b/adapter/sqs_redrive.go index 3fb13b8d3..cfc4437d0 100644 --- a/adapter/sqs_redrive.go +++ b/adapter/sqs_redrive.go @@ -237,7 +237,7 @@ func (s *SQSServer) redriveCandidateToDLQ( return false, err } if _, err := s.coordinator.Dispatch(ctx, req); err != nil { - if isRetryableTransactWriteError(err) { + if isIgnorableTransactRaceError(err) { return true, nil } return false, errors.WithStack(err) diff --git a/distribution/engine.go b/distribution/engine.go index 96d9a4c8e..16b4ad4b8 100644 --- a/distribution/engine.go +++ b/distribution/engine.go @@ -164,11 +164,16 @@ func (s RouteHistorySnapshot) Version() uint64 { return s.version } // scan from O(N) to "first non-covering gap" without changing the // resolution semantics). func (s RouteHistorySnapshot) OwnerOf(key []byte) (uint64, bool) { - r, ok := s.RouteOf(key) - if !ok { - return 0, false + for _, r := range s.routes { + if bytes.Compare(key, r.Start) < 0 { + break + } + if r.End != nil && bytes.Compare(key, r.End) >= 0 { + continue + } + return r.GroupID, true } - return r.GroupID, true + return 0, false } // RouteOf returns the route that covered key at this snapshot's version. diff --git a/kv/fsm.go b/kv/fsm.go index bc5da0b36..88c7b95cc 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -828,21 +828,27 @@ func (f *kvFSM) verifyWriteFence(r *pb.Request) error { return nil } observedVer := r.GetObservedRouteVersion() - if !f.writeFenceHistoryReady(observedVer) { + if !f.writeFenceHistoryReady() { return nil } - observedSnap, ok := f.routes.SnapshotAt(observedVer) + currentSnap, ok := f.routes.Current() if !ok { - return errors.WithStack(ErrComposed1VersionGCd) - } - if err := verifyWriteFenceFromSnapshot(r.GetMutations(), observedSnap, observedVer, "observed"); err != nil { - return err + return nil } - currentSnap, ok := f.routes.Current() - if !ok || currentSnap.Version() == observedSnap.Version() { - return nil + if observedVer != 0 { + observedSnap, ok := f.routes.SnapshotAt(observedVer) + if !ok { + return errors.WithStack(ErrComposed1VersionGCd) + } + if err := verifyWriteFenceFromSnapshot(r.GetMutations(), observedSnap, observedVer, "observed"); err != nil { + return err + } + if currentSnap.Version() == observedSnap.Version() { + return nil + } } + return verifyWriteFenceFromSnapshot(r.GetMutations(), currentSnap, currentSnap.Version(), "current") } @@ -859,8 +865,8 @@ func requestBypassesWriteFence(r *pb.Request) bool { return false } -func (f *kvFSM) writeFenceHistoryReady(observedVer uint64) bool { - return f.routes != nil && f.shardGroupID != 0 && observedVer != 0 +func (f *kvFSM) writeFenceHistoryReady() bool { + return f.routes != nil && f.shardGroupID != 0 } func verifyWriteFenceFromSnapshot(mutations []*pb.Mutation, snap RouteSnapshot, snapVer uint64, phase string) error { diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index 959358513..f1c027477 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -39,18 +39,14 @@ func newS3BucketAuxiliaryWriteFencedFSM(t *testing.T, bucket string) *kvFSM { return newComposed1FSM(t, engine, 1) } -func TestFSMAppliesRawPointWriteDespiteLocalWriteFencedRoute(t *testing.T) { +func TestFSMRejectsCurrentWriteFencedRawPointWrite(t *testing.T) { t.Parallel() fsm := newWriteFencedFSM(t) err := fsm.handleRawRequest(context.Background(), &pb.Request{ Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("v")}}, }, 10) - require.NoError(t, err) - - got, getErr := fsm.store.GetAt(context.Background(), []byte("z"), ^uint64(0)) - require.NoError(t, getErr) - require.Equal(t, []byte("v"), got) + require.ErrorIs(t, err, ErrRouteWriteFenced) } func TestFSMRejectsObservedWriteFencedRawPointWrite(t *testing.T) { @@ -84,7 +80,7 @@ func TestFSMRejectsCurrentWriteFenceAfterObservedActiveRawPointWrite(t *testing. require.ErrorIs(t, err, ErrRouteWriteFenced) } -func TestFSMAppliesS3BucketAuxiliaryPointWriteDespiteLocalWriteFencedRoute(t *testing.T) { +func TestFSMRejectsCurrentWriteFencedS3BucketAuxiliaryPointWrite(t *testing.T) { t.Parallel() ctx := context.Background() @@ -98,11 +94,7 @@ func TestFSMAppliesS3BucketAuxiliaryPointWriteDespiteLocalWriteFencedRoute(t *te err := fsm.handleRawRequest(ctx, &pb.Request{ Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: key, Value: []byte("v")}}, }, 10) - require.NoError(t, err) - - got, getErr := fsm.store.GetAt(ctx, key, ^uint64(0)) - require.NoError(t, getErr) - require.Equal(t, []byte("v"), got) + require.ErrorIs(t, err, ErrRouteWriteFenced) } } @@ -121,7 +113,7 @@ func TestFSMRejectsObservedWriteFencedS3BucketAuxiliaryPointWrite(t *testing.T) require.ErrorIs(t, err, ErrRouteWriteFenced) } -func TestFSMAppliesDelPrefixDespiteLocalWriteFencedRoute(t *testing.T) { +func TestFSMRejectsCurrentWriteFencedDelPrefix(t *testing.T) { t.Parallel() fsm := newWriteFencedFSM(t) @@ -130,10 +122,7 @@ func TestFSMAppliesDelPrefixDespiteLocalWriteFencedRoute(t *testing.T) { err := fsm.handleRawRequest(context.Background(), &pb.Request{ Mutations: []*pb.Mutation{{Op: pb.Op_DEL_PREFIX, Key: []byte("z")}}, }, 10) - require.NoError(t, err) - - _, getErr := fsm.store.GetAt(context.Background(), []byte("z"), ^uint64(0)) - require.Error(t, getErr) + require.ErrorIs(t, err, ErrRouteWriteFenced) } func TestFSMRejectsObservedWriteFencedDelPrefix(t *testing.T) { @@ -149,7 +138,7 @@ func TestFSMRejectsObservedWriteFencedDelPrefix(t *testing.T) { require.ErrorIs(t, err, ErrRouteWriteFenced) } -func TestFSMAppliesFullRangeDelPrefixDespiteLocalWriteFencedRoute(t *testing.T) { +func TestFSMRejectsCurrentWriteFencedFullRangeDelPrefix(t *testing.T) { t.Parallel() fsm := newWriteFencedFSM(t) @@ -158,13 +147,10 @@ func TestFSMAppliesFullRangeDelPrefixDespiteLocalWriteFencedRoute(t *testing.T) err := fsm.handleRawRequest(context.Background(), &pb.Request{ Mutations: []*pb.Mutation{{Op: pb.Op_DEL_PREFIX, Key: nil}}, }, 10) - require.NoError(t, err) - - _, getErr := fsm.store.GetAt(context.Background(), []byte("z"), ^uint64(0)) - require.Error(t, getErr) + require.ErrorIs(t, err, ErrRouteWriteFenced) } -func TestFSMAppliesBroadInternalDelPrefixDespiteLocalWriteFencedRoute(t *testing.T) { +func TestFSMRejectsCurrentWriteFencedBroadInternalDelPrefix(t *testing.T) { t.Parallel() fsm := newWriteFencedFSM(t) @@ -174,13 +160,10 @@ func TestFSMAppliesBroadInternalDelPrefixDespiteLocalWriteFencedRoute(t *testing err := fsm.handleRawRequest(context.Background(), &pb.Request{ Mutations: []*pb.Mutation{{Op: pb.Op_DEL_PREFIX, Key: []byte("!redis|")}}, }, 10) - require.NoError(t, err) - - _, getErr := fsm.store.GetAt(context.Background(), key, ^uint64(0)) - require.Error(t, getErr) + require.ErrorIs(t, err, ErrRouteWriteFenced) } -func TestFSMAppliesPrepareAndAbortDespiteLocalWriteFencedRoute(t *testing.T) { +func TestFSMRejectsCurrentWriteFencedPrepareButAllowsAbort(t *testing.T) { t.Parallel() ctx := context.Background() @@ -194,7 +177,7 @@ func TestFSMAppliesPrepareAndAbortDespiteLocalWriteFencedRoute(t *testing.T) { {Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("v")}, }, } - require.NoError(t, fsm.handleTxnRequest(ctx, prepare, 10)) + require.ErrorIs(t, fsm.handleTxnRequest(ctx, prepare, 10), ErrRouteWriteFenced) abort := &pb.Request{ IsTxn: true, diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index e8bfb5e80..147abc067 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -700,11 +700,6 @@ func (c *ShardedCoordinator) Dispatch(ctx context.Context, reqs *OperationGroup[ return nil, err } - c.maybeAutoPinRawObservedRouteVersion(reqs) - if resp, handled, err := c.dispatchBeforeShardRouting(ctx, reqs); handled { - return resp, err - } - // Capture whether the caller supplied a non-zero StartTS BEFORE // the coordinator-allocates-on-zero branch below mutates the // field. A caller-supplied StartTS names a specific snapshot @@ -736,10 +731,13 @@ func (c *ShardedCoordinator) Dispatch(ctx context.Context, reqs *OperationGroup[ } if reqs.IsTxn { + if resp, handled, err := c.dispatchBeforeShardRouting(ctx, reqs); handled { + return resp, err + } return c.dispatchTxnWithComposed1Retry(ctx, reqs, callerSuppliedStartTS) } - return c.dispatchNonTxn(ctx, reqs) + return c.dispatchRawWithComposed1Retry(ctx, reqs) } func (c *ShardedCoordinator) dispatchBeforeShardRouting(ctx context.Context, reqs *OperationGroup[OP]) (*CoordinateResponse, bool, error) { @@ -815,16 +813,17 @@ func (c *ShardedCoordinator) dispatchBeforeShardRouting(ctx context.Context, req // gate spuriously rejects resolver-routed commits even when // the resolver picked the correct gid. Skip the auto-pin // for any resolver-recognised key; the request flows with -// ObservedRouteVersion=0 and the M3 gate short-circuits — -// restoring the pre-auto-pin behaviour for resolver-routed -// txns. Resolver-aware M3 is M5+ work (codex P1 on -// 6a458a28, PR #900). +// ObservedRouteVersion=0 and the Composed-1 owner gate +// short-circuits — restoring the pre-auto-pin behaviour for +// resolver-routed txns. Resolver-aware M3 is M5+ work (PR #900). // // The non-auto-pin case (request flows with ObservedRouteVersion=0, -// M3 gate short-circuits) is the safe non-regressing posture for -// non-migrated callers — the gate cannot retroactively pin reads -// it was not present for. Adapters that want M3 protection must -// migrate to pin at BeginTxn per §4.1. +// Composed-1 owner gate short-circuits) is the safe non-regressing +// posture for non-migrated callers — the owner gate cannot +// retroactively pin reads it was not present for. The write-fence +// gate still checks the current route snapshot at apply time. +// Adapters that want full M3 owner protection must migrate to pin at +// BeginTxn per §4.1. // // Extracted from dispatchTxnWithComposed1Retry to keep its // cyclomatic complexity in the cyclop budget. @@ -838,16 +837,6 @@ func (c *ShardedCoordinator) maybeAutoPinObservedRouteVersion(reqs *OperationGro reqs.ObservedRouteVersion = c.engine.Version() } -func (c *ShardedCoordinator) maybeAutoPinRawObservedRouteVersion(reqs *OperationGroup[OP]) { - if c.engine == nil || reqs == nil || reqs.IsTxn || reqs.ObservedRouteVersion != 0 { - return - } - if c.anyResolverClaimedKey(reqs.Elems) { - return - } - reqs.ObservedRouteVersion = c.engine.Version() -} - // anyResolverClaimedKey reports whether any element's key is // claimed by the partition resolver. Returns false when the // router has no resolver installed (the most common case) so the @@ -1006,6 +995,23 @@ func (c *ShardedCoordinator) dispatchNonTxn(ctx context.Context, reqs *Operation return &CoordinateResponse{CommitIndex: r.CommitIndex}, nil } +func (c *ShardedCoordinator) dispatchRawWithComposed1Retry(ctx context.Context, reqs *OperationGroup[OP]) (*CoordinateResponse, error) { + for attempt := 0; attempt <= composed1RetryAttempts; attempt++ { + resp, handled, err := c.dispatchBeforeShardRouting(ctx, reqs) + if !handled { + resp, err = c.dispatchNonTxn(ctx, reqs) + } + if err == nil { + return resp, nil + } + if !errors.Is(err, ErrComposed1VersionGCd) || attempt == composed1RetryAttempts || c.engine == nil { + return resp, err + } + reqs.ObservedRouteVersion = c.engine.Version() + } + return nil, errors.WithStack(ErrInvalidRequest) +} + // hasDelPrefixElem returns true if any element is a DelPrefix operation. func hasDelPrefixElem(elems []*Elem[OP]) bool { for _, e := range elems { diff --git a/kv/sharded_coordinator_del_prefix_test.go b/kv/sharded_coordinator_del_prefix_test.go index 7b52fe105..d5f8632d3 100644 --- a/kv/sharded_coordinator_del_prefix_test.go +++ b/kv/sharded_coordinator_del_prefix_test.go @@ -122,7 +122,7 @@ func TestShardedCoordinator_DelPrefixBroadcastsToAllGroups(t *testing.T) { "same DEL_PREFIX element must use the same timestamp across shards") } -func TestShardedCoordinator_DelPrefixCarriesObservedRouteVersion(t *testing.T) { +func TestShardedCoordinator_DelPrefixDoesNotAutoPinObservedRouteVersion(t *testing.T) { t.Parallel() engine := distribution.NewEngine() @@ -142,10 +142,10 @@ func TestShardedCoordinator_DelPrefixCarriesObservedRouteVersion(t *testing.T) { }) require.NoError(t, err) require.Len(t, txn.requests, 1) - require.Equal(t, uint64(7), txn.requests[0].GetObservedRouteVersion()) + require.Zero(t, txn.requests[0].GetObservedRouteVersion()) } -func TestShardedCoordinator_RawWriteCarriesObservedRouteVersion(t *testing.T) { +func TestShardedCoordinator_RawWriteDoesNotAutoPinObservedRouteVersion(t *testing.T) { t.Parallel() engine := distribution.NewEngine() @@ -165,7 +165,61 @@ func TestShardedCoordinator_RawWriteCarriesObservedRouteVersion(t *testing.T) { }) require.NoError(t, err) require.Len(t, txn.requests, 1) - require.Equal(t, uint64(9), txn.requests[0].GetObservedRouteVersion()) + require.Zero(t, txn.requests[0].GetObservedRouteVersion()) +} + +func TestShardedCoordinator_RetriesRawWriteWhenObservedRouteVersionGCd(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 9, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: nil, End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }, + })) + txn := &recordingTransactional{ + errs: []error{ErrComposed1VersionGCd}, + } + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + ObservedRouteVersion: 3, + Elems: []*Elem[OP]{{Op: Put, Key: []byte("k"), Value: []byte("v")}}, + }) + require.NoError(t, err) + require.Len(t, txn.requests, 2) + require.Equal(t, uint64(3), txn.requests[0].GetObservedRouteVersion()) + require.Equal(t, uint64(9), txn.requests[1].GetObservedRouteVersion()) +} + +func TestShardedCoordinator_RetriesDelPrefixWhenObservedRouteVersionGCd(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 7, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: nil, End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }, + })) + txn := &recordingTransactional{ + errs: []error{ErrComposed1VersionGCd}, + } + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + ObservedRouteVersion: 2, + Elems: []*Elem[OP]{{Op: DelPrefix, Key: []byte("user:")}}, + }) + require.NoError(t, err) + require.Len(t, txn.requests, 2) + require.Equal(t, uint64(2), txn.requests[0].GetObservedRouteVersion()) + require.Equal(t, uint64(7), txn.requests[1].GetObservedRouteVersion()) } func TestShardedCoordinatorRejectsPointWriteOnWriteFencedRoute(t *testing.T) { From 04e8053bcb6649612ba8388f98c524c0268031a1 Mon Sep 17 00:00:00 2001 From: bootjp Date: Wed, 15 Jul 2026 07:52:00 +0900 Subject: [PATCH 20/26] Prioritize received raft snapshots --- internal/raftengine/etcd/engine.go | 40 +++++++++++++++++++++---- internal/raftengine/etcd/engine_test.go | 22 ++++++++++++++ 2 files changed, 56 insertions(+), 6 deletions(-) diff --git a/internal/raftengine/etcd/engine.go b/internal/raftengine/etcd/engine.go index b87a00d13..5316967db 100644 --- a/internal/raftengine/etcd/engine.go +++ b/internal/raftengine/etcd/engine.go @@ -60,9 +60,10 @@ const ( // from shrinking the shared inbound queue enough to drop heartbeats. minInboundQueueCapacity = 128 // priorityStepQueueCapacity is the inbound control-plane queue size. - // Heartbeats, votes, read-index responses, and timeout-now messages are - // tiny but time-sensitive; keeping them off the bulk stepCh prevents a - // MsgApp burst from forcing followers into avoidable elections. + // Heartbeats, votes, read-index responses, timeout-now messages, and + // received snapshot tokens are tiny but time-sensitive; keeping them off the + // bulk stepCh prevents a MsgApp burst from forcing followers into avoidable + // elections or rejecting a catch-up snapshot after its payload was streamed. priorityStepQueueCapacity = 1024 // priorityStepBurstLimit bounds consecutive non-blocking priority drains // so a sustained control-message stream cannot starve Tick, proposals, or @@ -2493,6 +2494,14 @@ func isPriorityMsg(t raftpb.MessageType) bool { t == raftpb.MsgTimeoutNow } +func isInboundPriorityMsg(t raftpb.MessageType) bool { + return isPriorityMsg(t) || t == raftpb.MsgSnap +} + +func isBlockingInboundStepMsg(t raftpb.MessageType) bool { + return t == raftpb.MsgSnap +} + // selectDispatchLane picks the per-peer channel for msgType. In the legacy // layout it returns pd.heartbeatResp for MsgHeartbeatResp, pd.heartbeat for // other priority control traffic, and pd.normal for everything else. In the @@ -4346,9 +4355,10 @@ func maxAppliedIndex(snapshot raftpb.Snapshot) uint64 { } func (e *Engine) enqueueStep(ctx context.Context, msg raftpb.Message) error { - ch := e.stepCh - if isPriorityMsg(msg.GetType()) && e.priorityStepCh != nil { - ch = e.priorityStepCh + ch := e.stepChannelFor(msg.GetType()) + + if isBlockingInboundStepMsg(msg.GetType()) { + return e.enqueueBlockingStep(ctx, ch, msg) } select { @@ -4364,6 +4374,24 @@ func (e *Engine) enqueueStep(ctx context.Context, msg raftpb.Message) error { } } +func (e *Engine) stepChannelFor(msgType raftpb.MessageType) chan raftpb.Message { + if isInboundPriorityMsg(msgType) && e.priorityStepCh != nil { + return e.priorityStepCh + } + return e.stepCh +} + +func (e *Engine) enqueueBlockingStep(ctx context.Context, ch chan raftpb.Message, msg raftpb.Message) error { + select { + case <-ctx.Done(): + return errors.WithStack(ctx.Err()) + case <-e.doneCh: + return e.currentErrorOrClosed() + case ch <- msg: + return nil + } +} + func (e *Engine) handleTransportMessage(ctx context.Context, msg raftpb.Message) error { select { case <-ctx.Done(): diff --git a/internal/raftengine/etcd/engine_test.go b/internal/raftengine/etcd/engine_test.go index ca92e1ff6..6bc836704 100644 --- a/internal/raftengine/etcd/engine_test.go +++ b/internal/raftengine/etcd/engine_test.go @@ -647,6 +647,28 @@ func TestEnqueueStepPriorityBypassesFullBulkQueue(t *testing.T) { require.Equal(t, uint64(1), engine.StepQueueFullCount()) } +func TestEnqueueStepSnapshotBypassesFullBulkQueue(t *testing.T) { + engine := &Engine{ + doneCh: make(chan struct{}), + stepCh: make(chan raftpb.Message, 1), + priorityStepCh: make(chan raftpb.Message, 1), + } + engine.stepCh <- raftpb.Message{Type: messageTypePtr(raftpb.MsgApp)} + + err := engine.enqueueStep(context.Background(), raftpb.Message{Type: messageTypePtr(raftpb.MsgSnap)}) + require.NoError(t, err) + require.Equal(t, uint64(0), engine.StepQueueFullCount()) + + select { + case msg := <-engine.priorityStepCh: + require.Equal(t, raftpb.MsgSnap, msg.GetType()) + default: + t.Fatal("snapshot step was not enqueued on the priority queue") + } + + require.Len(t, engine.stepCh, 1) +} + func TestHandleEventDrainsPriorityStepBeforeBulkStep(t *testing.T) { engine := &Engine{ closeCh: make(chan struct{}), From d818602593249f899bad1ce6d0b11da4f23508dc Mon Sep 17 00:00:00 2001 From: bootjp Date: Wed, 15 Jul 2026 08:23:05 +0900 Subject: [PATCH 21/26] Stabilize snapshot catch-up and fence checks --- adapter/dynamodb_transact.go | 2 + internal/raftengine/etcd/engine.go | 11 +++++ internal/raftengine/etcd/engine_test.go | 34 ++++++++++++++ internal/raftengine/etcd/fsm_snapshot_file.go | 31 ++++++++++++- .../raftengine/etcd/fsm_snapshot_file_test.go | 24 ++++++++++ kv/fsm.go | 3 -- kv/fsm_migration_fence_test.go | 46 +++++++++++++++++++ 7 files changed, 147 insertions(+), 4 deletions(-) diff --git a/adapter/dynamodb_transact.go b/adapter/dynamodb_transact.go index e1f6ab221..52fd98a42 100644 --- a/adapter/dynamodb_transact.go +++ b/adapter/dynamodb_transact.go @@ -1109,6 +1109,8 @@ func isRetryableTransactWriteError(err error) bool { } func isIgnorableTransactRaceError(err error) bool { + // Route fences reject before applying the write. They must be retried or + // propagated, not swallowed as "another worker already completed it". return errors.Is(err, store.ErrWriteConflict) || errors.Is(err, kv.ErrTxnLocked) } diff --git a/internal/raftengine/etcd/engine.go b/internal/raftengine/etcd/engine.go index 5316967db..d36e7cf8a 100644 --- a/internal/raftengine/etcd/engine.go +++ b/internal/raftengine/etcd/engine.go @@ -927,6 +927,17 @@ func newRawNode(cfg OpenConfig, storage *etcdraft.MemoryStorage, applied uint64) func waitForOpen(ctx context.Context, engine *Engine, waitForLeader bool) (*Engine, error) { if !waitForLeader { + select { + case <-ctx.Done(): + _ = engine.Close() + return nil, errors.WithStack(ctx.Err()) + case <-engine.doneCh: + if err := engine.currentError(); err != nil { + return nil, err + } + return nil, errors.WithStack(errClosed) + default: + } return engine, nil } if err := waitForOpenSignal(ctx, engine, engine.startedCh); err != nil { diff --git a/internal/raftengine/etcd/engine_test.go b/internal/raftengine/etcd/engine_test.go index 6bc836704..148f1d149 100644 --- a/internal/raftengine/etcd/engine_test.go +++ b/internal/raftengine/etcd/engine_test.go @@ -1748,6 +1748,40 @@ func TestWaitStartedTimeoutDoesNotCloseEngine(t *testing.T) { } } +func TestWaitForOpenCanceledContextClosesMultiNodeEngine(t *testing.T) { + engine := &Engine{ + doneCh: make(chan struct{}), + closeCh: make(chan struct{}), + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + type result struct { + engine *Engine + err error + } + resultCh := make(chan result, 1) + go func() { + opened, err := waitForOpen(ctx, engine, false) + resultCh <- result{engine: opened, err: err} + }() + + select { + case <-engine.closeCh: + case <-time.After(time.Second): + t.Fatal("canceled multi-node Open must close the engine before returning") + } + close(engine.doneCh) + + select { + case got := <-resultCh: + require.Nil(t, got.engine) + require.ErrorIs(t, got.err, context.Canceled) + case <-time.After(time.Second): + t.Fatal("waitForOpen did not return after Close completed") + } +} + func TestOpenMultiNodeReplicatesOverGRPCTransport(t *testing.T) { nodes, peers := newTransportTestNodes(t, 3) startTransportTestServers(nodes, peers) diff --git a/internal/raftengine/etcd/fsm_snapshot_file.go b/internal/raftengine/etcd/fsm_snapshot_file.go index 574724cae..78ca12ff4 100644 --- a/internal/raftengine/etcd/fsm_snapshot_file.go +++ b/internal/raftengine/etcd/fsm_snapshot_file.go @@ -831,7 +831,36 @@ func fsmSnapshotPairRestorable(snapDir, fsmSnapDir, snapName string, term, index if !ok { return false } - return verifyFSMSnapshotFileWithToken(fsmSnapPath(fsmSnapDir, index), tok.CRC32C, true) == nil + // Prewrite cleanup runs on the snapshot receive hot path before gRPC starts + // draining payload chunks. Only do a footer/token check here; full-payload + // CRC remains in the actual restore/open paths. + return fsmSnapshotFooterMatchesToken(fsmSnapPath(fsmSnapDir, index), tok.CRC32C) == nil +} + +func fsmSnapshotFooterMatchesToken(path string, tokenCRC uint32) error { + f, err := os.Open(path) + if err != nil { + return statFSMFileError(err) + } + defer f.Close() + + info, err := f.Stat() + if err != nil { + return errors.WithStack(err) + } + if info.Size() < fsmMinFileSize { + return errors.Wrapf(ErrFSMSnapshotTooSmall, + "file too small: %d bytes (minimum %d)", info.Size(), fsmMinFileSize) + } + footer, err := readFSMFooter(f, info.Size()) + if err != nil { + return err + } + if footer != tokenCRC { + return errors.Wrapf(ErrFSMSnapshotTokenCRC, + "path=%s footer=%08x token=%08x", path, footer, tokenCRC) + } + return nil } func snapshotTokenFromSnapFile(snapDir, snapName string, term, index uint64) (snapshotToken, bool) { diff --git a/internal/raftengine/etcd/fsm_snapshot_file_test.go b/internal/raftengine/etcd/fsm_snapshot_file_test.go index a7d509b44..f5b93f83a 100644 --- a/internal/raftengine/etcd/fsm_snapshot_file_test.go +++ b/internal/raftengine/etcd/fsm_snapshot_file_test.go @@ -446,6 +446,30 @@ func TestPrepareFSMSnapshotWriteKeepsTokenMatchingFallbackPair(t *testing.T) { require.FileExists(t, fsmSnapPath(fsmSnapDir, 100)) } +func TestPrewriteRestorableCheckUsesSnapshotFooterOnly(t *testing.T) { + snapDir := t.TempDir() + fsmSnapDir := t.TempDir() + payload := []byte("payload") + + crc, path := writeFSMFileForTest(t, fsmSnapDir, 100, payload) + createTokenSnapFileWithTerm(t, snapDir, 1, 100, crc) + + f, err := os.OpenFile(path, os.O_WRONLY, 0) + require.NoError(t, err) + _, err = f.WriteAt([]byte("X"), 0) + require.NoError(t, err) + require.NoError(t, f.Close()) + + require.True(t, fsmSnapshotPairRestorable( + snapDir, + fsmSnapDir, + "0000000000000001-0000000000000064.snap", + 1, + 100, + )) + require.ErrorIs(t, verifyFSMSnapshotFile(path, crc), ErrFSMSnapshotFileCRC) +} + func TestPrepareFSMSnapshotWriteKeepsWALValidAndRestorableFallbacksWhenWALValidCandidateIsBroken(t *testing.T) { snapDir := t.TempDir() fsmSnapDir := t.TempDir() diff --git a/kv/fsm.go b/kv/fsm.go index 88c7b95cc..1e1b3fb17 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -886,9 +886,6 @@ func verifyWriteFenceFromSnapshot(mutations []*pb.Mutation, snap RouteSnapshot, } continue } - if len(mut.Key) == 0 { - continue - } rKey := routeKey(mut.Key) if snap.WriteFencedForKey(rKey) { return errors.Wrapf(ErrRouteWriteFenced, diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index f1c027477..f668821c8 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -21,6 +21,17 @@ func newWriteFencedFSM(t *testing.T) *kvFSM { return newComposed1FSM(t, engine, 1) } +func newFirstRouteWriteFencedFSM(t *testing.T) *kvFSM { + t.Helper() + + engine := distribution.NewEngine() + applyComposed1Snapshot(t, engine, 1, []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: []byte("m"), GroupID: 1, State: distribution.RouteStateWriteFenced}, + {RouteID: 2, Start: []byte("m"), End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }) + return newComposed1FSM(t, engine, 1) +} + func s3BucketAuxiliaryFenceRoutes(bucket string, rawGroupID, fencedGroupID uint64) []distribution.RouteDescriptor { start := s3keys.RoutePrefixForBucketAnyGeneration(bucket) end := prefixScanEnd(start) @@ -49,6 +60,16 @@ func TestFSMRejectsCurrentWriteFencedRawPointWrite(t *testing.T) { require.ErrorIs(t, err, ErrRouteWriteFenced) } +func TestFSMRejectsCurrentWriteFencedEmptyRawPointWrite(t *testing.T) { + t.Parallel() + + fsm := newFirstRouteWriteFencedFSM(t) + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte(""), Value: []byte("v")}}, + }, 10) + require.ErrorIs(t, err, ErrRouteWriteFenced) +} + func TestFSMRejectsObservedWriteFencedRawPointWrite(t *testing.T) { t.Parallel() @@ -80,6 +101,31 @@ func TestFSMRejectsCurrentWriteFenceAfterObservedActiveRawPointWrite(t *testing. require.ErrorIs(t, err, ErrRouteWriteFenced) } +func TestFSMRejectsCurrentWriteFencedUnpinnedPrepare(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + applyComposed1Snapshot(t, engine, 1, []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }) + fsm := newComposed1FSM(t, engine, 1) + applyComposed1Snapshot(t, engine, 2, []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: []byte("m"), GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: []byte("m"), End: nil, GroupID: 1, State: distribution.RouteStateWriteFenced}, + }) + + err := fsm.handleTxnRequest(context.Background(), &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("z"), LockTTLms: defaultTxnLockTTLms})}, + {Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("v")}, + }, + }, 10) + require.ErrorIs(t, err, ErrRouteWriteFenced) +} + func TestFSMRejectsCurrentWriteFencedS3BucketAuxiliaryPointWrite(t *testing.T) { t.Parallel() From db9843abde7515833eb43025f93af213d4567aa6 Mon Sep 17 00:00:00 2001 From: bootjp Date: Wed, 15 Jul 2026 09:05:04 +0900 Subject: [PATCH 22/26] Stabilize startup reads and snapshot recovery --- adapter/distribution_server.go | 27 ++++++ adapter/distribution_server_test.go | 29 +++++++ adapter/grpc.go | 30 +++++++ adapter/grpc_test.go | 30 +++++++ internal/raftengine/etcd/engine.go | 57 ++++++++----- internal/raftengine/etcd/engine_test.go | 56 +++++++++++- internal/raftengine/etcd/fsm_snapshot_file.go | 78 +++++++++++++++-- .../raftengine/etcd/fsm_snapshot_file_test.go | 28 +++++- internal/raftengine/etcd/snapshot_spool.go | 14 +-- .../raftengine/etcd/snapshot_spool_test.go | 4 +- kv/fsm.go | 26 ++++++ kv/shard_key_test.go | 28 +++++- kv/sharded_coordinator.go | 31 +++++++ kv/sharded_coordinator_del_prefix_test.go | 85 ++++++++++++++++++- kv/sharded_coordinator_txn_test.go | 4 + main.go | 9 +- 16 files changed, 494 insertions(+), 42 deletions(-) diff --git a/adapter/distribution_server.go b/adapter/distribution_server.go index 3d3276f5c..f8b973f3d 100644 --- a/adapter/distribution_server.go +++ b/adapter/distribution_server.go @@ -23,6 +23,7 @@ type DistributionServer struct { catalog *distribution.CatalogStore coordinator kv.Coordinator readTracker *kv.ActiveTimestampTracker + readBlocked func() bool reloadRetry struct { attempts int interval time.Duration @@ -47,6 +48,12 @@ func WithDistributionActiveTimestampTracker(tracker *kv.ActiveTimestampTracker) } } +func WithDistributionReadGate(blocked func() bool) DistributionServerOption { + return func(s *DistributionServer) { + s.readBlocked = blocked + } +} + // 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 { @@ -96,6 +103,20 @@ func NewDistributionServer(e *distribution.Engine, catalog *distribution.Catalog return s } +func (s *DistributionServer) SetReadGate(blocked func() bool) { + if s != nil { + s.readBlocked = blocked + } +} + +func (s *DistributionServer) requireReadReady() error { + if s != nil && s.readBlocked != nil && s.readBlocked() { + //nolint:wrapcheck // Preserve the gRPC status code for startup readers. + return status.Error(codes.Unavailable, "distribution startup has not completed") + } + return nil +} + // UpdateRoute allows updating route information. func (s *DistributionServer) UpdateRoute(start, end []byte, group uint64) { s.engine.UpdateRoute(start, end, group) @@ -103,6 +124,9 @@ func (s *DistributionServer) UpdateRoute(start, end []byte, group uint64) { // GetRoute returns route for a key. func (s *DistributionServer) GetRoute(ctx context.Context, req *pb.GetRouteRequest) (*pb.GetRouteResponse, error) { + if err := s.requireReadReady(); err != nil { + return nil, err + } r, ok := s.engine.GetRoute(req.Key) if !ok { return &pb.GetRouteResponse{}, nil @@ -122,6 +146,9 @@ func (s *DistributionServer) GetTimestamp(ctx context.Context, req *pb.GetTimest // ListRoutes returns all durable routes from catalog storage. func (s *DistributionServer) ListRoutes(ctx context.Context, req *pb.ListRoutesRequest) (*pb.ListRoutesResponse, error) { + if err := s.requireReadReady(); err != nil { + return nil, err + } snapshot, err := s.loadCatalogSnapshot(ctx) if err != nil { return nil, err diff --git a/adapter/distribution_server_test.go b/adapter/distribution_server_test.go index f7cc59f17..876ef40a5 100644 --- a/adapter/distribution_server_test.go +++ b/adapter/distribution_server_test.go @@ -38,6 +38,35 @@ func TestDistributionServerGetRoute_HitAndMiss(t *testing.T) { require.Nil(t, miss.End) } +func TestDistributionServerRouteReadsHonorStartupGate(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + engine.UpdateRoute([]byte("a"), nil, 1) + catalog := distribution.NewCatalogStore(store.NewMVCCStore()) + _, err := catalog.Save(context.Background(), 0, []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte("a"), End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }) + require.NoError(t, err) + + blocked := true + s := NewDistributionServer(engine, catalog, WithDistributionReadGate(func() bool { return blocked })) + + _, err = s.GetRoute(context.Background(), &pb.GetRouteRequest{Key: []byte("a")}) + require.Error(t, err) + require.Equal(t, codes.Unavailable, status.Code(err)) + + _, err = s.ListRoutes(context.Background(), &pb.ListRoutesRequest{}) + require.Error(t, err) + require.Equal(t, codes.Unavailable, status.Code(err)) + + blocked = false + _, err = s.GetRoute(context.Background(), &pb.GetRouteRequest{Key: []byte("a")}) + require.NoError(t, err) + _, err = s.ListRoutes(context.Background(), &pb.ListRoutesRequest{}) + require.NoError(t, err) +} + func TestDistributionServerGetTimestamp_IsMonotonic(t *testing.T) { t.Parallel() diff --git a/adapter/grpc.go b/adapter/grpc.go index 773dacb93..21f3e1761 100644 --- a/adapter/grpc.go +++ b/adapter/grpc.go @@ -25,6 +25,7 @@ type GRPCServer struct { grpcTranscoder *grpcTranscoder coordinator kv.Coordinator store store.MVCCStore + readBlocked func() bool closeStore bool closeOnce sync.Once @@ -66,6 +67,12 @@ func WithCloseStore() GRPCServerOption { } } +func WithGRPCReadGate(blocked func() bool) GRPCServerOption { + return func(s *GRPCServer) { + s.readBlocked = blocked + } +} + func NewGRPCServer(store store.MVCCStore, coordinate kv.Coordinator, opts ...GRPCServerOption) *GRPCServer { s := &GRPCServer{ log: slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{ @@ -84,6 +91,14 @@ func NewGRPCServer(store store.MVCCStore, coordinate kv.Coordinator, opts ...GRP return s } +func (r *GRPCServer) requireReadReady() error { + if r != nil && r.readBlocked != nil && r.readBlocked() { + //nolint:wrapcheck // Preserve the gRPC status code for startup readers. + return status.Error(codes.Unavailable, "startup rotation has not completed") + } + return nil +} + func (r *GRPCServer) Close() error { if r == nil { return nil @@ -107,6 +122,9 @@ func (r *GRPCServer) clock() *kv.HLC { } func (r *GRPCServer) RawGet(ctx context.Context, req *pb.RawGetRequest) (*pb.RawGetResponse, error) { + if err := r.requireReadReady(); err != nil { + return nil, err + } readTS := req.GetTs() if readTS == 0 { readTS = globalSnapshotTS(ctx, r.clock(), r.store) @@ -143,6 +161,9 @@ func (r *GRPCServer) RawGet(ctx context.Context, req *pb.RawGetRequest) (*pb.Raw } func (r *GRPCServer) RawLatestCommitTS(ctx context.Context, req *pb.RawLatestCommitTSRequest) (*pb.RawLatestCommitTSResponse, error) { + if err := r.requireReadReady(); err != nil { + return nil, err + } key := req.GetKey() if len(key) == 0 { // No key: return the store's global last-committed watermark. @@ -173,6 +194,9 @@ func (r *GRPCServer) RawLatestCommitTS(ctx context.Context, req *pb.RawLatestCom } func (r *GRPCServer) RawScanAt(ctx context.Context, req *pb.RawScanAtRequest) (*pb.RawScanAtResponse, error) { + if err := r.requireReadReady(); err != nil { + return nil, err + } limit64 := req.GetLimit() limit, err := rawScanLimit(limit64) if err != nil { @@ -366,6 +390,9 @@ func (r *GRPCServer) Put(ctx context.Context, req *pb.PutRequest) (*pb.PutRespon } func (r *GRPCServer) Get(ctx context.Context, req *pb.GetRequest) (*pb.GetResponse, error) { + if err := r.requireReadReady(); err != nil { + return nil, err + } h := murmur3.New64() if _, err := h.Write(req.Key); err != nil { return nil, errors.WithStack(err) @@ -413,6 +440,9 @@ func (r *GRPCServer) Delete(ctx context.Context, req *pb.DeleteRequest) (*pb.Del } func (r *GRPCServer) Scan(ctx context.Context, req *pb.ScanRequest) (*pb.ScanResponse, error) { + if err := r.requireReadReady(); err != nil { + return nil, err + } limit, err := internal.Uint64ToInt(req.Limit) if err != nil { return &pb.ScanResponse{ diff --git a/adapter/grpc_test.go b/adapter/grpc_test.go index c70818adb..db5b8ce68 100644 --- a/adapter/grpc_test.go +++ b/adapter/grpc_test.go @@ -290,6 +290,36 @@ func TestGRPCServer_RawReadFenceHelpersStampCurrentRouteVersion(t *testing.T) { require.Equal(t, uint64(55), st.scanReadRouteVersion) } +func TestGRPCServer_ReadsHonorStartupGate(t *testing.T) { + t.Parallel() + + blocked := true + s := NewGRPCServer(store.NewMVCCStore(), nil, WithGRPCReadGate(func() bool { return blocked })) + ctx := context.Background() + + _, err := s.RawGet(ctx, &pb.RawGetRequest{Key: []byte("k"), Ts: 10}) + require.Error(t, err) + require.Equal(t, codes.Unavailable, status.Code(err)) + _, err = s.RawLatestCommitTS(ctx, &pb.RawLatestCommitTSRequest{}) + require.Error(t, err) + require.Equal(t, codes.Unavailable, status.Code(err)) + _, err = s.RawScanAt(ctx, &pb.RawScanAtRequest{Limit: 10, Ts: 10}) + require.Error(t, err) + require.Equal(t, codes.Unavailable, status.Code(err)) + _, err = s.Get(ctx, &pb.GetRequest{Key: []byte("k")}) + require.Error(t, err) + require.Equal(t, codes.Unavailable, status.Code(err)) + _, err = s.Scan(ctx, &pb.ScanRequest{Limit: 10}) + require.Error(t, err) + require.Equal(t, codes.Unavailable, status.Code(err)) + + blocked = false + _, err = s.RawGet(ctx, &pb.RawGetRequest{Key: []byte("k"), Ts: 10}) + require.NoError(t, err) + _, err = s.Scan(ctx, &pb.ScanRequest{Limit: 10}) + require.NoError(t, err) +} + func TestGRPCServer_RawReadFenceHelpersKeepCallerRouteVersion(t *testing.T) { t.Parallel() diff --git a/internal/raftengine/etcd/engine.go b/internal/raftengine/etcd/engine.go index d36e7cf8a..65ef0903f 100644 --- a/internal/raftengine/etcd/engine.go +++ b/internal/raftengine/etcd/engine.go @@ -85,6 +85,12 @@ const ( // upside is that a ~5 s transient pause (election-timeout scale) // no longer drops heartbeats and forces the peers' lease to expire. defaultHeartbeatBufPerPeer = 512 + // defaultReadIndexRespBufPerPeer sizes the dedicated follower-to-leader + // ReadIndex heartbeat response lane. etcd/raft encodes ReadIndex + // completions as MsgHeartbeatResp messages with Context set; unlike plain + // heartbeat acks, each context is tied to a caller waiting in handleRead + // and must not be coalesced or dropped behind superseded empty acks. + defaultReadIndexRespBufPerPeer = 512 // defaultHeartbeatRespBufPerPeer sizes the dedicated follower-to-leader // heartbeat response lane. When a follower is receiving a large snapshot, // the leader may continue to send heartbeats while the follower's outbound @@ -577,21 +583,23 @@ type dispatchRequest struct { // peerQueues holds separate dispatch channels per peer so that heartbeats // are never blocked behind large log-entry RPCs. // -// Legacy 3-lane layout (default): heartbeat + heartbeatResp + normal. +// Legacy 4-lane layout (default): heartbeat + heartbeatResp + +// readIndexResp + normal. // -// 5-lane layout (opt-in via ELASTICKV_RAFT_DISPATCHER_LANES=1): heartbeat + -// heartbeatResp + replication (MsgApp/MsgAppResp) + snapshot (MsgSnap) + -// other. Each lane gets its own goroutine so a bulky MsgSnap transfer cannot -// stall MsgApp replication and vice versa. Per-peer ordering within a given -// message type is preserved because a single peer's MsgApp stream all share -// one lane and one worker. +// 6-lane layout (opt-in via ELASTICKV_RAFT_DISPATCHER_LANES=1): heartbeat + +// heartbeatResp + readIndexResp + replication (MsgApp/MsgAppResp) + snapshot +// (MsgSnap) + other. Each lane gets its own goroutine so a bulky MsgSnap +// transfer cannot stall MsgApp replication and vice versa. Per-peer ordering +// within a given message type is preserved because a single peer's MsgApp +// stream all share one lane and one worker. type peerQueues struct { normal chan dispatchRequest heartbeat chan dispatchRequest heartbeatResp chan dispatchRequest - replication chan dispatchRequest // 5-lane mode only; nil otherwise - snapshot chan dispatchRequest // 5-lane mode only; nil otherwise - other chan dispatchRequest // 5-lane mode only; nil otherwise + readIndexResp chan dispatchRequest + replication chan dispatchRequest // 6-lane mode only; nil otherwise + snapshot chan dispatchRequest // 6-lane mode only; nil otherwise + other chan dispatchRequest // 6-lane mode only; nil otherwise ctx context.Context cancel context.CancelFunc } @@ -2416,7 +2424,7 @@ func (e *Engine) enqueueDispatchMessage(msg raftpb.Message) error { e.recordDroppedDispatch(msg) return nil } - ch := e.selectDispatchLane(pd, msg.GetType()) + ch := e.selectDispatchLaneForMessage(pd, msg) // Avoid the expensive deep-clone in prepareDispatchRequest when the channel // is already full. The len/cap check is safe here because this function is // only ever called from the single engine event-loop goroutine. @@ -2519,6 +2527,14 @@ func isBlockingInboundStepMsg(t raftpb.MessageType) bool { // opt-in multi-lane layout it additionally partitions the non-heartbeat traffic // so that MsgApp/MsgAppResp and MsgSnap do not share a goroutine and cannot // block each other. +func (e *Engine) selectDispatchLaneForMessage(pd *peerQueues, msg raftpb.Message) chan dispatchRequest { + msgType := msg.GetType() + if msgType == raftpb.MsgHeartbeatResp && len(msg.GetContext()) > 0 && pd.readIndexResp != nil { + return pd.readIndexResp + } + return e.selectDispatchLane(pd, msgType) +} + func (e *Engine) selectDispatchLane(pd *peerQueues, msgType raftpb.MessageType) chan dispatchRequest { if msgType == raftpb.MsgHeartbeatResp && pd.heartbeatResp != nil { return pd.heartbeatResp @@ -4457,24 +4473,25 @@ func (e *Engine) startPeerDispatcher(nodeID uint64) { pd := &peerQueues{ heartbeat: make(chan dispatchRequest, defaultHeartbeatBufPerPeer), heartbeatResp: make(chan dispatchRequest, defaultHeartbeatRespBufPerPeer), + readIndexResp: make(chan dispatchRequest, defaultReadIndexRespBufPerPeer), ctx: ctx, cancel: cancel, } var workers []chan dispatchRequest if e.dispatcherLanesEnabled { - // 5-lane layout: split MsgHeartbeatResp, MsgApp/MsgAppResp - // (replication), MsgSnap (snapshot), and misc (other) onto independent - // goroutines so a bulky snapshot transfer cannot stall replication or - // follower heartbeat responses. Each channel still serves a single - // peer, so within-type ordering (the raft invariant we care about for - // MsgApp) is preserved. + // 6-lane layout: split MsgHeartbeatResp, ReadIndex heartbeat + // responses, MsgApp/MsgAppResp (replication), MsgSnap (snapshot), and + // misc (other) onto independent goroutines so a bulky snapshot transfer + // cannot stall replication or follower heartbeat responses. Each channel + // still serves a single peer, so within-type ordering (the raft invariant + // we care about for MsgApp) is preserved. pd.replication = make(chan dispatchRequest, size) pd.snapshot = make(chan dispatchRequest, defaultSnapshotLaneBufPerPeer) pd.other = make(chan dispatchRequest, defaultOtherLaneBufPerPeer) - workers = []chan dispatchRequest{pd.heartbeat, pd.heartbeatResp, pd.replication, pd.snapshot, pd.other} + workers = []chan dispatchRequest{pd.heartbeat, pd.heartbeatResp, pd.readIndexResp, pd.replication, pd.snapshot, pd.other} } else { pd.normal = make(chan dispatchRequest, size) - workers = []chan dispatchRequest{pd.normal, pd.heartbeat, pd.heartbeatResp} + workers = []chan dispatchRequest{pd.normal, pd.heartbeat, pd.heartbeatResp, pd.readIndexResp} } e.peerDispatchers[nodeID] = pd e.dispatchWG.Add(len(workers)) @@ -4549,7 +4566,7 @@ func dispatcherLanesEnabledFromEnv() bool { // drain loops in runDispatchWorker exit. It is safe to call with any dispatch // layout because unused lanes are nil. func closePeerLanes(pd *peerQueues) { - for _, ch := range []chan dispatchRequest{pd.heartbeat, pd.heartbeatResp, pd.normal, pd.replication, pd.snapshot, pd.other} { + for _, ch := range []chan dispatchRequest{pd.heartbeat, pd.heartbeatResp, pd.readIndexResp, pd.normal, pd.replication, pd.snapshot, pd.other} { if ch != nil { close(ch) } diff --git a/internal/raftengine/etcd/engine_test.go b/internal/raftengine/etcd/engine_test.go index 148f1d149..ce81a539d 100644 --- a/internal/raftengine/etcd/engine_test.go +++ b/internal/raftengine/etcd/engine_test.go @@ -940,6 +940,7 @@ func TestUpsertPeerStartsDispatcherAndAcceptsMessages(t *testing.T) { require.True(t, ok, "dispatcher must be created on upsert") require.Equal(t, defaultHeartbeatBufPerPeer, cap(pd.heartbeat)) require.Equal(t, defaultHeartbeatRespBufPerPeer, cap(pd.heartbeatResp)) + require.Equal(t, defaultReadIndexRespBufPerPeer, cap(pd.readIndexResp)) require.Equal(t, 4, cap(pd.normal)) require.NoError(t, engine.enqueueDispatchMessage(raftpb.Message{Type: messageTypePtr(raftpb.MsgHeartbeat), To: uint64Ptr(2)})) @@ -962,6 +963,7 @@ func TestRemovePeerClosesDispatcherAndDropsSubsequentMessages(t *testing.T) { normal: make(chan dispatchRequest, 4), heartbeat: make(chan dispatchRequest, 4), heartbeatResp: make(chan dispatchRequest, 4), + readIndexResp: make(chan dispatchRequest, 4), ctx: ctx, cancel: cancel, } @@ -971,10 +973,11 @@ func TestRemovePeerClosesDispatcherAndDropsSubsequentMessages(t *testing.T) { peerDispatchers: map[uint64]*peerQueues{2: pd}, dispatchStopCh: stopCh, } - engine.dispatchWG.Add(3) + engine.dispatchWG.Add(4) go engine.runDispatchWorker(ctx, pd.normal) go engine.runDispatchWorker(ctx, pd.heartbeat) go engine.runDispatchWorker(ctx, pd.heartbeatResp) + go engine.runDispatchWorker(ctx, pd.readIndexResp) engine.removePeer(2) @@ -1583,6 +1586,44 @@ func TestEnqueueDispatchMessageCoalescesPlainHeartbeatBehindReadIndex(t *testing require.Equal(t, uint64(11), req.msg.GetIndex()) } +func TestEnqueueDispatchMessagePreservesIncomingReadIndexHeartbeatWhenRespLaneFull(t *testing.T) { + t.Parallel() + pd := &peerQueues{ + heartbeat: make(chan dispatchRequest, 1), + heartbeatResp: make(chan dispatchRequest, 2), + readIndexResp: make(chan dispatchRequest, 1), + } + engine := &Engine{ + nodeID: 1, + peerDispatchers: map[uint64]*peerQueues{ + 2: pd, + }, + } + pd.heartbeatResp <- prepareDispatchRequest(raftpb.Message{ + Type: messageTypePtr(raftpb.MsgHeartbeatResp), + To: uint64Ptr(2), + Context: []byte("read-index-a"), + }) + pd.heartbeatResp <- prepareDispatchRequest(raftpb.Message{ + Type: messageTypePtr(raftpb.MsgHeartbeatResp), + To: uint64Ptr(2), + Context: []byte("read-index-b"), + }) + + require.NoError(t, engine.enqueueDispatchMessage(raftpb.Message{ + Type: messageTypePtr(raftpb.MsgHeartbeatResp), + To: uint64Ptr(2), + Context: []byte("read-index-c"), + })) + + require.Zero(t, engine.DispatchDropCount()) + require.Len(t, pd.heartbeatResp, 2) + require.Len(t, pd.readIndexResp, 1) + req := <-pd.readIndexResp + require.Equal(t, raftpb.MsgHeartbeatResp, req.msg.GetType()) + require.Equal(t, []byte("read-index-c"), req.msg.Context) +} + func TestMaxAppliedIndexStartsFromSnapshotIndex(t *testing.T) { storage := etcdraft.NewMemoryStorage() snap := raftTestSnapshot(5, 2, []uint64{1}, nil) @@ -2335,6 +2376,7 @@ func TestSelectDispatchLane_LegacyThreeLane(t *testing.T) { normal: make(chan dispatchRequest, 1), heartbeat: make(chan dispatchRequest, 1), heartbeatResp: make(chan dispatchRequest, 1), + readIndexResp: make(chan dispatchRequest, 1), } cases := map[raftpb.MessageType]chan dispatchRequest{ @@ -2355,6 +2397,11 @@ func TestSelectDispatchLane_LegacyThreeLane(t *testing.T) { got := engine.selectDispatchLane(pd, mt) require.Equalf(t, want, got, "legacy mode routing for %s", mt) } + got := engine.selectDispatchLaneForMessage(pd, raftpb.Message{ + Type: messageTypePtr(raftpb.MsgHeartbeatResp), + Context: []byte("read-index"), + }) + require.Equal(t, pd.readIndexResp, got) } // TestSelectDispatchLane_FiveLane verifies that, when ELASTICKV_RAFT_DISPATCHER_LANES @@ -2367,6 +2414,7 @@ func TestSelectDispatchLane_FiveLane(t *testing.T) { pd := &peerQueues{ heartbeat: make(chan dispatchRequest, 1), heartbeatResp: make(chan dispatchRequest, 1), + readIndexResp: make(chan dispatchRequest, 1), replication: make(chan dispatchRequest, 1), snapshot: make(chan dispatchRequest, 1), other: make(chan dispatchRequest, 1), @@ -2390,6 +2438,11 @@ func TestSelectDispatchLane_FiveLane(t *testing.T) { got := engine.selectDispatchLane(pd, mt) require.Equalf(t, want, got, "multi-lane mode routing for %s", mt) } + got := engine.selectDispatchLaneForMessage(pd, raftpb.Message{ + Type: messageTypePtr(raftpb.MsgHeartbeatResp), + Context: []byte("read-index"), + }) + require.Equal(t, pd.readIndexResp, got) } // TestSelectDispatchLane_MsgPropReachesDefaultFallback verifies that MsgProp, @@ -2403,6 +2456,7 @@ func TestSelectDispatchLane_MsgPropReachesDefaultFallback(t *testing.T) { pd := &peerQueues{ heartbeat: make(chan dispatchRequest, 1), heartbeatResp: make(chan dispatchRequest, 1), + readIndexResp: make(chan dispatchRequest, 1), replication: make(chan dispatchRequest, 1), snapshot: make(chan dispatchRequest, 1), other: make(chan dispatchRequest, 1), diff --git a/internal/raftengine/etcd/fsm_snapshot_file.go b/internal/raftengine/etcd/fsm_snapshot_file.go index 78ca12ff4..8e31a05d3 100644 --- a/internal/raftengine/etcd/fsm_snapshot_file.go +++ b/internal/raftengine/etcd/fsm_snapshot_file.go @@ -619,6 +619,7 @@ type snapFileCandidate struct { index uint64 restorable bool walValid bool + tokenCRC uint32 } type prewriteSnapshotRetention struct { @@ -655,7 +656,7 @@ func purgeOlderSnapshotPairsBeforeWrite( return candidates[i].index < candidates[j].index }) - retention := keepRestorablePrewriteSnapshots(candidates) + retention := keepVerifiedPrewriteSnapshots(fsmSnapDir, candidates) var combined error combined = errors.CombineErrors(combined, purgeUnretainedPrewriteSnapshots(snapDir, fsmSnapDir, candidates, retention)) combined = errors.CombineErrors(combined, removePrewriteFSMOrphansBeforeIndex( @@ -746,6 +747,69 @@ func keepRestorablePrewriteSnapshots(candidates []snapFileCandidate) prewriteSna return retention } +func keepVerifiedPrewriteSnapshots(fsmSnapDir string, candidates []snapFileCandidate) prewriteSnapshotRetention { + retention := keepRestorablePrewriteSnapshots(candidates) + if fsmSnapDir == "" || len(candidates) <= prewriteSnapKeep { + return retention + } + for { + if !invalidateUnverifiedRetainedPrewriteSnapshots(fsmSnapDir, candidates, retention) { + return retention + } + retention = keepRestorablePrewriteSnapshots(candidates) + if !retentionHasRestorable(candidates, retention) { + return keepAllPrewriteSnapshots(candidates) + } + } +} + +func invalidateUnverifiedRetainedPrewriteSnapshots( + fsmSnapDir string, + candidates []snapFileCandidate, + retention prewriteSnapshotRetention, +) bool { + invalidated := false + for i := range candidates { + candidate := &candidates[i] + if !candidate.restorable || !retention.keep[candidate.name] || !retainedPrewriteSnapshotPrunesOlder(candidates, retention, *candidate) { + continue + } + if verifyFSMSnapshotFileWithToken(fsmSnapPath(fsmSnapDir, candidate.index), candidate.tokenCRC, true) == nil { + continue + } + candidate.restorable = false + invalidated = true + } + return invalidated +} + +func retainedPrewriteSnapshotPrunesOlder(candidates []snapFileCandidate, retention prewriteSnapshotRetention, retained snapFileCandidate) bool { + for _, candidate := range candidates { + if candidate.index >= retained.index || retention.keep[candidate.name] { + continue + } + return true + } + return false +} + +func retentionHasRestorable(candidates []snapFileCandidate, retention prewriteSnapshotRetention) bool { + for _, candidate := range candidates { + if candidate.restorable && retention.keep[candidate.name] { + return true + } + } + return false +} + +func keepAllPrewriteSnapshots(candidates []snapFileCandidate) prewriteSnapshotRetention { + retention := prewriteSnapshotRetention{keep: make(map[string]bool, len(candidates))} + for _, candidate := range candidates { + retention.keep[candidate.name] = true + } + return retention +} + func keepNewestMatchingPrewriteSnapshots( candidates []snapFileCandidate, retention *prewriteSnapshotRetention, @@ -813,28 +877,30 @@ func collectPrewriteSnapCandidates( if index == 0 || index >= nextIndex { continue } + restorable, tokenCRC := fsmSnapshotPairRestorable(snapDir, fsmSnapDir, e.Name(), term, index) candidates = append(candidates, snapFileCandidate{ name: e.Name(), index: index, - restorable: fsmSnapshotPairRestorable(snapDir, fsmSnapDir, e.Name(), term, index), + restorable: restorable, walValid: walValidIndexes == nil || walValidIndexes[walSnapshotKey{term: term, index: index}], + tokenCRC: tokenCRC, }) } return candidates } -func fsmSnapshotPairRestorable(snapDir, fsmSnapDir, snapName string, term, index uint64) bool { +func fsmSnapshotPairRestorable(snapDir, fsmSnapDir, snapName string, term, index uint64) (bool, uint32) { if fsmSnapDir == "" { - return false + return false, 0 } tok, ok := snapshotTokenFromSnapFile(snapDir, snapName, term, index) if !ok { - return false + return false, 0 } // Prewrite cleanup runs on the snapshot receive hot path before gRPC starts // draining payload chunks. Only do a footer/token check here; full-payload // CRC remains in the actual restore/open paths. - return fsmSnapshotFooterMatchesToken(fsmSnapPath(fsmSnapDir, index), tok.CRC32C) == nil + return fsmSnapshotFooterMatchesToken(fsmSnapPath(fsmSnapDir, index), tok.CRC32C) == nil, tok.CRC32C } func fsmSnapshotFooterMatchesToken(path string, tokenCRC uint32) error { diff --git a/internal/raftengine/etcd/fsm_snapshot_file_test.go b/internal/raftengine/etcd/fsm_snapshot_file_test.go index f5b93f83a..60cc33344 100644 --- a/internal/raftengine/etcd/fsm_snapshot_file_test.go +++ b/internal/raftengine/etcd/fsm_snapshot_file_test.go @@ -460,16 +460,40 @@ func TestPrewriteRestorableCheckUsesSnapshotFooterOnly(t *testing.T) { require.NoError(t, err) require.NoError(t, f.Close()) - require.True(t, fsmSnapshotPairRestorable( + restorable, _ := fsmSnapshotPairRestorable( snapDir, fsmSnapDir, "0000000000000001-0000000000000064.snap", 1, 100, - )) + ) + require.True(t, restorable) require.ErrorIs(t, verifyFSMSnapshotFile(path, crc), ErrFSMSnapshotFileCRC) } +func TestPrepareFSMSnapshotWriteVerifiesRetainedFooterOnlyFallbackBeforePruning(t *testing.T) { + snapDir := t.TempDir() + fsmSnapDir := t.TempDir() + + crc100, _ := writeFSMFileForTest(t, fsmSnapDir, 100, []byte("valid fallback")) + createTokenSnapFileWithTerm(t, snapDir, 1, 100, crc100) + crc200, path200 := writeFSMFileForTest(t, fsmSnapDir, 200, []byte("corrupt newest fallback")) + createTokenSnapFileWithTerm(t, snapDir, 1, 200, crc200) + + f, err := os.OpenFile(path200, os.O_WRONLY, 0) + require.NoError(t, err) + _, err = f.WriteAt([]byte("X"), 0) + require.NoError(t, err) + require.NoError(t, f.Close()) + + require.NoError(t, prepareFSMSnapshotWrite(snapDir, fsmSnapDir, 300)) + + require.FileExists(t, filepath.Join(snapDir, "0000000000000001-0000000000000064.snap")) + require.FileExists(t, fsmSnapPath(fsmSnapDir, 100)) + require.NoFileExists(t, filepath.Join(snapDir, "0000000000000001-00000000000000c8.snap")) + require.NoFileExists(t, fsmSnapPath(fsmSnapDir, 200)) +} + func TestPrepareFSMSnapshotWriteKeepsWALValidAndRestorableFallbacksWhenWALValidCandidateIsBroken(t *testing.T) { snapDir := t.TempDir() fsmSnapDir := t.TempDir() diff --git a/internal/raftengine/etcd/snapshot_spool.go b/internal/raftengine/etcd/snapshot_spool.go index cb1b3bc8e..678d7857a 100644 --- a/internal/raftengine/etcd/snapshot_spool.go +++ b/internal/raftengine/etcd/snapshot_spool.go @@ -33,6 +33,8 @@ const maxSnapshotPayloadBytesEnvVar = "ELASTICKV_RAFT_MAX_SNAPSHOT_PAYLOAD_BYTES const snapshotSpoolMinFreeBytesEnvVar = "ELASTICKV_RAFT_SNAPSHOT_SPOOL_MIN_FREE_BYTES" +const defaultReceiveSnapshotSpoolMinFreeBytes int64 = 1 << 30 // 1 GiB + // resolveMaxSnapshotPayloadBytes evaluates the env override once per spool // creation. Snapshots are infrequent enough that one Getenv + ParseInt per // spool is invisible in profiles, and resolving at construction means tests @@ -94,11 +96,13 @@ func newSnapshotSpool(dir string) (*snapshotSpool, error) { func newReceiveSnapshotSpool(dir string) (*snapshotSpool, error) { maxSize := resolveMaxSnapshotPayloadBytes() - // Keep one full max-size snapshot worth of headroom after receive-side - // spooling. Applying a token snapshot restores the FSM from the completed - // .fsm into a new local store directory, so a node can transiently need old - // snapshot + incoming .fsm + restored store bytes. - return newSnapshotSpoolWithLimits(dir, maxSize, resolveSnapshotSpoolMinFreeBytes(maxSize)) + // Keep a fixed emergency reserve after receive-side spooling. Tying this + // reserve to the configured maximum snapshot size made the default 16 GiB + // cap also require 16 GiB of free space after every chunk; production nodes + // with enough space for a real 13 GiB FSM snapshot were rejecting the stream + // around 4 GiB and retrying forever. Operators that restore into a layout + // needing a larger reserve can still raise it with the env knob. + return newSnapshotSpoolWithLimits(dir, maxSize, resolveSnapshotSpoolMinFreeBytes(defaultReceiveSnapshotSpoolMinFreeBytes)) } func newSnapshotSpoolWithLimits(dir string, maxSize, minFreeBytes int64) (*snapshotSpool, error) { diff --git a/internal/raftengine/etcd/snapshot_spool_test.go b/internal/raftengine/etcd/snapshot_spool_test.go index bed5fd223..b0facd96d 100644 --- a/internal/raftengine/etcd/snapshot_spool_test.go +++ b/internal/raftengine/etcd/snapshot_spool_test.go @@ -97,7 +97,7 @@ func TestSnapshotSpool_OverrideInvalidFallsBack(t *testing.T) { require.Equal(t, defaultMaxSnapshotPayloadBytes, spool.maxSize) } -func TestSnapshotSpoolDefaultReserveTracksPayloadCap(t *testing.T) { +func TestSnapshotSpoolDefaultReserveUsesFixedEmergencyHeadroom(t *testing.T) { const spoolCap = int64(4096) t.Setenv(maxSnapshotPayloadBytesEnvVar, strconv.FormatInt(spoolCap, 10)) @@ -106,7 +106,7 @@ func TestSnapshotSpoolDefaultReserveTracksPayloadCap(t *testing.T) { t.Cleanup(func() { _ = spool.Close() }) require.Equal(t, spoolCap, spool.maxSize) - require.Equal(t, spoolCap, spool.minFreeBytes) + require.Equal(t, defaultReceiveSnapshotSpoolMinFreeBytes, spool.minFreeBytes) } func TestSnapshotSpoolMaterializeDoesNotReserveDiskHeadroom(t *testing.T) { diff --git a/kv/fsm.go b/kv/fsm.go index 1e1b3fb17..b2e74b1db 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -546,6 +546,9 @@ func routePrefixRange(prefix []byte) ([]byte, []byte) { if start, ok := s3keys.BucketGenerationRoutePrefixForCleanupPrefix(prefix); ok { return start, prefixScanEnd(start) } + if start, ok := dynamoExactCleanupRouteKey(prefix); ok { + return start, routePointRangeEnd(start) + } if routeKeyspaceWideRawPrefix(prefix) { return []byte(""), nil } @@ -553,6 +556,29 @@ func routePrefixRange(prefix []byte) ([]byte, []byte) { return start, prefixScanEnd(start) } +func dynamoExactCleanupRouteKey(prefix []byte) ([]byte, bool) { + switch { + case bytes.HasPrefix(prefix, dynamoTableMetaPrefixBytes), + bytes.HasPrefix(prefix, dynamoTableGenerationPrefixBytes), + bytes.HasPrefix(prefix, dynamoItemPrefixBytes), + bytes.HasPrefix(prefix, dynamoGSIPrefixBytes): + default: + return nil, false + } + start := routeKey(prefix) + if len(start) == 0 || bytes.Equal(start, prefix) || !bytes.HasPrefix(start, dynamoRoutePrefixBytes) { + return nil, false + } + return start, true +} + +func routePointRangeEnd(start []byte) []byte { + end := make([]byte, 0, len(start)+1) + end = append(end, start...) + end = append(end, 0) + return end +} + func routeKeyspaceWideRawPrefix(prefix []byte) bool { if !rawPrefixMayContainRouteMappedKeys(prefix) { return false diff --git a/kv/shard_key_test.go b/kv/shard_key_test.go index 817a1d809..6f6f2a653 100644 --- a/kv/shard_key_test.go +++ b/kv/shard_key_test.go @@ -306,7 +306,7 @@ func TestRoutePrefixRangeTreatsBroadMappedPrefixesAsFullKeyspace(t *testing.T) { name: "dynamo table cleanup prefix", prefix: []byte(DynamoItemPrefix + tableSegment + "|7|"), wantStart: dynamoTableRoute, - wantEnd: prefixScanEnd(dynamoTableRoute), + wantEnd: routePointRangeEnd(dynamoTableRoute), }, { name: "broad redis namespace", @@ -343,6 +343,32 @@ func TestRoutePrefixRangeTreatsBroadMappedPrefixesAsFullKeyspace(t *testing.T) { } } +func TestRoutePrefixRangeTreatsDynamoCleanupAsExactRouteKey(t *testing.T) { + t.Parallel() + + fooSegment := base64.RawURLEncoding.EncodeToString([]byte("foo")) + foobarSegment := base64.RawURLEncoding.EncodeToString([]byte("foobar")) + fooRoute := dynamoRouteTableKey([]byte(fooSegment)) + foobarRoute := dynamoRouteTableKey([]byte(foobarSegment)) + + start, end := routePrefixRange([]byte(DynamoItemPrefix + fooSegment + "|7|")) + + require.Equal(t, fooRoute, start) + require.Equal(t, routePointRangeEnd(fooRoute), end) + require.False(t, rangesIntersectForTest(start, end, foobarRoute, prefixScanEnd(foobarRoute)), + "cleanup for table foo must not intersect table foobar's route") +} + +func rangesIntersectForTest(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 TestRouteKeyFilterTreatsNilAndEmptyEndAsInfinity(t *testing.T) { t.Parallel() diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index 147abc067..73420bb86 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -1007,11 +1007,42 @@ func (c *ShardedCoordinator) dispatchRawWithComposed1Retry(ctx context.Context, if !errors.Is(err, ErrComposed1VersionGCd) || attempt == composed1RetryAttempts || c.engine == nil { return resp, err } + if !c.canRetryRawVersionGC(reqs) { + return resp, err + } reqs.ObservedRouteVersion = c.engine.Version() } return nil, errors.WithStack(ErrInvalidRequest) } +func (c *ShardedCoordinator) canRetryRawVersionGC(reqs *OperationGroup[OP]) bool { + if c == nil || c.router == nil || reqs == nil || hasDelPrefixElem(reqs.Elems) { + return false + } + var ( + firstGID uint64 + seen bool + ) + for _, elem := range reqs.Elems { + if elem == nil { + return false + } + gid, ok := c.router.ResolveGroup(elem.Key) + if !ok { + return false + } + if !seen { + firstGID = gid + seen = true + continue + } + if gid != firstGID { + return false + } + } + return seen +} + // hasDelPrefixElem returns true if any element is a DelPrefix operation. func hasDelPrefixElem(elems []*Elem[OP]) bool { for _, e := range elems { diff --git a/kv/sharded_coordinator_del_prefix_test.go b/kv/sharded_coordinator_del_prefix_test.go index d5f8632d3..cd314eef6 100644 --- a/kv/sharded_coordinator_del_prefix_test.go +++ b/kv/sharded_coordinator_del_prefix_test.go @@ -195,7 +195,7 @@ func TestShardedCoordinator_RetriesRawWriteWhenObservedRouteVersionGCd(t *testin require.Equal(t, uint64(9), txn.requests[1].GetObservedRouteVersion()) } -func TestShardedCoordinator_RetriesDelPrefixWhenObservedRouteVersionGCd(t *testing.T) { +func TestShardedCoordinator_DoesNotRetryDelPrefixWhenObservedRouteVersionGCd(t *testing.T) { t.Parallel() engine := distribution.NewEngine() @@ -216,10 +216,87 @@ func TestShardedCoordinator_RetriesDelPrefixWhenObservedRouteVersionGCd(t *testi ObservedRouteVersion: 2, Elems: []*Elem[OP]{{Op: DelPrefix, Key: []byte("user:")}}, }) - require.NoError(t, err) - require.Len(t, txn.requests, 2) + require.ErrorIs(t, err, ErrComposed1VersionGCd) + require.Len(t, txn.requests, 1) require.Equal(t, uint64(2), txn.requests[0].GetObservedRouteVersion()) - require.Equal(t, uint64(7), txn.requests[1].GetObservedRouteVersion()) +} + +func TestShardedCoordinator_DoesNotRetryMultiShardRawWriteWhenObservedRouteVersionGCd(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 11, + 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{errs: []error{ErrComposed1VersionGCd}} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + ObservedRouteVersion: 3, + Elems: []*Elem[OP]{ + {Op: Put, Key: []byte("a"), Value: []byte("v1")}, + {Op: Put, Key: []byte("z"), Value: []byte("v2")}, + }, + }) + + require.ErrorIs(t, err, ErrComposed1VersionGCd) + require.Len(t, g1Txn.requests, 1) + require.Len(t, g2Txn.requests, 1) + require.Equal(t, uint64(3), g1Txn.requests[0].GetObservedRouteVersion()) + require.Equal(t, uint64(3), g2Txn.requests[0].GetObservedRouteVersion()) +} + +func TestShardedCoordinator_DoesNotRetryRawWriteWhenRetryRouteBecomesMultiShard(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 11, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }, + })) + g1Txn := &recordingTransactional{ + errs: []error{ErrComposed1VersionGCd}, + onCommit: func(call int, _ *pb.Request) { + if call != 0 { + return + } + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 12, + 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}, + }, + })) + }, + } + g2Txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + ObservedRouteVersion: 3, + Elems: []*Elem[OP]{ + {Op: Put, Key: []byte("a"), Value: []byte("v1")}, + {Op: Put, Key: []byte("z"), Value: []byte("v2")}, + }, + }) + + require.ErrorIs(t, err, ErrComposed1VersionGCd) + require.Len(t, g1Txn.requests, 1) + require.Len(t, g2Txn.requests, 0) + require.Equal(t, uint64(3), g1Txn.requests[0].GetObservedRouteVersion()) } func TestShardedCoordinatorRejectsPointWriteOnWriteFencedRoute(t *testing.T) { diff --git a/kv/sharded_coordinator_txn_test.go b/kv/sharded_coordinator_txn_test.go index 964f2713f..f3e29b377 100644 --- a/kv/sharded_coordinator_txn_test.go +++ b/kv/sharded_coordinator_txn_test.go @@ -19,6 +19,7 @@ type recordingTransactional struct { requests []*pb.Request responses []*TransactionResponse errs []error + onCommit func(call int, req *pb.Request) } func (s *recordingTransactional) Commit(_ context.Context, reqs []*pb.Request) (*TransactionResponse, error) { @@ -30,6 +31,9 @@ func (s *recordingTransactional) Commit(_ context.Context, reqs []*pb.Request) ( } s.requests = append(s.requests, cloneTxnRequest(reqs[0])) call := len(s.requests) - 1 + if s.onCommit != nil { + s.onCommit(call, s.requests[call]) + } if call < len(s.errs) && s.errs[call] != nil { return nil, s.errs[call] } diff --git a/main.go b/main.go index a630690d0..734436134 100644 --- a/main.go +++ b/main.go @@ -1564,6 +1564,9 @@ func prepareRuntimeServerRunner(waitRotateOnStartup startupRotationWaiter, in se }) } publicKVGate := &startupPublicKVGate{} + if in.distServer != nil { + in.distServer.SetReadGate(publicKVGate.blocked) + } installHLCLeaseRenewalBlocker(in.coordinate, waitRotateOnStartup.BlockMutators) adapterCoordinate := startupGatedCoordinator{ inner: in.coordinate, @@ -2185,6 +2188,7 @@ func startRaftServers( shardGroups map[uint64]*kv.ShardGroup, shardStore *kv.ShardStore, coordinate kv.Coordinator, + readGate func() bool, distServer *adapter.DistributionServer, relay *adapter.RedisPubSubRelay, proposalObserverForGroup func(uint64) kv.ProposalObserver, @@ -2218,7 +2222,7 @@ func startRaftServers( } gs := grpc.NewServer(opts...) trx := kv.NewTransactionWithProposer(proposerForGroup(rt, shardGroups), kv.WithProposalObserver(observerForGroup(proposalObserverForGroup, rt.spec.id))) - grpcSvc := adapter.NewGRPCServer(shardStore, coordinate) + grpcSvc := adapter.NewGRPCServer(shardStore, coordinate, adapter.WithGRPCReadGate(readGate)) pb.RegisterRawKVServer(gs, grpcSvc) pb.RegisterTransactionalKVServer(gs, grpcSvc) pb.RegisterInternalServer(gs, adapter.NewInternalWithEngine( @@ -2658,8 +2662,10 @@ func (r *runtimeServerRunner) startRaftTransport() error { return r.startupFailure(err) } adminGRPCOpts := r.adminGRPCOpts + var readGate func() bool if r.publicKVGate != nil { adminGRPCOpts.unary = append(adminGRPCOpts.unary, r.publicKVGate.unaryInterceptor) + readGate = r.publicKVGate.blocked } forwardDeps := adminForwardServerDeps{ tables: newDynamoTablesSource(r.dynamoServer), @@ -2674,6 +2680,7 @@ func (r *runtimeServerRunner) startRaftTransport() error { r.shardGroups, r.shardStore, r.coordinate, + readGate, r.distServer, r.pubsubRelay, func(groupID uint64) kv.ProposalObserver { From a6d7d7609c8cb502773ca293ea4964f6640f373b Mon Sep 17 00:00:00 2001 From: bootjp Date: Wed, 15 Jul 2026 09:26:39 +0900 Subject: [PATCH 23/26] Reset stale raft peer connections --- internal/raftengine/etcd/grpc_transport.go | 22 ++++++++++++-- .../raftengine/etcd/grpc_transport_test.go | 29 +++++++++++++++++++ 2 files changed, 48 insertions(+), 3 deletions(-) diff --git a/internal/raftengine/etcd/grpc_transport.go b/internal/raftengine/etcd/grpc_transport.go index 32d82e83e..50ea58460 100644 --- a/internal/raftengine/etcd/grpc_transport.go +++ b/internal/raftengine/etcd/grpc_transport.go @@ -533,7 +533,9 @@ func (t *GRPCTransport) dispatchRegular(ctx context.Context, msg raftpb.Message) } req := &pb.EtcdRaftMessage{Message: raw} if isPriorityMsg(msg.GetType()) || !t.sendStreamEnabledNow() || !t.allowPeerStreamProbe(peer.Address, time.Now()) { - return t.dispatchRegularUnary(ctx, client, req) + err := t.dispatchRegularUnary(ctx, client, req) + t.closePeerConnOnRetryableDialError(peer.Address, err) + return err } err = t.dispatchRegularStream(ctx, peer.Address, client, req) if err == nil { @@ -541,11 +543,16 @@ func (t *GRPCTransport) dispatchRegular(ctx context.Context, msg raftpb.Message) } if grpcStatusCode(err) == codes.Unimplemented { t.markPeerStreamUnsupported(peer.Address) - return t.dispatchRegularUnary(ctx, client, req) + err := t.dispatchRegularUnary(ctx, client, req) + t.closePeerConnOnRetryableDialError(peer.Address, err) + return err } if isSendStreamDisabled(err) { - return t.dispatchRegularUnary(ctx, client, req) + err := t.dispatchRegularUnary(ctx, client, req) + t.closePeerConnOnRetryableDialError(peer.Address, err) + return err } + t.closePeerConnOnRetryableDialError(peer.Address, err) return errors.WithStack(err) } @@ -554,6 +561,15 @@ func (t *GRPCTransport) dispatchRegularUnary(ctx context.Context, client pb.Etcd return errors.WithStack(err) } +func (t *GRPCTransport) closePeerConnOnRetryableDialError(address string, err error) { + if grpcStatusCode(err) != codes.Unavailable { + return + } + t.mu.Lock() + defer t.mu.Unlock() + t.closePeerConnLocked(address) +} + func (t *GRPCTransport) dispatchRegularStream(ctx context.Context, address string, client pb.EtcdRaftClient, req *pb.EtcdRaftMessage) error { stream, err := t.streamFor(ctx, address, client) if err != nil { diff --git a/internal/raftengine/etcd/grpc_transport_test.go b/internal/raftengine/etcd/grpc_transport_test.go index 02ffcb97f..fac1feb9b 100644 --- a/internal/raftengine/etcd/grpc_transport_test.go +++ b/internal/raftengine/etcd/grpc_transport_test.go @@ -19,6 +19,7 @@ import ( "google.golang.org/grpc" "google.golang.org/grpc/codes" "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" "google.golang.org/protobuf/proto" ) @@ -908,6 +909,30 @@ func TestDispatchRegularUsesUnaryForPriorityMessages(t *testing.T) { require.Zero(t, client.sendStreamCalls.Load()) } +func TestDispatchRegularDropsCachedPeerConnAfterUnaryUnavailable(t *testing.T) { + const addr = "host:2" + transport := NewGRPCTransport([]Peer{{NodeID: 2, Address: addr}}) + t.Cleanup(func() { require.NoError(t, transport.Close()) }) + client := &testEtcdRaftClient{sendErr: status.Error(codes.Unavailable, "connection refused")} + injectClient(t, transport, addr, client) + + err := transport.dispatchRegular(context.Background(), raftpb.Message{ + Type: messageTypePtr(raftpb.MsgHeartbeat), + From: uint64Ptr(1), + To: uint64Ptr(2), + Term: uint64Ptr(4), + Commit: uint64Ptr(22), + }) + require.Error(t, err) + require.Equal(t, codes.Unavailable, grpcStatusCode(err)) + + transport.mu.RLock() + _, cached := transport.clients[addr] + transport.mu.RUnlock() + require.False(t, cached) + require.Equal(t, int32(1), client.sendCalls.Load()) +} + func TestDispatchRegularUsesUnaryWhenSendStreamDisabled(t *testing.T) { t.Setenv(sendStreamEnabledEnvVar, "false") const addr = "host:2" @@ -1522,12 +1547,16 @@ type testEtcdRaftClient struct { blockSendStreamUntilContext bool sendStreamStarted chan struct{} releaseSendStream chan struct{} + sendErr error sendCalls atomic.Int32 sendStreamCalls atomic.Int32 } func (c *testEtcdRaftClient) Send(_ context.Context, _ *pb.EtcdRaftMessage, _ ...grpc.CallOption) (*pb.EtcdRaftAck, error) { c.sendCalls.Add(1) + if c.sendErr != nil { + return nil, c.sendErr + } return &pb.EtcdRaftAck{}, nil } From c8f865ea84ca03de6d96a604494b5902e974af19 Mon Sep 17 00:00:00 2001 From: bootjp Date: Wed, 15 Jul 2026 10:04:49 +0900 Subject: [PATCH 24/26] Stabilize snapshot catch-up and fence routing --- internal/raftengine/etcd/engine.go | 3 + .../etcd/engine_applied_index_test.go | 12 ++ internal/raftengine/etcd/grpc_transport.go | 121 ++++++++++++++++-- .../raftengine/etcd/grpc_transport_test.go | 52 ++++++++ kv/coordinator.go | 7 +- kv/fsm.go | 24 +++- kv/fsm_migration_fence_test.go | 28 ++++ kv/sharded_coordinator.go | 52 +++++++- kv/sharded_coordinator_partition_test.go | 43 +++++++ proto/internal.pb.go | 17 ++- proto/internal.proto | 5 + 11 files changed, 341 insertions(+), 23 deletions(-) diff --git a/internal/raftengine/etcd/engine.go b/internal/raftengine/etcd/engine.go index 65ef0903f..a273f5e8d 100644 --- a/internal/raftengine/etcd/engine.go +++ b/internal/raftengine/etcd/engine.go @@ -3386,6 +3386,9 @@ func (e *Engine) protectReceivedFSMSnapshot(index uint64) bool { if e.protectedReceivedFSMSnaps == nil { e.protectedReceivedFSMSnaps = make(map[uint64]int, 1) } + if e.protectedReceivedFSMSnaps[index] > 0 { + return false + } e.protectedReceivedFSMSnaps[index]++ return true } diff --git a/internal/raftengine/etcd/engine_applied_index_test.go b/internal/raftengine/etcd/engine_applied_index_test.go index 8ebb2ecfe..cf4596306 100644 --- a/internal/raftengine/etcd/engine_applied_index_test.go +++ b/internal/raftengine/etcd/engine_applied_index_test.go @@ -208,6 +208,18 @@ func TestProtectReceivedFSMSnapshotWaitsOnSnapshotMuForAlreadyAppliedIndex(t *te require.Empty(t, e.protectedReceivedFSMSnaps) } +func TestProtectReceivedFSMSnapshotRejectsDuplicateInFlightIndex(t *testing.T) { + e := &Engine{} + + require.True(t, e.protectReceivedFSMSnapshot(9)) + require.False(t, e.protectReceivedFSMSnapshot(9)) + require.Equal(t, map[uint64]int{9: 1}, e.protectedReceivedFSMSnaps) + + e.unprotectReceivedFSMSnapshot(9) + require.Empty(t, e.protectedReceivedFSMSnaps) + require.True(t, e.protectReceivedFSMSnapshot(9)) +} + func TestUnprotectReceivedFSMSnapshotTokenIfApplied(t *testing.T) { e := &Engine{ protectedReceivedFSMSnaps: map[uint64]int{9: 1}, diff --git a/internal/raftengine/etcd/grpc_transport.go b/internal/raftengine/etcd/grpc_transport.go index 50ea58460..29f67129f 100644 --- a/internal/raftengine/etcd/grpc_transport.go +++ b/internal/raftengine/etcd/grpc_transport.go @@ -41,6 +41,7 @@ var ( errSnapshotMetadataDuplicate = errors.New("etcd raft snapshot metadata was sent more than once") errSnapshotMessageNil = errors.New("etcd raft snapshot message is required") errSnapshotStreamShort = errors.New("etcd raft snapshot stream closed before final chunk") + errSnapshotDispatchBusy = errors.New("etcd raft snapshot dispatch already in progress") errReceivedFSMSnapshotStale = errors.New("etcd raft received fsm snapshot is stale") errPeerStreamClosed = errors.New("etcd raft SendStream closed") errSendStreamDisabled = errors.New("etcd raft SendStream is disabled") @@ -132,7 +133,8 @@ type GRPCTransport struct { // bridgeSem limits concurrent bridge-mode snapshot materializations so // that aggregate in-memory allocation stays bounded even when multiple // dispatch workers run simultaneously. - bridgeSem chan struct{} + bridgeSem chan struct{} + snapshotSendSem chan struct{} } func NewGRPCTransport(peers []Peer) *GRPCTransport { @@ -167,6 +169,7 @@ func NewGRPCTransport(peers []Peer) *GRPCTransport { sendStreamCancel: sendStreamCancel, snapshotChunkSize: defaultSnapshotChunkSize, bridgeSem: make(chan struct{}, defaultBridgeMaterializeLimit), + snapshotSendSem: make(chan struct{}, 1), } } @@ -389,6 +392,10 @@ func isSnapshotMsg(msg raftpb.Message) bool { func (t *GRPCTransport) dispatchSnapshot(ctx context.Context, msg raftpb.Message) error { ctx, cancel := transportContext(ctx, defaultSnapshotDispatchTimeout) defer cancel() + if !t.tryAcquireSnapshotSend() { + return errors.WithStack(errors.Mark(status.Error(codes.ResourceExhausted, errSnapshotDispatchBusy.Error()), errSnapshotDispatchBusy)) + } + defer t.releaseSnapshotSend() // Prefer streaming when the snapshot holds a token and an opener is wired. // This avoids materialising the full FSM payload in memory on the sender. @@ -413,6 +420,28 @@ func (t *GRPCTransport) dispatchSnapshot(ctx context.Context, msg raftpb.Message return t.sendSnapshot(ctx, patched) } +func (t *GRPCTransport) tryAcquireSnapshotSend() bool { + if t == nil || t.snapshotSendSem == nil { + return true + } + select { + case t.snapshotSendSem <- struct{}{}: + return true + default: + return false + } +} + +func (t *GRPCTransport) releaseSnapshotSend() { + if t == nil || t.snapshotSendSem == nil { + return + } + select { + case <-t.snapshotSendSem: + default: + } +} + // streamFSMSnapshot streams the .fsm payload file directly to the peer using // chunked gRPC without loading the full content into memory. func (t *GRPCTransport) streamFSMSnapshot(ctx context.Context, msg raftpb.Message, index uint64, openFn func(uint64) (io.ReadCloser, error)) error { @@ -1247,6 +1276,11 @@ func (t *GRPCTransport) receiveSnapshotStream(stream pb.EtcdRaft_SendSnapshotSer if err != nil { return raftpb.Message{}, err } + protection, err := protectReceivedSnapshotMetadata(metadata, fsmSnapDir, protectFn, unprotectFn) + if err != nil { + return raftpb.Message{}, err + } + defer protection.releaseUnlessHandedOff() spool, err := newReceiveSnapshotSpool(spoolPlacement) if err != nil { return raftpb.Message{}, err @@ -1275,10 +1309,12 @@ func (t *GRPCTransport) receiveSnapshotStream(stream pb.EtcdRaft_SendSnapshotSer metadata, firstPayloadChunk, preparedFSMWrite, + protection.protected, ) if err != nil { return raftpb.Message{}, err } + protection.handoff() index := uint64(0) if msg.Snapshot != nil { index = msg.Snapshot.GetMetadata().GetIndex() @@ -1292,6 +1328,50 @@ func (t *GRPCTransport) receiveSnapshotStream(stream pb.EtcdRaft_SendSnapshotSer return msg, nil } +type receivedSnapshotProtection struct { + index uint64 + protected bool + handedOff bool + unprotectFn func(uint64) +} + +func protectReceivedSnapshotMetadata( + metadata raftpb.Message, + fsmSnapDir string, + protectFn func(uint64) bool, + unprotectFn func(uint64), +) (receivedSnapshotProtection, error) { + if fsmSnapDir == "" || metadata.Snapshot == nil { + return receivedSnapshotProtection{}, nil + } + index := metadata.Snapshot.GetMetadata().GetIndex() + if index == 0 || protectFn == nil { + return receivedSnapshotProtection{}, nil + } + if !protectFn(index) { + return receivedSnapshotProtection{}, errors.WithStack(errReceivedFSMSnapshotStale) + } + return receivedSnapshotProtection{ + index: index, + protected: true, + unprotectFn: unprotectFn, + }, nil +} + +func (p *receivedSnapshotProtection) handoff() { + if p == nil { + return + } + p.handedOff = true +} + +func (p *receivedSnapshotProtection) releaseUnlessHandedOff() { + if p == nil || !p.protected || p.handedOff || p.unprotectFn == nil { + return + } + p.unprotectFn(p.index) +} + // drainSnapshotChunks consumes the SendSnapshot stream into spool, computes // CRC32C over the payload bytes as they hit disk, and on the final chunk // hands off to finalizeReceivedSnapshot — which decides between the @@ -1308,7 +1388,7 @@ func drainSnapshotChunks( unprotectFn func(uint64), ) (raftpb.Message, int64, error) { var metadata raftpb.Message - return drainSnapshotChunksFrom(stream, spool, fsmSnapDir, prepareFn, protectFn, unprotectFn, metadata, nil, false) + return drainSnapshotChunksFrom(stream, spool, fsmSnapDir, prepareFn, protectFn, unprotectFn, metadata, nil, false, false) } func drainSnapshotChunksFrom( @@ -1321,6 +1401,7 @@ func drainSnapshotChunksFrom( metadata raftpb.Message, firstPayloadChunk *pb.EtcdRaftSnapshotChunk, preparedFSMWrite bool, + preprotected bool, ) (raftpb.Message, int64, error) { seenMetadata := metadata.Snapshot != nil // Wrap spool with crc32CWriter so the CRC accumulates as bytes hit @@ -1353,7 +1434,7 @@ func drainSnapshotChunksFrom( } payloadBytes += int64(len(chunk.Chunk)) if chunk.Final { - msg, err := finalizeReceivedSnapshot(metadata, spool, crcWriter.Sum32(), fsmSnapDir, protectFn, unprotectFn, seenMetadata) + msg, err := finalizeReceivedSnapshot(metadata, spool, crcWriter.Sum32(), fsmSnapDir, protectFn, unprotectFn, seenMetadata, preprotected) if err != nil { return raftpb.Message{}, 0, err } @@ -1436,6 +1517,7 @@ func finalizeReceivedSnapshot( protectFn func(uint64) bool, unprotectFn func(uint64), seenMetadata bool, + preprotected bool, ) (raftpb.Message, error) { if !seenMetadata || metadata.Snapshot == nil { return raftpb.Message{}, errors.WithStack(errSnapshotMetadataNil) @@ -1447,23 +1529,38 @@ func finalizeReceivedSnapshot( // rename to). return buildSnapshotMessage(metadata, spool, seenMetadata) } - protected := false - if protectFn != nil { - if !protectFn(index) { - return raftpb.Message{}, errors.WithStack(errReceivedFSMSnapshotStale) - } - protected = true + protected, err := protectReceivedSnapshotForFinalize(index, preprotected, protectFn) + if err != nil { + return raftpb.Message{}, err } if err := spool.FinalizeAsFSMFile(fsmSnapDir, index, crc32c); err != nil { - if protected && unprotectFn != nil { - unprotectFn(index) - } + releaseFinalizeSnapshotProtection(index, protected, preprotected, unprotectFn) return raftpb.Message{}, err } metadata.Snapshot.Data = encodeSnapshotToken(index, crc32c) return metadata, nil } +func protectReceivedSnapshotForFinalize(index uint64, preprotected bool, protectFn func(uint64) bool) (bool, error) { + if preprotected { + return true, nil + } + if protectFn == nil { + return false, nil + } + if !protectFn(index) { + return false, errors.WithStack(errReceivedFSMSnapshotStale) + } + return true, nil +} + +func releaseFinalizeSnapshotProtection(index uint64, protected bool, preprotected bool, unprotectFn func(uint64)) { + if !protected || preprotected || unprotectFn == nil { + return + } + unprotectFn(index) +} + func maybePrepareReceivedFSMSnapshotWrite( metadata raftpb.Message, fsmSnapDir string, diff --git a/internal/raftengine/etcd/grpc_transport_test.go b/internal/raftengine/etcd/grpc_transport_test.go index fac1feb9b..b7e43c703 100644 --- a/internal/raftengine/etcd/grpc_transport_test.go +++ b/internal/raftengine/etcd/grpc_transport_test.go @@ -125,6 +125,58 @@ func TestReceiveSnapshotStreamRejectsDuplicateMetadata(t *testing.T) { require.True(t, errors.Is(err, errSnapshotMetadataDuplicate)) } +func TestSnapshotSendGateAllowsOnlyOneInFlightSnapshot(t *testing.T) { + transport := NewGRPCTransport(nil) + + require.True(t, transport.tryAcquireSnapshotSend()) + require.False(t, transport.tryAcquireSnapshotSend()) + + transport.releaseSnapshotSend() + require.True(t, transport.tryAcquireSnapshotSend()) + transport.releaseSnapshotSend() +} + +func TestReceiveSnapshotStreamRejectsProtectedIndexBeforePayload(t *testing.T) { + const index = uint64(127) + metadata := raftpb.Message{ + Type: messageTypePtr(raftpb.MsgSnap), + From: uint64Ptr(1), + To: uint64Ptr(2), + Snapshot: &raftpb.Snapshot{ + Metadata: testSnapshotMetadata(index, 1, nil), + }, + } + raw, err := proto.Marshal(&metadata) + require.NoError(t, err) + + transport := NewGRPCTransport(nil) + fsmSnapDir := t.TempDir() + transport.SetFSMSnapDir(fsmSnapDir) + var protected []uint64 + transport.SetFSMSnapshotProtection( + func(got uint64) bool { + protected = append(protected, got) + return false + }, + func(uint64) { + t.Fatal("stale snapshot rejection must not unprotect an index it did not protect") + }, + ) + stream := &testSendSnapshotServer{ + chunks: []*pb.EtcdRaftSnapshotChunk{ + {Metadata: raw}, + {Chunk: []byte("payload that must not be read"), Final: true}, + }, + } + + _, err = transport.receiveSnapshotStream(stream) + require.Error(t, err) + require.True(t, errors.Is(err, errReceivedFSMSnapshotStale)) + require.Equal(t, []uint64{index}, protected) + require.Equal(t, 1, stream.index, "receiver must reject after metadata before reading payload chunks") + require.NoFileExists(t, fsmSnapPath(fsmSnapDir, index)) +} + // TestReceiveSnapshotStream_StreamingTokenWhenFSMSnapDirSet pins the // memory-safety win: when the receive transport has fsmSnapDir wired, // the spool file is renamed to fsmSnapPath(...) and Snapshot.Data is diff --git a/kv/coordinator.go b/kv/coordinator.go index 174c3b636..0c1567a26 100644 --- a/kv/coordinator.go +++ b/kv/coordinator.go @@ -1088,7 +1088,7 @@ func (c *Coordinate) dispatchTxn(ctx context.Context, reqs []*Elem[OP], readKeys // carries the option-2 one-phase dedup probe key for a retry that reuses // a failed attempt's write set. r, err := c.transactionManager.Commit(ctx, []*pb.Request{ - onePhaseTxnRequestWithPrevCommit(startTS, commitTS, prevCommitTS, primary, reqs, readKeys, observedRouteVersion), + onePhaseTxnRequestWithPrevCommit(startTS, commitTS, prevCommitTS, primary, reqs, readKeys, observedRouteVersion, nil), }) if err != nil { return nil, errors.WithStack(err) @@ -1265,7 +1265,7 @@ func (c *Coordinate) buildRedirectRequests(reqs *OperationGroup[OP]) ([]*pb.Requ commitTS = 0 } return []*pb.Request{ - onePhaseTxnRequestWithPrevCommit(reqs.StartTS, commitTS, reqs.PrevCommitTS, primary, reqs.Elems, reqs.ReadKeys, reqs.ObservedRouteVersion), + onePhaseTxnRequestWithPrevCommit(reqs.StartTS, commitTS, reqs.PrevCommitTS, primary, reqs.Elems, reqs.ReadKeys, reqs.ObservedRouteVersion, nil), }, nil } @@ -1319,7 +1319,7 @@ func elemToMutation(req *Elem[OP]) *pb.Mutation { // route catalog snapshot at txn-begin (M1 plumbing, see // docs/design/2026_05_29_implemented_composed1_cross_group_commit_guard.md). // Zero is the legacy "unpinned" sentinel. -func onePhaseTxnRequestWithPrevCommit(startTS, commitTS, prevCommitTS uint64, primaryKey []byte, reqs []*Elem[OP], readKeys [][]byte, observedRouteVersion uint64) *pb.Request { +func onePhaseTxnRequestWithPrevCommit(startTS, commitTS, prevCommitTS uint64, primaryKey []byte, reqs []*Elem[OP], readKeys [][]byte, observedRouteVersion uint64, writeFenceBypassKeys [][]byte) *pb.Request { muts := make([]*pb.Mutation, 0, len(reqs)+1) muts = append(muts, &pb.Mutation{ Op: pb.Op_PUT, @@ -1336,6 +1336,7 @@ func onePhaseTxnRequestWithPrevCommit(startTS, commitTS, prevCommitTS uint64, pr Mutations: muts, ReadKeys: readKeys, ObservedRouteVersion: observedRouteVersion, + WriteFenceBypassKeys: writeFenceBypassKeys, } } diff --git a/kv/fsm.go b/kv/fsm.go index b2e74b1db..dedc7829a 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -867,7 +867,7 @@ func (f *kvFSM) verifyWriteFence(r *pb.Request) error { if !ok { return errors.WithStack(ErrComposed1VersionGCd) } - if err := verifyWriteFenceFromSnapshot(r.GetMutations(), observedSnap, observedVer, "observed"); err != nil { + if err := verifyWriteFenceFromSnapshot(r.GetMutations(), r.GetWriteFenceBypassKeys(), observedSnap, observedVer, "observed"); err != nil { return err } if currentSnap.Version() == observedSnap.Version() { @@ -875,7 +875,7 @@ func (f *kvFSM) verifyWriteFence(r *pb.Request) error { } } - return verifyWriteFenceFromSnapshot(r.GetMutations(), currentSnap, currentSnap.Version(), "current") + return verifyWriteFenceFromSnapshot(r.GetMutations(), r.GetWriteFenceBypassKeys(), currentSnap, currentSnap.Version(), "current") } func requestBypassesWriteFence(r *pb.Request) bool { @@ -895,7 +895,8 @@ func (f *kvFSM) writeFenceHistoryReady() bool { return f.routes != nil && f.shardGroupID != 0 } -func verifyWriteFenceFromSnapshot(mutations []*pb.Mutation, snap RouteSnapshot, snapVer uint64, phase string) error { +func verifyWriteFenceFromSnapshot(mutations []*pb.Mutation, writeFenceBypassKeys [][]byte, snap RouteSnapshot, snapVer uint64, phase string) error { + bypassKeys := writeFenceBypassKeySet(writeFenceBypassKeys) for _, mut := range mutations { if mut == nil { continue @@ -912,6 +913,9 @@ func verifyWriteFenceFromSnapshot(mutations []*pb.Mutation, snap RouteSnapshot, } continue } + if _, ok := bypassKeys[string(mut.Key)]; ok { + continue + } rKey := routeKey(mut.Key) if snap.WriteFencedForKey(rKey) { return errors.Wrapf(ErrRouteWriteFenced, @@ -928,6 +932,20 @@ func verifyWriteFenceFromSnapshot(mutations []*pb.Mutation, snap RouteSnapshot, return nil } +func writeFenceBypassKeySet(keys [][]byte) map[string]struct{} { + if len(keys) == 0 { + return nil + } + out := make(map[string]struct{}, len(keys)) + for _, key := range keys { + if len(key) == 0 { + continue + } + out[string(key)] = struct{}{} + } + return out +} + // verifyOwnerFromSnapshot is the shared per-mutation owner-check // loop used by verifyComposed1's observed-version and current- // version passes. `phase` is the diagnostic label ("observed" / diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index f668821c8..ecb771b38 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -81,6 +81,34 @@ func TestFSMRejectsObservedWriteFencedRawPointWrite(t *testing.T) { require.ErrorIs(t, err, ErrRouteWriteFenced) } +func TestFSMWriteFenceBypassAllowsMarkedRawPointWrite(t *testing.T) { + t.Parallel() + + fsm := newFirstRouteWriteFencedFSM(t) + key := []byte("!sqs|msg|data|p|partitioned-key") + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + WriteFenceBypassKeys: [][]byte{key}, + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: key, Value: []byte("v")}}, + }, 10) + require.NoError(t, err) + + got, err := fsm.store.GetAt(context.Background(), key, 10) + require.NoError(t, err) + require.Equal(t, []byte("v"), got) +} + +func TestFSMWriteFenceBypassDoesNotAllowDelPrefix(t *testing.T) { + t.Parallel() + + fsm := newFirstRouteWriteFencedFSM(t) + prefix := []byte("!sqs|msg|data|p|") + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + WriteFenceBypassKeys: [][]byte{prefix}, + Mutations: []*pb.Mutation{{Op: pb.Op_DEL_PREFIX, Key: prefix}}, + }, 10) + require.ErrorIs(t, err, ErrRouteWriteFenced) +} + func TestFSMRejectsCurrentWriteFenceAfterObservedActiveRawPointWrite(t *testing.T) { t.Parallel() diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index 73420bb86..ec7c43843 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -1143,6 +1143,42 @@ func (c *ShardedCoordinator) partitionResolverRecognisesPointKey(key []byte) boo return c.router.partitionResolver.RecognisesPartitionedKey(key) } +func (c *ShardedCoordinator) resolverClaimedWriteFenceBypassKeysForElems(elems []*Elem[OP]) [][]byte { + if len(elems) == 0 { + return nil + } + out := make([][]byte, 0, len(elems)) + for _, elem := range elems { + if elem == nil || elem.Op == DelPrefix || !c.partitionResolverClaimsPointKey(elem.Key) { + continue + } + out = append(out, bytes.Clone(elem.Key)) + } + return out +} + +func (c *ShardedCoordinator) resolverClaimedWriteFenceBypassKeys(muts []*pb.Mutation) [][]byte { + if len(muts) == 0 { + return nil + } + out := make([][]byte, 0, len(muts)) + for _, mut := range muts { + if mut == nil || mut.GetOp() == pb.Op_DEL_PREFIX || !c.partitionResolverClaimsPointKey(mut.Key) { + continue + } + out = append(out, bytes.Clone(mut.Key)) + } + return out +} + +func (c *ShardedCoordinator) partitionResolverClaimsPointKey(key []byte) bool { + if c == nil || c.router == nil || c.router.partitionResolver == nil || len(key) == 0 { + return false + } + _, ok := c.router.partitionResolver.ResolveGroup(key) + return ok +} + func (c *ShardedCoordinator) rejectWriteFencedDelPrefixes(elems []*Elem[OP]) error { if c == nil || c.engine == nil { return nil @@ -1337,7 +1373,7 @@ func (c *ShardedCoordinator) dispatchSingleShardTxn(ctx context.Context, startTS // carries the one-phase dedup probe key for a retry that reuses a failed // attempt's write set. resp, err := g.Txn.Commit(ctx, []*pb.Request{ - onePhaseTxnRequestWithPrevCommit(startTS, commitTS, prevCommitTS, primaryKey, elems, readKeys, observedRouteVersion), + onePhaseTxnRequestWithPrevCommit(startTS, commitTS, prevCommitTS, primaryKey, elems, readKeys, observedRouteVersion, c.resolverClaimedWriteFenceBypassKeysForElems(elems)), }) if err != nil { return nil, errors.WithStack(err) @@ -1369,6 +1405,7 @@ func (c *ShardedCoordinator) prewriteTxn(ctx context.Context, startTS, commitTS Mutations: append([]*pb.Mutation{prepareMeta}, grouped[gid]...), ReadKeys: groupedReadKeys[gid], ObservedRouteVersion: observedRouteVersion, + WriteFenceBypassKeys: c.resolverClaimedWriteFenceBypassKeys(grouped[gid]), } if _, err := g.Txn.Commit(ctx, []*pb.Request{req}); err != nil { // Same WithoutCancel pattern as dispatchTxn's @@ -1421,6 +1458,7 @@ func (c *ShardedCoordinator) commitPrimaryTxn(ctx context.Context, startTS uint6 Ts: startTS, Mutations: append([]*pb.Mutation{meta}, keys...), ObservedRouteVersion: observedRouteVersion, + WriteFenceBypassKeys: c.resolverClaimedWriteFenceBypassKeys(keys), } r, err := g.Txn.Commit(ctx, []*pb.Request{req}) @@ -1470,6 +1508,7 @@ func (c *ShardedCoordinator) commitSecondaryTxns(ctx context.Context, startTS ui Ts: startTS, Mutations: append([]*pb.Mutation{meta}, keyMutations(grouped[gid])...), ObservedRouteVersion: observedRouteVersion, + WriteFenceBypassKeys: c.resolverClaimedWriteFenceBypassKeys(grouped[gid]), } r, err := commitSecondaryWithRetry(ctx, g, req) if err != nil { @@ -2067,6 +2106,7 @@ func (c *ShardedCoordinator) rawLogs(ctx context.Context, reqs *OperationGroup[O Ts: ts, Mutations: grouped[gid], ObservedRouteVersion: reqs.ObservedRouteVersion, + WriteFenceBypassKeys: c.resolverClaimedWriteFenceBypassKeys(grouped[gid]), }) } return logs, nil @@ -2107,7 +2147,11 @@ func (c *ShardedCoordinator) txnLogs(ctx context.Context, reqs *OperationGroup[O if err != nil { return nil, err } - return buildTxnLogs(reqs.StartTS, commitTS, grouped, gids, reqs.ObservedRouteVersion) + bypassKeysByGroup := make(map[uint64][][]byte, len(grouped)) + for gid, muts := range grouped { + bypassKeysByGroup[gid] = c.resolverClaimedWriteFenceBypassKeys(muts) + } + return buildTxnLogs(reqs.StartTS, commitTS, grouped, gids, reqs.ObservedRouteVersion, bypassKeysByGroup) } // observeMutation: counted pre-commit, so a mutation that subsequently @@ -2183,7 +2227,7 @@ func (c *ShardedCoordinator) groupMutations(reqs []*Elem[OP], label keyviz.Label return grouped, gids, nil } -func buildTxnLogs(startTS uint64, commitTS uint64, grouped map[uint64][]*pb.Mutation, gids []uint64, observedRouteVersion uint64) ([]*pb.Request, error) { +func buildTxnLogs(startTS uint64, commitTS uint64, grouped map[uint64][]*pb.Mutation, gids []uint64, observedRouteVersion uint64, writeFenceBypassKeysByGroup map[uint64][][]byte) ([]*pb.Request, error) { logs := make([]*pb.Request, 0, len(gids)*txnPhaseCount) for _, gid := range gids { muts := grouped[gid] @@ -2200,6 +2244,7 @@ func buildTxnLogs(startTS uint64, commitTS uint64, grouped map[uint64][]*pb.Muta {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: primaryKey, LockTTLms: defaultTxnLockTTLms, CommitTS: 0})}, }, muts...), ObservedRouteVersion: observedRouteVersion, + WriteFenceBypassKeys: writeFenceBypassKeysByGroup[gid], }, &pb.Request{ IsTxn: true, @@ -2209,6 +2254,7 @@ func buildTxnLogs(startTS uint64, commitTS uint64, grouped map[uint64][]*pb.Muta {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: primaryKey, LockTTLms: 0, CommitTS: commitTS})}, }, keys...), ObservedRouteVersion: observedRouteVersion, + WriteFenceBypassKeys: writeFenceBypassKeysByGroup[gid], }, ) } diff --git a/kv/sharded_coordinator_partition_test.go b/kv/sharded_coordinator_partition_test.go index e9b5c16bb..7b224db8a 100644 --- a/kv/sharded_coordinator_partition_test.go +++ b/kv/sharded_coordinator_partition_test.go @@ -151,6 +151,8 @@ func TestShardedCoordinatorWriteFencePrecheckHonoursPartitionResolver(t *testing require.Equal(t, uint64(42), resp.CommitIndex) require.Empty(t, g1.requests, "engine route fence must not preempt resolver-owned keys") require.Len(t, g42.requests, 1) + require.Equal(t, [][]byte{key}, g42.requests[0].GetWriteFenceBypassKeys(), + "resolver-owned point writes must carry the FSM write-fence bypass marker") } func TestShardedCoordinatorWriteFencePrecheckSkipsUnresolvedPartitionKey(t *testing.T) { @@ -185,6 +187,47 @@ func TestShardedCoordinatorWriteFencePrecheckSkipsUnresolvedPartitionKey(t *test require.Empty(t, g1.requests) } +func TestShardedCoordinatorTxnCarriesResolverWriteFenceBypassKeys(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: nil, GroupID: 1, State: distribution.RouteStateWriteFenced}, + }, + })) + + g1 := &recordingTransactional{ + responses: []*TransactionResponse{{CommitIndex: 1}}, + } + g42 := &recordingTransactional{ + responses: []*TransactionResponse{{CommitIndex: 42}}, + } + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1, Store: store.NewMVCCStore()}, + 42: {Txn: g42, Store: store.NewMVCCStore()}, + }, 1, NewHLC(), nil) + + key := []byte("!sqs|msg|data|p|txn-partitioned-key") + coord.WithPartitionResolver(&stubResolver{claim: map[string]uint64{ + string(key): 42, + }}) + + resp, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + IsTxn: true, + StartTS: 100, + CommitTS: 200, + Elems: []*Elem[OP]{{Op: Put, Key: key, Value: []byte("v")}}, + }) + require.NoError(t, err) + require.NotNil(t, resp) + require.Empty(t, g1.requests) + require.Len(t, g42.requests, 1) + require.Equal(t, [][]byte{key}, g42.requests[0].GetWriteFenceBypassKeys(), + "single-shard txn prepare/one-phase request must preserve the resolver-owned point key for the FSM") +} + // TestShardedCoordinator_DispatchSplitsMutationsByResolverGroup is // the genuine regression for the Gemini-HIGH groupMutations // bypass: a Dispatch with mutations belonging to TWO different diff --git a/proto/internal.pb.go b/proto/internal.pb.go index 0e2bd4664..7490809ce 100644 --- a/proto/internal.pb.go +++ b/proto/internal.pb.go @@ -208,6 +208,11 @@ type Request struct { // is plumbing only — the FSM ignores the value, so all existing // callers see no behaviour change. ObservedRouteVersion uint64 `protobuf:"varint,6,opt,name=observed_route_version,json=observedRouteVersion,proto3" json:"observed_route_version,omitempty"` + // write_fence_bypass_keys carries point keys whose owner was resolved by a + // higher-level partition resolver before the request was appended to Raft. + // The FSM uses it only to skip the byte-route write fence for those exact + // point mutations; prefix deletes and unmarked keys still fail closed. + WriteFenceBypassKeys [][]byte `protobuf:"bytes,7,rep,name=write_fence_bypass_keys,json=writeFenceBypassKeys,proto3" json:"write_fence_bypass_keys,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -284,6 +289,13 @@ func (x *Request) GetObservedRouteVersion() uint64 { return 0 } +func (x *Request) GetWriteFenceBypassKeys() [][]byte { + if x != nil { + return x.WriteFenceBypassKeys + } + return nil +} + type RaftCommand struct { state protoimpl.MessageState `protogen:"open.v1"` Requests []*Request `protobuf:"bytes,1,rep,name=requests,proto3" json:"requests,omitempty"` @@ -909,14 +921,15 @@ const file_internal_proto_rawDesc = "" + "\bMutation\x12\x13\n" + "\x02op\x18\x01 \x01(\x0e2\x03.OpR\x02op\x12\x10\n" + "\x03key\x18\x02 \x01(\fR\x03key\x12\x14\n" + - "\x05value\x18\x03 \x01(\fR\x05value\"\xca\x01\n" + + "\x05value\x18\x03 \x01(\fR\x05value\"\x81\x02\n" + "\aRequest\x12\x15\n" + "\x06is_txn\x18\x01 \x01(\bR\x05isTxn\x12\x1c\n" + "\x05phase\x18\x02 \x01(\x0e2\x06.PhaseR\x05phase\x12\x0e\n" + "\x02ts\x18\x03 \x01(\x04R\x02ts\x12'\n" + "\tmutations\x18\x04 \x03(\v2\t.MutationR\tmutations\x12\x1b\n" + "\tread_keys\x18\x05 \x03(\fR\breadKeys\x124\n" + - "\x16observed_route_version\x18\x06 \x01(\x04R\x14observedRouteVersion\"3\n" + + "\x16observed_route_version\x18\x06 \x01(\x04R\x14observedRouteVersion\x125\n" + + "\x17write_fence_bypass_keys\x18\a \x03(\fR\x14writeFenceBypassKeys\"3\n" + "\vRaftCommand\x12$\n" + "\brequests\x18\x01 \x03(\v2\b.RequestR\brequests\"M\n" + "\x0eForwardRequest\x12\x15\n" + diff --git a/proto/internal.proto b/proto/internal.proto index 1f7c77712..f58257cec 100644 --- a/proto/internal.proto +++ b/proto/internal.proto @@ -55,6 +55,11 @@ message Request { // is plumbing only — the FSM ignores the value, so all existing // callers see no behaviour change. uint64 observed_route_version = 6; + // write_fence_bypass_keys carries point keys whose owner was resolved by a + // higher-level partition resolver before the request was appended to Raft. + // The FSM uses it only to skip the byte-route write fence for those exact + // point mutations; prefix deletes and unmarked keys still fail closed. + repeated bytes write_fence_bypass_keys = 7; } message RaftCommand { From b553090f299c5243ed2307b67ece1bf1f90d9930 Mon Sep 17 00:00:00 2001 From: bootjp Date: Wed, 15 Jul 2026 10:40:25 +0900 Subject: [PATCH 25/26] Keep SQS receive route fences visible --- adapter/sqs_messages.go | 7 +- adapter/sqs_receive_route_fence_test.go | 86 +++++++++++++++++++++++++ 2 files changed, 90 insertions(+), 3 deletions(-) create mode 100644 adapter/sqs_receive_route_fence_test.go diff --git a/adapter/sqs_messages.go b/adapter/sqs_messages.go index 83bfe6009..a75e0c96d 100644 --- a/adapter/sqs_messages.go +++ b/adapter/sqs_messages.go @@ -1340,8 +1340,9 @@ func (s *SQSServer) rotateMessagesForDelivery( // - (msg, false, nil) → delivered, caller appends. // - (nil, true, nil) → expected race; skip this candidate only. // Covers ErrKeyNotFound (someone deleted the record between the -// vis-index scan and our GetAt) and ErrWriteConflict on dispatch -// (another receive rotated the same record). +// vis-index scan and our GetAt), plus non-fence dispatch races +// like ErrWriteConflict / ErrTxnLocked (another receive rotated +// the same record). // - (nil, false, err) → non-retryable failure; propagate up the // stack so ReceiveMessage returns an actionable 5xx instead of // a false-empty 200. @@ -1441,7 +1442,7 @@ func (s *SQSServer) commitReceiveRotation(ctx context.Context, queueName string, return nil, false, err } if _, err := s.coordinator.Dispatch(ctx, req); err != nil { - if isRetryableTransactWriteError(err) { + if isIgnorableTransactRaceError(err) { return nil, true, nil } return nil, false, errors.WithStack(err) diff --git a/adapter/sqs_receive_route_fence_test.go b/adapter/sqs_receive_route_fence_test.go new file mode 100644 index 000000000..9c584f791 --- /dev/null +++ b/adapter/sqs_receive_route_fence_test.go @@ -0,0 +1,86 @@ +package adapter + +import ( + "context" + "testing" + + "github.com/bootjp/elastickv/kv" + "github.com/bootjp/elastickv/store" + "github.com/stretchr/testify/require" +) + +type receiveRotationErrorCoordinator struct { + stubAdapterCoordinator + err error +} + +func (c *receiveRotationErrorCoordinator) Dispatch(context.Context, *kv.OperationGroup[kv.OP]) (*kv.CoordinateResponse, error) { + return nil, c.err +} + +func TestSQSReceiveRotationClassifiesDispatchErrors(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + err error + wantSkip bool + wantErr error + }{ + {name: "write conflict", err: store.ErrWriteConflict, wantSkip: true}, + {name: "txn locked", err: kv.ErrTxnLocked, wantSkip: true}, + {name: "route write fenced", err: kv.ErrRouteWriteFenced, wantErr: kv.ErrRouteWriteFenced}, + } + + for _, tt := range tests { + + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + const queueName = "receive-route-fence" + const messageID = "msg-1" + const gen = uint64(1) + const visibleAt = int64(1000) + + meta := &sqsQueueMeta{Name: queueName, Generation: gen, PartitionCount: 1} + rec := &sqsMessageRecord{ + MessageID: messageID, + Body: []byte("body"), + MD5OfBody: sqsMD5Hex([]byte("body")), + SendTimestampMillis: visibleAt, + VisibleAtMillis: visibleAt, + QueueGeneration: gen, + } + cand := sqsMsgCandidate{ + visKey: sqsMsgVisKeyDispatch(meta, queueName, 0, gen, visibleAt, messageID), + messageID: messageID, + partition: 0, + } + dataKey := sqsMsgDataKeyDispatch(meta, queueName, 0, gen, messageID) + srv := &SQSServer{ + coordinator: &receiveRotationErrorCoordinator{err: tt.err}, + } + + msg, skip, err := srv.commitReceiveRotation( + context.Background(), + queueName, + meta, + cand, + dataKey, + rec, + uint64(visibleAt), + sqsReceiveOptions{VisibilityTimeout: 30}, + nil, + fifoLockAcquire, + ) + + require.Nil(t, msg) + require.Equal(t, tt.wantSkip, skip) + if tt.wantErr != nil { + require.ErrorIs(t, err, tt.wantErr) + } else { + require.NoError(t, err) + } + }) + } +} From d2b22acc76035ec00d4f79c6d83bed7b232e16b7 Mon Sep 17 00:00:00 2001 From: bootjp Date: Wed, 15 Jul 2026 19:45:09 +0900 Subject: [PATCH 26/26] Stabilize redis proxy and raft recovery --- adapter/redis_txn.go | 96 ++++++++-- adapter/redis_txn_test.go | 86 +++++++++ cmd/redis-proxy/main.go | 122 ++++++++----- ..._implemented_etcd_snapshot_disk_offload.md | 4 +- ...29_implemented_snapshot_logical_decoder.md | 4 +- .../2026_04_29_proposed_logical_backup.md | 12 +- .../2026_05_28_implemented_tla_safety_spec.md | 8 +- ...implemented_idempotent_snapshot_restore.md | 166 +++++++----------- .../2026_06_12_proposed_scaling_roadmap.md | 2 +- internal/raftengine/etcd/engine.go | 47 +++-- .../etcd/engine_applied_index_test.go | 22 +++ internal/raftengine/etcd/engine_test.go | 62 ++++++- internal/raftengine/etcd/fsm_snapshot_file.go | 70 +++++++- internal/raftengine/etcd/grpc_transport.go | 18 +- .../raftengine/etcd/grpc_transport_test.go | 54 ++++-- .../raftengine/etcd/snapshot_spool_test.go | 42 +++++ internal/raftengine/etcd/wal_store.go | 162 ++++++----------- .../etcd/wal_store_skip_gate_test.go | 98 +++++------ internal/raftengine/statemachine.go | 26 ++- kv/coordinator.go | 22 ++- kv/fsm.go | 31 ++-- kv/lease_read_test.go | 8 +- kv/lease_warmup_test.go | 64 +++++++ kv/sharded_coordinator.go | 2 +- .../dashboards/elastickv-redis-summary.json | 2 +- monitoring/hotpath.go | 2 +- monitoring/hotpath_test.go | 2 +- proxy/config.go | 3 + proxy/dualwrite.go | 122 ++++++++++--- proxy/metrics.go | 11 +- proxy/noop_backend.go | 38 ++++ proxy/proxy.go | 15 +- proxy/proxy_test.go | 105 +++++++++-- proxy/pubsub.go | 10 +- proxy/raw_redis_proxy.go | 156 ++++++++++++++++ proxy/raw_redis_proxy_test.go | 77 ++++++++ 36 files changed, 1282 insertions(+), 489 deletions(-) create mode 100644 proxy/noop_backend.go create mode 100644 proxy/raw_redis_proxy.go create mode 100644 proxy/raw_redis_proxy_test.go diff --git a/adapter/redis_txn.go b/adapter/redis_txn.go index 53b937ed2..dcbcc3b48 100644 --- a/adapter/redis_txn.go +++ b/adapter/redis_txn.go @@ -98,9 +98,11 @@ type txnValue struct { } type stringReplacement struct { - key []byte - value []byte - ttl *time.Time + key []byte + value []byte + ttl *time.Time + rawTyp redisValueType + rawTypKnown bool } type txnContext struct { @@ -617,18 +619,48 @@ func (t *txnContext) loadTTLState(key []byte) (*ttlTxnState, error) { } func (t *txnContext) stagedKeyType(key []byte) (redisValueType, error) { + view, err := t.stagedKeyTypeView(key) + if err != nil { + return redisTypeNone, err + } + return view.typ, nil +} + +type txnKeyTypeView struct { + typ redisValueType + rawTyp redisValueType + rawTypKnown bool +} + +func (t *txnContext) stagedKeyTypeView(key []byte) (txnKeyTypeView, error) { k := string(key) - if _, ok := t.replacers[k]; ok { - return redisTypeString, nil + if repl, ok := t.replacers[k]; ok { + return txnKeyTypeView{ + typ: redisTypeString, + rawTyp: repl.rawTyp, + rawTypKnown: repl.rawTypKnown, + }, nil } if typ, ok := t.stagedPositiveKeyType(k); ok { - return typ, nil + return txnKeyTypeView{typ: typ}, nil } if t.hasStagedTypeDeletion(k) { - return redisTypeNone, nil + return txnKeyTypeView{typ: redisTypeNone}, nil } t.trackTypeReadKeys(key) - return t.server.keyTypeAt(t.ctxOrBackground(), key, t.startTS) + rawTyp, err := t.server.rawKeyTypeAt(t.ctxOrBackground(), key, t.startTS) + if err != nil { + return txnKeyTypeView{}, err + } + typ, err := t.server.applyTTLFilter(t.ctxOrBackground(), key, t.startTS, rawTyp) + if err != nil { + return txnKeyTypeView{}, err + } + return txnKeyTypeView{ + typ: typ, + rawTyp: rawTyp, + rawTypKnown: true, + }, nil } func (t *txnContext) stagedPositiveKeyType(key string) (redisValueType, bool) { @@ -721,10 +753,11 @@ func (t *txnContext) applySet(cmd redcon.Command) (redisResult, error) { if err != nil { return redisResult{}, err } - typ, err := t.stagedKeyType(cmd.Args[1]) + typeView, err := t.stagedKeyTypeView(cmd.Args[1]) if err != nil { return redisResult{}, err } + typ := typeView.typ // NX/XX: skip the write if the key-existence condition is not met. exists := typ != redisTypeNone @@ -739,7 +772,8 @@ func (t *txnContext) applySet(cmd redcon.Command) (redisResult, error) { if err != nil { return redisResult{}, err } - t.stageStringReplacement(cmd.Args[1], cmd.Args[2], opts.ttl) + t.trackWideCollectionFenceReads(cmd.Args[1]) + t.stageStringReplacementWithRawType(cmd.Args[1], cmd.Args[2], opts.ttl, typeView.rawTyp, typeView.rawTypKnown) return applySetResult(opts, oldValue), nil } @@ -777,15 +811,32 @@ func cloneTimePtr(in *time.Time) *time.Time { } func (t *txnContext) stageStringReplacement(key, value []byte, ttl *time.Time) { + t.stageStringReplacementWithRawType(key, value, ttl, redisTypeNone, false) +} + +func (t *txnContext) stageStringReplacementWithRawType(key, value []byte, ttl *time.Time, rawTyp redisValueType, rawTypKnown bool) { if t.replacers == nil { t.replacers = map[string]*stringReplacement{} } k := string(key) - t.replacers[k] = &stringReplacement{ - key: bytes.Clone(key), - value: bytes.Clone(value), - ttl: cloneTimePtr(ttl), + if repl, ok := t.replacers[k]; ok { + repl.value = bytes.Clone(value) + repl.ttl = cloneTimePtr(ttl) + if rawTypKnown { + repl.rawTyp = rawTyp + repl.rawTypKnown = true + } + delete(t.deletedKeys, k) + return + } + repl := &stringReplacement{ + key: bytes.Clone(key), + value: bytes.Clone(value), + ttl: cloneTimePtr(ttl), + rawTyp: rawTyp, + rawTypKnown: rawTypKnown, } + t.replacers[k] = repl delete(t.deletedKeys, k) } @@ -1576,11 +1627,13 @@ func (t *txnContext) buildReplacementElems(ctx context.Context) ([]*kv.Elem[kv.O for _, k := range keys { repl := t.replacers[k] t.trackWideCollectionFenceReads(repl.key) - deleteElems, _, err := t.server.deleteLogicalKeyElems(ctx, repl.key, t.startTS) - if err != nil { - return nil, err + if repl.needsFullLogicalDelete() { + deleteElems, _, err := t.server.deleteLogicalKeyElems(ctx, repl.key, t.startTS) + if err != nil { + return nil, err + } + elems = append(elems, deleteElems...) } - elems = append(elems, deleteElems...) elems = append(elems, redisTxnWideCollectionFenceElems(repl.key)...) elems = append(elems, &kv.Elem[kv.OP]{ Op: kv.Put, @@ -1596,6 +1649,13 @@ func (t *txnContext) buildReplacementElems(ctx context.Context) ([]*kv.Elem[kv.O return elems, nil } +func (r *stringReplacement) needsFullLogicalDelete() bool { + if !r.rawTypKnown { + return true + } + return isNonStringCollectionType(r.rawTyp) +} + func (t *txnContext) buildLogicalDeletionElems(ctx context.Context) ([]*kv.Elem[kv.OP], error) { if len(t.logicalDeletes) == 0 { return nil, nil diff --git a/adapter/redis_txn_test.go b/adapter/redis_txn_test.go index 3349c4b8b..c8e49c1d9 100644 --- a/adapter/redis_txn_test.go +++ b/adapter/redis_txn_test.go @@ -684,6 +684,92 @@ func TestRedisTxnSetReplacementConflictsWithConcurrentWideHashWrite(t *testing.T "SET replacement in MULTI must conflict with concurrent HSET of a new field") } +func TestRedisTxnSetReplacementTracksWideFencesBeforeBuild(t *testing.T) { + t.Parallel() + + ctx := context.Background() + server, st := newRedisStorageMigrationTestServer(t) + key := []byte("set-replace:fence-read") + + txn := newRedisTxnTestContext(server) + res, err := txn.applySet(redcon.Command{Args: [][]byte{[]byte(cmdSet), key, []byte("string")}}) + require.NoError(t, err) + require.Equal(t, "OK", res.str) + for _, fenceKey := range redisTxnWideCollectionFenceKeys(key) { + require.Contains(t, txn.readKeys, string(fenceKey)) + } + + require.NoError(t, st.PutAt(ctx, redisTxnWideHashFenceKey(key), []byte{}, redisTxnTestStartTS+1, 0)) + require.ErrorIs(t, txn.validateReadSet(ctx), store.ErrWriteConflict) +} + +func TestRedisTxnSetReplacementSkipsWideCleanupForRawStringOrMissing(t *testing.T) { + t.Parallel() + + ctx := context.Background() + cases := []struct { + name string + seed func(store.MVCCStore, []byte) + }{ + {name: "missing"}, + { + name: "string", + seed: func(st store.MVCCStore, key []byte) { + require.NoError(t, st.PutAt(ctx, redisStrKey(key), encodeRedisStr([]byte("old"), nil), redisTxnTestStartTS, 0)) + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + server, st := newRedisStorageMigrationTestServer(t) + key := []byte("set-replace:no-wide-cleanup:" + tc.name) + if tc.seed != nil { + tc.seed(st, key) + } + + txn := newRedisTxnTestContext(server) + res, err := txn.applySet(redcon.Command{Args: [][]byte{[]byte(cmdSet), key, []byte("string")}}) + require.NoError(t, err) + require.Equal(t, "OK", res.str) + + elems, err := txn.buildReplacementElems(ctx) + require.NoError(t, err) + require.False(t, elemKeysContain(elems, store.HashMetaKey(key))) + require.False(t, elemKeysContain(elems, store.SetMetaKey(key))) + require.False(t, elemKeysContain(elems, store.ZSetMetaKey(key))) + require.False(t, elemKeysContain(elems, store.ListMetaKey(key))) + require.False(t, elemKeysContain(elems, store.StreamMetaKey(key))) + require.True(t, elemKeysContain(elems, redisStrKey(key))) + }) + } +} + +func TestRedisTxnSetReplacementDeletesExpiredRawHash(t *testing.T) { + t.Parallel() + + ctx := context.Background() + server, st := newRedisStorageMigrationTestServer(t) + key := []byte("set-replace:expired-hash") + expired := time.Now().Add(-time.Hour) + require.NoError(t, st.PutAt(ctx, store.HashFieldKey(key, []byte("old")), []byte("v"), redisTxnTestStartTS, 0)) + require.NoError(t, st.PutAt(ctx, store.HashMetaKey(key), store.MarshalHashMeta(store.HashMeta{Len: 1}), redisTxnTestStartTS, 0)) + require.NoError(t, st.PutAt(ctx, redisTTLKey(key), encodeRedisTTL(expired), redisTxnTestStartTS, 0)) + + txn := newRedisTxnTestContext(server) + res, err := txn.applySet(redcon.Command{Args: [][]byte{[]byte(cmdSet), key, []byte("string")}}) + require.NoError(t, err) + require.Equal(t, "OK", res.str) + + elems, err := txn.buildReplacementElems(ctx) + require.NoError(t, err) + require.True(t, elemKeysContain(elems, store.HashFieldKey(key, []byte("old")))) + require.True(t, elemKeysContain(elems, store.HashMetaKey(key))) + require.True(t, elemKeysContain(elems, redisStrKey(key))) +} + func TestRedisTxnSetReplacementConflictsWithConcurrentListPush(t *testing.T) { t.Parallel() diff --git a/cmd/redis-proxy/main.go b/cmd/redis-proxy/main.go index ec2622eb3..4ad067987 100644 --- a/cmd/redis-proxy/main.go +++ b/cmd/redis-proxy/main.go @@ -7,6 +7,7 @@ import ( "log/slog" "net" "net/http" + "net/http/pprof" "os" "os/signal" "strings" @@ -57,6 +58,8 @@ func run() error { flag.StringVar(&cfg.SentryEnv, "sentry-env", cfg.SentryEnv, "Sentry environment") flag.Float64Var(&cfg.SentrySampleRate, "sentry-sample", cfg.SentrySampleRate, "Sentry sample rate") flag.StringVar(&cfg.MetricsAddr, "metrics", cfg.MetricsAddr, "Prometheus metrics address") + flag.StringVar(&cfg.PProfAddr, "pprof", cfg.PProfAddr, "pprof listen address (empty = disabled)") + flag.BoolVar(&cfg.RedisOnlyRaw, "redis-only-raw", cfg.RedisOnlyRaw, "Use raw TCP bridging in redis-only mode") flag.Parse() mode, err := parseRuntimeOptions(modeStr, primaryPoolSize, elasticKVPoolSize, secondaryWriteConcurrency, secondaryScriptConcurrency) @@ -78,11 +81,38 @@ func run() error { sentryReporter := proxy.NewSentryReporter(cfg.SentryDSN, cfg.SentryEnv, cfg.SentrySampleRate, logger) defer sentryReporter.Flush(sentryFlushTimeout) - // Prometheus reg := prometheus.NewRegistry() metrics := proxy.NewProxyMetrics(reg) - // Backends + primary, secondary, err := newBackends(cfg, primaryPoolSize, elasticKVPoolSize, logger) + if err != nil { + return err + } + defer primary.Close() + defer secondary.Close() + + dual := proxy.NewDualWriter(primary, secondary, cfg, metrics, sentryReporter, logger) + defer dual.Close() // wait for in-flight async goroutines + srv := proxy.NewProxyServer(cfg, dual, metrics, sentryReporter, logger) + + // Context for graceful shutdown + ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer cancel() + + go serveMetrics(ctx, cfg.MetricsAddr, reg, logger) + + if cfg.PProfAddr != "" { + go servePProf(ctx, cfg.PProfAddr, logger) + } + + // Start proxy + if err := srv.ListenAndServe(ctx); err != nil { + return fmt.Errorf("proxy server: %w", err) + } + return nil +} + +func newBackends(cfg proxy.ProxyConfig, primaryPoolSize, elasticKVPoolSize int, logger *slog.Logger) (proxy.Backend, proxy.Backend, error) { primaryOpts := proxy.DefaultBackendOptions() primaryOpts.DB = cfg.PrimaryDB primaryOpts.Password = cfg.PrimaryPassword @@ -94,59 +124,63 @@ func run() error { secondarySeeds := parseAddrList(cfg.SecondaryAddr) if len(secondarySeeds) == 0 { - return fmt.Errorf("at least one secondary address is required") + return nil, nil, fmt.Errorf("at least one secondary address is required") } - var primary, secondary proxy.Backend switch cfg.Mode { - case proxy.ModeElasticKVPrimary, proxy.ModeElasticKVOnly: - primary = proxy.NewLeaderAwareRedisBackend(secondarySeeds, "elastickv", secondaryOpts, logger) - secondary = proxy.NewRedisBackendWithOptions(cfg.PrimaryAddr, "redis", primaryOpts) - case proxy.ModeRedisOnly, proxy.ModeDualWrite, proxy.ModeDualWriteShadow: - primary = proxy.NewRedisBackendWithOptions(cfg.PrimaryAddr, "redis", primaryOpts) - secondary = proxy.NewLeaderAwareRedisBackend(secondarySeeds, "elastickv", secondaryOpts, logger) + case proxy.ModeElasticKVPrimary: + return proxy.NewLeaderAwareRedisBackend(secondarySeeds, "elastickv", secondaryOpts, logger), + proxy.NewRedisBackendWithOptions(cfg.PrimaryAddr, "redis", primaryOpts), nil + case proxy.ModeElasticKVOnly: + return proxy.NewLeaderAwareRedisBackend(secondarySeeds, "elastickv", secondaryOpts, logger), + proxy.NewNoopBackend("redis"), nil + case proxy.ModeRedisOnly: + return proxy.NewRedisBackendWithOptions(cfg.PrimaryAddr, "redis", primaryOpts), + proxy.NewNoopBackend("elastickv"), nil + case proxy.ModeDualWrite, proxy.ModeDualWriteShadow: + return proxy.NewRedisBackendWithOptions(cfg.PrimaryAddr, "redis", primaryOpts), + proxy.NewLeaderAwareRedisBackend(secondarySeeds, "elastickv", secondaryOpts, logger), nil + default: + return nil, nil, fmt.Errorf("unsupported mode: %s", cfg.Mode.String()) } - defer primary.Close() - defer secondary.Close() +} - dual := proxy.NewDualWriter(primary, secondary, cfg, metrics, sentryReporter, logger) - defer dual.Close() // wait for in-flight async goroutines - srv := proxy.NewProxyServer(cfg, dual, metrics, sentryReporter, logger) +func serveMetrics(ctx context.Context, addr string, reg *prometheus.Registry, logger *slog.Logger) { + mux := http.NewServeMux() + mux.Handle("/metrics", promhttp.HandlerFor(reg, promhttp.HandlerOpts{})) + serveHTTP(ctx, addr, mux, "metrics", logger) +} - // Context for graceful shutdown - ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) - defer cancel() +func servePProf(ctx context.Context, addr string, logger *slog.Logger) { + mux := http.NewServeMux() + mux.HandleFunc("/debug/pprof/", pprof.Index) + mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline) + mux.HandleFunc("/debug/pprof/profile", pprof.Profile) + mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol) + mux.HandleFunc("/debug/pprof/trace", pprof.Trace) + serveHTTP(ctx, addr, mux, "pprof", logger) +} - // Start metrics server +func serveHTTP(ctx context.Context, addr string, handler http.Handler, name string, logger *slog.Logger) { + var lc net.ListenConfig + ln, err := lc.Listen(ctx, "tcp", addr) + if err != nil { + logger.Error(name+" listen failed", "addr", addr, "err", err) + return + } + srv := &http.Server{Handler: handler, ReadHeaderTimeout: time.Second} go func() { - mux := http.NewServeMux() - mux.Handle("/metrics", promhttp.HandlerFor(reg, promhttp.HandlerOpts{})) - var lc net.ListenConfig - ln, err := lc.Listen(ctx, "tcp", cfg.MetricsAddr) - if err != nil { - logger.Error("metrics listen failed", "addr", cfg.MetricsAddr, "err", err) - return - } - metricsSrv := &http.Server{Handler: mux, ReadHeaderTimeout: time.Second} - go func() { - <-ctx.Done() - shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), metricsShutdownTimeout) - defer shutdownCancel() - if err := metricsSrv.Shutdown(shutdownCtx); err != nil { - logger.Warn("metrics server shutdown error", "err", err) - } - }() - logger.Info("metrics server starting", "addr", cfg.MetricsAddr) - if err := metricsSrv.Serve(ln); err != nil && err != http.ErrServerClosed { - logger.Error("metrics server error", "err", err) + <-ctx.Done() + shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), metricsShutdownTimeout) + defer shutdownCancel() + if err := srv.Shutdown(shutdownCtx); err != nil { + logger.Warn(name+" server shutdown error", "err", err) } }() - - // Start proxy - if err := srv.ListenAndServe(ctx); err != nil { - return fmt.Errorf("proxy server: %w", err) + logger.Info(name+" server starting", "addr", addr) + if err := srv.Serve(ln); err != nil && err != http.ErrServerClosed { + logger.Error(name+" server error", "err", err) } - return nil } func parseRuntimeOptions(modeStr string, primaryPoolSize, elasticKVPoolSize, secondaryWriteConcurrency, secondaryScriptConcurrency int) (proxy.ProxyMode, error) { diff --git a/docs/design/2026_04_14_implemented_etcd_snapshot_disk_offload.md b/docs/design/2026_04_14_implemented_etcd_snapshot_disk_offload.md index 5312c3784..b800f2906 100644 --- a/docs/design/2026_04_14_implemented_etcd_snapshot_disk_offload.md +++ b/docs/design/2026_04_14_implemented_etcd_snapshot_disk_offload.md @@ -5,7 +5,7 @@ ### Observed Symptoms Clusters using the etcd engine exhibit large memory spikes across all nodes at snapshot -creation intervals (`defaultSnapshotEvery = 10,000` entries). The spike magnitude scales +creation intervals (`defaultSnapshotEvery = 100,000` entries). The spike magnitude scales with FSM data size, and simultaneous spikes across multiple nodes compound the pressure. ### Root Cause @@ -34,7 +34,7 @@ storage.CreateSnapshot(applied, &confState, payload) ``` `MemoryStorage` holds `raftpb.Snapshot.Data = payload` until the next snapshot is created -(i.e., until another 10,000 entries are processed), keeping the full FSM export resident +(i.e., until another 100,000 entries are processed), keeping the full FSM export resident in memory the entire time. #### Problem 3: Re-allocation when sending to followers diff --git a/docs/design/2026_04_29_implemented_snapshot_logical_decoder.md b/docs/design/2026_04_29_implemented_snapshot_logical_decoder.md index 07205fc7c..f4f721ba4 100644 --- a/docs/design/2026_04_29_implemented_snapshot_logical_decoder.md +++ b/docs/design/2026_04_29_implemented_snapshot_logical_decoder.md @@ -37,7 +37,7 @@ which is a different framing the decoder/encoder do not touch. See `2026_05_25_implemented_snapshot_logical_encoder.md` §"Why a separate design doc" item 3. -Snapshots are taken automatically every `defaultSnapshotEvery = 10000` +Snapshots are taken automatically every `defaultSnapshotEvery = 100000` log entries (`internal/raftengine/etcd/engine.go:92`) and stored under `{dataDir}/fsm-snap/.fsm`. They are crash-consistent by construction — the writer takes a Pebble snapshot at the FSM's @@ -608,7 +608,7 @@ bespoke parser, the format has failed its goal. - **Staleness.** Whatever was written between the snapshot's `applied_index` and "now" is not in the snapshot, so it is not in the decoded output. The gap is bounded by `SnapshotEvery × write_rate` - (default 10000 entries; for a write-heavy cluster, seconds; for a + (default 100000 entries; for a write-heavy cluster, minutes; for a quiet one, hours). - **Cadence is not user-controlled.** "Snapshot now" requires a Raft trigger; the decoder cannot create a fresh snapshot, only consume diff --git a/docs/design/2026_04_29_proposed_logical_backup.md b/docs/design/2026_04_29_proposed_logical_backup.md index 39350d188..1e307c832 100644 --- a/docs/design/2026_04_29_proposed_logical_backup.md +++ b/docs/design/2026_04_29_proposed_logical_backup.md @@ -884,7 +884,7 @@ refuse if: remaining_headroom < --snapshot-headroom-entries ``` `SnapshotEvery` is the per-engine snapshot trigger (default -`defaultSnapshotEvery = 10000` from `internal/raftengine/etcd/engine.go:92`, +`defaultSnapshotEvery = 100000` from `internal/raftengine/etcd/engine.go:92`, overridden via the `ELASTICKV_RAFT_SNAPSHOT_COUNT` env var — there is no `--raftSnapshotEvery` CLI flag). The value is currently a private field on the etcd `Engine` struct (`internal/raftengine/etcd/engine.go:224`) @@ -906,9 +906,9 @@ implemented on the etcd backend by returning the place, `BeginBackup` reads each group's `SnapshotEvery` rather than hardcoding `defaultSnapshotEvery`, so an operator who tuned `ELASTICKV_RAFT_SNAPSHOT_COUNT` sees consistent behavior. With -`SnapshotEvery = 10000` and -`--snapshot-headroom-entries = 1000` (default; one-tenth of -SnapshotEvery), the check refuses backups when fewer than 1000 entries +`SnapshotEvery = 100000` and +`--snapshot-headroom-entries = 10000` (default; one-tenth of +SnapshotEvery), the check refuses backups when fewer than 10000 entries remain before the next snapshot fires — i.e. when an in-flight backup is at risk of triggering the snapshot-installation corner case. A freshly-snapshotted cluster has the *largest* remaining headroom and @@ -1341,7 +1341,7 @@ written. `s.groups[id]` map without a typed assertion fallback. Tests that mock `AdminGroup` (e.g. `adapter/admin_grpc_test.go`) gain one extra method to implement; they can return - `defaultSnapshotEvery = 10000` for parity with production. + `defaultSnapshotEvery = 100000` for parity with production. - Extend `kv/active_timestamp_tracker.go` with `PinWithDeadline`, `Extend`, and the per-second sweeper goroutine that reaps expired pins and emits the `backup_pin_expired` structured warning. @@ -1601,7 +1601,7 @@ Scope: out of this proposal; mentioned only to draw the boundary. | `TestAdminGRPCConnCacheReuse` | Two consecutive `BeginBackup` calls dialing the same peer share one underlying `*grpc.ClientConn` (verified via the cache size); shutdown of a peer evicts only that entry, not the whole cache; admin cache is independent of `ShardStore.connCache` (no cross-cache eviction) | | `TestBeginBackupPropagatesAdminAuthToken` | A cluster booted with `--adminToken` accepts a `BeginBackup` call carrying `authorization: Bearer `; the handler propagates the same metadata via `metadata.NewOutgoingContext` to every `GetNodeVersion` fan-out dial, so peers return their version (not `Unauthenticated`). A `BeginBackup` call without the token is itself rejected before any fan-out happens | | `TestVersionCacheRaceUnderLoad` | `go test -race` with 50 concurrent `GetRaftGroups` callers and async `GetNodeVersion` probe goroutines writing the cache simultaneously emits no data-race report; the `sync.Map` choice is enforced by the lack of a separate `versionCacheMu` field on `AdminServer` | -| `TestSnapshotEveryReadsFromEngine` | A node started with `ELASTICKV_RAFT_SNAPSHOT_COUNT=5000` reports `Engine.SnapshotEvery() == 5000`; `BeginBackup` uses 5000 (not the default 10000) when computing remaining headroom | +| `TestSnapshotEveryReadsFromEngine` | A node started with `ELASTICKV_RAFT_SNAPSHOT_COUNT=5000` reports `Engine.SnapshotEvery() == 5000`; `BeginBackup` uses 5000 (not the default 100000) when computing remaining headroom | | `TestRenewBackupRetriesLeaderElection` | Force a leader election mid-`RenewBackup`; the admin server retries `BackupExtend` up to 3 times with 500ms backoff and succeeds once the new leader is established, without aborting the dump | | `TestPinWithDeadlineExpiry` | `PinWithDeadline(ts, now+100ms)` is auto-released by the sweeper after the deadline; compactor unblocked; `backup_pin_expired` log emitted | | `TestBeginBackupWaitsForLaggingShard` | Force shard B's `applied_index` to lag; `BeginBackup` polls until it catches up or times out with `FailedPrecondition`; no scan starts in the timeout case | diff --git a/docs/design/2026_05_28_implemented_tla_safety_spec.md b/docs/design/2026_05_28_implemented_tla_safety_spec.md index cd94bb5d0..463fa16ba 100644 --- a/docs/design/2026_05_28_implemented_tla_safety_spec.md +++ b/docs/design/2026_05_28_implemented_tla_safety_spec.md @@ -199,8 +199,8 @@ cannot check them as state invariants. independently reach `ceiling + 1` before a fresh ceiling is renewed, the new leader's first `Next()` can tie or undercut the old leader's last commit. Bounding inter-node skew to less than - one ceiling window (< 3s with the current `hlcPhysicalWindowMs = - 3000ms`) keeps the window wide enough that the new leader cannot + one ceiling window (< 15s with the current `hlcPhysicalWindowMs = + 15000ms`) keeps the window wide enough that the new leader cannot independently reach the overflow value before a renewal applies. Should be surfaced in operator docs as a cluster prerequisite. - **(ii) Logical-counter handoff.** The 16-bit logical half of the HLC @@ -270,7 +270,7 @@ cannot check them as state invariants. `hlcPhysicalWindowMs` cannot serve any persistence timestamp — every client commit is rejected until renewal succeeds. This is a CP, not AP, trade-off and operators must size - `hlcPhysicalWindowMs` (currently 3s) relative to expected + `hlcPhysicalWindowMs` (currently 15s) relative to expected partition duration; see §9 risk 7. ### 5.2 OCC @@ -676,7 +676,7 @@ does not keep this document in `partial`. 7. **Fail-closed availability under partition.** HLC-4 precondition (iii) makes the ceiling-fence behaviour normative: a leader partitioned from the default group's quorum for longer than - `hlcPhysicalWindowMs` (currently 3s) cannot serve any persistence + `hlcPhysicalWindowMs` (currently 15s) cannot serve any persistence timestamp, so client commits are rejected until renewal succeeds. This is a CP, not AP, trade-off and is a stricter regime than the current implementation (which silently keeps issuing). Mitigation: diff --git a/docs/design/2026_06_02_implemented_idempotent_snapshot_restore.md b/docs/design/2026_06_02_implemented_idempotent_snapshot_restore.md index 2317fb9db..a3db4b2f2 100644 --- a/docs/design/2026_06_02_implemented_idempotent_snapshot_restore.md +++ b/docs/design/2026_06_02_implemented_idempotent_snapshot_restore.md @@ -297,9 +297,9 @@ func restoreSnapshotState(fsm StateMachine, snapshot raftpb.Snapshot, fsmSnapDir // The body restore is skipped, but we MUST still consume // the v1/v2 snapshot header so the FSM picks up the HLC // ceiling AND the Stage 8a cutover (see §5). Thread - // tok.CRC32C through so the skip path verifies the file - // before mutating FSM state, matching the existing - // openAndRestoreFSMSnapshot safety contract. + // tok.CRC32C through so the skip path verifies the cheap + // snapshot envelope before mutating FSM state. Full-body CRC + // remains on the restore path where body bytes are consumed. return applyHeaderStateOnSkip(fsm, fsmSnapPath(fsmSnapDir, tok.Index), tok.CRC32C) } return openAndRestoreFSMSnapshot(fsm, fsmSnapPath(fsmSnapDir, tok.Index), tok.CRC32C) @@ -423,42 +423,35 @@ or `internal/raftengine/etcd → kv` edges (test-only on both sides), and adding either to satisfy the round-6 design either duplicates the CRC verifier in `kv` or breaks the layering. -Round-7 keeps the CRC verifier in its existing package and splits the -seam into two phases — a **parse** phase that reads the header from a -caller-supplied reader (and drains the rest, for CRC coverage), and an -**apply** phase that is pure assignment. The engine orchestrates -size + footer + tee'd CRC computation around the parse phase, then -calls apply only after all three pass: +Round-7 keeps the envelope checks in the engine package and splits the +seam into two phases — a **parse** phase that reads only the header from +a caller-supplied reader, and an **apply** phase that is pure assignment. +The engine validates size + footer-vs-tokenCRC before parsing and calls +apply only after those checks and the header parse pass: ```go // internal/raftengine/statemachine.go (new, sibling to ApplyIndexAware) type SnapshotHeaderApplier interface { - // ParseSnapshotHeader reads the v1/v2 header from r, drains the - // remaining bytes (so a wrapping crc32 TeeReader covers the full - // payload), and returns the parsed (ceiling, cutover) pair WITHOUT - // mutating FSM state. Implementations MUST NOT touch any FSM - // fields here; the engine calls ApplySnapshotHeader separately - // only after the wrapping CRC verification passes. + // ParseSnapshotHeader reads the v1/v2 header from r and returns the + // parsed (ceiling, cutover) pair WITHOUT mutating FSM state. + // Implementations MUST NOT touch any FSM fields here; the engine + // calls ApplySnapshotHeader separately only after the snapshot file + // footer matches the raft token. // // Errors propagate from the underlying header parser - // (ErrSnapshotHeaderUnknownMagic / InvalidLength) or from the - // drain pass (I/O errors). FSM state stays untouched on error. + // (ErrSnapshotHeaderUnknownMagic / InvalidLength). FSM state stays + // untouched on error. ParseSnapshotHeader(r io.Reader) (ceiling, cutover uint64, err error) // ApplySnapshotHeader is pure assignment of the verified header // state. The engine calls this only after ParseSnapshotHeader - // returned and the wrapping crc32 hash matched the file footer. + // returned and the snapshot file's footer matched the raft token. ApplySnapshotHeader(ceiling, cutover uint64) } // internal/raftengine/etcd/wal_store.go -- never imports kv; -// CRC verification stays here where the helpers live. +// cheap snapshot envelope checks stay here where the helpers live. func applyHeaderStateOnSkip(fsm StateMachine, snapPath string, tokenCRC uint32) error { - setter, ok := fsm.(SnapshotHeaderApplier) - if !ok { - return nil // FSM has no header state; skip is harmless. - } - file, err := os.Open(snapPath) if err != nil { return statFSMFileError(err) } defer file.Close() @@ -480,29 +473,24 @@ func applyHeaderStateOnSkip(fsm StateMachine, snapPath string, tokenCRC uint32) "path=%s footer=%08x token=%08x", snapPath, footer, tokenCRC) } - // Step 3: full-body CRC. Wrap the payload in a crc32 TeeReader and - // hand it to the FSM's ParseSnapshotHeader for header parse + drain. - // The header bytes are included in the computed CRC because the - // FSM reads them from the tee'd reader. + setter, ok := fsm.(SnapshotHeaderApplier) + if !ok { + return nil // FSM has no header state; skip is harmless. + } + if _, err := file.Seek(0, io.SeekStart); err != nil { return errors.WithStack(err) } payloadSize := info.Size() - fsmFooterSize - h := crc32.New(crc32cTable) - tee := io.TeeReader(io.LimitReader(file, payloadSize), h) - ceiling, cutover, perr := setter.ParseSnapshotHeader(tee) + ceiling, cutover, perr := setter.ParseSnapshotHeader(io.LimitReader(file, payloadSize)) if perr != nil { - // ErrSnapshotHeaderUnknownMagic / InvalidLength / I/O error - // surfaced from the FSM's parse pass. State unchanged. + // ErrSnapshotHeaderUnknownMagic / InvalidLength surfaced from + // the FSM's parse pass. State unchanged. return errors.WithStack(perr) } - if h.Sum32() != footer { - return errors.Wrapf(ErrFSMSnapshotFileCRC, - "path=%s footer=%08x computed=%08x", snapPath, footer, h.Sum32()) - } - // All three checks passed; apply side-effects. + // Envelope checks and header parse passed; apply side-effects. setter.ApplySnapshotHeader(ceiling, cutover) return nil } @@ -510,16 +498,10 @@ func applyHeaderStateOnSkip(fsm StateMachine, snapPath string, tokenCRC uint32) // kv/fsm.go (new methods on kvFSM) -- kv.ReadSnapshotHeader stays inside kv; // no imports of internal/raftengine/etcd or its private helpers. func (f *kvFSM) ParseSnapshotHeader(r io.Reader) (uint64, uint64, error) { - // The engine has already wrapped r in a crc32 TeeReader sized at - // the body payload (file size minus 4-byte footer). We read the - // header, then drain the rest of the body so the engine's CRC - // covers every byte (matching restoreAndComputeCRC's behaviour). - br := bufio.NewReaderSize(r, 1<<20) //nolint:mnd // 1 MiB, local to kv - ceiling, cutover, err := ReadSnapshotHeader(br) + // The skip path already has the FSM body locally, so read only + // the snapshot header needed for HLC/cutover state. + ceiling, cutover, err := ReadSnapshotHeader(bufio.NewReaderSize(r, 4<<10)) if err != nil { return 0, 0, errors.WithStack(err) } - if _, err := io.Copy(io.Discard, br); err != nil { - return 0, 0, errors.WithStack(err) - } return ceiling, cutover, nil } @@ -531,27 +513,16 @@ func (f *kvFSM) ApplySnapshotHeader(ceiling, cutover uint64) { } ``` -**Cost note**. Step 3 reads the full snapshot file once (through the -crc32 TeeReader). For multi-GiB FSMs this is a non-trivial I/O cost -— but it is **strictly cheaper** than the restore path it replaces -(which also reads the file once via `restoreAndComputeCRC` AND -additionally writes a temp Pebble database with sstable / WAL output). -Observed restore wall-clock is dominated by Pebble writes, not reads; -eliding the writes preserves the bulk of the win. A future -optimisation could persist the HLC ceiling + cutover durably -(analogous to `metaAppliedIndex`) and elide the file read entirely — -out of scope here, flagged under Open Questions. - -**Why this seam shape**. The two-phase split lets the CRC verifier -stay co-located with its private helpers in -`internal/raftengine/etcd/fsm_snapshot_file.go`'s package, **and** -keeps the v1/v2 header parser inside `kv` where it already lives. -Neither package imports the other in production. The "do CRC on -engine side, side-effects on FSM side after verify" contract is -exactly the inversion of `openAndRestoreFSMSnapshot` (which inlines -`fsm.Restore` inside the CRC tee for performance reasons): for the -skip path we don't need a single-pass restore, so splitting the -phases costs nothing and buys layer hygiene. +**Cost note**. The skip path does not read the multi-GiB body. Startup +cost is bounded by opening the snapshot file, reading the footer, and +parsing the fixed-size header. Full-body CRC verification remains on +the execute/full-restore path, where the body bytes are actually used. + +**Why this seam shape**. The two-phase split keeps the snapshot +envelope checks in `internal/raftengine/etcd`, and keeps the v1/v2 +header parser inside `kv` where it already lives. Neither package +imports the other in production. The skip path deliberately avoids the +restore path's body CRC work because it does not consume body bytes. ### 6. Crash-safety argument @@ -617,7 +588,7 @@ window. By forcing `pebble.Sync` on `SetDurableAppliedIndex` we make the checkpoint at least as durable as the snapshot pointer that follows. Cost: +1 extra fsync per snapshot persist (rare; default -`SnapshotCount=10000`). Negligible vs. the savings, and the only +`SnapshotCount=100000`). Negligible vs. the savings, and the only way to keep the round-4 ordering proof intact under nosync mode. #### Encryption opcodes (`OpRegistration`/`OpBootstrap`/`OpRotation`) @@ -768,7 +739,7 @@ func (s *pebbleStore) SetDurableAppliedIndex(idx uint64) error { // because WAL compaction starts at the snapshot index, no future // replay can re-bump metaAppliedIndex from the lost lease/data // applies, and the skip permanently falls back. The +1 extra fsync - // per snapshot persist (rare; default SnapshotCount=10000) is the + // per snapshot persist (rare; default SnapshotCount=100000) is the // right price. return errors.WithStack(b.Commit(pebble.Sync)) } @@ -801,12 +772,12 @@ permanent fallback case is closed. **Cost**: one extra pebble `Batch.Commit` (Sync per `ELASTICKV_FSM_SYNC_MODE`) per snapshot persist. Snapshots fire on the -etcd raft `SnapshotCount` cadence (default 10000 entries), so this is -~one extra fsync per ~10000 entries — negligible. +etcd raft `SnapshotCount` cadence (default 100000 entries), so this is +~one extra fsync per ~100000 entries — negligible. **Why not bump on every HLC lease apply**. Option A (1 pebble batch per lease tick) costs ~1 fsync/sec/group continuously. Option B (the -snapshot-persist hook) costs ~1 fsync per 10000 entries. Both close +snapshot-persist hook) costs ~1 fsync per 100000 entries. Both close the skip gap; B costs ~10⁴× less and aligns with the natural durability boundary the engine already maintains. @@ -903,7 +874,7 @@ restoreSnapshotState skipped (FSM at index %d, snapshot at %d, ceiling=%d, cutov |---|---|---| | **B1** (this PR) | Design doc | None | | **B2** | `ApplyMutationsRaftAt` / `DeletePrefixAtRaftAt` overloads + meta-key bundling in both leaves + `pebbleStore.LastAppliedIndex()` (under `dbMu.RLock()`) + `pebbleStore.SetDurableAppliedIndex()` (under `dbMu.RLock()` + `applyMu.Lock()` RMW monotonic guard, **`pebble.Sync` unconditionally**) + `kvFSM.LastAppliedIndex()` directly satisfies `raftengine.AppliedIndexReader` (compile-time guard in `kv/fsm_applied_index_iface_check.go`) + `kvFSM.SetDurableAppliedIndex` forwarding + thread `f.pendingApplyIdx` into the data-Apply leaves + BOTH `persistCreatedSnapshot` (`engine.go:2679`) AND `e.persistLocalSnapshotPayload` (`engine.go:4032`, the SnapshotCount-triggered hot path) call `SetDurableAppliedIndex` BEFORE the corresponding `persist.SaveSnap` | Meta key starts being written on every data Apply AND at every snapshot persist (both config-snapshot and steady-state local-snapshot paths). Skip is still disabled. Soak in production for one release. | -| **B3** | `restoreSnapshotState` skip gate + `applyHeaderStateOnSkip(snapPath, tok.CRC32C)` orchestrating size + footer-vs-tokenCRC + full-body-CRC verification using `internal/raftengine/etcd`'s existing helpers (matching `openAndRestoreFSMSnapshot`'s safety contract) + two-phase `SnapshotHeaderApplier` seam on `kvFSM` (`ParseSnapshotHeader(r io.Reader) (ceiling, cutover, err)` + pure `ApplySnapshotHeader(ceiling, cutover)`) + metrics + INFO log | **User-visible cold-start win.** | +| **B3** | `restoreSnapshotState` skip gate + `applyHeaderStateOnSkip(snapPath, tok.CRC32C)` orchestrating size + footer-vs-tokenCRC checks using `internal/raftengine/etcd`'s existing helpers + two-phase `SnapshotHeaderApplier` seam on `kvFSM` (`ParseSnapshotHeader(r io.Reader) (ceiling, cutover, err)` + pure `ApplySnapshotHeader(ceiling, cutover)`) + metrics + INFO log. Full-body CRC remains on the execute/full-restore path. | **User-visible cold-start win.** | | **B4** | Lower `HEALTH_TIMEOUT_SECONDS` default once production data shows steady-state skip rate ≥ 90 % | Tighter ceiling; the env override remains honoured. | Each of B2–B3 ships behind tests: @@ -927,8 +898,8 @@ Each of B2–B3 ships behind tests: test asserts that `applyHeaderStateOnSkip` sets `f.hlc.PhysicalCeiling()` **and** `f.restoredCutover` for both v1 and v2 snapshot headers — the ceiling+cutover are invariant under - the optimisation. Three additional CRC-corruption tests (round-6, - one per failure mode) inject the corruption and drive the skip path + the optimisation. Additional envelope/header tests inject corruption + and drive the skip path through `applyHeaderStateOnSkip` (either directly or via `restoreSnapshotState` with `fsmAlreadyAtIndex` returning true), asserting the specific typed error surfaces and that the FSM did @@ -937,16 +908,15 @@ Each of B2–B3 ships behind tests: `ErrFSMSnapshotTooSmall`. - Pair the file with a wrong-token CRC → `ErrFSMSnapshotTokenCRC`. - - Flip one body byte (post-header) → - `ErrFSMSnapshotFileCRC` (round-6's full-body CRC pass catches - this; without round-6 the skip would silently install state - from a corrupt file). + - Flip one body byte (post-header) and assert the skip path still + succeeds without scanning the body. Body corruption is caught by + the full restore path if that path is needed. An idle-cluster integration test runs a 3-node cluster with `ELASTICKV_RAFT_SNAPSHOT_COUNT=10` (overriding the default - `defaultSnapshotEvery = 10000` at `engine.go:93` so the scenario is + `defaultSnapshotEvery = 100000` at `engine.go:93` so the scenario is tractable — at default + `hlcRenewalInterval = 1 s` an idle period - would need ≥ 20 000 s). With the override, the test issues no data + would need ≥ 200 000 s). With the override, the test issues no data writes for `2 × 10 × hlcRenewalInterval = 20 s`, takes a snapshot, restarts a node, and asserts the skip fires — proving the codex round-3 P2 scenario is closed end-to-end through the @@ -1061,7 +1031,7 @@ subsection + B2 row + B2 test list) closes the gap by bumping successful snapshot persist, `LastAppliedIndex >= snapshot.Index` holds unconditionally, so the skip fires reliably on the next restart. Cost: one extra pebble `Batch.Commit` per snapshot persist -(~one extra fsync per `SnapshotCount` entries, default 10000) versus +(~one extra fsync per `SnapshotCount` entries, default 100000) versus Option A's continuous ~1 fsync/sec/group. Lesson: "rare" should be a quantitative claim against the actual @@ -1138,16 +1108,13 @@ and silently apply `ceiling=0, cutover=0` for a too-short file. Round 6 threads `tok.CRC32C` through `applyHeaderStateOnSkip` and `SnapshotHeaderApplier.ApplySnapshotHeaderFromFile`, and the §5 -pseudocode runs the same three-step verification before applying any -side-effect. Cost: one extra full-file read on the skip path — still -strictly cheaper than the restore path it replaces (the same read -happens during `restoreAndComputeCRC`, plus restore additionally -writes a temp Pebble database via `restoreBatchLoopInto`). - -A follow-up optimisation that persists the HLC ceiling + cutover as -durable meta keys (analogous to `metaAppliedIndex`) would let the -skip path elide the file read entirely; flagged under Open Questions -but not part of this proposal. +pseudocode originally ran the same three-step verification before +applying any side-effect. Production profiling later showed that the +extra full-file read is too expensive for multi-GiB snapshots during +restart. The final implementation keeps the cheap fail-closed checks +that matter for the skipped header side-effect (size, footer-vs-token, +header parse), and leaves full-body CRC on the full restore path where +body bytes are consumed. Lesson: when a §X seam takes over a side-effect previously gated by existing fail-closed checks, **inventory those checks first and @@ -1176,12 +1143,11 @@ or force a layering change. Round 7 splits `SnapshotHeaderApplier` into two methods: `ParseSnapshotHeader(r io.Reader) (ceiling, cutover, err)` and the -pure-assignment `ApplySnapshotHeader(ceiling, cutover)`. The CRC -verification orchestration stays in `internal/raftengine/etcd/wal_store.go` -where the helpers already live; the engine wraps the file in a crc32 -TeeReader and hands the reader to `setter.ParseSnapshotHeader`, which -calls the still-in-`kv`-package `ReadSnapshotHeader(*bufio.Reader)` -and drains. No package imports change; no helpers need to be +pure-assignment `ApplySnapshotHeader(ceiling, cutover)`. The envelope +checks stay in `internal/raftengine/etcd/wal_store.go` where the helpers +already live; the engine hands a payload-limited reader to +`setter.ParseSnapshotHeader`, which calls the still-in-`kv`-package +`ReadSnapshotHeader`. No package imports change; no helpers need to be exported. **P2 line 660 — `SetDurableAppliedIndex` honoured nosync mode.** @@ -1201,7 +1167,7 @@ Round 7 pins `pebbleStore.SetDurableAppliedIndex` to `pebble.Sync` unconditionally. The checkpoint must be at least as durable as the snapshot pointer that immediately follows; there is no raft log entry to replay it from after WAL compaction. Cost: +1 extra fsync per -snapshot persist (rare, default `SnapshotCount=10000`). +snapshot persist (rare, default `SnapshotCount=100000`). Lesson: - (P2 line 410) **Before pseudocoding cross-package helper use, diff --git a/docs/design/2026_06_12_proposed_scaling_roadmap.md b/docs/design/2026_06_12_proposed_scaling_roadmap.md index 6da71b736..4002f2209 100644 --- a/docs/design/2026_06_12_proposed_scaling_roadmap.md +++ b/docs/design/2026_06_12_proposed_scaling_roadmap.md @@ -124,7 +124,7 @@ Storage breakage at 1–10 TB/shard: - Snapshot transfer is a single-stream full-iter scan; at 1 TB it is hours, pins SSTs (blocks compaction reclamation), and any raft snapshot transfer serializes through it. -- WAL replay between `defaultSnapshotEvery = 10 000` raft entries +- WAL replay between `defaultSnapshotEvery = 100 000` raft entries can be multi-GiB; with `pebble.Sync` durability path the replay is tens of minutes. - Pebble L0CompactionThreshold / LBaseMaxBytes / compaction diff --git a/internal/raftengine/etcd/engine.go b/internal/raftengine/etcd/engine.go index a273f5e8d..b508b2f00 100644 --- a/internal/raftengine/etcd/engine.go +++ b/internal/raftengine/etcd/engine.go @@ -116,10 +116,10 @@ const ( // defaultSnapshotEvery is the fallback trigger threshold: take an FSM // snapshot once the applied index has advanced this many entries past // the last snapshot's index. etcd/raft itself uses 10_000 as a default, - // but with fat proposal payloads (e.g. Lua scripts) this can produce a - // multi-GiB WAL between snapshots. Operators can lower via + // but multi-GiB FSM snapshots can take tens of seconds to persist and + // contend with the Raft hot path. Operators can lower or raise via // ELASTICKV_RAFT_SNAPSHOT_COUNT without a rebuild. - defaultSnapshotEvery = 10_000 + defaultSnapshotEvery = 100_000 snapshotEveryEnvVar = "ELASTICKV_RAFT_SNAPSHOT_COUNT" defaultSnapshotQueueSize = 1 defaultAdminPollInterval = 10 * time.Millisecond @@ -1250,12 +1250,11 @@ func (e *Engine) recordDispatchErrorCode(code string) uint64 { } // StepQueueFullCount returns the total number of inbound raft messages -// that could not be enqueued into the selected inbound step queue -// because the channel was at capacity. This is the "etcd raft inbound -// step queue is full" signal from the task description: a spike -// indicates the local raft loop is starved, usually by something -// blocking the apply path such as -// the pre-#560 rawKeyTypeAt seek storm. +// that found the selected inbound step queue at capacity. Blocking inbound +// message classes wait for space after incrementing this counter; best-effort +// classes still return errStepQueueFull. A spike indicates the local raft loop +// is starved, usually by something blocking the apply path such as the +// pre-#560 rawKeyTypeAt seek storm. func (e *Engine) StepQueueFullCount() uint64 { if e == nil { return 0 @@ -2298,6 +2297,7 @@ func (e *Engine) handleStep(msg raftpb.Message) { commitBeforeStep := e.rawNode.Status().GetCommit() if err := e.rawNode.Step(&msg); err != nil { if errors.Is(err, etcdraft.ErrStepPeerNotFound) { + e.removeReceivedFSMSnapshotToken(msg) e.unprotectReceivedFSMSnapshotToken(msg) return } @@ -2305,9 +2305,11 @@ func (e *Engine) handleStep(msg raftpb.Message) { return } if e.unprotectReceivedFSMSnapshotTokenIfCommitted(msg, commitBeforeStep) { + e.removeReceivedFSMSnapshotToken(msg) return } if !e.rawNode.HasReady() { + e.removeReceivedFSMSnapshotToken(msg) e.unprotectReceivedFSMSnapshotToken(msg) return } @@ -2518,7 +2520,7 @@ func isInboundPriorityMsg(t raftpb.MessageType) bool { } func isBlockingInboundStepMsg(t raftpb.MessageType) bool { - return t == raftpb.MsgSnap + return t == raftpb.MsgSnap || t == raftpb.MsgApp || t == raftpb.MsgAppResp } // selectDispatchLane picks the per-peer channel for msgType. In the legacy @@ -3475,6 +3477,7 @@ func (e *Engine) releaseIgnoredReceivedFSMSnapshotSteps(rd etcdraft.Ready) { if index == readySnapshotIndex { continue } + e.removeReceivedFSMSnapshotIndex(index) for i := 0; i < count; i++ { e.unprotectReceivedFSMSnapshot(index) } @@ -3489,6 +3492,23 @@ func (e *Engine) unprotectReceivedFSMSnapshotToken(msg raftpb.Message) { e.unprotectReceivedFSMSnapshot(index) } +func (e *Engine) removeReceivedFSMSnapshotToken(msg raftpb.Message) { + index, ok := receivedFSMSnapshotTokenIndex(msg) + if !ok { + return + } + e.removeReceivedFSMSnapshotIndex(index) +} + +func (e *Engine) removeReceivedFSMSnapshotIndex(index uint64) { + if e == nil || e.fsmSnapDir == "" || index == 0 { + return + } + e.snapshotMu.Lock() + defer e.snapshotMu.Unlock() + removeWithWarn(fsmSnapPath(e.fsmSnapDir, index), "ignored received fsm snapshot") +} + func receivedFSMSnapshotTokenIndex(msg raftpb.Message) (uint64, bool) { if msg.GetType() != raftpb.MsgSnap || msg.GetSnapshot() == nil || !isSnapshotToken(msg.GetSnapshot().GetData()) { return 0, false @@ -4412,6 +4432,13 @@ func (e *Engine) stepChannelFor(msgType raftpb.MessageType) chan raftpb.Message } func (e *Engine) enqueueBlockingStep(ctx context.Context, ch chan raftpb.Message, msg raftpb.Message) error { + select { + case ch <- msg: + return nil + default: + e.stepQueueFullCount.Add(1) + } + select { case <-ctx.Done(): return errors.WithStack(ctx.Err()) diff --git a/internal/raftengine/etcd/engine_applied_index_test.go b/internal/raftengine/etcd/engine_applied_index_test.go index cf4596306..6d8d3e48f 100644 --- a/internal/raftengine/etcd/engine_applied_index_test.go +++ b/internal/raftengine/etcd/engine_applied_index_test.go @@ -268,8 +268,29 @@ func TestReleaseIgnoredReceivedFSMSnapshotStepsUnprotectsNonSnapshotReady(t *tes require.Empty(t, e.pendingReceivedFSMSnapshotStep) } +func TestReleaseIgnoredReceivedFSMSnapshotStepsRemovesIgnoredFSMFile(t *testing.T) { + fsmSnapDir := t.TempDir() + writeFSMFileForTest(t, fsmSnapDir, 10, []byte("ignored snapshot")) + e := &Engine{ + fsmSnapDir: fsmSnapDir, + protectedReceivedFSMSnaps: map[uint64]int{10: 1}, + pendingReceivedFSMSnapshotStep: map[uint64]int{ + 10: 1, + }, + } + + e.releaseIgnoredReceivedFSMSnapshotSteps(etcdraft.Ready{}) + + require.NoFileExists(t, fsmSnapPath(fsmSnapDir, 10)) + require.Empty(t, e.protectedReceivedFSMSnaps) + require.Empty(t, e.pendingReceivedFSMSnapshotStep) +} + func TestReleaseIgnoredReceivedFSMSnapshotStepsKeepsSnapshotReadyProtected(t *testing.T) { + fsmSnapDir := t.TempDir() + writeFSMFileForTest(t, fsmSnapDir, 10, []byte("accepted snapshot")) e := &Engine{ + fsmSnapDir: fsmSnapDir, protectedReceivedFSMSnaps: map[uint64]int{10: 1}, pendingReceivedFSMSnapshotStep: map[uint64]int{ 10: 1, @@ -283,6 +304,7 @@ func TestReleaseIgnoredReceivedFSMSnapshotStepsKeepsSnapshotReadyProtected(t *te require.Equal(t, map[uint64]int{10: 1}, e.protectedReceivedFSMSnaps) require.Empty(t, e.pendingReceivedFSMSnapshotStep) + require.FileExists(t, fsmSnapPath(fsmSnapDir, 10)) } // TestRecordingFSM_SatisfiesAppliedIndexWriter is a compile-time- diff --git a/internal/raftengine/etcd/engine_test.go b/internal/raftengine/etcd/engine_test.go index ce81a539d..713862a9c 100644 --- a/internal/raftengine/etcd/engine_test.go +++ b/internal/raftengine/etcd/engine_test.go @@ -599,7 +599,7 @@ func TestHandleTransportMessageWaitsForStartup(t *testing.T) { require.NoError(t, <-errCh) } -func TestEnqueueStepReturnsQueueFull(t *testing.T) { +func TestEnqueueStepBestEffortReturnsQueueFull(t *testing.T) { engine := &Engine{ doneCh: make(chan struct{}), stepCh: make(chan raftpb.Message, 1), @@ -608,20 +608,66 @@ func TestEnqueueStepReturnsQueueFull(t *testing.T) { require.Equal(t, uint64(0), engine.StepQueueFullCount()) - err := engine.enqueueStep(context.Background(), raftpb.Message{Type: messageTypePtr(raftpb.MsgApp)}) + err := engine.enqueueStep(context.Background(), raftpb.Message{Type: messageTypePtr(raftpb.MsgStorageAppend)}) require.Error(t, err) require.True(t, errors.Is(err, errStepQueueFull)) // The Prometheus hot-path dashboard relies on StepQueueFullCount - // advancing exactly once per rejected enqueue so the scraped rate - // equals the true drop rate, not a multiple of it. + // advancing exactly once per full enqueue attempt so the scraped + // rate equals the true congestion rate, not a multiple of it. require.Equal(t, uint64(1), engine.StepQueueFullCount()) - err = engine.enqueueStep(context.Background(), raftpb.Message{Type: messageTypePtr(raftpb.MsgApp)}) + err = engine.enqueueStep(context.Background(), raftpb.Message{Type: messageTypePtr(raftpb.MsgStorageAppend)}) require.Error(t, err) require.Equal(t, uint64(2), engine.StepQueueFullCount()) } +func TestEnqueueStepMsgAppWaitsForQueueSlot(t *testing.T) { + engine := &Engine{ + doneCh: make(chan struct{}), + stepCh: make(chan raftpb.Message, 1), + } + engine.stepCh <- raftpb.Message{Type: messageTypePtr(raftpb.MsgHeartbeat)} + + errCh := make(chan error, 1) + go func() { + errCh <- engine.enqueueStep(context.Background(), raftpb.Message{Type: messageTypePtr(raftpb.MsgApp)}) + }() + + select { + case err := <-errCh: + t.Fatalf("MsgApp enqueue returned before the queue had room: %v", err) + case <-time.After(20 * time.Millisecond): + } + require.Equal(t, uint64(1), engine.StepQueueFullCount()) + + firstMsg := <-engine.stepCh + require.Equal(t, raftpb.MsgHeartbeat, firstMsg.GetType()) + select { + case err := <-errCh: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("MsgApp enqueue did not resume after the queue had room") + } + nextMsg := <-engine.stepCh + require.Equal(t, raftpb.MsgApp, nextMsg.GetType()) +} + +func TestEnqueueStepMsgAppReturnsContextDeadlineWhenQueueStaysFull(t *testing.T) { + engine := &Engine{ + doneCh: make(chan struct{}), + stepCh: make(chan raftpb.Message, 1), + } + engine.stepCh <- raftpb.Message{Type: messageTypePtr(raftpb.MsgHeartbeat)} + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) + defer cancel() + err := engine.enqueueStep(ctx, raftpb.Message{Type: messageTypePtr(raftpb.MsgApp)}) + require.Error(t, err) + require.True(t, errors.Is(err, context.DeadlineExceeded)) + require.Equal(t, uint64(1), engine.StepQueueFullCount()) +} + func TestEnqueueStepPriorityBypassesFullBulkQueue(t *testing.T) { engine := &Engine{ doneCh: make(chan struct{}), @@ -641,9 +687,11 @@ func TestEnqueueStepPriorityBypassesFullBulkQueue(t *testing.T) { t.Fatal("priority heartbeat was not enqueued") } - err = engine.enqueueStep(context.Background(), raftpb.Message{Type: messageTypePtr(raftpb.MsgApp)}) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) + defer cancel() + err = engine.enqueueStep(ctx, raftpb.Message{Type: messageTypePtr(raftpb.MsgApp)}) require.Error(t, err) - require.True(t, errors.Is(err, errStepQueueFull)) + require.True(t, errors.Is(err, context.DeadlineExceeded)) require.Equal(t, uint64(1), engine.StepQueueFullCount()) } diff --git a/internal/raftengine/etcd/fsm_snapshot_file.go b/internal/raftengine/etcd/fsm_snapshot_file.go index 8e31a05d3..e27bf0b0b 100644 --- a/internal/raftengine/etcd/fsm_snapshot_file.go +++ b/internal/raftengine/etcd/fsm_snapshot_file.go @@ -46,6 +46,12 @@ const ( // fsmWriteBufSize is the bufio.Writer buffer size used when writing .fsm files. fsmWriteBufSize = 1 << 20 // 1 MiB + // defaultMaxRetainedFSMSnapshotBytes bounds retained .fsm payload bytes + // after successful snapshot publication. Large production FSM snapshots can + // be tens of GiB; keeping defaultMaxSnapFiles full copies leaves too little + // headroom for the next receive-side spool. + defaultMaxRetainedFSMSnapshotBytes = int64(16 << 30) // 16 GiB + // fsmMaxInMemPayload is the maximum payload size that readFSMSnapshotPayload // will materialise into memory. Larger snapshots must use the streaming path // (openFSMSnapshotPayloadReader) to avoid OOM. 1 GiB is chosen as a generous @@ -53,6 +59,8 @@ const ( fsmMaxInMemPayload = int64(1 << 30) // 1 GiB ) +const maxRetainedFSMSnapshotBytesEnvVar = "ELASTICKV_RAFT_MAX_RETAINED_FSM_SNAPSHOT_BYTES" + var ( snapshotTokenMagic = [snapshotTokenMagicLen]byte{'E', 'K', 'V', 'T'} crc32cTable = crc32.MakeTable(crc32.Castagnoli) @@ -1081,7 +1089,7 @@ func purgeOldSnapshotFiles(snapDir, fsmSnapDir string) error { } snaps := collectSnapNames(entries) - if len(snaps) <= defaultMaxSnapFiles { + if len(snaps) == 0 { return nil } // Sort explicitly: os.ReadDir returns lexicographic order on most systems, @@ -1089,8 +1097,12 @@ func purgeOldSnapshotFiles(snapDir, fsmSnapDir string) error { // hex, so lexicographic == chronological order (oldest first). sort.Strings(snaps) + maxKeep := retainedSnapshotFileLimit(snaps, fsmSnapDir) + if len(snaps) <= maxKeep { + return nil + } var combined error - for _, name := range snaps[:len(snaps)-defaultMaxSnapFiles] { + for _, name := range snaps[:len(snaps)-maxKeep] { if err := purgeSnapPair(snapDir, fsmSnapDir, name); err != nil { combined = errors.CombineErrors(combined, err) } @@ -1101,6 +1113,60 @@ func purgeOldSnapshotFiles(snapDir, fsmSnapDir string) error { return errors.WithStack(combined) } +func retainedSnapshotFileLimit(snaps []string, fsmSnapDir string) int { + maxKeep := defaultMaxSnapFiles + if maxKeep < 1 { + maxKeep = 1 + } + if maxKeep > len(snaps) { + maxKeep = len(snaps) + } + if fsmSnapDir == "" { + return maxKeep + } + budget := maxRetainedFSMSnapshotBytes() + if budget <= 0 { + return maxKeep + } + for maxKeep > 1 && retainedFSMSnapshotBytes(snaps[len(snaps)-maxKeep:], fsmSnapDir) > budget { + maxKeep-- + } + return maxKeep +} + +func maxRetainedFSMSnapshotBytes() int64 { + raw := strings.TrimSpace(os.Getenv(maxRetainedFSMSnapshotBytesEnvVar)) + if raw == "" { + return defaultMaxRetainedFSMSnapshotBytes + } + n, err := strconv.ParseInt(raw, 10, 64) + if err != nil { + slog.Warn("invalid max retained FSM snapshot bytes; using default", + "env", maxRetainedFSMSnapshotBytesEnvVar, + "value", raw, + "error", err, + ) + return defaultMaxRetainedFSMSnapshotBytes + } + return n +} + +func retainedFSMSnapshotBytes(snaps []string, fsmSnapDir string) int64 { + var total int64 + for _, name := range snaps { + idx := parseSnapFileIndex(name) + if idx == 0 { + continue + } + info, err := os.Stat(fsmSnapPath(fsmSnapDir, idx)) + if err != nil { + continue + } + total += info.Size() + } + return total +} + func collectSnapNames(entries []os.DirEntry) []string { var snaps []string for _, e := range entries { diff --git a/internal/raftengine/etcd/grpc_transport.go b/internal/raftengine/etcd/grpc_transport.go index 29f67129f..ab7fdf335 100644 --- a/internal/raftengine/etcd/grpc_transport.go +++ b/internal/raftengine/etcd/grpc_transport.go @@ -1087,8 +1087,11 @@ func snapshotMessageHeader(msg raftpb.Message) ([]byte, error) { } func sendSnapshotChunks(stream pb.EtcdRaft_SendSnapshotClient, header []byte, payload []byte, chunkSize int) error { + if err := sendSnapshotChunk(stream, &pb.EtcdRaftSnapshotChunk{Metadata: header}); err != nil { + return err + } if len(payload) == 0 { - return sendSnapshotChunk(stream, &pb.EtcdRaftSnapshotChunk{Metadata: header, Final: true}) + return sendSnapshotChunk(stream, &pb.EtcdRaftSnapshotChunk{Final: true}) } for offset := 0; offset < len(payload); offset += chunkSize { end := offset + chunkSize @@ -1099,9 +1102,6 @@ func sendSnapshotChunks(stream pb.EtcdRaft_SendSnapshotClient, header []byte, pa Chunk: payload[offset:end], Final: end == len(payload), } - if offset == 0 { - chunk.Metadata = header - } if err := sendSnapshotChunk(stream, chunk); err != nil { return err } @@ -1120,6 +1120,9 @@ func sendSnapshotReaderChunks(stream pb.EtcdRaft_SendSnapshotClient, header []by if chunkSize <= 0 { chunkSize = defaultSnapshotChunkSize } + if err := sendSnapshotChunk(stream, &pb.EtcdRaftSnapshotChunk{Metadata: header}); err != nil { + return err + } buffered := bufio.NewReaderSize(reader, chunkSize) current, err := readSnapshotChunk(buffered, chunkSize) if err != nil { @@ -1129,14 +1132,13 @@ func sendSnapshotReaderChunks(stream pb.EtcdRaft_SendSnapshotClient, header []by // a small snapshot. Include it so the receiver does not get an // empty snapshot.Data. return sendSnapshotChunk(stream, &pb.EtcdRaftSnapshotChunk{ - Metadata: header, - Chunk: current, - Final: true, + Chunk: current, + Final: true, }) } return errors.WithStack(err) } - return streamReaderChunks(stream, header, buffered, current, chunkSize) + return streamReaderChunks(stream, nil, buffered, current, chunkSize) } // streamReaderChunks drains buffered starting from `current` (the first full diff --git a/internal/raftengine/etcd/grpc_transport_test.go b/internal/raftengine/etcd/grpc_transport_test.go index b7e43c703..a57fc43bb 100644 --- a/internal/raftengine/etcd/grpc_transport_test.go +++ b/internal/raftengine/etcd/grpc_transport_test.go @@ -1549,10 +1549,13 @@ func TestSendSnapshotReaderChunksSmallPayloadPreservesData(t *testing.T) { err := sendSnapshotReaderChunks(client, header, bytes.NewReader(payload), defaultSnapshotChunkSize) require.NoError(t, err) - require.Len(t, client.chunks, 1) + require.Len(t, client.chunks, 2) require.Equal(t, header, client.chunks[0].Metadata) - require.Equal(t, payload, client.chunks[0].Chunk) - require.True(t, client.chunks[0].Final) + require.Empty(t, client.chunks[0].Chunk) + require.False(t, client.chunks[0].Final) + require.Empty(t, client.chunks[1].Metadata) + require.Equal(t, payload, client.chunks[1].Chunk) + require.True(t, client.chunks[1].Final) } func TestSendSnapshotReaderChunksEmptyPayloadSendsHeaderOnly(t *testing.T) { @@ -1562,10 +1565,13 @@ func TestSendSnapshotReaderChunksEmptyPayloadSendsHeaderOnly(t *testing.T) { err := sendSnapshotReaderChunks(client, header, bytes.NewReader(nil), defaultSnapshotChunkSize) require.NoError(t, err) - require.Len(t, client.chunks, 1) + require.Len(t, client.chunks, 2) require.Equal(t, header, client.chunks[0].Metadata) require.Empty(t, client.chunks[0].Chunk) - require.True(t, client.chunks[0].Final) + require.False(t, client.chunks[0].Final) + require.Empty(t, client.chunks[1].Metadata) + require.Empty(t, client.chunks[1].Chunk) + require.True(t, client.chunks[1].Final) } // testSnapshotSendClient captures chunks sent via sendSnapshotReaderChunks / sendSnapshotChunks. @@ -1718,19 +1724,23 @@ func TestSendSnapshotReaderChunksMultiChunk(t *testing.T) { err := sendSnapshotReaderChunks(client, header, bytes.NewReader(payload), chunkSize) require.NoError(t, err) - require.Len(t, client.chunks, 3) + require.Len(t, client.chunks, 4) require.Equal(t, header, client.chunks[0].Metadata) - require.Equal(t, []byte("1234"), client.chunks[0].Chunk) + require.Empty(t, client.chunks[0].Chunk) require.False(t, client.chunks[0].Final) require.Empty(t, client.chunks[1].Metadata) - require.Equal(t, []byte("abcd"), client.chunks[1].Chunk) + require.Equal(t, []byte("1234"), client.chunks[1].Chunk) require.False(t, client.chunks[1].Final) require.Empty(t, client.chunks[2].Metadata) - require.Equal(t, []byte("5678"), client.chunks[2].Chunk) - require.True(t, client.chunks[2].Final) + require.Equal(t, []byte("abcd"), client.chunks[2].Chunk) + require.False(t, client.chunks[2].Final) + + require.Empty(t, client.chunks[3].Metadata) + require.Equal(t, []byte("5678"), client.chunks[3].Chunk) + require.True(t, client.chunks[3].Final) } func TestSendSnapshotReaderChunksExactBoundary(t *testing.T) { @@ -1743,13 +1753,16 @@ func TestSendSnapshotReaderChunksExactBoundary(t *testing.T) { err := sendSnapshotReaderChunks(client, header, bytes.NewReader(payload), chunkSize) require.NoError(t, err) - require.Len(t, client.chunks, 2) + require.Len(t, client.chunks, 3) require.Equal(t, header, client.chunks[0].Metadata) - require.Equal(t, []byte("1234"), client.chunks[0].Chunk) + require.Empty(t, client.chunks[0].Chunk) require.False(t, client.chunks[0].Final) - require.Equal(t, []byte("5678"), client.chunks[1].Chunk) - require.True(t, client.chunks[1].Final) + require.Equal(t, []byte("1234"), client.chunks[1].Chunk) + require.False(t, client.chunks[1].Final) + + require.Equal(t, []byte("5678"), client.chunks[2].Chunk) + require.True(t, client.chunks[2].Final) } // TestSendSnapshotReaderChunksTrailingPartialChunk regressions a production @@ -1772,17 +1785,20 @@ func TestSendSnapshotReaderChunksTrailingPartialChunk(t *testing.T) { err := sendSnapshotReaderChunks(client, header, bytes.NewReader(payload), chunkSize) require.NoError(t, err) - require.Len(t, client.chunks, 3, "expected two full chunks plus a trailing partial") + require.Len(t, client.chunks, 4, "expected metadata, two full chunks, and a trailing partial") require.Equal(t, header, client.chunks[0].Metadata) - require.Equal(t, []byte("1234"), client.chunks[0].Chunk) + require.Empty(t, client.chunks[0].Chunk) require.False(t, client.chunks[0].Final) - require.Equal(t, []byte("5678"), client.chunks[1].Chunk) + require.Equal(t, []byte("1234"), client.chunks[1].Chunk) require.False(t, client.chunks[1].Final) - require.Equal(t, []byte("X"), client.chunks[2].Chunk) - require.True(t, client.chunks[2].Final) + require.Equal(t, []byte("5678"), client.chunks[2].Chunk) + require.False(t, client.chunks[2].Final) + + require.Equal(t, []byte("X"), client.chunks[3].Chunk) + require.True(t, client.chunks[3].Final) var delivered []byte for _, c := range client.chunks { diff --git a/internal/raftengine/etcd/snapshot_spool_test.go b/internal/raftengine/etcd/snapshot_spool_test.go index b0facd96d..211077312 100644 --- a/internal/raftengine/etcd/snapshot_spool_test.go +++ b/internal/raftengine/etcd/snapshot_spool_test.go @@ -374,6 +374,48 @@ func TestPurgeOldSnapFiles(t *testing.T) { require.NoError(t, err) } +func TestPurgeOldSnapFilesHonorsFSMByteBudget(t *testing.T) { + t.Setenv(maxRetainedFSMSnapshotBytesEnvVar, "100") + + snapDir := t.TempDir() + fsmSnapDir := t.TempDir() + + for i := uint64(1); i <= 3; i++ { + index := i * 10000 + createSnapFile(t, snapDir, index) + require.NoError(t, os.WriteFile(fsmSnapPath(fsmSnapDir, index), bytes.Repeat([]byte("x"), 60), 0o600)) + } + + require.NoError(t, purgeOldSnapshotFiles(snapDir, fsmSnapDir)) + + require.NoFileExists(t, filepath.Join(snapDir, fmt.Sprintf("%016x-%016x.snap", 1, uint64(10000)))) + require.NoFileExists(t, filepath.Join(snapDir, fmt.Sprintf("%016x-%016x.snap", 1, uint64(20000)))) + require.FileExists(t, filepath.Join(snapDir, fmt.Sprintf("%016x-%016x.snap", 1, uint64(30000)))) + require.NoFileExists(t, fsmSnapPath(fsmSnapDir, 10000)) + require.NoFileExists(t, fsmSnapPath(fsmSnapDir, 20000)) + require.FileExists(t, fsmSnapPath(fsmSnapDir, 30000)) +} + +func TestPurgeOldSnapFilesBudgetCanBeDisabled(t *testing.T) { + t.Setenv(maxRetainedFSMSnapshotBytesEnvVar, "0") + + snapDir := t.TempDir() + fsmSnapDir := t.TempDir() + + for i := uint64(1); i <= 4; i++ { + index := i * 10000 + createSnapFile(t, snapDir, index) + require.NoError(t, os.WriteFile(fsmSnapPath(fsmSnapDir, index), bytes.Repeat([]byte("x"), 60), 0o600)) + } + + require.NoError(t, purgeOldSnapshotFiles(snapDir, fsmSnapDir)) + + require.NoFileExists(t, filepath.Join(snapDir, fmt.Sprintf("%016x-%016x.snap", 1, uint64(10000)))) + require.FileExists(t, filepath.Join(snapDir, fmt.Sprintf("%016x-%016x.snap", 1, uint64(20000)))) + require.FileExists(t, filepath.Join(snapDir, fmt.Sprintf("%016x-%016x.snap", 1, uint64(30000)))) + require.FileExists(t, filepath.Join(snapDir, fmt.Sprintf("%016x-%016x.snap", 1, uint64(40000)))) +} + func TestPurgeOldSnapFilesUnderLimit(t *testing.T) { snapDir := t.TempDir() fsmSnapDir := t.TempDir() diff --git a/internal/raftengine/etcd/wal_store.go b/internal/raftengine/etcd/wal_store.go index f6fcc8a44..2a38ac3d0 100644 --- a/internal/raftengine/etcd/wal_store.go +++ b/internal/raftengine/etcd/wal_store.go @@ -2,7 +2,6 @@ package etcd import ( "bytes" - "hash/crc32" "io" "os" "path/filepath" @@ -142,22 +141,20 @@ func loadWalState(logger *zap.Logger, walDir, snapDir, fsmSnapDir string, fsm St return nil, err } - // Codex P1 #934: open the WAL BEFORE the skip-gate decision so we - // know the post-snapshot entry tail. The skip path is only safe - // when the FSM is at least as advanced as the last WAL entry; if - // the FSM is past `tok.Index` but the WAL still carries entries - // `tok.Index+1 .. have` (the normal interval between snapshots, - // since metaAppliedIndex advances on each Apply), those entries - // would re-apply onto a Pebble store that already contains them, - // hitting OCC conflicts and leaving the HLC below timestamps - // already on disk. Compute the WAL tail's last index and gate - // the skip on `have >= lastWalIndex`. + // Open the WAL before the skip-gate decision so we know the committed + // replay target for metrics and raw-node initialization. The skip + // gate itself only requires the FSM to be at least as fresh as the + // persisted snapshot index: Engine.Open seeds e.applied with the + // FSM's durable applied index, rawNodeAppliedForOpen trims the RawNode + // applied pointer when volatile/conf-change entries must replay, and + // applyNormalCommitted drops data-mutating duplicates while applying + // the remaining committed tail. w, hardState, entries, err := openAndReadWALWithRepair(logger, walDir, walSnapshotFor(snapshot)) if err != nil { return nil, err } - lastCommittedIndex := coldStartSkipThreshold(snapshot, hardState) - effectiveApplied, err := restoreSnapshotState(fsm, snapshot, lastCommittedIndex, fsmSnapDir, obs, logger) + committedReplayTarget := coldStartReplayTarget(snapshot, hardState) + effectiveApplied, err := restoreSnapshotState(fsm, snapshot, committedReplayTarget, fsmSnapDir, obs, logger) if err != nil { if closeErr := w.Close(); closeErr != nil { logger.Warn("WAL close failed after restoreSnapshotState error", @@ -221,25 +218,26 @@ func reportColdStartExecute(obs raftengine.ColdStartObserver, logger *zap.Logger ) } -// coldStartSkipThreshold returns the maximum log index the cold- -// start replay can deliver via Ready.CommittedEntries on this -// node: max(snapshot.Metadata.Index, hardState.Commit). The skip -// gate compares the FSM's durable applied index against this -// value; skip is only safe when the FSM is at least this fresh. +// coldStartReplayTarget returns the maximum log index cold-start replay +// can deliver via Ready.CommittedEntries on this node: +// max(snapshot.Metadata.Index, hardState.Commit). The skip gate uses the +// snapshot index for safety; this target is still reported in metrics and +// used by RawNode initialization to replay only the entries still needed +// after e.applied is seeded from the FSM's durable applied index. // // Followers can carry an UNCOMMITTED WAL suffix // (entries[n-1].Index > hardState.Commit). Raft does NOT surface // those entries in CommittedEntries until the leader confirms -// them. The previous gate used the WAL tail (entries[n-1].Index) +// them. The previous target used the WAL tail (entries[n-1].Index) // which forced a multi-GiB restore on every restart of any // follower with an uncommitted suffix, defeating the cold-start -// optimization. Codex P2 #934 round 3. +// optimization. // // The lower bound stays at the snapshot pointer because an empty // WAL still requires the FSM to be at least at the snapshot // index. Raft's invariant guarantees hardState.Commit >= 0; we do // not need to bound from below explicitly beyond snap.Index. -func coldStartSkipThreshold(snapshot raftpb.Snapshot, hardState raftpb.HardState) uint64 { +func coldStartReplayTarget(snapshot raftpb.Snapshot, hardState raftpb.HardState) uint64 { threshold := snapshot.GetMetadata().GetIndex() if hardState.GetCommit() > threshold { threshold = hardState.GetCommit() @@ -356,12 +354,12 @@ func loadPersistedSnapshot(logger *zap.Logger, walDir string, snapshotter *snap. // <= have or the Pebble store would observe them twice (OCC // conflicts; HLC ceiling inversion). The execute path returns // snapshot.Metadata.Index to leave engine behaviour unchanged. -func restoreSnapshotState(fsm StateMachine, snapshot raftpb.Snapshot, lastWalIndex uint64, fsmSnapDir string, obs raftengine.ColdStartObserver, logger *zap.Logger) (uint64, error) { +func restoreSnapshotState(fsm StateMachine, snapshot raftpb.Snapshot, replayTarget uint64, fsmSnapDir string, obs raftengine.ColdStartObserver, logger *zap.Logger) (uint64, error) { if etcdraft.IsEmptySnap(&snapshot) || len(snapshot.Data) == 0 || fsm == nil { return 0, nil } if isSnapshotToken(snapshot.Data) { - return restoreSnapshotStateFromToken(fsm, snapshot, lastWalIndex, fsmSnapDir, obs, logger) + return restoreSnapshotStateFromToken(fsm, snapshot, replayTarget, fsmSnapDir, obs, logger) } // Legacy format: full FSM payload embedded in snapshot.Data. if err := fsm.Restore(bytes.NewReader(snapshot.Data)); err != nil { @@ -376,32 +374,33 @@ func restoreSnapshotState(fsm StateMachine, snapshot raftpb.Snapshot, lastWalInd // - skip path: `have` (FSM is already past snapshot.Metadata.Index) // - execute path: snapshot.Metadata.Index (restored from snapshot) // -// The skip threshold is lastWalIndex (NOT tok.Index): the FSM must be -// at least as fresh as the last WAL entry the cold-start replay would -// deliver, otherwise entries between tok.Index and have would re-apply -// onto a Pebble store that already contains them. Codex P1 #934. +// The skip threshold is tok.Index: if the FSM already contains the +// snapshot state, the engine can seed e.applied with the FSM's durable +// applied index and let normal replay apply any committed tail after +// that point. Entries at or below e.applied are handled by the existing +// duplicate-replay seam in applyCommitted. // // Metrics + log fire AFTER the restore-side work succeeds (coderabbit // Major #934): a header/CRC failure must not register a "successful" // outcome in the soak metrics. -func restoreSnapshotStateFromToken(fsm StateMachine, snapshot raftpb.Snapshot, lastWalIndex uint64, fsmSnapDir string, obs raftengine.ColdStartObserver, logger *zap.Logger) (uint64, error) { +func restoreSnapshotStateFromToken(fsm StateMachine, snapshot raftpb.Snapshot, replayTarget uint64, fsmSnapDir string, obs raftengine.ColdStartObserver, logger *zap.Logger) (uint64, error) { tok, err := decodeSnapshotToken(snapshot.Data) if err != nil { return 0, err } - decision, have := decideSkipOutcome(fsm, lastWalIndex) + decision, have := decideSkipOutcome(fsm, tok.Index) snapPath := fsmSnapPath(fsmSnapDir, tok.Index) if decision == coldStartSkip { if err := applyHeaderStateOnSkip(fsm, snapPath, tok.CRC32C); err != nil { return 0, err } - reportColdStart(obs, logger, decision, tok.Index, lastWalIndex, have) + reportColdStart(obs, logger, decision, tok.Index, replayTarget, have) return have, nil } if err := openAndRestoreFSMSnapshot(fsm, snapPath, tok.CRC32C); err != nil { return 0, err } - reportColdStart(obs, logger, decision, tok.Index, lastWalIndex, have) + reportColdStart(obs, logger, decision, tok.Index, replayTarget, have) return snapshot.GetMetadata().GetIndex(), nil } @@ -461,8 +460,8 @@ func reportColdStart(obs raftengine.ColdStartObserver, logger *zap.Logger, d col case coldStartSkip: // Observer contract (cold_start.go + monitoring/cold_start.go // Prometheus impl): args are (snapshotIndex, haveAppliedIndex); - // gauges compute have-snapIndex. Codex P2 + coderabbit Major - // #934: do NOT pass target/lastWalIndex here or the exported + // gauges compute have-snapIndex. Do NOT pass the replay target + // here or the exported // gauge measures the wrong baseline. if obs != nil { obs.RestoreSkipped(snapIndex, have) @@ -501,21 +500,18 @@ func reportColdStart(obs raftengine.ColdStartObserver, logger *zap.Logger, d col } } -// applyHeaderStateOnSkip mirrors openAndRestoreFSMSnapshot's safety -// contract (size + footer-vs-tokenCRC + full-body-CRC) but applies -// only the header side-effects (HLC ceiling + Stage 8a cutover) -// instead of running the body restore. The body bytes are read for -// CRC coverage but discarded -- fsm.db already holds equivalent -// state, which is precisely the reason we're skipping the restore. +// applyHeaderStateOnSkip validates the cheap snapshot envelope checks +// (size + footer-vs-tokenCRC) and applies only the header side-effects +// (HLC ceiling + Stage 8a cutover) instead of running the body restore. +// The body bytes are not read here: fsm.db already holds equivalent state, +// which is precisely the reason we're skipping the restore. Full-body CRC +// verification still runs on the execute/full-restore path where body +// bytes are actually consumed. // // FSMs that do not implement raftengine.SnapshotHeaderApplier // silently no-op the apply phase -- the FSM has no header state to -// carry forward, and the CRC verification still runs (with no -// observable side-effect on success). On any verification failure -// the typed error propagates and FSM state stays untouched. -// -// See PR #910 design §5 round-7 (two-phase seam) + round-6 -// (three-step CRC mirroring openAndRestoreFSMSnapshot). +// carry forward. On any envelope or header parse failure the typed error +// propagates and FSM state stays untouched. func applyHeaderStateOnSkip(fsm StateMachine, snapPath string, tokenCRC uint32) error { file, err := os.Open(snapPath) if err != nil { @@ -527,61 +523,30 @@ func applyHeaderStateOnSkip(fsm StateMachine, snapPath string, tokenCRC uint32) if err != nil { return errors.WithStack(err) } - footer, err := verifyFSMSnapshotPrefix(file, info.Size(), snapPath, tokenCRC) - if err != nil { + if _, err := verifyFSMSnapshotPrefix(file, info.Size(), snapPath, tokenCRC); err != nil { return err } - // Step 3: full-body CRC. Wrap the payload in a crc32 TeeReader - // and hand it to the FSM's ParseSnapshotHeader for header parse - // + drain. Every payload byte flows through h, matching - // restoreAndComputeCRC's boundary in openAndRestoreFSMSnapshot. - // - // Error-ordering contract (claude #934 R1-F1): header parse - // errors surface BEFORE the body-CRC compare runs, so callers - // (the skip-gate fallback in restoreSnapshotState) may observe - // either an ErrSnapshotHeaderUnknownMagic / InvalidLength chain or - // an ErrFSMSnapshotFileCRC chain depending on which check fails - // first. This is the same ordering openAndRestoreFSMSnapshot has - // — both errors are equally fatal for the skip path (they signal - // snapshot file corruption) and both must propagate without ever - // calling ApplySnapshotHeader. The CRC check stays AFTER the - // header parse so the TeeReader has actually been drained before - // we read h.Sum32(); inverting the order would let a CRC mismatch - // surface on a truncated body even when the header was valid, - // muddying the operator-facing diagnostic. + setter, hasSetter := fsm.(raftengine.SnapshotHeaderApplier) + if !hasSetter { + return nil + } + if _, err := file.Seek(0, io.SeekStart); err != nil { return errors.WithStack(err) } payloadSize := info.Size() - fsmFooterSize - h := crc32.New(crc32cTable) - tee := io.TeeReader(io.LimitReader(file, payloadSize), h) - - setter, hasSetter := fsm.(raftengine.SnapshotHeaderApplier) - ceiling, cutover, err := readSnapshotHeaderOrDrain(setter, hasSetter, tee) + ceiling, cutover, err := setter.ParseSnapshotHeader(io.LimitReader(file, payloadSize)) if err != nil { - return err - } - - if h.Sum32() != footer { - return errors.Wrapf(ErrFSMSnapshotFileCRC, - "path=%s footer=%08x computed=%08x", snapPath, footer, h.Sum32()) - } - - // All three checks passed; apply side-effects (pure assignment - // in the FSM). Skipped silently when the FSM does not expose - // the seam. - if hasSetter { - setter.ApplySnapshotHeader(ceiling, cutover) + return errors.WithStack(err) } + setter.ApplySnapshotHeader(ceiling, cutover) return nil } -// verifyFSMSnapshotPrefix runs the first two cheap checks of -// openAndRestoreFSMSnapshot's three-step contract: size and -// footer-vs-tokenCRC. Returns the on-disk footer value (caller -// reuses it for the step-3 full-body CRC compare). Typed errors -// surface unchanged. +// verifyFSMSnapshotPrefix runs the cheap checks shared with +// openAndRestoreFSMSnapshot: size and footer-vs-tokenCRC. Returns the +// on-disk footer value. Typed errors surface unchanged. func verifyFSMSnapshotPrefix(file *os.File, fileSize int64, snapPath string, tokenCRC uint32) (uint32, error) { if fileSize < fsmMinFileSize { return 0, errors.Wrapf(ErrFSMSnapshotTooSmall, @@ -598,27 +563,6 @@ func verifyFSMSnapshotPrefix(file *os.File, fileSize int64, snapPath string, tok return footer, nil } -// readSnapshotHeaderOrDrain branches on whether the FSM exposes the -// SnapshotHeaderApplier seam: when present, delegate to -// ParseSnapshotHeader (which parses the header AND drains the rest); -// otherwise drain the entire payload through the tee'd reader so the -// CRC pass covers every byte. The (ceiling, cutover) tuple is zero -// in the no-seam case -- the caller's ApplySnapshotHeader branch -// short-circuits on hasSetter, so the zero values are inert. -func readSnapshotHeaderOrDrain(setter raftengine.SnapshotHeaderApplier, hasSetter bool, tee io.Reader) (uint64, uint64, error) { - if hasSetter { - ceiling, cutover, err := setter.ParseSnapshotHeader(tee) - if err != nil { - return 0, 0, errors.WithStack(err) - } - return ceiling, cutover, nil - } - if _, err := io.Copy(io.Discard, tee); err != nil { - return 0, 0, errors.WithStack(err) - } - return 0, 0, nil -} - func walSnapshotFor(snapshot raftpb.Snapshot) walpb.Snapshot { return walpb.Snapshot{ Index: proto.Uint64(snapshot.GetMetadata().GetIndex()), diff --git a/internal/raftengine/etcd/wal_store_skip_gate_test.go b/internal/raftengine/etcd/wal_store_skip_gate_test.go index 212a47ec6..ffde0a0b3 100644 --- a/internal/raftengine/etcd/wal_store_skip_gate_test.go +++ b/internal/raftengine/etcd/wal_store_skip_gate_test.go @@ -53,17 +53,13 @@ func (f *skipGateFSM) ParseSnapshotHeader(r io.Reader) (uint64, uint64, error) { if f.parseErr != nil { return 0, 0, f.parseErr } - // Mimic the real kvFSM contract: parse + drain. We don't actually - // parse a header here; the test fixtures embed magic+ceiling but - // for the gate-level tests we just drain so the CRC matches. + // Mimic the real kvFSM contract: parse only the header. The skip path + // must not drain multi-GiB snapshot bodies just to seed header state. hdrLen := 16 hdr := make([]byte, hdrLen) if n, _ := io.ReadFull(r, hdr); n == hdrLen && bytes.HasPrefix(hdr, []byte("EKVTHLC1")) { f.parsedCeiling = binary.BigEndian.Uint64(hdr[8:16]) } - if _, err := io.Copy(io.Discard, r); err != nil { - return 0, 0, err - } return f.parsedCeiling, 0, nil } @@ -230,13 +226,10 @@ func TestSkipGate_EmitsAfterSuccess(t *testing.T) { require.Empty(t, obs.fallbacks) } -// TestColdStartSkipThreshold pins codex P2 #934 round 3. The -// threshold caps at hardState.Commit so a follower carrying an -// uncommitted WAL suffix is NOT forced to run the full restore -// every restart (the original gate used the WAL tail, which can -// exceed Commit, and raft would not deliver those entries until -// the leader confirmed them). -func TestColdStartSkipThreshold(t *testing.T) { +// TestColdStartReplayTarget verifies the committed replay target caps at +// hardState.Commit so a follower carrying an uncommitted WAL suffix does not +// report or initialize from entries raft cannot deliver yet. +func TestColdStartReplayTarget(t *testing.T) { t.Parallel() mkSnap := func(idx uint64) raftpb.Snapshot { return raftTestSnapshot(idx, 0, nil, nil) @@ -253,36 +246,30 @@ func TestColdStartSkipThreshold(t *testing.T) { {"hardState.Commit zero", mkSnap(100), testHardState(0, 0), 100}, } for _, c := range cases { - got := coldStartSkipThreshold(c.snap, c.hs) + got := coldStartReplayTarget(c.snap, c.hs) if got != c.expected { t.Errorf("%s: got %d, want %d", c.name, got, c.expected) } } } -// TestSkipGate_ExecutesWhenWALCarriesPostSnapshotEntries pins -// codex P1 #934. When the FSM is past tok.Index but the WAL still -// carries entries tok.Index+1 .. have (the normal interval between -// snapshots — metaAppliedIndex advances on each Apply), the skip -// path MUST NOT fire even though have > tok.Index. Those WAL -// entries would re-apply onto a Pebble store that already contains -// them, hitting OCC conflicts and leaving the HLC below timestamps -// already on disk. -// -// Fixture: snap.Index=100, fsm.applied=150, lastWalIndex=150 (the -// WAL has entries 101..150 mirroring the applied tail). Gate -// criterion is have >= lastWalIndex, which holds; that's the -// happy-skip case. To exercise the bug, set lastWalIndex=200 (the -// WAL still has entries 151..200 that have NOT been applied yet); -// have=150 < lastWalIndex=200 must trigger execute, not skip. -func TestSkipGate_ExecutesWhenWALCarriesPostSnapshotEntries(t *testing.T) { +// TestSkipGate_SkipsWhenWALCarriesPostSnapshotTail verifies a node can skip +// the multi-GiB snapshot restore even when the committed WAL has entries past +// the FSM's durable applied index. Engine.Open seeds e.applied with `have`, +// then RawNode/applyCommitted replays only the needed tail and drops +// data-mutating duplicates at or below `have`. +func TestSkipGate_SkipsWhenWALCarriesPostSnapshotTail(t *testing.T) { dir := t.TempDir() const ( snapIndex uint64 = 100 appliedIdx uint64 = 150 - lastWalIndex uint64 = 200 + replayTarget uint64 = 200 + ceilingMs uint64 = 1700_000_000_002 ) - payload := []byte("body-bytes-for-execute") + payload := make([]byte, 16, 16+len("body-bytes-after-header")) + copy(payload[:8], "EKVTHLC1") + binary.BigEndian.PutUint64(payload[8:], ceilingMs) + payload = append(payload, []byte("body-bytes-after-header")...) crc, _ := writeFSMFileForTest(t, dir, snapIndex, payload) fsm := &skipGateFSM{applied: appliedIdx, appliedPresent: true} @@ -291,19 +278,15 @@ func TestSkipGate_ExecutesWhenWALCarriesPostSnapshotEntries(t *testing.T) { Metadata: testSnapshotMetadata(snapIndex, 0, nil), } obs := &recordingObs{} - _, gateErr := restoreSnapshotState(fsm, snap, lastWalIndex, dir, obs, nil) + effective, gateErr := restoreSnapshotState(fsm, snap, replayTarget, dir, obs, nil) require.NoError(t, gateErr) - require.Equal(t, payload, fsm.bodyBytes, - "have(150) < lastWalIndex(200) MUST execute full restore so the WAL replay does not duplicate-apply") - require.False(t, fsm.restoredHeader, "execute path MUST NOT use ApplySnapshotHeader") - require.Empty(t, obs.skipped) - // recordingObs now stores |snapIndex - have| (round-5 fix; mirrors - // monitoring.ColdStartObserver semantics so the FSM-ahead-of- - // snapshot case doesn't underflow). For have(150) > snapIndex(100) - // the absolute gap is 50. - require.Equal(t, []uint64{appliedIdx - snapIndex}, obs.executed, - "observer MUST record the absolute snapshot-relative gap") + require.Empty(t, fsm.bodyBytes, "skip path MUST NOT call fsm.Restore") + require.True(t, fsm.restoredHeader) + require.Equal(t, ceilingMs, fsm.appliedCeiling) + require.Equal(t, appliedIdx, effective) + require.Equal(t, []uint64{appliedIdx - snapIndex}, obs.skipped) + require.Empty(t, obs.executed) require.Empty(t, obs.fallbacks) } @@ -392,30 +375,37 @@ func TestApplyHeaderStateOnSkip_WrongTokenCRC(t *testing.T) { require.False(t, fsm.restoredHeader, "FSM state MUST NOT mutate on verification failure") } -// TestApplyHeaderStateOnSkip_BodyCorruption asserts step 3 catches a -// flipped body byte (CRC mismatch). -func TestApplyHeaderStateOnSkip_BodyCorruption(t *testing.T) { +// TestApplyHeaderStateOnSkip_DoesNotScanBody asserts skip startup cost is +// bounded by the header, not the snapshot body. Body corruption is still +// caught by the full restore path, where the body is actually consumed. +func TestApplyHeaderStateOnSkip_DoesNotScanBody(t *testing.T) { dir := t.TempDir() - crc, path := writeFSMFileForTest(t, dir, 1, []byte("payload-bytes")) + const ceilingMs uint64 = 1700_000_000_001 + payload := make([]byte, 16, 16+len("payload-bytes")) + copy(payload[:8], "EKVTHLC1") + binary.BigEndian.PutUint64(payload[8:], ceilingMs) + payload = append(payload, []byte("payload-bytes")...) + crc, path := writeFSMFileForTest(t, dir, 1, payload) - // Flip the first byte of the body in-place. The footer still + // Flip a byte after the fixed 16-byte header. The footer still // reads as `crc`, but the on-the-wire content no longer matches - // it. Step 2 (footer-vs-token) passes (we pass the same `crc` - // as tokenCRC), step 3 (full-body CRC) fails. + // it. The skip path should still succeed because it never uses body + // bytes; the full restore path remains responsible for body CRC. f, err := os.OpenFile(path, os.O_RDWR, 0) require.NoError(t, err) defer f.Close() var b [1]byte - _, err = f.ReadAt(b[:], 0) + _, err = f.ReadAt(b[:], 16) require.NoError(t, err) b[0] ^= 0x01 - _, err = f.WriteAt(b[:], 0) + _, err = f.WriteAt(b[:], 16) require.NoError(t, err) fsm := &skipGateFSM{} err = applyHeaderStateOnSkip(fsm, path, crc) - require.ErrorIs(t, err, ErrFSMSnapshotFileCRC) - require.False(t, fsm.restoredHeader, "FSM state MUST NOT mutate on verification failure") + require.NoError(t, err) + require.True(t, fsm.restoredHeader) + require.Equal(t, ceilingMs, fsm.appliedCeiling) } // --- kvFSM header preservation contract --- diff --git a/internal/raftengine/statemachine.go b/internal/raftengine/statemachine.go index 9c796f4f8..ec14dc517 100644 --- a/internal/raftengine/statemachine.go +++ b/internal/raftengine/statemachine.go @@ -100,24 +100,22 @@ type AppliedIndexWriter interface { // // The interface is two-phase by design: // -// - ParseSnapshotHeader reads the v1/v2 header from a caller- -// supplied io.Reader (wrapped in a crc32 TeeReader by the -// engine) and drains the remaining bytes so the wrapping CRC -// covers the full payload. It returns the parsed (ceiling, -// cutover) pair WITHOUT mutating FSM state. Errors propagate -// from the underlying header parser -// (ErrSnapshotHeaderUnknownMagic / InvalidLength) or from the -// drain pass (I/O errors); FSM state stays untouched on error. +// - ParseSnapshotHeader reads only the v1/v2 header from a caller- +// supplied io.Reader. It returns the parsed (ceiling, cutover) pair +// WITHOUT mutating FSM state. Errors propagate from the underlying +// header parser (ErrSnapshotHeaderUnknownMagic / InvalidLength); +// FSM state stays untouched on error. // // - ApplySnapshotHeader is pure assignment of the verified header // state. The engine calls this only after ParseSnapshotHeader -// returned successfully AND the wrapping crc32 hash matched -// the file footer. +// returned successfully and the snapshot file's footer matches the +// raft token. // -// Splitting parse from apply lets the CRC verifier stay co-located -// with its private helpers in internal/raftengine/etcd (matching -// the openAndRestoreFSMSnapshot safety contract) while the v1/v2 -// header parser stays inside the kv package where it already lives. +// Splitting parse from apply keeps the v1/v2 header parser inside the +// kv package where it already lives, while the engine remains responsible +// for deciding whether the snapshot body is actually needed. Full-body CRC +// verification still happens on the restore path; the skip path only reads +// the header because the FSM body state is already present locally. // Neither package imports the other in production. type SnapshotHeaderApplier interface { ParseSnapshotHeader(r io.Reader) (ceiling, cutover uint64, err error) diff --git a/kv/coordinator.go b/kv/coordinator.go index 0c1567a26..b32c46716 100644 --- a/kv/coordinator.go +++ b/kv/coordinator.go @@ -36,15 +36,20 @@ const dispatchLeaderRetryInterval = 25 * time.Millisecond // hlcPhysicalWindowMs is the duration in milliseconds that the Raft-agreed // physical ceiling extends ahead of the current wall clock. Modelled after -// TiDB's TSO 3-second window: the leader commits ceiling = now + window, and +// TiDB's TSO window strategy: the leader commits ceiling = now + window, and // renews before the window expires. A new leader inherits the committed ceiling // so it never issues timestamps that collide with the previous leader's window. -const hlcPhysicalWindowMs int64 = 3_000 +const hlcPhysicalWindowMs int64 = 15_000 // hlcRenewalInterval controls how often the leader proposes a new ceiling. // Must be less than hlcPhysicalWindowMs to guarantee the window never expires. const hlcRenewalInterval = 1 * time.Second +// hlcRenewalTimeout bounds a single renewal proposal. It is intentionally +// longer than hlcRenewalInterval so transient Raft write backlog does not +// cancel the renewal before the physical ceiling has real risk of expiring. +const hlcRenewalTimeout = 5 * time.Second + // CoordinatorOption is a functional option for Coordinate constructors. type CoordinatorOption func(*Coordinate) @@ -813,10 +818,10 @@ func (c *Coordinate) extendLeaseAfterRenewal(dispatchStart monoclock.Instant, ex // RunHLCLeaseRenewal runs a background loop that periodically proposes a new // physical ceiling to the Raft cluster while this node is the leader. // -// The ceiling is set to now + hlcPhysicalWindowMs (3 s) and is renewed every -// hlcRenewalInterval (1 s), mirroring TiDB's TSO window strategy. Because the -// window is always at least 2 s ahead of any real timestamp, a new leader will -// never issue timestamps that overlap with the previous leader's window. +// The ceiling is set to now + hlcPhysicalWindowMs and is renewed every +// hlcRenewalInterval, mirroring TiDB's TSO window strategy. Because the window +// stays ahead of real timestamps, a new leader will never issue timestamps that +// overlap with the previous leader's window. // // RunHLCLeaseRenewal blocks until ctx is cancelled; call it in a goroutine. func (c *Coordinate) RunHLCLeaseRenewal(ctx context.Context) { @@ -835,7 +840,10 @@ func (c *Coordinate) RunHLCLeaseRenewal(ctx context.Context) { } if c.IsLeaderAcceptingWrites() { ceilingMs := time.Now().UnixMilli() + hlcPhysicalWindowMs - if err := c.ProposeHLCLease(ctx, ceilingMs); err != nil { + pctx, cancel := context.WithTimeout(ctx, hlcRenewalTimeout) + err := c.ProposeHLCLease(pctx, ceilingMs) + cancel() + if err != nil { c.log.WarnContext(ctx, "hlc lease renewal failed", slog.Int64("ceiling_ms", ceilingMs), slog.Any("err", err), diff --git a/kv/fsm.go b/kv/fsm.go index dedc7829a..34450dab7 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -687,37 +687,30 @@ func (f *kvFSM) RestoredCutover() uint64 { // ParseSnapshotHeader implements raftengine.SnapshotHeaderApplier // phase 1 — the cold-start skip path's parse-without-side-effect -// step. The engine has wrapped `r` in a crc32 TeeReader sized at -// the body payload (file size minus 4-byte footer), so every byte -// pulled from `r` flows through the engine's hash. We read the -// v1/v2 header via ReadSnapshotHeader, then drain the rest of the -// body so the wrapping hash covers every payload byte — matching -// restoreAndComputeCRC's behaviour in openAndRestoreFSMSnapshot. +// step. The skip path only needs the header state because the FSM body +// is already present locally, so this reads the v1/v2 header and leaves +// the remainder untouched. Full-body CRC verification still happens on +// the restore path where the body bytes are consumed. // // IMPORTANT: this method MUST NOT touch f.hlc or f.restoredCutover. // The engine calls ApplySnapshotHeader separately, only after the -// wrapping CRC verification passes. Mutating FSM state here would -// defeat the "no side-effect on CRC failure" contract that the -// PR #910 design §5 round-7 split is designed to preserve. +// snapshot envelope checks pass. Mutating FSM state here would defeat +// the "no side-effect on parse failure" contract that the PR #910 +// design §5 round-7 split is designed to preserve. func (f *kvFSM) ParseSnapshotHeader(r io.Reader) (uint64, uint64, error) { - br := bufio.NewReaderSize(r, 1<<20) //nolint:mnd // 1 MiB, local to kv - ceiling, cutover, err := ReadSnapshotHeader(br) + const headerReadBufferSize = 4 << 10 + + ceiling, cutover, err := ReadSnapshotHeader(bufio.NewReaderSize(r, headerReadBufferSize)) if err != nil { return 0, 0, errors.WithStack(err) } - // Drain the remainder so the engine's TeeReader-wrapped CRC - // covers every byte of the body (LimitReader exhaustion - // signals "full payload consumed" to the caller). - if _, err := io.Copy(io.Discard, br); err != nil { - return 0, 0, errors.WithStack(err) - } return ceiling, cutover, nil } // ApplySnapshotHeader implements raftengine.SnapshotHeaderApplier // phase 2 — pure assignment of the verified header state. Called -// only after ParseSnapshotHeader returned successfully AND the -// engine's wrapping crc32 hash matched the file footer. Mirrors +// only after ParseSnapshotHeader returned successfully and the +// snapshot file's footer matched the raft token. Mirrors // the two side-effects Restore would have applied for the header // portion (HLC physical ceiling + restoredCutover). See PR #910 // design §5 round-7. diff --git a/kv/lease_read_test.go b/kv/lease_read_test.go index d62c103d2..2fc83a299 100644 --- a/kv/lease_read_test.go +++ b/kv/lease_read_test.go @@ -23,7 +23,8 @@ type fakeLeaseEngine struct { linearizableCalls atomic.Int32 proposeErr error // when set, Propose returns it (warm-up failure tests) proposeCalls atomic.Int32 - proposeHook func() // invoked inside Propose before returning (race injection) + proposeHook func() // invoked inside Propose before returning (race injection) + proposeCtxHook func(context.Context) state atomic.Value // stores raftengine.State; default Leader lastQuorumAckMonoNs atomic.Int64 // 0 = no ack yet. Updated by setQuorumAck(). leaderLossCallbacksMu sync.Mutex @@ -62,8 +63,11 @@ func (e *fakeLeaseEngine) Status() raftengine.Status { func (e *fakeLeaseEngine) Configuration(context.Context) (raftengine.Configuration, error) { return raftengine.Configuration{}, nil } -func (e *fakeLeaseEngine) Propose(context.Context, []byte) (*raftengine.ProposalResult, error) { +func (e *fakeLeaseEngine) Propose(ctx context.Context, _ []byte) (*raftengine.ProposalResult, error) { e.proposeCalls.Add(1) + if e.proposeCtxHook != nil { + e.proposeCtxHook(ctx) + } if e.proposeHook != nil { e.proposeHook() } diff --git a/kv/lease_warmup_test.go b/kv/lease_warmup_test.go index 5558a913c..18d64d2d6 100644 --- a/kv/lease_warmup_test.go +++ b/kv/lease_warmup_test.go @@ -169,6 +169,41 @@ func TestCoordinate_RunHLCLeaseRenewal_BlockerSuppressesProposals(t *testing.T) "HLC renewal should resume after startup rotation blocker clears") } +func TestCoordinate_RunHLCLeaseRenewal_UsesRenewalTimeout(t *testing.T) { + eng := &fakeLeaseEngine{applied: 11, leaseDur: time.Hour} + c := NewCoordinatorWithEngine(nil, eng) + deadline := make(chan time.Duration, 1) + eng.proposeCtxHook = func(ctx context.Context) { + d, ok := ctx.Deadline() + if !ok { + deadline <- 0 + return + } + deadline <- time.Until(d) + } + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + c.RunHLCLeaseRenewal(ctx) + close(done) + }() + t.Cleanup(func() { + cancel() + <-done + }) + + select { + case got := <-deadline: + require.Greater(t, got, hlcRenewalInterval, + "HLC renewal proposal deadline must outlive the renewal cadence") + require.LessOrEqual(t, got, hlcRenewalTimeout, + "HLC renewal proposal deadline must remain bounded") + case <-time.After(2 * hlcRenewalInterval): + t.Fatal("timed out waiting for HLC renewal proposal") + } +} + // TestShardedCoordinator_RenewHLCLease_WarmsGroupLease proves the // sharded renewal path warms the target group's lease on a successful // propose, so LeaseReadForKey on a key owned by that group serves from the @@ -294,6 +329,35 @@ func TestShardedCoordinator_RenewHLCLeases_ProposesToEveryLedGroup(t *testing.T) "the non-default group lease must be warmed by all-group renewal") } +func TestShardedCoordinator_RenewHLCLeases_UsesRenewalTimeout(t *testing.T) { + t.Parallel() + eng1 := newShardedLeaseEngine(100) + eng2 := newShardedLeaseEngine(200) + deadline := make(chan time.Duration, 1) + eng1.proposeCtxHook = func(ctx context.Context) { + d, ok := ctx.Deadline() + if !ok { + deadline <- 0 + return + } + deadline <- time.Until(d) + } + coord := mustShardedLeaseCoord(t, eng1, eng2) + + done := coord.renewHLCLeases(context.Background()) + requireRenewalDone(t, done) + + select { + case got := <-deadline: + require.Greater(t, got, hlcRenewalInterval, + "HLC renewal proposal deadline must outlive the renewal cadence") + require.LessOrEqual(t, got, hlcRenewalTimeout, + "HLC renewal proposal deadline must remain bounded") + default: + t.Fatal("missing HLC renewal proposal deadline sample") + } +} + func TestShardedCoordinator_RenewHLCLeases_SkipsNonLeaders(t *testing.T) { t.Parallel() eng1 := newShardedLeaseEngine(100) diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index ec7c43843..124acfb7b 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -2342,7 +2342,7 @@ func (c *ShardedCoordinator) renewHLCLeases(ctx context.Context) <-chan struct{} go func(gid uint64, group *ShardGroup) { defer wg.Done() defer c.finishHLCLeaseRenewal(gid) - pctx, cancel := context.WithTimeout(ctx, hlcRenewalInterval) + pctx, cancel := context.WithTimeout(ctx, hlcRenewalTimeout) defer cancel() c.renewHLCLease(pctx, gid, group) }(gid, group) diff --git a/monitoring/grafana/dashboards/elastickv-redis-summary.json b/monitoring/grafana/dashboards/elastickv-redis-summary.json index 3d05d6176..ddfce4f7c 100644 --- a/monitoring/grafana/dashboards/elastickv-redis-summary.json +++ b/monitoring/grafana/dashboards/elastickv-redis-summary.json @@ -1683,7 +1683,7 @@ ], "title": "Raft Queue Saturation (stepCh full / outbound drops / errors)", "type": "timeseries", - "description": "Counter rates from the etcd raft engine. step-queue-full means inbound messages from remote peers were dropped because the local raft loop was too slow to consume the selected step queue (the 'etcd raft inbound step queue is full' log line). dispatch-dropped means outbound messages were discarded before transport because the per-peer channel was full. dispatch-errors means transport delivery failed. The pre-#560 seek storm caused all three to spike together; watch for them to fall after the rollout and stay flat." + "description": "Counter rates from the etcd raft engine. step-queue-full means inbound messages from remote peers found the selected step queue full; blocking replication messages now wait for space, while best-effort message classes may still be rejected. dispatch-dropped means outbound messages were discarded before transport because the per-peer channel was full. dispatch-errors means transport delivery failed. The pre-#560 seek storm caused all three to spike together; watch for them to fall after the rollout and stay flat." }, { "datasource": "$datasource", diff --git a/monitoring/hotpath.go b/monitoring/hotpath.go index 98b645b26..3fe7a3b51 100644 --- a/monitoring/hotpath.go +++ b/monitoring/hotpath.go @@ -98,7 +98,7 @@ func newHotPathMetrics(registerer prometheus.Registerer) *HotPathMetrics { stepQueueFullTotal: prometheus.NewCounterVec( prometheus.CounterOpts{ Name: "elastickv_raft_step_queue_full_total", - Help: "Inbound raft messages that could not be enqueued because the selected step queue was full; indicates the raft loop is starved (classic pre-#560 seek-storm symptom).", + Help: "Inbound raft messages that found the selected step queue full. Blocking replication messages wait for space; best-effort messages may still be rejected. Indicates the raft loop is starved.", }, []string{"group"}, ), diff --git a/monitoring/hotpath_test.go b/monitoring/hotpath_test.go index fc10cda8a..d865c5983 100644 --- a/monitoring/hotpath_test.go +++ b/monitoring/hotpath_test.go @@ -147,7 +147,7 @@ elastickv_raft_dispatch_dropped_total{group="1",node_address="10.0.0.1:50051",no # HELP elastickv_raft_dispatch_errors_total Outbound raft dispatches that reached the transport but failed. Mirrors etcd raft Engine.dispatchErrorCount. # TYPE elastickv_raft_dispatch_errors_total counter elastickv_raft_dispatch_errors_total{group="1",node_address="10.0.0.1:50051",node_id="n1"} 2 -# HELP elastickv_raft_step_queue_full_total Inbound raft messages that could not be enqueued because the selected step queue was full; indicates the raft loop is starved (classic pre-#560 seek-storm symptom). +# HELP elastickv_raft_step_queue_full_total Inbound raft messages that found the selected step queue full. Blocking replication messages wait for space; best-effort messages may still be rejected. Indicates the raft loop is starved. # TYPE elastickv_raft_step_queue_full_total counter elastickv_raft_step_queue_full_total{group="1",node_address="10.0.0.1:50051",node_id="n1"} 1 `), diff --git a/proxy/config.go b/proxy/config.go index 25cfd2689..f25b240b3 100644 --- a/proxy/config.go +++ b/proxy/config.go @@ -71,7 +71,9 @@ type ProxyConfig struct { SentryEnv string SentrySampleRate float64 MetricsAddr string + PProfAddr string PubSubCompareWindow time.Duration + RedisOnlyRaw bool } // DefaultConfig returns a ProxyConfig with sensible defaults. @@ -86,5 +88,6 @@ func DefaultConfig() ProxyConfig { SentrySampleRate: 1.0, MetricsAddr: ":9191", PubSubCompareWindow: defaultPubSubCompareWindow, + RedisOnlyRaw: true, } } diff --git a/proxy/dualwrite.go b/proxy/dualwrite.go index 474f7d89e..1c940e557 100644 --- a/proxy/dualwrite.go +++ b/proxy/dualwrite.go @@ -25,16 +25,14 @@ const ( // (EVAL / EVALSHA). Lua scripts under high load cause write conflicts in the Raft // layer, and each conflict triggers a full script re-execution. Capping the // concurrency reduces contention so individual scripts complete within - // SecondaryTimeout. Excess secondary script writes may be dropped to keep - // contention bounded; this is only tolerable in modes where the script write - // is targeting the non-authoritative backend. + // SecondaryTimeout. Strict dual-write script replays wait for capacity instead + // of being dropped; best-effort users of goScript may still drop. maxScriptWriteGoroutines = 64 - // maxCompactedRetries caps retries when the secondary returns - // "read timestamp has been compacted". Each attempt re-sends the command so - // the secondary re-selects a fresh read snapshot; a small bound is enough - // because the compaction waterline advances slowly relative to SecondaryTimeout. - maxCompactedRetries = 3 + // maxSecondaryTransientRetries caps proxy-level retries when the secondary + // returns a transient OCC/read-snapshot error after exhausting its own retry + // loop. SecondaryTimeout still bounds the whole replay. + maxSecondaryTransientRetries = 3 // compactedRetryInitialBackoff is the first delay before retrying a secondary // command that failed with a compacted-read error. compactedRetryInitialBackoff = 10 * time.Millisecond @@ -67,6 +65,21 @@ func isReadTSCompactedError(err error) bool { return strings.Contains(err.Error(), readTSCompactedMarker) } +func isRetryableSecondaryWriteError(err error) bool { + if err == nil { + return false + } + if isReadTSCompactedError(err) { + return true + } + switch classifySecondaryWriteError(err) { + case "retry_limit", "write_conflict", "txn_locked": + return true + default: + return false + } +} + // DualWriter routes commands to primary and secondary backends based on mode. type DualWriter struct { primary Backend @@ -148,7 +161,9 @@ func (d *DualWriter) Close() { d.wg.Wait() } -// Write sends a write command to the primary synchronously, then to the secondary asynchronously. +// Write sends a write command to the primary synchronously, then to the secondary. +// The secondary write uses a bounded async slot when possible, but applies +// caller backpressure instead of dropping when the slot pool is saturated. // cmd must be the pre-uppercased command name. func (d *DualWriter) Write(ctx context.Context, cmd string, args [][]byte) (any, error) { iArgs := bytesArgsToInterfaces(args) @@ -166,9 +181,8 @@ func (d *DualWriter) Write(ctx context.Context, cmd string, args [][]byte) (any, } d.metrics.CommandTotal.WithLabelValues(cmd, d.primary.Name(), "ok").Inc() - // Secondary: async fire-and-forget (bounded) if d.hasSecondaryWrite() { - d.goWrite(func() { d.writeSecondary(cmd, iArgs) }) + d.runSecondaryWrite(func() { d.writeSecondary(cmd, iArgs) }) } return resp, err //nolint:wrapcheck // redis.Nil must pass through unwrapped for callers to detect nil replies @@ -257,7 +271,9 @@ func (d *DualWriter) Admin(ctx context.Context, cmd string, args [][]byte) (any, return resp, err //nolint:wrapcheck // redis.Nil must pass through unwrapped for callers to detect nil replies } -// Script forwards EVAL/EVALSHA to the primary, and async replays to secondary. +// Script forwards EVAL/EVALSHA to the primary, and replays to secondary. +// Secondary script replays are concurrency-limited and apply caller +// backpressure rather than dropping when the script slot pool is saturated. // cmd must be the pre-uppercased command name. func (d *DualWriter) Script(ctx context.Context, cmd string, args [][]byte) (any, error) { iArgs := bytesArgsToInterfaces(args) @@ -275,23 +291,21 @@ func (d *DualWriter) Script(ctx context.Context, cmd string, args [][]byte) (any d.rememberScript(cmd, args) if d.hasSecondaryWrite() { - d.goScript(func() { d.writeSecondary(cmd, iArgs) }) + d.runSecondaryScript(func() { d.writeSecondary(cmd, iArgs) }) } return resp, err //nolint:wrapcheck // redis.Nil must pass through unwrapped for callers to detect nil replies } // writeSecondary sends the command to the secondary, handling the NOSCRIPT -// → EVAL fallback and transparently retrying when the secondary reports that -// the read snapshot has been compacted. A re-sent command causes the backend -// to re-select a fresh read timestamp, which is the only way to recover once -// the original startTS has fallen behind MinRetainedTS on a peer node. +// → EVAL fallback and transparently retrying transient secondary errors. A +// re-sent command causes the backend to re-select a fresh timestamp and can +// also recover from hot-key OCC retry exhaustion in the secondary Redis adapter. // // The secondary's raw redis error is kept in sErr (not wrapped) so that // writeSecondary can classify it via errors.Is(sErr, redis.Nil), attach the // original message to Sentry and the structured log, and so the retry -// predicate isReadTSCompactedError matches the exact substring coming back -// from gRPC. +// predicate can match the exact substring coming back from gRPC. func (d *DualWriter) writeSecondary(cmd string, iArgs []any) { sCtx, cancel := context.WithTimeout(context.Background(), d.cfg.SecondaryTimeout) defer cancel() @@ -317,14 +331,14 @@ func (d *DualWriter) writeSecondary(cmd string, iArgs []any) { _, sErr = result.Result() } } - if !isReadTSCompactedError(sErr) { + if !isRetryableSecondaryWriteError(sErr) { break } - if attempt >= maxCompactedRetries { + if attempt >= maxSecondaryTransientRetries { break } - d.logger.Debug("retrying secondary write on compacted snapshot", - "cmd", cmd, "attempt", attempt+1, "backoff", backoff, "err", sErr) + d.logger.Debug("retrying secondary write after transient error", + "cmd", cmd, "attempt", attempt+1, "backoff", backoff, "reason", classifySecondaryWriteError(sErr), "err", sErr) if !waitCompactedRetryBackoff(sCtx, backoff) { break } @@ -367,6 +381,19 @@ func (d *DualWriter) recordSecondaryWriteFailure(cmd string, iArgs []any, elapse d.logger.Warn("secondary write failed", warnArgs...) } +func (d *DualWriter) replaySecondaryPipeline(cmds [][]any) { + d.runSecondaryWrite(func() { + sCtx, cancel := context.WithTimeout(context.Background(), d.cfg.SecondaryTimeout) + defer cancel() + _, pErr := d.secondary.Pipeline(sCtx, cmds) + if pErr != nil { + d.logger.Warn("secondary txn replay failed", "err", pErr) + d.metrics.SecondaryWriteErrors.Inc() + d.metrics.SecondaryWriteErrorsByReason.WithLabelValues("PIPELINE", classifySecondaryWriteError(pErr)).Inc() + } + }) +} + // waitCompactedRetryBackoff sleeps for a jittered interval or returns early // when the context is cancelled. Returns false if the caller should abort // the retry loop (context done). @@ -416,13 +443,15 @@ func nextCompactedRetryBackoff(current time.Duration) time.Duration { } // goWrite launches fn in a bounded write goroutine. +// It is best-effort: when the pool is saturated the work is dropped. +// Strict secondary writes should use runSecondaryWrite. func (d *DualWriter) goWrite(fn func()) { d.goAsyncWithSem(d.writeSem, fn) } // goScript launches fn in a bounded Lua-script write goroutine. -// It uses a smaller semaphore than goWrite to cap the number of concurrent -// EVAL/EVALSHA secondary writes. When the cap is reached the write is dropped. +// It is best-effort: when the pool is saturated the work is dropped. +// Strict secondary scripts should use runSecondaryScript. func (d *DualWriter) goScript(fn func()) { d.goAsyncWithSem(d.scriptSem, fn) } @@ -465,6 +494,49 @@ func (d *DualWriter) goAsyncWithSem(sem chan struct{}, fn func()) { } } +// runSecondaryWrite launches a strict secondary write. It uses the async write +// pool while capacity is available, and otherwise blocks the caller until a +// slot is available. This preserves dual-write consistency while still bounding +// secondary backend concurrency. +func (d *DualWriter) runSecondaryWrite(fn func()) { + d.runSecondaryWithBackpressure(d.writeSem, fn) +} + +// runSecondaryScript is the strict variant of goScript for EVAL/EVALSHA replay. +func (d *DualWriter) runSecondaryScript(fn func()) { + d.runSecondaryWithBackpressure(d.scriptSem, fn) +} + +func (d *DualWriter) runSecondaryWithBackpressure(sem chan struct{}, fn func()) { + d.mu.Lock() + if d.closed { + d.mu.Unlock() + return + } + select { + case sem <- struct{}{}: + d.wg.Add(1) + d.mu.Unlock() + go func() { + defer func() { + <-sem + d.wg.Done() + }() + fn() + }() + return + default: + d.metrics.AsyncBackpressure.Inc() + d.wg.Add(1) + d.mu.Unlock() + } + + defer d.wg.Done() + sem <- struct{}{} + defer func() { <-sem }() + fn() +} + func (d *DualWriter) logAsyncDrop() { nowNano := time.Now().UnixNano() diff --git a/proxy/metrics.go b/proxy/metrics.go index cd8b3837e..1436a6c23 100644 --- a/proxy/metrics.go +++ b/proxy/metrics.go @@ -17,7 +17,8 @@ type ProxyMetrics struct { ActiveConnections prometheus.Gauge - AsyncDrops prometheus.Counter + AsyncDrops prometheus.Counter + AsyncBackpressure prometheus.Counter PubSubShadowDivergences *prometheus.CounterVec PubSubShadowErrors prometheus.Counter @@ -78,7 +79,12 @@ func NewProxyMetrics(reg prometheus.Registerer) *ProxyMetrics { AsyncDrops: prometheus.NewCounter(prometheus.CounterOpts{ Namespace: "proxy", Name: "async_drops_total", - Help: "Total async operations dropped due to semaphore backpressure.", + Help: "Total best-effort async operations dropped due to semaphore backpressure.", + }), + AsyncBackpressure: prometheus.NewCounter(prometheus.CounterOpts{ + Namespace: "proxy", + Name: "async_backpressure_total", + Help: "Total strict secondary writes that waited for an async semaphore slot instead of being dropped.", }), ActiveConnections: prometheus.NewGauge(prometheus.GaugeOpts{ @@ -110,6 +116,7 @@ func NewProxyMetrics(reg prometheus.Registerer) *ProxyMetrics { m.Divergences, m.MigrationGaps, m.AsyncDrops, + m.AsyncBackpressure, m.ActiveConnections, m.PubSubShadowDivergences, m.PubSubShadowErrors, diff --git a/proxy/noop_backend.go b/proxy/noop_backend.go new file mode 100644 index 000000000..fbe5dc811 --- /dev/null +++ b/proxy/noop_backend.go @@ -0,0 +1,38 @@ +package proxy + +import ( + "context" + + "github.com/redis/go-redis/v9" +) + +// noopBackend is used for modes with no secondary backend. Keeping a concrete +// Backend avoids starting unnecessary leader-discovery goroutines in redis-only +// and elastickv-only modes while preserving DualWriter's simple shape. +type noopBackend struct { + name string +} + +// NewNoopBackend returns a Backend placeholder for an intentionally unused +// side of the proxy. +func NewNoopBackend(name string) Backend { + return noopBackend{name: name} +} + +func (n noopBackend) Do(ctx context.Context, args ...any) *redis.Cmd { + cmd := redis.NewCmd(ctx, args...) + cmd.SetErr(ErrNoLeaderBackend) + return cmd +} + +func (n noopBackend) Pipeline(context.Context, [][]any) ([]*redis.Cmd, error) { + return nil, ErrNoLeaderBackend +} + +func (n noopBackend) Close() error { + return nil +} + +func (n noopBackend) Name() string { + return n.name +} diff --git a/proxy/proxy.go b/proxy/proxy.go index 9ced8559a..9b2add2b1 100644 --- a/proxy/proxy.go +++ b/proxy/proxy.go @@ -59,6 +59,10 @@ func NewProxyServer(cfg ProxyConfig, dual *DualWriter, metrics *ProxyMetrics, se // ListenAndServe starts the redcon proxy server. func (p *ProxyServer) ListenAndServe(ctx context.Context) error { + if p.cfg.Mode == ModeRedisOnly && p.cfg.RedisOnlyRaw { + return p.listenAndServeRawRedis(ctx) + } + p.shutdownCtx = ctx var lc net.ListenConfig @@ -430,17 +434,8 @@ func (p *ProxyServer) execTxn(conn redcon.Conn, state *proxyConnState) { conn.WriteError("ERR empty transaction response") } - // Async replay to secondary (bounded) if p.dual.hasSecondaryWrite() { - p.dual.goAsync(func() { - sCtx, cancel := context.WithTimeout(context.Background(), p.cfg.SecondaryTimeout) - defer cancel() - _, pErr := p.dual.Secondary().Pipeline(sCtx, cmds) - if pErr != nil { - p.logger.Warn("secondary txn replay failed", "err", pErr) - p.metrics.SecondaryWriteErrors.Inc() - } - }) + p.dual.replaySecondaryPipeline(cmds) } } diff --git a/proxy/proxy_test.go b/proxy/proxy_test.go index 768e7ee9d..4398aca72 100644 --- a/proxy/proxy_test.go +++ b/proxy/proxy_test.go @@ -688,7 +688,53 @@ func TestDualWriter_GoAsync_DropLogsAreRateLimited(t *testing.T) { d.Close() } -func TestDualWriter_Script_DropsWhenScriptSemFull(t *testing.T) { +func TestDualWriter_Write_WaitsWhenWriteSemFull(t *testing.T) { + primary := newMockBackend("primary") + primary.doFunc = makeCmd("OK", nil) + secondary := newMockBackend("secondary") + secondary.doFunc = makeCmd("OK", nil) + + metrics := newTestMetrics() + cfg := ProxyConfig{Mode: ModeDualWrite, SecondaryWriteConcurrency: 1, SecondaryTimeout: 10 * time.Second} + d := NewDualWriter(primary, secondary, cfg, metrics, newTestSentry(), testLogger) + + blocker := make(chan struct{}) + d.goAsync(func() { + <-blocker + }) + + done := make(chan struct{}) + go func() { + _, err := d.Write(context.Background(), "SET", [][]byte{ + []byte("SET"), []byte("key"), []byte("value"), + }) + assert.NoError(t, err) + close(done) + }() + + select { + case <-done: + t.Fatal("Write returned while the strict secondary write slot was full") + case <-time.After(100 * time.Millisecond): + // good: the primary succeeded, but strict secondary replay is applying backpressure. + } + + assert.Equal(t, 0, secondary.CallCount()) + assert.InDelta(t, 1, testutil.ToFloat64(metrics.AsyncBackpressure), 0.001) + assert.InDelta(t, 0, testutil.ToFloat64(metrics.AsyncDrops), 0.001) + + close(blocker) + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("Write did not complete after a secondary write slot was released") + } + assert.Equal(t, 1, secondary.CallCount()) + d.Close() +} + +func TestDualWriter_Script_WaitsWhenScriptSemFull(t *testing.T) { primary := newMockBackend("primary") primary.doFunc = makeCmd("OK", nil) secondary := newMockBackend("secondary") @@ -707,7 +753,6 @@ func TestDualWriter_Script_DropsWhenScriptSemFull(t *testing.T) { }) } - // Script should return promptly even when scriptSem is full. done := make(chan struct{}) go func() { _, err := d.Script(context.Background(), "EVALSHA", [][]byte{ @@ -719,16 +764,23 @@ func TestDualWriter_Script_DropsWhenScriptSemFull(t *testing.T) { select { case <-done: - // good — Script returned without blocking on a full scriptSem - case <-time.After(time.Second): - t.Fatal("Script blocked when script semaphore was full") + t.Fatal("Script returned while the strict script replay slot was full") + case <-time.After(100 * time.Millisecond): + // good: strict EVALSHA replay is applying backpressure instead of dropping. } - // The async replay to secondary must be dropped. assert.Equal(t, 0, secondary.CallCount()) - assert.InDelta(t, 1, testutil.ToFloat64(metrics.AsyncDrops), 0.001) + assert.InDelta(t, 1, testutil.ToFloat64(metrics.AsyncBackpressure), 0.001) + assert.InDelta(t, 0, testutil.ToFloat64(metrics.AsyncDrops), 0.001) close(blocker) + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("Script did not complete after a secondary script slot was released") + } + assert.Equal(t, 1, secondary.CallCount()) d.Close() } @@ -1024,8 +1076,8 @@ func TestDualWriter_writeSecondary_RetriesReadTSCompacted(t *testing.T) { } func TestDualWriter_writeSecondary_ReadTSCompactedRetriesAreBounded(t *testing.T) { - // When the compacted error is persistent, the retry loop must stop after - // maxCompactedRetries+1 attempts so the secondary goroutine returns + // When the transient error is persistent, the retry loop must stop after + // maxSecondaryTransientRetries+1 attempts so the secondary goroutine returns // instead of burning a scriptSem slot indefinitely. primary := newMockBackend("primary") primary.doFunc = makeCmd("OK", nil) @@ -1039,12 +1091,43 @@ func TestDualWriter_writeSecondary_ReadTSCompactedRetriesAreBounded(t *testing.T d.writeSecondary("EVALSHA", []any{[]byte("EVALSHA"), []byte("deadbeef"), []byte("0")}) - assert.Equal(t, maxCompactedRetries+1, secondary.CallCount(), - "secondary must stop after maxCompactedRetries+1 attempts") + assert.Equal(t, maxSecondaryTransientRetries+1, secondary.CallCount(), + "secondary must stop after maxSecondaryTransientRetries+1 attempts") assert.InDelta(t, 1, testutil.ToFloat64(metrics.SecondaryWriteErrors), 0.001, "a persistent compacted error must still be reported as a secondary write error") } +func TestDualWriter_writeSecondary_RetriesRetryLimitWriteConflict(t *testing.T) { + // A hot key can exhaust the secondary Redis adapter's internal OCC retry loop. + // Re-sending the command lets the backend choose a fresh timestamp and keeps + // dual-write traffic from permanently diverging on a transient conflict. + primary := newMockBackend("primary") + primary.doFunc = makeCmd("OK", nil) + + secondary := newMockBackend("secondary") + retryLimitErr := testRedisErr("redis txn retry limit exceeded: key: myzset: write conflict") + var calls int + secondary.doFunc = func(ctx context.Context, args ...any) *redis.Cmd { + calls++ + cmd := redis.NewCmd(ctx, args...) + if calls < 3 { + cmd.SetErr(retryLimitErr) + return cmd + } + cmd.SetVal("OK") + return cmd + } + + metrics := newTestMetrics() + d := NewDualWriter(primary, secondary, ProxyConfig{Mode: ModeDualWrite, SecondaryTimeout: time.Second}, metrics, newTestSentry(), testLogger) + + d.writeSecondary("ZADD", []any{[]byte("ZADD"), []byte("myzset"), []byte("1"), []byte("member")}) + + assert.Equal(t, 3, calls, "secondary must retry transient retry-limit write conflicts") + assert.InDelta(t, 0, testutil.ToFloat64(metrics.SecondaryWriteErrors), 0.001, + "a retried success must not count as a secondary write error") +} + func TestDualWriter_writeSecondary_RetriesDoNotRepeatNoScriptProbe(t *testing.T) { // After the EVAL fallback resolves a NOSCRIPT, a compacted-retry must // re-send the resolved EVAL form directly — never the known-missing diff --git a/proxy/pubsub.go b/proxy/pubsub.go index 1ecd9e284..221d085c1 100644 --- a/proxy/pubsub.go +++ b/proxy/pubsub.go @@ -470,15 +470,7 @@ func (s *pubsubSession) execTxn() { s.writeMu.Unlock() if s.proxy.dual.hasSecondaryWrite() { - s.proxy.dual.goAsync(func() { - sCtx, cancel := context.WithTimeout(context.Background(), s.proxy.cfg.SecondaryTimeout) - defer cancel() - _, pErr := s.proxy.dual.Secondary().Pipeline(sCtx, cmds) - if pErr != nil { - s.proxy.logger.Warn("secondary txn replay failed", "err", pErr) - s.proxy.metrics.SecondaryWriteErrors.Inc() - } - }) + s.proxy.dual.replaySecondaryPipeline(cmds) } } diff --git a/proxy/raw_redis_proxy.go b/proxy/raw_redis_proxy.go new file mode 100644 index 000000000..f478d5480 --- /dev/null +++ b/proxy/raw_redis_proxy.go @@ -0,0 +1,156 @@ +package proxy + +import ( + "bufio" + "context" + "fmt" + "io" + "net" + "strconv" + "strings" + "sync" + "time" +) + +const ( + rawRedisCopyDirections = 2 + rawRedisHandshakeTimeout = 5 * time.Second +) + +func (p *ProxyServer) listenAndServeRawRedis(ctx context.Context) error { + p.shutdownCtx = ctx + + var lc net.ListenConfig + ln, err := lc.Listen(ctx, "tcp", p.cfg.ListenAddr) + if err != nil { + return fmt.Errorf("raw redis proxy listen: %w", err) + } + defer ln.Close() + + p.logger.Info("raw redis proxy starting", + "addr", p.cfg.ListenAddr, + "mode", p.cfg.Mode.String(), + "primary", p.cfg.PrimaryAddr, + "primary_db", p.cfg.PrimaryDB, + ) + + var wg sync.WaitGroup + defer wg.Wait() + + go func() { + <-ctx.Done() + p.logger.Info("shutting down raw redis proxy") + _ = ln.Close() + }() + + for { + client, err := ln.Accept() + if err != nil { + if ctx.Err() != nil { + return nil + } + return fmt.Errorf("raw redis proxy accept: %w", err) + } + wg.Add(1) + go func() { + defer wg.Done() + p.handleRawRedisConn(ctx, client) + }() + } +} + +func (p *ProxyServer) handleRawRedisConn(ctx context.Context, client net.Conn) { + p.metrics.ActiveConnections.Inc() + defer p.metrics.ActiveConnections.Dec() + defer client.Close() + + var dialer net.Dialer + upstream, err := dialer.DialContext(ctx, "tcp", p.cfg.PrimaryAddr) + if err != nil { + p.logger.Warn("raw redis upstream dial failed", "addr", p.cfg.PrimaryAddr, "err", err) + return + } + defer upstream.Close() + + upstreamReader := bufio.NewReader(upstream) + if err := p.prepareRawRedisUpstream(upstream, upstreamReader); err != nil { + p.logger.Warn("raw redis upstream prepare failed", "addr", p.cfg.PrimaryAddr, "err", err) + return + } + + stop := make(chan struct{}) + defer close(stop) + go func() { + select { + case <-ctx.Done(): + _ = client.Close() + _ = upstream.Close() + case <-stop: + } + }() + + errCh := make(chan struct{}, rawRedisCopyDirections) + go rawCopy(upstream, client, errCh) + go rawCopy(client, upstreamReader, errCh) + <-errCh +} + +func rawCopy(dst net.Conn, src io.Reader, done chan<- struct{}) { + _, _ = io.Copy(dst, src) + _ = dst.Close() + done <- struct{}{} +} + +func (p *ProxyServer) prepareRawRedisUpstream(upstream net.Conn, upstreamReader *bufio.Reader) error { + if err := upstream.SetDeadline(time.Now().Add(rawRedisHandshakeTimeout)); err != nil { + return fmt.Errorf("set handshake deadline: %w", err) + } + if p.cfg.PrimaryPassword != "" { + if err := rawRedisRoundTrip(upstream, upstreamReader, "AUTH", p.cfg.PrimaryPassword); err != nil { + return fmt.Errorf("auth: %w", err) + } + } + if p.cfg.PrimaryDB != 0 { + if err := rawRedisRoundTrip(upstream, upstreamReader, "SELECT", strconv.Itoa(p.cfg.PrimaryDB)); err != nil { + return fmt.Errorf("select db %d: %w", p.cfg.PrimaryDB, err) + } + } + if err := upstream.SetDeadline(time.Time{}); err != nil { + return fmt.Errorf("clear handshake deadline: %w", err) + } + return nil +} + +func rawRedisRoundTrip(w io.Writer, r *bufio.Reader, args ...string) error { + if _, err := w.Write(rawRedisCommand(args...)); err != nil { + return fmt.Errorf("write command: %w", err) + } + prefix, err := r.ReadByte() + if err != nil { + return fmt.Errorf("read reply prefix: %w", err) + } + line, err := r.ReadString('\n') + if err != nil { + return fmt.Errorf("read reply line: %w", err) + } + line = strings.TrimRight(line, "\r\n") + if prefix == '-' { + return fmt.Errorf("%s", line) + } + return nil +} + +func rawRedisCommand(args ...string) []byte { + var b strings.Builder + b.WriteByte('*') + b.WriteString(strconv.Itoa(len(args))) + b.WriteString("\r\n") + for _, arg := range args { + b.WriteByte('$') + b.WriteString(strconv.Itoa(len(arg))) + b.WriteString("\r\n") + b.WriteString(arg) + b.WriteString("\r\n") + } + return []byte(b.String()) +} diff --git a/proxy/raw_redis_proxy_test.go b/proxy/raw_redis_proxy_test.go new file mode 100644 index 000000000..285c1bd3d --- /dev/null +++ b/proxy/raw_redis_proxy_test.go @@ -0,0 +1,77 @@ +package proxy + +import ( + "bufio" + "io" + "net" + "strconv" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestRawRedisCommand(t *testing.T) { + t.Parallel() + + require.Equal(t, + []byte("*2\r\n$4\r\nAUTH\r\n$6\r\nsecret\r\n"), + rawRedisCommand("AUTH", "secret"), + ) +} + +func TestPrepareRawRedisUpstreamAuthenticatesAndSelectsDB(t *testing.T) { + t.Parallel() + + client, server := net.Pipe() + defer client.Close() + defer server.Close() + + seen := make(chan []string, 2) + go func() { + reader := bufio.NewReader(server) + for i := 0; i < 2; i++ { + cmd, err := readRawRedisTestCommand(reader) + if err != nil { + return + } + seen <- cmd + _, _ = server.Write([]byte("+OK\r\n")) + } + }() + + p := &ProxyServer{cfg: ProxyConfig{PrimaryPassword: "secret", PrimaryDB: 1}} + err := p.prepareRawRedisUpstream(client, bufio.NewReader(client)) + require.NoError(t, err) + require.Equal(t, []string{"AUTH", "secret"}, <-seen) + require.Equal(t, []string{"SELECT", "1"}, <-seen) +} + +func readRawRedisTestCommand(r *bufio.Reader) ([]string, error) { + line, err := r.ReadString('\n') + if err != nil { + return nil, err + } + line = strings.TrimRight(line, "\r\n") + n, err := strconv.Atoi(strings.TrimPrefix(line, "*")) + if err != nil { + return nil, err + } + out := make([]string, 0, n) + for range n { + line, err = r.ReadString('\n') + if err != nil { + return nil, err + } + size, err := strconv.Atoi(strings.TrimPrefix(strings.TrimRight(line, "\r\n"), "$")) + if err != nil { + return nil, err + } + arg := make([]byte, size+2) + if _, err := io.ReadFull(r, arg); err != nil { + return nil, err + } + out = append(out, string(arg[:size])) + } + return out, nil +}