diff --git a/internal/raftengine/etcd/engine.go b/internal/raftengine/etcd/engine.go index df65236b7..9d210122b 100644 --- a/internal/raftengine/etcd/engine.go +++ b/internal/raftengine/etcd/engine.go @@ -99,6 +99,9 @@ const ( // raft hot path is high blast radius and a regression here can cause // cluster-wide elections. dispatcherLanesEnvVar = "ELASTICKV_RAFT_DISPATCHER_LANES" + // preVoteEnvVar permits an operator to temporarily disable raft + // pre-vote during manual quorum recovery. It defaults to enabled. + preVoteEnvVar = "ELASTICKV_RAFT_PRE_VOTE" // 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, @@ -817,7 +820,13 @@ func rawNodeAppliedForOpen(storage *etcdraft.MemoryStorage, applied uint64, cfg if applied <= baseApplied { return applied, nil } - return trimRawNodeAppliedForReplay(storage, baseApplied, applied, cfg), nil + rawApplied := trimRawNodeAppliedForReplay(storage, baseApplied, applied, cfg) + if rawApplied > baseApplied { + if err := replayColdStartVolatileEntries(storage, baseApplied+1, rawApplied+1, cfg); err != nil { + return 0, err + } + } + return rawApplied, nil } func rawNodeAppliedBounds(storage *etcdraft.MemoryStorage, applied uint64) (uint64, uint64, error) { @@ -867,23 +876,58 @@ func coldStartEntryRequiresReplay(entry *raftpb.Entry, cfg OpenConfig) bool { default: return false } + _, _, err := coldStartNormalPayload(entry, cfg) + return err != nil +} + +func replayColdStartVolatileEntries(storage *etcdraft.MemoryStorage, lo, hi uint64, cfg OpenConfig) error { + entries, err := storage.Entries(lo, hi, math.MaxUint64) + if err != nil { + return errors.WithStack(err) + } + for _, entry := range entries { + if entry == nil || !entryTypeIsNormal(entry) { + continue + } + payload, ok, err := coldStartNormalPayload(entry, cfg) + if err != nil { + return err + } + if !ok || !coldStartPayloadIsVolatile(payload, cfg) { + continue + } + cfg.StateMachine.Apply(payload) + } + return nil +} + +func entryTypeIsNormal(entry *raftpb.Entry) bool { + return entry.GetType() == raftpb.EntryNormal +} + +func coldStartNormalPayload(entry *raftpb.Entry, cfg OpenConfig) ([]byte, bool, error) { if len(entry.GetData()) == 0 { - return false + return nil, false, nil } _, payload, ok := decodeProposalEnvelope(entry.GetData()) if !ok { - return false + return nil, false, nil } - if entry.GetIndex() > orInertCutover(cfg.RaftCutoverIndex)() { - if cfg.RaftCipher == nil { - return true - } - plain, err := unwrapRaftPayload(cfg.RaftCipher, payload) - if err != nil { - return true - } - payload = plain + if entry.GetIndex() <= orInertCutover(cfg.RaftCutoverIndex)() { + return payload, true, nil + } + if cfg.RaftCipher == nil { + return nil, false, errors.Wrap(ErrRaftUnwrapFailed, + "raftengine/etcd: cold-start entry past raft envelope cutover but no raft cipher wired") } + plain, err := unwrapRaftPayload(cfg.RaftCipher, payload) + if err != nil { + return nil, false, err + } + return plain, true, nil +} + +func coldStartPayloadIsVolatile(payload []byte, cfg OpenConfig) bool { classifier, ok := cfg.StateMachine.(raftengine.VolatileEntryClassifier) if !ok { return false @@ -902,7 +946,7 @@ func newRawNode(cfg OpenConfig, storage *etcdraft.MemoryStorage, applied uint64) MaxCommittedSizePerReady: cfg.MaxSizePerMsg, MaxInflightMsgs: cfg.MaxInflightMsg, CheckQuorum: true, - PreVote: true, + PreVote: preVoteEnabledFromEnv(), ReadOnlyOption: etcdraft.ReadOnlySafe, DisableProposalForwarding: true, }) @@ -4363,6 +4407,20 @@ func dispatcherLanesEnabledFromEnv() bool { return enabled } +func preVoteEnabledFromEnv() bool { + v := strings.TrimSpace(os.Getenv(preVoteEnvVar)) + if v == "" { + return true + } + enabled, err := strconv.ParseBool(v) + if err != nil { + slog.Warn("invalid ELASTICKV_RAFT_PRE_VOTE; using default", + "value", v, "default", true) + return true + } + return enabled +} + // 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. diff --git a/internal/raftengine/etcd/engine_test.go b/internal/raftengine/etcd/engine_test.go index 00aa807c2..724b031ed 100644 --- a/internal/raftengine/etcd/engine_test.go +++ b/internal/raftengine/etcd/engine_test.go @@ -1905,19 +1905,19 @@ func TestRawNodeAppliedForOpenPreservesVolatileReplay(t *testing.T) { }, }) cfg := rawNodeTestConfig() - cfg.StateMachine = &volatileTagFakeFSM{} + fsm := &volatileTagFakeFSM{} + cfg.StateMachine = fsm rawApplied, err := rawNodeAppliedForOpen(storage, 150, cfg) require.NoError(t, err) - require.Equal(t, uint64(124), rawApplied, - "RawNode must still deliver volatile duplicate entries for in-memory replay") + require.Equal(t, uint64(150), rawApplied, + "volatile duplicate entries should replay without forcing RawNode to redeliver the committed tail") + require.Equal(t, int32(1), fsm.calls.Load()) + require.Equal(t, []byte{0x02, 0xaa}, fsm.lastPayload) rawNode, err := newRawNode(cfg, storage, rawApplied) require.NoError(t, err) - require.True(t, rawNode.HasReady()) - ready := rawNode.Ready() - require.NotEmpty(t, ready.CommittedEntries) - require.Equal(t, uint64(125), ready.CommittedEntries[0].GetIndex()) + require.False(t, rawNode.HasReady()) } func TestRawNodeAppliedForOpenPreservesConfChangeReplay(t *testing.T) { @@ -1935,6 +1935,20 @@ func TestRawNodeAppliedForOpenPreservesConfChangeReplay(t *testing.T) { "RawNode must deliver committed config changes so ApplyConfChange rebuilds membership") } +func TestPreVoteEnabledFromEnv(t *testing.T) { + t.Setenv(preVoteEnvVar, "") + require.True(t, preVoteEnabledFromEnv()) + + t.Setenv(preVoteEnvVar, "false") + require.False(t, preVoteEnabledFromEnv()) + + t.Setenv(preVoteEnvVar, "true") + require.True(t, preVoteEnabledFromEnv()) + + t.Setenv(preVoteEnvVar, "not-bool") + require.True(t, preVoteEnabledFromEnv()) +} + // TestErrNotLeaderMatchesRaftEngineSentinel pins the invariant that the // etcd engine's internal leadership-loss errors are marked against the // shared raftengine sentinels. The lease-read fast path in package kv diff --git a/internal/raftengine/etcd/wal_store.go b/internal/raftengine/etcd/wal_store.go index f6fcc8a44..c69982f90 100644 --- a/internal/raftengine/etcd/wal_store.go +++ b/internal/raftengine/etcd/wal_store.go @@ -142,16 +142,11 @@ 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 restore decision so we know the committed + // replay target. The FSM restore can be skipped as soon as the + // durable FSM is at least at the snapshot pointer: Open seeds the + // engine's applied counter from EffectiveApplied, and WAL replay + // then delivers only committed entries above that durable point. w, hardState, entries, err := openAndReadWALWithRepair(logger, walDir, walSnapshotFor(snapshot)) if err != nil { return nil, err @@ -202,30 +197,36 @@ func reportColdStartExecute(obs raftengine.ColdStartObserver, logger *zap.Logger if logger == nil { return } - // gap_to_snapshot uses absolute value because the gate now - // permits have > snapIndex (FSM ahead of snapshot but behind - // committed tail). gap_behind_committed is target-have; can be - // 0 when have==target. + // gap_to_snapshot uses absolute value because stale metadata can + // still report an FSM ahead of the snapshot while the execute path + // runs due to an earlier fallback. gap_behind_committed is clamped + // to avoid underflow when tests call this helper with a lower + // synthetic target. var gapToSnapshot uint64 if have >= snapIndex { gapToSnapshot = have - snapIndex } else { gapToSnapshot = snapIndex - have } + var gapBehindCommitted uint64 + if target > have { + gapBehindCommitted = target - have + } logger.Info("restoreSnapshotState executed (FSM behind WAL committed tail)", zap.Uint64("fsm_applied", have), zap.Uint64("snapshot_index", snapIndex), zap.Uint64("last_committed_index", target), zap.Uint64("gap_to_snapshot", gapToSnapshot), - zap.Uint64("gap_behind_committed", target-have), + zap.Uint64("gap_behind_committed", gapBehindCommitted), ) } // 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. +// node: max(snapshot.Metadata.Index, hardState.Commit). This is the +// committed replay target used for observability; the snapshot body +// restore decision itself only requires the durable FSM to be at +// least at snapshot.Metadata.Index. // // Followers can carry an UNCOMMITTED WAL suffix // (entries[n-1].Index > hardState.Commit). Raft does NOT surface @@ -252,8 +253,7 @@ func coldStartSkipThreshold(snapshot raftpb.Snapshot, hardState raftpb.HardState // skip gate fires with the FSM at `have > snapshot.Metadata.Index`, // EffectiveApplied carries `have`; without this seed the engine // would deliver entries snapshot.Index+1..have to applyCommitted -// and re-apply them onto a Pebble store already containing them -// (codex P1 #934 root cause). +// and re-apply them onto a Pebble store already containing them. func coldStartApplied(disk *diskState) uint64 { base := maxAppliedIndex(disk.LocalSnap) if disk.EffectiveApplied > base { @@ -356,12 +356,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, committedTailIndex 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, committedTailIndex, fsmSnapDir, obs, logger) } // Legacy format: full FSM payload embedded in snapshot.Data. if err := fsm.Restore(bytes.NewReader(snapshot.Data)); err != nil { @@ -376,32 +376,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 snapshot.Metadata.Index: once the FSM is at +// the snapshot pointer, the engine can seed its applied counter from +// the durable FSM index and replay only the committed WAL suffix above +// that point. // // 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, committedTailIndex 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) + snapIndex := snapshot.GetMetadata().GetIndex() + decision, have := decideSkipOutcome(fsm, snapIndex) 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, snapIndex, committedTailIndex, 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, snapIndex, committedTailIndex, have) return snapshot.GetMetadata().GetIndex(), nil } @@ -468,6 +469,13 @@ func reportColdStart(obs raftengine.ColdStartObserver, logger *zap.Logger, d col obs.RestoreSkipped(snapIndex, have) } if logger != nil { + var gapAheadCommitted uint64 + var gapBehindCommitted uint64 + if have >= target { + gapAheadCommitted = have - target + } else { + gapBehindCommitted = target - have + } // Two named gap fields so an operator correlating the // log against the Prometheus gauge sees consistent // magnitudes (claude #934 round 5): @@ -480,7 +488,8 @@ func reportColdStart(obs raftengine.ColdStartObserver, logger *zap.Logger, d col zap.Uint64("snapshot_index", snapIndex), zap.Uint64("last_committed_index", target), zap.Uint64("gap_ahead_snapshot", have-snapIndex), - zap.Uint64("gap_ahead_committed", have-target), + zap.Uint64("gap_ahead_committed", gapAheadCommitted), + zap.Uint64("gap_behind_committed", gapBehindCommitted), ) } case coldStartExecute: diff --git a/internal/raftengine/etcd/wal_store_skip_gate_test.go b/internal/raftengine/etcd/wal_store_skip_gate_test.go index 212a47ec6..6348a46a0 100644 --- a/internal/raftengine/etcd/wal_store_skip_gate_test.go +++ b/internal/raftengine/etcd/wal_store_skip_gate_test.go @@ -177,10 +177,10 @@ func TestSkipGate_ExecutesWhenFSMStale(t *testing.T) { require.Empty(t, obs.fallbacks) } -// TestSkipGate_ReturnsEffectiveAppliedOnSkip pins codex P1 #934 -// round 2. The skip path MUST return `have` so Engine.Open can seed -// e.applied above snapshot.Index, preventing applyCommitted from -// re-delivering the snapshot..have tail. +// TestSkipGate_ReturnsEffectiveAppliedOnSkip verifies that the skip +// path returns `have` so Engine.Open can seed e.applied above +// snapshot.Index, preventing applyCommitted from re-delivering the +// snapshot..have tail. func TestSkipGate_ReturnsEffectiveAppliedOnSkip(t *testing.T) { dir := t.TempDir() const ( @@ -230,8 +230,8 @@ 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 +// TestColdStartSkipThreshold verifies that 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 @@ -260,29 +260,20 @@ func TestColdStartSkipThreshold(t *testing.T) { } } -// 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_SkipsWhenFSMAtSnapshotButBehindCommitTail verifies +// the cold-start fast path for the normal crash window: the FSM is +// already at or beyond the snapshot pointer, but durable Raft commit +// is ahead. The snapshot body must still be skipped; Engine.Open uses +// the returned EffectiveApplied to replay only the committed WAL suffix +// above the FSM's durable applied index. +func TestSkipGate_SkipsWhenFSMAtSnapshotButBehindCommitTail(t *testing.T) { dir := t.TempDir() const ( - snapIndex uint64 = 100 - appliedIdx uint64 = 150 - lastWalIndex uint64 = 200 + snapIndex uint64 = 100 + appliedIdx uint64 = 150 + committedTailIndex uint64 = 200 ) - payload := []byte("body-bytes-for-execute") + payload := []byte("body-bytes-for-skip") crc, _ := writeFSMFileForTest(t, dir, snapIndex, payload) fsm := &skipGateFSM{applied: appliedIdx, appliedPresent: true} @@ -291,19 +282,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, committedTailIndex, 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.Equal(t, appliedIdx, effective, + "skip path MUST return the durable FSM applied index so WAL replay starts above it") + require.Empty(t, fsm.bodyBytes, "skip path MUST NOT call fsm.Restore") + require.True(t, fsm.restoredHeader, "skip path MUST still apply snapshot header state") + require.Equal(t, []uint64{appliedIdx - snapIndex}, obs.skipped) + require.Empty(t, obs.executed) require.Empty(t, obs.fallbacks) } diff --git a/kv/fsm.go b/kv/fsm.go index 421f5fc3b..aa2c75352 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -485,9 +485,9 @@ func (f *kvFSM) handleRawRequest(ctx context.Context, r *pb.Request, commitTS ui if isTxnInternalKey(mut.Key) { return errors.WithStack(ErrInvalidRequest) } - if err := f.assertNoConflictingTxnLock(ctx, mut.Key, nil, 0); err != nil { - return err - } + } + if err := f.assertNoConflictingTxnLocks(ctx, r.Mutations, nil, 0); err != nil { + return err } muts, err := toStoreMutations(r.Mutations) @@ -1111,9 +1111,9 @@ func (f *kvFSM) buildOnePhaseStoreMutations(ctx context.Context, muts []*pb.Muta if isTxnInternalKey(mut.Key) { return nil, errors.WithStack(ErrInvalidRequest) } - if err := f.assertNoConflictingTxnLock(ctx, mut.Key, nil, 0); err != nil { - return nil, err - } + } + if err := f.assertNoConflictingTxnLocks(ctx, muts, nil, 0); err != nil { + return nil, err } storeMuts, err := toStoreMutations(muts) if err != nil { @@ -1398,6 +1398,131 @@ func (f *kvFSM) assertNoConflictingTxnLock(ctx context.Context, key, primaryKey return NewTxnLockedError(key) } +type txnLockCheckTarget struct { + key []byte + lockKey []byte + index int +} + +type txnLockCheckResult struct { + err error + index int +} + +const txnLockBatchScanPageSize = 1024 + +type mvccBatchGetter interface { + GetAtBatch(ctx context.Context, keys [][]byte, ts uint64) (map[string][]byte, error) +} + +func (f *kvFSM) assertNoConflictingTxnLocks(ctx context.Context, muts []*pb.Mutation, primaryKey []byte, startTS uint64) error { + if len(muts) == 0 { + return nil + } + targets, scanStart, scanEnd := txnLockCheckTargets(muts) + if len(targets) == 0 { + return nil + } + if getter, ok := f.store.(mvccBatchGetter); ok { + return errors.WithStack(assertNoConflictingTxnLocksWithBatchGet(ctx, getter, targets, primaryKey, startTS, len(muts))) + } + return f.assertNoConflictingTxnLocksWithScan(ctx, targets, scanStart, scanEnd, primaryKey, startTS, len(muts)) +} + +func (f *kvFSM) assertNoConflictingTxnLocksWithScan(ctx context.Context, targets map[string]txnLockCheckTarget, scanStart, scanEnd, primaryKey []byte, startTS uint64, targetCount int) error { + best := txnLockCheckResult{index: targetCount} + for start := scanStart; len(start) > 0; { + kvs, err := f.store.ScanAt(ctx, start, scanEnd, txnLockBatchScanPageSize, ^uint64(0)) + if err != nil { + return errors.WithStack(err) + } + if len(kvs) == 0 { + break + } + updateTxnLockCheckResult(kvs, targets, primaryKey, startTS, &best) + next := nextScanCursor(kvs[len(kvs)-1].Key) + if scanEnd != nil && bytes.Compare(next, scanEnd) >= 0 { + break + } + start = next + } + return best.err +} + +func assertNoConflictingTxnLocksWithBatchGet(ctx context.Context, getter mvccBatchGetter, targets map[string]txnLockCheckTarget, primaryKey []byte, startTS uint64, targetCount int) error { + values, err := getter.GetAtBatch(ctx, txnLockTargetKeys(targets), ^uint64(0)) + if err != nil { + return errors.WithStack(err) + } + best := txnLockCheckResult{index: targetCount} + for lockKey, value := range values { + target, ok := targets[lockKey] + if !ok || target.index >= best.index { + continue + } + if err := txnLockConflictError(value, target.key, primaryKey, startTS); err != nil { + best.err = err + best.index = target.index + } + } + return best.err +} + +func txnLockTargetKeys(targets map[string]txnLockCheckTarget) [][]byte { + keys := make([][]byte, 0, len(targets)) + for _, target := range targets { + keys = append(keys, target.lockKey) + } + return keys +} + +func updateTxnLockCheckResult(kvs []*store.KVPair, targets map[string]txnLockCheckTarget, primaryKey []byte, startTS uint64, best *txnLockCheckResult) { + for _, kvp := range kvs { + target, ok := targets[string(kvp.Key)] + if !ok || target.index >= best.index { + continue + } + if err := txnLockConflictError(kvp.Value, target.key, primaryKey, startTS); err != nil { + best.err = err + best.index = target.index + } + } +} + +func txnLockCheckTargets(muts []*pb.Mutation) (map[string]txnLockCheckTarget, []byte, []byte) { + targets := make(map[string]txnLockCheckTarget, len(muts)) + var minLockKey []byte + var maxLockKey []byte + for i, mut := range muts { + lockKey := txnLockKey(mut.Key) + if _, ok := targets[string(lockKey)]; ok { + continue + } + targets[string(lockKey)] = txnLockCheckTarget{key: mut.Key, lockKey: lockKey, index: i} + if minLockKey == nil || bytes.Compare(lockKey, minLockKey) < 0 { + minLockKey = lockKey + } + if maxLockKey == nil || bytes.Compare(lockKey, maxLockKey) > 0 { + maxLockKey = lockKey + } + } + if minLockKey == nil { + return nil, nil, nil + } + return targets, minLockKey, nextScanCursor(maxLockKey) +} + +func txnLockConflictError(lockBytes []byte, key, primaryKey []byte, startTS uint64) error { + lock, err := decodeTxnLock(lockBytes) + if err != nil { + return errors.WithStack(err) + } + if startTS != 0 && lock.StartTS == startTS && bytes.Equal(lock.PrimaryKey, primaryKey) { + return nil + } + return NewTxnLockedError(key) +} + func toStoreMutations(muts []*pb.Mutation) ([]*store.KVPairMutation, error) { out := make([]*store.KVPairMutation, 0, len(muts)) for _, mut := range muts { diff --git a/kv/fsm_txn_test.go b/kv/fsm_txn_test.go index 1a6a4a808..db35beab5 100644 --- a/kv/fsm_txn_test.go +++ b/kv/fsm_txn_test.go @@ -126,6 +126,42 @@ func TestPrepareRejectsSameStartTSDifferentPrimary(t *testing.T) { require.ErrorIs(t, err, ErrTxnLocked) } +func TestOnePhaseBatchLockCheckPreservesMutationOrder(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + fsm, ok := NewKvFSMWithHLC(st, NewHLC()).(*kvFSM) + require.True(t, ok) + + 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("p"), LockTTLms: defaultTxnLockTTLms})}, + {Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("vz")}, + {Op: pb.Op_PUT, Key: []byte("a"), Value: []byte("va")}, + }, + } + require.NoError(t, applyFSMRequest(t, fsm, prepare)) + + onePhase := &pb.Request{ + IsTxn: true, + Phase: pb.Phase_NONE, + Ts: 20, + Mutations: []*pb.Mutation{ + {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: []byte("z"), CommitTS: 30})}, + {Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("new-z")}, + {Op: pb.Op_PUT, Key: []byte("a"), Value: []byte("new-a")}, + }, + } + err := applyFSMRequest(t, fsm, onePhase) + require.ErrorIs(t, err, ErrTxnLocked) + key, _, ok := TxnLockedDetails(err) + require.True(t, ok) + require.Equal(t, []byte("z"), key) +} + func TestCommitIsIdempotentAfterCommitRecordExists(t *testing.T) { t.Parallel() diff --git a/main.go b/main.go index 530dbf4b7..c66963025 100644 --- a/main.go +++ b/main.go @@ -270,6 +270,11 @@ const memoryShutdownThresholdEnvVar = "ELASTICKV_MEMORY_SHUTDOWN_THRESHOLD_MB" // warning and fall through to the default. const memoryShutdownPollIntervalEnvVar = "ELASTICKV_MEMORY_SHUTDOWN_POLL_INTERVAL" +const ( + lockResolverEnabledEnvVar = "ELASTICKV_LOCK_RESOLVER_ENABLED" + fsmCompactorEnabledEnvVar = "ELASTICKV_FSM_COMPACTOR_ENABLED" +) + const bytesPerMiB = 1024 * 1024 func main() { @@ -337,6 +342,45 @@ func memwatchConfigFromEnv() (memwatch.Config, bool) { return cfg, true } +func enabledEnv(name string) bool { + raw := strings.TrimSpace(os.Getenv(name)) + if raw == "" { + return true + } + enabled, err := strconv.ParseBool(raw) + if err != nil { + slog.Warn("invalid "+name+"; using default", + "value", raw, + "default", true, + ) + return true + } + return enabled +} + +func startLockResolverIfEnabled(shardStore *kv.ShardStore, shardGroups map[uint64]*kv.ShardGroup, cleanup *internalutil.CleanupStack) { + if enabledEnv(lockResolverEnabledEnvVar) { + lockResolver := kv.NewLockResolver(shardStore, shardGroups, nil) + cleanup.Add(func() { lockResolver.Close() }) + return + } + slog.Info("background lock resolver disabled", "env", lockResolverEnabledEnvVar) +} + +func startFSMCompactorIfEnabled(ctx context.Context, eg *errgroup.Group, runtimes []*raftGroupRuntime, readTracker *kv.ActiveTimestampTracker) { + if enabledEnv(fsmCompactorEnabledEnvVar) { + compactor := kv.NewFSMCompactor( + fsmCompactionRuntimes(runtimes), + kv.WithFSMCompactorActiveTimestampTracker(readTracker), + ) + eg.Go(func() error { + return compactor.Run(ctx) + }) + return + } + slog.Info("fsm compactor disabled", "env", fsmCompactorEnabledEnvVar) +} + func run() error { cfg, engineType, bootstrapCfg, bootstrap, err := resolveRuntimeInputs() if err != nil { @@ -445,8 +489,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 +548,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 diff --git a/main_background_env_test.go b/main_background_env_test.go new file mode 100644 index 000000000..2a38c7169 --- /dev/null +++ b/main_background_env_test.go @@ -0,0 +1,26 @@ +package main + +import "testing" + +func TestEnabledEnv(t *testing.T) { + t.Run("unset returns default", func(t *testing.T) { + t.Setenv("ELASTICKV_TEST_BOOL_ENV", "") + if !enabledEnv("ELASTICKV_TEST_BOOL_ENV") { + t.Fatal("enabledEnv() = false, want true") + } + }) + + t.Run("false disables default true", func(t *testing.T) { + t.Setenv("ELASTICKV_TEST_BOOL_ENV", "false") + if enabledEnv("ELASTICKV_TEST_BOOL_ENV") { + t.Fatal("enabledEnv() = true, want false") + } + }) + + t.Run("invalid returns default", func(t *testing.T) { + t.Setenv("ELASTICKV_TEST_BOOL_ENV", "maybe") + if !enabledEnv("ELASTICKV_TEST_BOOL_ENV") { + t.Fatal("enabledEnv() = false, want true") + } + }) +} diff --git a/store/lsm_store.go b/store/lsm_store.go index d818b7bde..0b14632e5 100644 --- a/store/lsm_store.go +++ b/store/lsm_store.go @@ -12,6 +12,7 @@ import ( "math" "os" "path/filepath" + "sort" "strconv" "strings" "sync" @@ -109,6 +110,18 @@ const ( // fsmSyncModeEnv. Any other value falls back to the default. fsmSyncModeSync = "sync" fsmSyncModeNoSync = "nosync" + + // conflictCheckSmallLinearAdvanceLimit bounds the number of physical Pebble + // keys we scan with Next before falling back to SeekGE during small batched + // OCC checks. + conflictCheckSmallLinearAdvanceLimit = 64 + + // conflictCheckLargeLinearAdvanceLimit stays intentionally small: large + // replay batches often touch sparse key ranges, where broad Next scans burn + // CPU before the server starts listening. Fallback seeks are bounded to the + // target key range with SeekGEWithLimit. + conflictCheckLargeLinearAdvanceLimit = 16 + conflictCheckLargeBatchThreshold = 1024 ) // pebbleCacheBytes is the effective per-store Pebble block-cache capacity, @@ -365,6 +378,17 @@ func encodeKey(key []byte, ts uint64) []byte { return k } +func appendEncodedKey(dst []byte, key []byte, ts uint64) []byte { + needed := encodedKeyLen(key) + if cap(dst) < needed { + dst = make([]byte, needed) + } else { + dst = dst[:needed] + } + fillEncodedKey(dst, key, ts) + return dst +} + func encodedKeyLen(key []byte) int { return len(key) + timestampSize } @@ -392,6 +416,22 @@ func keyUpperBound(key []byte) []byte { return nil // key is all 0xFF; no finite upper bound } +func appendKeyUpperBound(dst []byte, key []byte) []byte { + if cap(dst) < len(key) { + dst = make([]byte, len(key)) + } else { + dst = dst[:len(key)] + } + copy(dst, key) + for i := len(dst) - 1; i >= 0; i-- { + dst[i]++ + if dst[i] != 0 { + return dst[:i+1] + } + } + return nil +} + func decodeKey(k []byte) ([]byte, uint64) { if len(k) < timestampSize { return nil, 0 @@ -903,6 +943,75 @@ func (s *pebbleStore) GetAt(ctx context.Context, key []byte, ts uint64) ([]byte, return s.getAt(ctx, key, ts) } +func (s *pebbleStore) GetAtBatch(ctx context.Context, keys [][]byte, ts uint64) (map[string][]byte, error) { + if len(keys) == 0 { + return map[string][]byte{}, nil + } + if err := ctx.Err(); err != nil { + return nil, errors.WithStack(err) + } + + s.dbMu.RLock() + defer s.dbMu.RUnlock() + if err := ctx.Err(); err != nil { + return nil, errors.WithStack(err) + } + if readTSCompacted(ts, s.effectiveMinRetainedTS()) { + return nil, ErrReadTSCompacted + } + + sortedKeys := uniqueSortedKeys(keys) + iter, err := s.db.NewIter(&pebble.IterOptions{}) + if err != nil { + return nil, errors.WithStack(err) + } + defer iter.Close() + + return s.getAtBatchWithIter(ctx, iter, sortedKeys, ts) +} + +func (s *pebbleStore) getAtBatchWithIter(ctx context.Context, iter *pebble.Iterator, sortedKeys [][]byte, ts uint64) (map[string][]byte, error) { + values := make(map[string][]byte, len(sortedKeys)) + var seekKey []byte + var upperBound []byte + for _, key := range sortedKeys { + if err := ctx.Err(); err != nil { + return nil, errors.WithStack(err) + } + seekKey = appendEncodedKey(seekKey, key, ts) + upperBound = appendKeyUpperBound(upperBound, key) + if iter.SeekGEWithLimit(seekKey, upperBound) != pebble.IterValid { + continue + } + value, err := s.readVisibleVersion(iter, key, ts) + if err != nil { + if errors.Is(err, ErrKeyNotFound) { + continue + } + return nil, err + } + values[string(key)] = bytes.Clone(value) + } + return values, nil +} + +func uniqueSortedKeys(keys [][]byte) [][]byte { + sortedKeys := make([][]byte, 0, len(keys)) + seen := make(map[string]struct{}, len(keys)) + for _, key := range keys { + keyString := string(key) + if _, ok := seen[keyString]; ok { + continue + } + seen[keyString] = struct{}{} + sortedKeys = append(sortedKeys, key) + } + sort.Slice(sortedKeys, func(i, j int) bool { + return bytes.Compare(sortedKeys[i], sortedKeys[j]) < 0 + }) + return sortedKeys +} + func (s *pebbleStore) ExistsAt(ctx context.Context, key []byte, ts uint64) (bool, error) { s.dbMu.RLock() defer s.dbMu.RUnlock() @@ -941,6 +1050,10 @@ func (s *pebbleStore) ExistsAt(ctx context.Context, key []byte, ts uint64) (bool func (s *pebbleStore) CommittedVersionAt(_ context.Context, key []byte, commitTS uint64) (bool, error) { s.dbMu.RLock() defer s.dbMu.RUnlock() + return s.committedVersionAtLocked(key, commitTS) +} + +func (s *pebbleStore) committedVersionAtLocked(key []byte, commitTS uint64) (bool, error) { _, closer, err := s.db.Get(encodeKey(key, commitTS)) if err != nil { if errors.Is(err, pebble.ErrNotFound) { @@ -1422,36 +1535,163 @@ func (s *pebbleStore) LatestCommitTS(ctx context.Context, key []byte) (uint64, b return s.latestCommitTS(ctx, key) } -func (s *pebbleStore) checkConflicts(ctx context.Context, mutations []*KVPairMutation, startTS uint64) error { - for _, mut := range mutations { - // Use the internal latestCommitTS to avoid re-acquiring dbMu (the caller - // – ApplyMutations – already holds dbMu.RLock()). - ts, exists, err := s.latestCommitTS(ctx, mut.Key) - if err != nil { - return err +type conflictCheckKey struct { + key []byte + index int +} + +func sortConflictCheckKeys(keys []conflictCheckKey) { + sort.Slice(keys, func(i, j int) bool { + if c := bytes.Compare(keys[i].key, keys[j].key); c != 0 { + return c < 0 } - if exists && ts > startTS { - // Record the first conflicting key's bucket only — a single - // failing ApplyMutations corresponds to one OCC conflict - // regardless of how many subsequent mutations would also - // collide, which matches the Prometheus rate semantics. - s.writeConflicts.record(WriteConflictKindWrite, classifyWriteConflictKey(mut.Key)) - return NewWriteConflictError(mut.Key) + return keys[i].index < keys[j].index + }) +} + +func sortedMutationConflictKeys(mutations []*KVPairMutation) []conflictCheckKey { + keys := make([]conflictCheckKey, len(mutations)) + for i, mut := range mutations { + keys[i] = conflictCheckKey{key: mut.Key, index: i} + } + sortConflictCheckKeys(keys) + return keys +} + +func sortedReadConflictKeys(readKeys [][]byte) []conflictCheckKey { + keys := make([]conflictCheckKey, len(readKeys)) + for i, key := range readKeys { + keys[i] = conflictCheckKey{key: key, index: i} + } + sortConflictCheckKeys(keys) + return keys +} + +func conflictTSFromCurrentIter(iter *pebble.Iterator, key []byte, limit int) (uint64, bool, bool, error) { + for steps := 0; steps < limit && iter.Valid(); steps++ { + userKey, version := decodeKeyView(iter.Key()) + switch c := bytes.Compare(userKey, key); { + case c < 0: + iter.Next() + continue + case c == 0: + return version, true, true, nil + default: + return 0, false, true, nil } } - return nil + if err := iter.Error(); err != nil { + return 0, false, true, errors.WithStack(err) + } + if !iter.Valid() { + return 0, false, true, nil + } + return 0, false, false, nil } -func (s *pebbleStore) checkReadConflicts(ctx context.Context, readKeys [][]byte, startTS uint64) error { - for _, key := range readKeys { - ts, exists, err := s.latestCommitTS(ctx, key) +func conflictCheckLinearAdvanceLimitFor(keyCount int) int { + if keyCount >= conflictCheckLargeBatchThreshold { + return conflictCheckLargeLinearAdvanceLimit + } + return conflictCheckSmallLinearAdvanceLimit +} + +func (s *pebbleStore) latestCommitTSWithIter(ctx context.Context, iter *pebble.Iterator, key []byte, seekKey, upperBound *[]byte, preferNext bool, linearAdvanceLimit int) (uint64, bool, error) { + if err := ctx.Err(); err != nil { + return 0, false, errors.WithStack(err) + } + if preferNext { + ts, exists, decided, err := conflictTSFromCurrentIter(iter, key, linearAdvanceLimit) + if err != nil || decided { + return ts, exists, err + } + } + *seekKey = appendEncodedKey(*seekKey, key, math.MaxUint64) + *upperBound = appendKeyUpperBound(*upperBound, key) + if iter.SeekGEWithLimit(*seekKey, *upperBound) == pebble.IterValid { + userKey, version := decodeKeyView(iter.Key()) + if bytes.Equal(userKey, key) { + return version, true, nil + } + } + if err := iter.Error(); err != nil { + return 0, false, errors.WithStack(err) + } + return 0, false, nil +} + +func nextConflictCheckKeyGroup(keys []conflictCheckKey, start int) ([]byte, int, int) { + current := keys[start].key + minOriginalIndex := keys[start].index + next := start + 1 + for next < len(keys) && bytes.Equal(keys[next].key, current) { + if keys[next].index < minOriginalIndex { + minOriginalIndex = keys[next].index + } + next++ + } + return current, minOriginalIndex, next +} + +func (s *pebbleStore) firstConflictKey(ctx context.Context, keys []conflictCheckKey, startTS uint64) ([]byte, bool, error) { + if len(keys) == 0 { + return nil, false, nil + } + iter, err := s.db.NewIter(&pebble.IterOptions{}) + if err != nil { + return nil, false, errors.WithStack(err) + } + defer iter.Close() + + firstConflictIndex := len(keys) + var firstConflictKey []byte + var seekKey []byte + var upperBound []byte + preferNext := false + linearAdvanceLimit := conflictCheckLinearAdvanceLimitFor(len(keys)) + for i := 0; i < len(keys); { + current, minOriginalIndex, next := nextConflictCheckKeyGroup(keys, i) + ts, exists, err := s.latestCommitTSWithIter(ctx, iter, current, &seekKey, &upperBound, preferNext, linearAdvanceLimit) if err != nil { - return err + return nil, false, err } - if exists && ts > startTS { - s.writeConflicts.record(WriteConflictKindRead, classifyWriteConflictKey(key)) - return NewWriteConflictError(key) + preferNext = iter.Valid() + if exists && ts > startTS && minOriginalIndex < firstConflictIndex { + firstConflictIndex = minOriginalIndex + firstConflictKey = current } + i = next + } + if firstConflictKey == nil { + return nil, false, nil + } + return firstConflictKey, true, nil +} + +func (s *pebbleStore) checkConflicts(ctx context.Context, mutations []*KVPairMutation, startTS uint64) error { + conflictKey, found, err := s.firstConflictKey(ctx, sortedMutationConflictKeys(mutations), startTS) + if err != nil { + return err + } + if found { + // Record the first conflicting key's bucket only — a single + // failing ApplyMutations corresponds to one OCC conflict + // regardless of how many subsequent mutations would also + // collide, which matches the Prometheus rate semantics. + s.writeConflicts.record(WriteConflictKindWrite, classifyWriteConflictKey(conflictKey)) + return NewWriteConflictError(conflictKey) + } + return nil +} + +func (s *pebbleStore) checkReadConflicts(ctx context.Context, readKeys [][]byte, startTS uint64) error { + conflictKey, found, err := s.firstConflictKey(ctx, sortedReadConflictKeys(readKeys), startTS) + if err != nil { + return err + } + if found { + s.writeConflicts.record(WriteConflictKindRead, classifyWriteConflictKey(conflictKey)) + return NewWriteConflictError(conflictKey) } return nil } @@ -1572,6 +1812,94 @@ func (s *pebbleStore) ApplyMutationsRaftAt(ctx context.Context, mutations []*KVP return s.applyMutationsWithOpts(ctx, mutations, readKeys, startTS, commitTS, s.raftApplyWriteOpts(), false, appliedIndex) } +func (s *pebbleStore) raftApplyAlreadyLandedLocked(mutations []*KVPairMutation, commitTS uint64) (bool, error) { + if len(mutations) == 0 { + return false, nil + } + for _, mut := range mutations { + if mut == nil || len(mut.Key) == 0 { + continue + } + return s.committedVersionAtLocked(mut.Key, commitTS) + } + return false, nil +} + +func (s *pebbleStore) shouldWriteAppliedIndexLocked(appliedIndex uint64) (bool, error) { + if appliedIndex == 0 { + return false, nil + } + existing, present, err := s.readAppliedIndexLocked() + if err != nil { + return false, err + } + return !present || existing < appliedIndex, nil +} + +func (s *pebbleStore) writeRaftApplyMetaLocked(newLastTS uint64, writeAppliedIndex bool, appliedIndex uint64, writeOpts *pebble.WriteOptions) error { + b := s.db.NewBatch() + defer func() { _ = b.Close() }() + if err := setPebbleUint64InBatch(b, metaLastCommitTSBytes, newLastTS); err != nil { + return err + } + if writeAppliedIndex { + if err := setPebbleUint64InBatch(b, metaAppliedIndexBytes, appliedIndex); err != nil { + return err + } + } + return errors.WithStack(b.Commit(writeOpts)) +} + +func (s *pebbleStore) markRaftApplyAlreadyLandedLocked(commitTS, appliedIndex uint64, writeOpts *pebble.WriteOptions) error { + writeAppliedIndex, err := s.shouldWriteAppliedIndexLocked(appliedIndex) + if err != nil { + return err + } + + s.mtx.Lock() + defer s.mtx.Unlock() + + newLastTS := s.lastCommitTS + if commitTS > newLastTS { + newLastTS = commitTS + } + if !writeAppliedIndex && newLastTS == s.lastCommitTS { + return nil + } + if err := s.writeRaftApplyMetaLocked(newLastTS, writeAppliedIndex, appliedIndex, writeOpts); err != nil { + return err + } + s.updateLastCommitTS(newLastTS) + return nil +} + +func (s *pebbleStore) staleRaftApplyFastPathLocked(mutations []*KVPairMutation, commitTS, appliedIndex uint64, writeOpts *pebble.WriteOptions) (bool, error) { + if appliedIndex == 0 { + return false, nil + } + landed, err := s.raftApplyAlreadyLandedLocked(mutations, commitTS) + if err != nil || !landed { + return false, err + } + return true, s.markRaftApplyAlreadyLandedLocked(commitTS, appliedIndex, writeOpts) +} + +func (s *pebbleStore) hasCommitsAfter(startTS uint64) bool { + s.mtx.RLock() + defer s.mtx.RUnlock() + return s.lastCommitTS > startTS +} + +func (s *pebbleStore) checkApplyConflicts(ctx context.Context, mutations []*KVPairMutation, readKeys [][]byte, startTS uint64) error { + if !s.hasCommitsAfter(startTS) { + return nil + } + if err := s.checkConflicts(ctx, mutations, startTS); err != nil { + return err + } + return s.checkReadConflicts(ctx, readKeys, startTS) +} + func (s *pebbleStore) applyMutationsWithOpts(ctx context.Context, mutations []*KVPairMutation, readKeys [][]byte, startTS, commitTS uint64, writeOpts *pebble.WriteOptions, gateRegistration bool, appliedIndex uint64) error { s.dbMu.RLock() defer s.dbMu.RUnlock() @@ -1581,13 +1909,21 @@ func (s *pebbleStore) applyMutationsWithOpts(ctx context.Context, mutations []*K s.applyMu.Lock() defer s.applyMu.Unlock() + fastPathDone, err := s.staleRaftApplyFastPathLocked(mutations, commitTS, appliedIndex, writeOpts) + if err != nil { + return err + } + if fastPathDone { + return nil + } + b := s.db.NewBatch() defer b.Close() - if err := s.checkConflicts(ctx, mutations, startTS); err != nil { - return err - } - if err := s.checkReadConflicts(ctx, readKeys, startTS); err != nil { + // Keep OCC at the store boundary: callers pass the read snapshot via + // startTS/readKeys, and leader-issued commit timestamps alone do not prove + // the read set is still current. + if err := s.checkApplyConflicts(ctx, mutations, readKeys, startTS); err != nil { return err } diff --git a/store/lsm_store_applied_index_test.go b/store/lsm_store_applied_index_test.go index ff00ea950..bc7d8f97a 100644 --- a/store/lsm_store_applied_index_test.go +++ b/store/lsm_store_applied_index_test.go @@ -96,6 +96,54 @@ func TestApplyMutationsRaftAt_BundlesMetaAppliedIndex(t *testing.T) { require.Equal(t, []byte("v1"), val) } +func TestApplyMutationsRaftAt_AlreadyLandedAdvancesStaleAppliedIndex(t *testing.T) { + ctx := context.Background() + st := newApplyIndexPebbleStore(t) + ps := pebbleStoreApplied(t, st) + + muts := []*KVPairMutation{ + {Op: OpTypePut, Key: []byte("k1"), Value: []byte("v1")}, + {Op: OpTypePut, Key: []byte("k2"), Value: []byte("v2")}, + } + const ( + startTS uint64 = 100 + commitTS uint64 = 200 + entryIdx uint64 = 42 + ) + + require.NoError(t, ps.ApplyMutations(ctx, muts, nil, commitTS, commitTS)) + require.NoError(t, ps.db.Set(metaAppliedIndexBytes, []byte{0x01}, pebble.Sync)) + + require.NoError(t, ps.ApplyMutationsRaftAt(ctx, muts, nil, startTS, commitTS, entryIdx)) + + got, present, err := ps.LastAppliedIndex() + require.NoError(t, err) + require.True(t, present) + require.Equal(t, entryIdx, got) + + val, err := ps.GetAt(ctx, []byte("k1"), commitTS) + require.NoError(t, err) + require.Equal(t, []byte("v1"), val) +} + +func TestApplyMutationsRaftAt_AlreadyLandedFastPathDoesNotHideConflict(t *testing.T) { + ctx := context.Background() + st := newApplyIndexPebbleStore(t) + ps := pebbleStoreApplied(t, st) + + require.NoError(t, ps.PutAt(ctx, []byte("k"), []byte("newer"), 200, 0)) + + err := ps.ApplyMutationsRaftAt(ctx, []*KVPairMutation{ + {Op: OpTypePut, Key: []byte("k"), Value: []byte("raft")}, + }, nil, 100, 300, 42) + require.ErrorIs(t, err, ErrWriteConflict) + + got, present, idxErr := ps.LastAppliedIndex() + require.NoError(t, idxErr) + require.False(t, present) + require.Equal(t, uint64(0), got) +} + // TestApplyMutationsRaftAt_ZeroIndexLeavesMetaKey covers the // appliedIndex==0 escape hatch — callers without a raft entry index // (test fakes, legacy ApplyMutationsRaft) MUST NOT bump the meta key. diff --git a/store/lsm_store_test.go b/store/lsm_store_test.go index 27b0fe306..509a3cd2f 100644 --- a/store/lsm_store_test.go +++ b/store/lsm_store_test.go @@ -146,6 +146,36 @@ func TestPebbleStore_LastCommitTSPersistedAcrossRestart(t *testing.T) { require.Equal(t, uint64(42), reopened.LastCommitTS()) } +func TestPebbleStore_GetAtBatch(t *testing.T) { + dir, err := os.MkdirTemp("", "pebble-get-at-batch-test") + require.NoError(t, err) + defer os.RemoveAll(dir) + + ctx := context.Background() + s, err := NewPebbleStore(dir) + require.NoError(t, err) + defer s.Close() + + require.NoError(t, s.PutAt(ctx, []byte("b"), []byte("v-b"), 10, 0)) + require.NoError(t, s.PutAt(ctx, []byte("a"), []byte("v-a"), 20, 0)) + require.NoError(t, s.DeleteAt(ctx, []byte("deleted"), 30)) + + ps, ok := s.(*pebbleStore) + require.True(t, ok) + got, err := ps.GetAtBatch(ctx, [][]byte{ + []byte("missing"), + []byte("b"), + []byte("a"), + []byte("b"), + []byte("deleted"), + }, ^uint64(0)) + require.NoError(t, err) + require.Equal(t, map[string][]byte{ + "b": []byte("v-b"), + "a": []byte("v-a"), + }, got) +} + func TestPebbleStore_MinRetainedTSPersistedAcrossRestart(t *testing.T) { dir, err := os.MkdirTemp("", "pebble-min-retained-ts-test") require.NoError(t, err) diff --git a/store/lsm_store_txn_test.go b/store/lsm_store_txn_test.go index be985cc8b..0eee5228d 100644 --- a/store/lsm_store_txn_test.go +++ b/store/lsm_store_txn_test.go @@ -2,6 +2,7 @@ package store import ( "context" + "fmt" "testing" "github.com/cockroachdb/errors" @@ -13,6 +14,62 @@ import ( // ApplyMutations // --------------------------------------------------------------------------- +func TestAppendEncodedKeyReusesBuffer(t *testing.T) { + var buf []byte + buf = appendEncodedKey(buf, []byte("alpha"), 10) + firstPtr := &buf[0] + firstCap := cap(buf) + + buf = appendEncodedKey(buf, []byte("bravo"), 20) + require.Equal(t, firstPtr, &buf[0]) + require.Equal(t, firstCap, cap(buf)) + + userKey, ts := decodeKeyView(buf) + require.Equal(t, []byte("bravo"), userKey) + require.EqualValues(t, 20, ts) +} + +func TestAppendKeyUpperBoundReusesBuffer(t *testing.T) { + var buf []byte + buf = appendKeyUpperBound(buf, []byte("alpha")) + firstPtr := &buf[0] + firstCap := cap(buf) + require.Equal(t, []byte("alphb"), buf) + + buf = appendKeyUpperBound(buf, []byte("bravo")) + require.Equal(t, firstPtr, &buf[0]) + require.Equal(t, firstCap, cap(buf)) + require.Equal(t, []byte("bravp"), buf) + + buf = appendKeyUpperBound(buf, []byte{'a', 0xff}) + require.Equal(t, []byte{'b'}, buf) + + buf = appendKeyUpperBound(buf, []byte{0xff}) + require.Nil(t, buf) +} + +func TestConflictCheckLinearAdvanceLimitFor(t *testing.T) { + require.Equal(t, conflictCheckSmallLinearAdvanceLimit, conflictCheckLinearAdvanceLimitFor(conflictCheckLargeBatchThreshold-1)) + require.Equal(t, conflictCheckLargeLinearAdvanceLimit, conflictCheckLinearAdvanceLimitFor(conflictCheckLargeBatchThreshold)) +} + +func TestPebbleStore_HasCommitsAfter(t *testing.T) { + dir := t.TempDir() + st, err := NewPebbleStore(dir) + require.NoError(t, err) + defer st.Close() + ps, ok := st.(*pebbleStore) + require.True(t, ok, "expected *pebbleStore, got %T", st) + + ctx := context.Background() + require.False(t, ps.hasCommitsAfter(0)) + + require.NoError(t, ps.PutAt(ctx, []byte("k"), []byte("v"), 10, 0)) + require.True(t, ps.hasCommitsAfter(9)) + require.False(t, ps.hasCommitsAfter(10)) + require.False(t, ps.hasCommitsAfter(11)) +} + func TestPebbleStore_ApplyMutations_BasicPut(t *testing.T) { dir := t.TempDir() s, err := NewPebbleStore(dir) @@ -128,6 +185,53 @@ func TestPebbleStore_ApplyMutations_WriteConflict(t *testing.T) { assert.Equal(t, []byte("v1"), val) } +func TestPebbleStore_ApplyMutations_BatchedConflictCheckPreservesInputOrder(t *testing.T) { + dir := t.TempDir() + s, err := NewPebbleStore(dir) + require.NoError(t, err) + defer s.Close() + + ctx := context.Background() + require.NoError(t, s.PutAt(ctx, []byte("z-conflict"), []byte("old-z"), 50, 0)) + require.NoError(t, s.PutAt(ctx, []byte("a-conflict"), []byte("old-a"), 50, 0)) + + mutations := []*KVPairMutation{ + {Op: OpTypePut, Key: []byte("z-conflict"), Value: []byte("new-z")}, + {Op: OpTypePut, Key: []byte("middle"), Value: []byte("middle")}, + {Op: OpTypePut, Key: []byte("a-conflict"), Value: []byte("new-a")}, + } + err = s.ApplyMutations(ctx, mutations, nil, 10, 60) + require.Error(t, err) + require.True(t, errors.Is(err, ErrWriteConflict)) + + conflictKey, ok := WriteConflictKey(err) + require.True(t, ok) + assert.Equal(t, []byte("z-conflict"), conflictKey) +} + +func TestPebbleStore_ApplyMutations_LargeBatchNoConflict(t *testing.T) { + dir := t.TempDir() + s, err := NewPebbleStore(dir) + require.NoError(t, err) + defer s.Close() + + ctx := context.Background() + const mutationCount = 4096 + mutations := make([]*KVPairMutation, 0, mutationCount) + for i := mutationCount - 1; i >= 0; i-- { + mutations = append(mutations, &KVPairMutation{ + Op: OpTypePut, + Key: []byte(fmt.Sprintf("batch-key-%04d", i)), + Value: []byte("value"), + }) + } + + require.NoError(t, s.ApplyMutations(ctx, mutations, nil, 10, 20)) + val, err := s.GetAt(ctx, []byte("batch-key-0000"), 20) + require.NoError(t, err) + assert.Equal(t, []byte("value"), val) +} + func TestPebbleStore_ApplyMutations_NoConflictWhenStartTSGECommit(t *testing.T) { dir := t.TempDir() s, err := NewPebbleStore(dir) @@ -498,6 +602,33 @@ func TestPebbleStore_ApplyMutations_ReadConflict(t *testing.T) { assert.ErrorIs(t, err, ErrKeyNotFound) } +func TestPebbleStore_ApplyMutations_ReadConflictPreservesInputOrder(t *testing.T) { + dir := t.TempDir() + s, err := NewPebbleStore(dir) + require.NoError(t, err) + defer s.Close() + + ctx := context.Background() + require.NoError(t, s.PutAt(ctx, []byte("z-read"), []byte("old-z"), 50, 0)) + require.NoError(t, s.PutAt(ctx, []byte("a-read"), []byte("old-a"), 50, 0)) + + mutations := []*KVPairMutation{ + {Op: OpTypePut, Key: []byte("new-key"), Value: []byte("value")}, + } + readKeys := [][]byte{ + []byte("z-read"), + []byte("middle-read"), + []byte("a-read"), + } + err = s.ApplyMutations(ctx, mutations, readKeys, 10, 60) + require.Error(t, err) + require.True(t, errors.Is(err, ErrWriteConflict)) + + conflictKey, ok := WriteConflictKey(err) + require.True(t, ok) + assert.Equal(t, []byte("z-read"), conflictKey) +} + func TestPebbleStore_ApplyMutations_ReadConflict_NoConflictWhenReadKeyUnchanged(t *testing.T) { dir := t.TempDir() s, err := NewPebbleStore(dir) diff --git a/store/mvcc_store.go b/store/mvcc_store.go index abbda3240..41a9e48ed 100644 --- a/store/mvcc_store.go +++ b/store/mvcc_store.go @@ -606,8 +606,10 @@ func (s *mvccStore) ApplyMutations(ctx context.Context, mutations []*KVPairMutat s.mtx.Lock() defer s.mtx.Unlock() - if err := s.checkConflictsLocked(mutations, readKeys, startTS); err != nil { - return err + if s.lastCommitTS > startTS { + if err := s.checkConflictsLocked(mutations, readKeys, startTS); err != nil { + return err + } } commitTS = s.alignCommitTS(commitTS)