From 97a42ba72c18205fcb84a43ad4ca0cedc48efa31 Mon Sep 17 00:00:00 2001 From: bootjp Date: Sun, 19 Jul 2026 21:15:25 +0900 Subject: [PATCH 01/14] tso: add dedicated ceiling fsm --- .../2026_04_16_partial_centralized_tso.md | 139 +++++-- kv/tso_fsm.go | 339 +++++++++++++++++ kv/tso_fsm_test.go | 345 ++++++++++++++++++ main.go | 228 +++++++----- main_encryption_admin.go | 43 ++- main_encryption_admin_test.go | 62 ++++ multiraft_runtime_test.go | 127 ++++++- 7 files changed, 1144 insertions(+), 139 deletions(-) create mode 100644 kv/tso_fsm.go create mode 100644 kv/tso_fsm_test.go diff --git a/docs/design/2026_04_16_partial_centralized_tso.md b/docs/design/2026_04_16_partial_centralized_tso.md index 631ae4dbf..d3d6732d8 100644 --- a/docs/design/2026_04_16_partial_centralized_tso.md +++ b/docs/design/2026_04_16_partial_centralized_tso.md @@ -2,10 +2,11 @@ - Status: Partial — M1 all-led-group HLC renewal, the M2 reserved-group bootstrap bridge, and the M3-M5 TSO allocator/batch cutover are implemented; - the minimal TSO-only FSM and follower redirect/admin exposure remain open + the minimal TSO-only FSM and runtime group-0 wiring are implemented, while + follower redirect/admin exposure and shadow validation remain open - Author: bootjp - Date: 2026-04-16 -- Updated: 2026-07-11 +- Updated: 2026-07-18 --- @@ -44,13 +45,27 @@ Implemented: bridge still issues timestamps from the locally led data shard through `LocalTSOAllocator`; pinning timestamp issuance to group 0 remains deferred until the TSO-leader redirect path exists. +9. `TSOStateMachine` implements the dedicated-group FSM contract: it accepts + HLC lease entries and explicit allocation-floor entries, halt-fails invalid + entries, snapshots/restores TSO-owned ceiling/floor state, and classifies + full ceiling/floor mirror entries as volatile-only so post-snapshot HLC + mirror raises are not lost on duplicate replay skip paths. +10. Runtime bootstrap instantiates `TSOStateMachine` for `groupID = 0` without + opening an MVCC store. Group 0 retains the normal wrap-aware proposer path + for Raft-envelope cutovers, while route validation keeps all user data out + of the control group. Restore accepts snapshots written by the earlier + compatibility `kvFSM`, imports their HLC ceiling, derives a safe allocation + floor, and drains the legacy MVCC payload for full CRC verification. The + group-0 `EncryptionAdmin` surface is capability-only: mutators are left + unwired and return `FailedPrecondition`. During upgrade replay, valid + encryption control entries committed by the earlier compatibility FSM are + decoded and deterministically rejected without halting the TSO apply loop; + malformed control entries still halt fail-closed. Remaining: -1. Replace the compatibility bridge FSM with the minimal TSO-only FSM that - persists only the HLC ceiling. -2. Add follower redirect/admin exposure for the dedicated TSO leader. -3. Add Phase B shadow-read validation before making dedicated TSO the only +1. Add follower redirect/admin exposure for the dedicated TSO leader. +2. Add Phase B shadow-read validation before making dedicated TSO the only production timestamp path. ### 1.1 Original Limitation @@ -183,7 +198,10 @@ Node-3 ──┘ │ - `groupID = 0` is reserved for the TSO group; all user-data shards use `groupID >= 1`. -- The TSO FSM applies only HLC lease entries. It carries no key-value storage, +- The TSO FSM applies HLC lease entries and explicit allocation-floor entries. + Allocation floors use a versioned TSO envelope whose first byte remains in + the old data-FSM fail-closed range; the full magic prevents an encryption + opcode from being interpreted as TSO state. It carries no key-value storage, so compaction and snapshot overhead are negligible. - Membership mirrors the full cluster so any node can become TSO leader. @@ -191,62 +209,108 @@ Node-3 ──┘ │ ```go // TSOStateMachine implements raftengine.StateMachine. It is a minimal FSM -// that only tracks the HLC physical ceiling; its entire persisted state is a -// single int64 ceiling value. +// that tracks TSO-owned HLC physical ceiling and allocation floor state. The +// shared HLC is a volatile mirror only; snapshots are sourced from the FSM's +// own fields so data-group lease renewals cannot advance group-0 state outside +// group-0 consensus. // // Interface mapping (raftengine.StateMachine vs raw raft.FSM): // Apply([]byte) ← engine strips log envelope, passes raw payload // Snapshot() (Snapshot, _) ← returns raftengine.Snapshot (WriteTo/Close) // Restore(io.Reader) ← caller owns and closes the reader type TSOStateMachine struct { - hlc *HLC + hlc *HLC + ceilingMs atomic.Int64 + allocationFloor atomic.Uint64 } -// Apply processes an HLC lease entry. The engine strips the raft.Log +// Apply processes TSO ceiling/floor entries. The engine strips the raft.Log // envelope and passes only the raw command bytes, matching the // raftengine.StateMachine.Apply([]byte) signature. func (f *TSOStateMachine) Apply(data []byte) any { - if len(data) == 0 || data[0] != raftEncodeHLCLease { - return nil + if len(data) == 0 { + return haltErr(errors.Wrap(ErrTSOStateMachineInvalidEntry, + "empty entry")) } - return f.applyHLCLease(data[1:]) + switch { + case data[0] == raftEncodeHLCLease: + if len(data) != hlcLeaseEntryLen { + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "expected %d bytes, got %d", hlcLeaseEntryLen, len(data))) + } + ceiling := int64(binary.BigEndian.Uint64(data[1:])) + if ceiling <= 0 { + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "non-positive HLC lease ceiling %d", ceiling)) + } + f.applyLeaseCeiling(ceiling) + case bytes.HasPrefix(data, []byte(tsoAllocationFloorEnvelope)): + if len(data) != len(tsoAllocationFloorEnvelope)+8 { + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "invalid allocation-floor envelope length %d", len(data))) + } + floor := binary.BigEndian.Uint64(data[len(tsoAllocationFloorEnvelope):]) + if floor == 0 { + return haltErr(errors.Wrap(ErrTSOStateMachineInvalidEntry, + "zero TSO allocation floor")) + } + f.applyAllocationFloor(floor) + default: + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "unexpected tag 0x%02x", data[0])) + } + return nil } -// Snapshot serialises the ceiling as 8 big-endian bytes. +// Snapshot serialises ceiling and allocationFloor as two 8-byte big-endian +// fields. The source is TSO-applied state, not the shared HLC mirror. // Returns raftengine.Snapshot (WriteTo/Close) rather than raft.FSMSnapshot. func (f *TSOStateMachine) Snapshot() (raftengine.Snapshot, error) { - var ceiling int64 - if f.hlc != nil { - ceiling = f.hlc.PhysicalCeiling() - } - return &tsoSnapshot{ceiling: ceiling}, nil + return &tsoSnapshot{ + ceiling: f.ceilingMs.Load(), + allocationFloor: f.allocationFloor.Load(), + }, nil } -// Restore deserialises the ceiling and updates the shared HLC. +// Restore deserialises ceiling/floor state and mirrors it to the shared HLC. +// Current snapshots are 16 bytes; legacy 8-byte ceiling snapshots derive the +// allocation floor from the ceiling used by the previous FSM format. // The caller owns the reader and is responsible for closing it. func (f *TSOStateMachine) Restore(r io.Reader) error { - var buf [8]byte - if _, err := io.ReadFull(r, buf[:]); err != nil { + payload, err := io.ReadAll(io.LimitReader(r, 17)) + if err != nil { return err } - ceiling := int64(binary.BigEndian.Uint64(buf[:])) - if f.hlc != nil && ceiling > 0 { - f.hlc.SetPhysicalCeiling(ceiling) + var ceiling int64 + var floor uint64 + switch len(payload) { + case 8: + ceiling = int64(binary.BigEndian.Uint64(payload[:8])) + floor = tsoLeaseAllocationFloor(ceiling) + case 16: + ceiling = int64(binary.BigEndian.Uint64(payload[:8])) + floor = binary.BigEndian.Uint64(payload[8:]) + default: + return errors.Newf("expected 8 or 16 snapshot bytes, got %d", + len(payload)) } + f.restoreSnapshotState(ceiling, floor) return nil } -// tsoSnapshot implements raftengine.Snapshot. It serialises the HLC physical -// ceiling as 8 big-endian bytes — the entire persisted state of the TSO group. +// tsoSnapshot implements raftengine.Snapshot. It serialises TSO-owned ceiling +// and floor state as 16 bytes — the entire persisted state of the TSO group. // WriteTo/Close matches the raftengine.Snapshot interface (not Persist/Release // from raft.FSMSnapshot). type tsoSnapshot struct { - ceiling int64 + ceiling int64 + allocationFloor uint64 } func (s *tsoSnapshot) WriteTo(w io.Writer) (int64, error) { - var buf [8]byte + var buf [16]byte binary.BigEndian.PutUint64(buf[:], uint64(s.ceiling)) + binary.BigEndian.PutUint64(buf[8:], s.allocationFloor) n, err := w.Write(buf[:]) return int64(n), err } @@ -633,7 +697,7 @@ least as large as the maximum shard ceiling. | M3 — shipped | Define `TSOAllocator` interface; implement backed by `defaultGroup` | Medium | | M4 — shipped | `BatchAllocator` with atomic counter for low-latency timestamp serving | Medium | | M5 — shipped for default-group bridge | Coordinator feature-flag cutover via `--tsoEnabled`; shadow validation against a dedicated group remains deferred to M6 | Medium | -| M6 — partial | Dedicated TSO Raft group (`groupID = 0`) is reserved/bootstrap-capable and warmed by the HLC renewal bridge; TSO-leader-only timestamp issuance and the minimal `TSOStateMachine` remain open | Low | +| M6 — partial | Dedicated TSO Raft group (`groupID = 0`) is reserved/bootstrap-capable, warmed by the HLC renewal bridge, and runs the minimal `TSOStateMachine` without an MVCC store; TSO-leader redirect/admin exposure and TSO-leader-only timestamp issuance remain open | Low | | M7 | Phase D legacy cleanup + cross-shard SSI read-timestamp validation via TSO | Low | --- @@ -654,6 +718,13 @@ least as large as the maximum shard ceiling. leader via gRPC, or support follower reads with a known-safe timestamp bound? -5. **Backward compatibility:** Existing snapshots and Raft logs store ceiling - values inline in shard FSMs. Migration to a dedicated TSO FSM requires a - coordinated snapshot + compaction pass. +5. **Backward compatibility (resolved for group 0):** Existing group-0 Raft + logs already carry compatible HLC lease entries. `TSOStateMachine.Restore` + recognizes legacy `kvFSM` snapshot headers, imports the ceiling, derives the + allocation floor, and drains the old store payload before Raft resumes log + replay. Valid historical encryption control entries are decoded and rejected + as obsolete ordinary apply responses, so replay advances without mutating + TSO state; new group-0 encryption mutators are disabled at the RPC boundary, + while the group-0 engine remains wired as a leader view so read-only sidecar + recovery cannot be served from a stale follower. Runtime wiring therefore + does not require deleting group-0 state. diff --git a/kv/tso_fsm.go b/kv/tso_fsm.go new file mode 100644 index 000000000..fa18e3176 --- /dev/null +++ b/kv/tso_fsm.go @@ -0,0 +1,339 @@ +package kv + +import ( + "bufio" + "bytes" + "encoding/binary" + "io" + "sync/atomic" + + "github.com/bootjp/elastickv/internal/encryption/fsmwire" + "github.com/bootjp/elastickv/internal/raftengine" + "github.com/cockroachdb/errors" +) + +var _ raftengine.StateMachine = (*TSOStateMachine)(nil) +var _ raftengine.VolatileEntryClassifier = (*TSOStateMachine)(nil) + +var ( + ErrTSOStateMachineInvalidEntry = errors.New("tso fsm: invalid entry") + ErrTSOLegacyEncryptionEntryRejected = errors.New("tso fsm: legacy encryption entry rejected") +) + +const ( + // tsoAllocationFloorEnvelope starts in kvFSM's fail-closed encryption + // range so old data-group FSMs halt on a misrouted entry. The remaining + // magic and version bytes keep it distinct from every encryption opcode. + tsoAllocationFloorEnvelope = "\x07TSOF\x01" + tsoSnapshotV1Len = hlcLeasePayloadLen + tsoSnapshotV2Len = hlcLeasePayloadLen * 2 +) + +// TSOStateMachine is the minimal state machine for the dedicated timestamp +// group. It accepts HLC lease-renewal entries plus explicit allocation-floor +// entries. The HLC is only a volatile mirror; snapshots are sourced from the +// TSO FSM's own applied state so unrelated shard-group lease renewals cannot +// advance group-0 state outside the group-0 log. +type TSOStateMachine struct { + hlc *HLC + ceilingMs atomic.Int64 + allocationFloor atomic.Uint64 +} + +func NewTSOStateMachine(hlc *HLC) *TSOStateMachine { + return &TSOStateMachine{hlc: hlc} +} + +func (f *TSOStateMachine) Apply(data []byte) any { + if len(data) == 0 { + return haltErr(errors.Wrap(ErrTSOStateMachineInvalidEntry, "empty entry")) + } + switch { + case data[0] == raftEncodeHLCLease: + return f.applyLeaseEntry(data) + case bytes.HasPrefix(data, []byte(tsoAllocationFloorEnvelope)): + return f.applyAllocationFloorEntry(data) + case data[0] >= fsmwire.OpEncryptionMin && data[0] <= fsmwire.OpEncryptionMax: + return rejectLegacyTSOEncryptionEntry(data) + default: + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, "unexpected tag 0x%02x", data[0])) + } +} + +// rejectLegacyTSOEncryptionEntry lets an upgraded group 0 advance past a +// control entry committed while it still ran kvFSM. New group-0 listeners do +// not expose encryption mutators, so these entries can only be historical. +// Validate the old wire shape before returning an ordinary apply response: +// malformed control bytes still halt, while a valid obsolete entry is +// deterministically rejected without mutating TSO state or halting replay. +func rejectLegacyTSOEncryptionEntry(data []byte) any { + var err error + switch data[0] { + case fsmwire.OpRegistration: + _, err = fsmwire.DecodeRegistration(data[1:]) + case fsmwire.OpBootstrap: + _, err = fsmwire.DecodeBootstrap(data[1:]) + case fsmwire.OpRotation: + _, err = fsmwire.DecodeRotation(data[1:]) + default: + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "reserved encryption entry tag 0x%02x", data[0])) + } + if err != nil { + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "malformed legacy encryption entry tag 0x%02x: %v", data[0], err)) + } + return errors.Wrapf(ErrTSOLegacyEncryptionEntryRejected, + "tag 0x%02x is not part of dedicated TSO state", data[0]) +} + +func (f *TSOStateMachine) applyLeaseEntry(data []byte) any { + if len(data) != hlcLeaseEntryLen { + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, "expected HLC lease entry length %d, got %d", hlcLeaseEntryLen, len(data))) + } + ceilingMs := int64(binary.BigEndian.Uint64(data[1:])) //nolint:gosec // value is a Unix ms timestamp encoded as uint64. + if ceilingMs <= 0 { + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, "non-positive HLC lease ceiling %d", ceilingMs)) + } + if f != nil { + f.applyLeaseCeiling(ceilingMs) + } + return nil +} + +func (f *TSOStateMachine) applyAllocationFloorEntry(data []byte) any { + expectedLen := len(tsoAllocationFloorEnvelope) + hlcLeasePayloadLen + if len(data) != expectedLen { + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, "expected TSO allocation floor entry length %d, got %d", expectedLen, len(data))) + } + floor := binary.BigEndian.Uint64(data[len(tsoAllocationFloorEnvelope):]) + if floor == 0 { + return haltErr(errors.Wrap(ErrTSOStateMachineInvalidEntry, "zero TSO allocation floor")) + } + if f != nil { + f.applyAllocationFloor(floor) + } + return nil +} + +func (f *TSOStateMachine) Snapshot() (raftengine.Snapshot, error) { + var ceilingMs int64 + var allocationFloor uint64 + if f != nil { + ceilingMs = f.ceilingMs.Load() + allocationFloor = f.allocationFloor.Load() + } + return &tsoFSMSnapshot{ceilingMs: ceilingMs, allocationFloor: allocationFloor}, nil +} + +func (f *TSOStateMachine) Restore(r io.Reader) error { + if r == nil { + return errors.New("tso fsm snapshot: reader is nil") + } + br := tsoSnapshotReader(r) + if legacy, err := restoreLegacyKVFSMSnapshot(f, br); legacy || err != nil { + return err + } + ceilingMs, allocationFloor, err := readTSOSnapshotState(br) + if err != nil { + return err + } + if f != nil { + f.restoreSnapshotState(ceilingMs, allocationFloor) + } + return nil +} + +func tsoSnapshotReader(r io.Reader) *bufio.Reader { + if br, ok := r.(*bufio.Reader); ok { + return br + } + return bufio.NewReader(r) +} + +func readTSOSnapshotState(br *bufio.Reader) (int64, uint64, error) { + payload, err := io.ReadAll(io.LimitReader(br, tsoSnapshotV2Len+1)) + if err != nil { + return 0, 0, errors.Wrap(err, "restore tso fsm snapshot") + } + var ceilingMs int64 + var allocationFloor uint64 + var legacySnapshot bool + switch len(payload) { + case tsoSnapshotV1Len: + legacySnapshot = true + ceilingMs = int64(binary.BigEndian.Uint64(payload[:hlcLeasePayloadLen])) //nolint:gosec // legacy snapshot value. + case tsoSnapshotV2Len: + ceilingMs = int64(binary.BigEndian.Uint64(payload[:hlcLeasePayloadLen])) //nolint:gosec // snapshot value. + allocationFloor = binary.BigEndian.Uint64(payload[hlcLeasePayloadLen:]) + default: + return 0, 0, errors.Wrapf(ErrTSOStateMachineInvalidEntry, "tso fsm snapshot: expected %d or %d bytes, got %d", tsoSnapshotV1Len, tsoSnapshotV2Len, len(payload)) + } + if ceilingMs < 0 { + return 0, 0, errors.Wrapf(ErrTSOStateMachineInvalidEntry, "tso fsm snapshot: negative ceiling %d", ceilingMs) + } + if legacySnapshot && ceilingMs > 0 { + allocationFloor = tsoLeaseAllocationFloor(ceilingMs) + } + return ceilingMs, allocationFloor, nil +} + +// restoreLegacyKVFSMSnapshot migrates snapshots produced while reserved group +// 0 still used kvFSM as a compatibility bridge. Only the HLC header is TSO +// state; the empty MVCC payload is drained so the raft engine can verify the +// complete snapshot CRC. Non-kvFSM snapshots are left untouched in br. +func restoreLegacyKVFSMSnapshot(f *TSOStateMachine, br *bufio.Reader) (bool, error) { + legacy, err := hasLegacyKVFSMSnapshotHeader(br) + if err != nil || !legacy { + return legacy, err + } + ceiling, _, err := ReadSnapshotHeader(br) + if err != nil { + return true, errors.Wrap(err, "tso fsm snapshot: read legacy kv fsm header") + } + if _, err := io.Copy(io.Discard, br); err != nil { + return true, errors.Wrap(err, "tso fsm snapshot: drain legacy kv fsm payload") + } + ceilingMs := int64(ceiling) //nolint:gosec // validated below before use. + if ceilingMs < 0 { + return true, errors.Wrapf(ErrTSOStateMachineInvalidEntry, "tso fsm snapshot: negative legacy ceiling %d", ceilingMs) + } + if f == nil || ceilingMs == 0 { + return true, nil + } + f.restoreSnapshotState(ceilingMs, tsoLeaseAllocationFloor(ceilingMs)) + return true, nil +} + +func hasLegacyKVFSMSnapshotHeader(br *bufio.Reader) (bool, error) { + peeked, err := br.Peek(len(hlcSnapshotMagic)) + if err != nil { + if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) { + return false, nil + } + return false, errors.Wrap(err, "tso fsm snapshot: peek legacy header") + } + switch { + case isV1Magic(peeked): + return true, nil + case isV2Magic(peeked): + return true, nil + case isUnknownEKVTHLC(peeked): + return true, nil + default: + return false, nil + } +} + +func (f *TSOStateMachine) IsVolatileOnlyPayload(payload []byte) bool { + return len(payload) == hlcLeaseEntryLen && payload[0] == raftEncodeHLCLease || + len(payload) == len(tsoAllocationFloorEnvelope)+hlcLeasePayloadLen && + bytes.HasPrefix(payload, []byte(tsoAllocationFloorEnvelope)) +} + +func (f *TSOStateMachine) applyLeaseCeiling(ceilingMs int64) { + if f == nil || ceilingMs <= 0 { + return + } + storeMaxInt64(&f.ceilingMs, ceilingMs) + if f.hlc != nil { + f.hlc.SetPhysicalCeiling(f.ceilingMs.Load()) + } +} + +func (f *TSOStateMachine) applyAllocationFloor(floor uint64) { + if f == nil || floor == 0 { + return + } + storeMaxUint64(&f.allocationFloor, floor) + if f.hlc != nil { + f.hlc.Observe(f.allocationFloor.Load()) + } +} + +func (f *TSOStateMachine) restoreSnapshotState(ceilingMs int64, allocationFloor uint64) { + if f == nil { + return + } + if ceilingMs > 0 { + storeMaxInt64(&f.ceilingMs, ceilingMs) + } + if allocationFloor > 0 { + storeMaxUint64(&f.allocationFloor, allocationFloor) + } + if f.hlc != nil { + if currentCeiling := f.ceilingMs.Load(); currentCeiling > 0 { + f.hlc.SetPhysicalCeiling(currentCeiling) + } + if currentFloor := f.allocationFloor.Load(); currentFloor > 0 { + f.hlc.Observe(currentFloor) + } + } +} + +func storeMaxInt64(value *atomic.Int64, candidate int64) { + for { + current := value.Load() + if candidate <= current { + return + } + if value.CompareAndSwap(current, candidate) { + return + } + } +} + +func storeMaxUint64(value *atomic.Uint64, candidate uint64) { + for { + current := value.Load() + if candidate <= current { + return + } + if value.CompareAndSwap(current, candidate) { + return + } + } +} + +func tsoLeaseAllocationFloor(ceilingMs int64) uint64 { + return (nonNegativeUint64(ceilingMs) << hlcLogicalBits) | hlcLogicalMask +} + +func marshalTSOAllocationFloor(floor uint64) []byte { + out := make([]byte, len(tsoAllocationFloorEnvelope)+hlcLeasePayloadLen) + copy(out, tsoAllocationFloorEnvelope) + binary.BigEndian.PutUint64(out[len(tsoAllocationFloorEnvelope):], floor) + return out +} + +type tsoFSMSnapshot struct { + ceilingMs int64 + allocationFloor uint64 +} + +func (s *tsoFSMSnapshot) WriteTo(w io.Writer) (int64, error) { + if w == nil { + return 0, errors.New("tso fsm snapshot: writer is nil") + } + var ceilingMs int64 + var allocationFloor uint64 + if s != nil { + ceilingMs = s.ceilingMs + allocationFloor = s.allocationFloor + } + var buf [tsoSnapshotV2Len]byte + binary.BigEndian.PutUint64(buf[:], uint64(ceilingMs)) //nolint:gosec // ceilingMs is a Unix ms timestamp. + binary.BigEndian.PutUint64(buf[hlcLeasePayloadLen:], allocationFloor) + n, err := w.Write(buf[:]) + if err != nil { + return int64(n), errors.Wrap(err, "write tso fsm snapshot") + } + if n != len(buf) { + return int64(n), errors.WithStack(io.ErrShortWrite) + } + return int64(n), nil +} + +func (s *tsoFSMSnapshot) Close() error { + return nil +} diff --git a/kv/tso_fsm_test.go b/kv/tso_fsm_test.go new file mode 100644 index 000000000..4c60bbba5 --- /dev/null +++ b/kv/tso_fsm_test.go @@ -0,0 +1,345 @@ +package kv + +import ( + "bufio" + "bytes" + "encoding/binary" + "testing" + "time" + + "github.com/bootjp/elastickv/internal/encryption/fsmwire" + "github.com/bootjp/elastickv/store" + "github.com/stretchr/testify/require" +) + +func TestTSOStateMachineApplyHLCLeaseUpdatesCeiling(t *testing.T) { + t.Parallel() + + const ceilingMs = int64(1_700_000_123_456) + hlc := NewHLC() + fsm := NewTSOStateMachine(hlc) + + result := fsm.Apply(marshalHLCLeaseRenew(ceilingMs)) + require.Nil(t, result) + require.Equal(t, ceilingMs, hlc.PhysicalCeiling()) + require.Zero(t, hlc.Current()) +} + +func TestTSOStateMachineApplyAllocationFloorAdvancesHLC(t *testing.T) { + t.Parallel() + + ceilingMs := time.Now().Add(time.Hour).UnixMilli() + hlc := NewHLC() + fsm := NewTSOStateMachine(hlc) + + require.Nil(t, fsm.Apply(marshalHLCLeaseRenew(ceilingMs))) + floor := tsoLeaseAllocationFloor(ceilingMs) + require.Zero(t, hlc.Current()) + + require.Nil(t, fsm.Apply(marshalTSOAllocationFloor(floor))) + require.Equal(t, floor, hlc.Current()) + + base, err := hlc.NextBatchFenced(1) + require.NoError(t, err) + require.Greater(t, base, floor) +} + +func TestTSOStateMachineRejectsNonLeaseEntry(t *testing.T) { + t.Parallel() + + payload := make([]byte, hlcLeaseEntryLen) + payload[0] = raftEncodeSingle + + err := requireTSOHaltError(t, NewTSOStateMachine(NewHLC()).Apply(payload)) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) +} + +func TestTSOStateMachineRejectsLegacyEncryptionEntriesWithoutHalting(t *testing.T) { + t.Parallel() + + registration := fsmwire.RegistrationPayload{DEKID: 1, FullNodeID: 2, LocalEpoch: 3} + tests := []struct { + name string + opcode byte + payload []byte + }{ + { + name: "registration", + opcode: fsmwire.OpRegistration, + payload: fsmwire.EncodeRegistration(registration), + }, + { + name: "bootstrap", + opcode: fsmwire.OpBootstrap, + payload: fsmwire.EncodeBootstrap(fsmwire.BootstrapPayload{ + StorageDEKID: 1, + WrappedStorage: []byte("storage"), + RaftDEKID: 2, + WrappedRaft: []byte("raft"), + BatchRegistry: []fsmwire.RegistrationPayload{registration}, + }), + }, + { + name: "rotation", + opcode: fsmwire.OpRotation, + payload: fsmwire.EncodeRotation(fsmwire.RotationPayload{ + SubTag: fsmwire.RotateSubRotateDEK, + DEKID: 4, + Purpose: fsmwire.PurposeStorage, + Wrapped: []byte("rotated"), + ProposerRegistration: registration, + }), + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + entry := append([]byte{tc.opcode}, tc.payload...) + result := NewTSOStateMachine(NewHLC()).Apply(entry) + err, ok := result.(error) + require.Truef(t, ok, "legacy control response type = %T, want ordinary error", result) + require.ErrorIs(t, err, ErrTSOLegacyEncryptionEntryRejected) + _, halts := result.(interface{ HaltApply() error }) + require.False(t, halts, "valid legacy control entry must not halt replay") + }) + } +} + +func TestTSOStateMachineHaltsMalformedLegacyEncryptionEntry(t *testing.T) { + t.Parallel() + + err := requireTSOHaltError(t, NewTSOStateMachine(NewHLC()).Apply([]byte{fsmwire.OpRegistration})) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) +} + +func TestTSOStateMachineRejectsMalformedLease(t *testing.T) { + t.Parallel() + + for _, payload := range [][]byte{ + {}, + {raftEncodeHLCLease}, + append([]byte{raftEncodeHLCLease}, make([]byte, hlcLeasePayloadLen+1)...), + } { + err := requireTSOHaltError(t, NewTSOStateMachine(NewHLC()).Apply(payload)) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) + } +} + +func TestTSOStateMachineRejectsNonPositiveLeaseCeiling(t *testing.T) { + t.Parallel() + + for _, ceilingMs := range []int64{0, -1} { + err := requireTSOHaltError(t, NewTSOStateMachine(NewHLC()).Apply(marshalHLCLeaseRenew(ceilingMs))) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) + } +} + +func TestTSOStateMachineRejectsMalformedAllocationFloor(t *testing.T) { + t.Parallel() + + for _, payload := range [][]byte{ + []byte(tsoAllocationFloorEnvelope), + append([]byte(tsoAllocationFloorEnvelope), make([]byte, hlcLeasePayloadLen+1)...), + marshalTSOAllocationFloor(0), + } { + err := requireTSOHaltError(t, NewTSOStateMachine(NewHLC()).Apply(payload)) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) + } +} + +func TestTSOStateMachineRejectsBareEncryptionReservedAllocationFloor(t *testing.T) { + t.Parallel() + + payload := make([]byte, hlcLeaseEntryLen) + payload[0] = fsmwire.OpEncryptionMax + binary.BigEndian.PutUint64(payload[1:], 42) + err := requireTSOHaltError(t, NewTSOStateMachine(NewHLC()).Apply(payload)) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) +} + +func TestTSOStateMachineNilHLCDoesNotPanic(t *testing.T) { + t.Parallel() + + require.Nil(t, NewTSOStateMachine(nil).Apply(marshalHLCLeaseRenew(1_700_000_123_456))) + require.Nil(t, NewTSOStateMachine(nil).Apply(marshalTSOAllocationFloor(1))) +} + +func TestTSOStateMachineSnapshotRestoreRoundTrip(t *testing.T) { + t.Parallel() + + const ceilingMs = int64(1_700_000_654_321) + sourceHLC := NewHLC() + source := NewTSOStateMachine(sourceHLC) + require.Nil(t, source.Apply(marshalHLCLeaseRenew(ceilingMs))) + floor := tsoLeaseAllocationFloor(ceilingMs) + require.Nil(t, source.Apply(marshalTSOAllocationFloor(floor))) + + snap, err := source.Snapshot() + require.NoError(t, err) + defer func() { require.NoError(t, snap.Close()) }() + + var buf bytes.Buffer + n, err := snap.WriteTo(&buf) + require.NoError(t, err) + require.EqualValues(t, tsoSnapshotV2Len, n) + require.Len(t, buf.Bytes(), tsoSnapshotV2Len) + + targetHLC := NewHLC() + target := NewTSOStateMachine(targetHLC) + require.NoError(t, target.Restore(bytes.NewReader(buf.Bytes()))) + require.Equal(t, ceilingMs, targetHLC.PhysicalCeiling()) + require.Equal(t, floor, targetHLC.Current()) +} + +func TestTSOStateMachineSnapshotUsesTSOOwnedCeiling(t *testing.T) { + t.Parallel() + + const ( + tsoCeiling = int64(1_000) + unrelatedCeiling = int64(2_000) + ) + sourceHLC := NewHLC() + source := NewTSOStateMachine(sourceHLC) + require.Nil(t, source.Apply(marshalHLCLeaseRenew(tsoCeiling))) + sourceHLC.SetPhysicalCeiling(unrelatedCeiling) + + snap, err := source.Snapshot() + require.NoError(t, err) + defer func() { require.NoError(t, snap.Close()) }() + + var buf bytes.Buffer + _, err = snap.WriteTo(&buf) + require.NoError(t, err) + + targetHLC := NewHLC() + require.NoError(t, NewTSOStateMachine(targetHLC).Restore(bytes.NewReader(buf.Bytes()))) + require.Equal(t, tsoCeiling, targetHLC.PhysicalCeiling()) + require.Zero(t, targetHLC.Current()) +} + +func TestTSOStateMachineRestoreRejectsTruncatedSnapshot(t *testing.T) { + t.Parallel() + + err := NewTSOStateMachine(NewHLC()).Restore(bytes.NewReader([]byte{0x01, 0x02})) + require.Error(t, err) +} + +func TestTSOStateMachineLegacySnapshotProbeRejectsEveryShortMagicPrefix(t *testing.T) { + t.Parallel() + + for size := range len(hlcSnapshotMagic) { + payload := append([]byte(nil), hlcSnapshotMagic[:size]...) + legacy, err := hasLegacyKVFSMSnapshotHeader(bufio.NewReader(bytes.NewReader(payload))) + require.NoError(t, err) + require.False(t, legacy) + } +} + +func TestTSOStateMachineRestoreLegacySnapshotDerivesAllocationFloor(t *testing.T) { + t.Parallel() + + const ceilingMs = int64(1_700_000_654_321) + var buf [hlcLeasePayloadLen]byte + binary.BigEndian.PutUint64(buf[:], uint64(ceilingMs)) + + hlc := NewHLC() + require.NoError(t, NewTSOStateMachine(hlc).Restore(bytes.NewReader(buf[:]))) + require.Equal(t, ceilingMs, hlc.PhysicalCeiling()) + require.Equal(t, tsoLeaseAllocationFloor(ceilingMs), hlc.Current()) +} + +func TestTSOStateMachineRestoreKeepsMonotonicCeiling(t *testing.T) { + t.Parallel() + + const ( + higherCeiling = int64(2_000) + lowerCeiling = int64(1_000) + ) + hlc := NewHLC() + fsm := NewTSOStateMachine(hlc) + require.Nil(t, fsm.Apply(marshalHLCLeaseRenew(higherCeiling))) + require.Nil(t, fsm.Apply(marshalTSOAllocationFloor(tsoLeaseAllocationFloor(higherCeiling)))) + + var buf [tsoSnapshotV2Len]byte + binary.BigEndian.PutUint64(buf[:], uint64(lowerCeiling)) + + require.NoError(t, fsm.Restore(bytes.NewReader(buf[:]))) + require.Equal(t, higherCeiling, hlc.PhysicalCeiling()) + require.Equal(t, tsoLeaseAllocationFloor(higherCeiling), hlc.Current()) + + snap, err := fsm.Snapshot() + require.NoError(t, err) + defer func() { require.NoError(t, snap.Close()) }() + + var snapBuf bytes.Buffer + n, err := snap.WriteTo(&snapBuf) + require.NoError(t, err) + require.EqualValues(t, tsoSnapshotV2Len, n) + require.Equal(t, uint64(higherCeiling), binary.BigEndian.Uint64(snapBuf.Bytes()[:hlcLeasePayloadLen])) + require.Equal(t, tsoLeaseAllocationFloor(higherCeiling), binary.BigEndian.Uint64(snapBuf.Bytes()[hlcLeasePayloadLen:])) +} + +func TestTSOStateMachineRestoresLegacyKVFSMSnapshot(t *testing.T) { + const ceilingMs = int64(1_700_000_123_456) + tests := []struct { + name string + opts []FSMOption + }{ + {name: "v1"}, + {name: "v2", opts: []FSMOption{WithCutoverSource(&staticCutoverSource{v: 42})}}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + legacyClock := NewHLC() + legacyClock.SetPhysicalCeiling(ceilingMs) + legacyFSM := NewKvFSMWithHLC(store.NewMVCCStore(), legacyClock, tc.opts...) + legacySnapshot, err := legacyFSM.Snapshot() + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, legacySnapshot.Close()) }) + + var legacy bytes.Buffer + _, err = legacySnapshot.WriteTo(&legacy) + require.NoError(t, err) + + targetClock := NewHLC() + target := NewTSOStateMachine(targetClock) + require.NoError(t, target.Restore(bytes.NewReader(legacy.Bytes()))) + require.Equal(t, ceilingMs, targetClock.PhysicalCeiling()) + require.Equal(t, tsoLeaseAllocationFloor(ceilingMs), targetClock.Current()) + + got, err := target.Snapshot() + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, got.Close()) }) + var payload bytes.Buffer + _, err = got.WriteTo(&payload) + require.NoError(t, err) + require.Len(t, payload.Bytes(), tsoSnapshotV2Len) + require.Equal(t, uint64(ceilingMs), binary.BigEndian.Uint64(payload.Bytes()[:hlcLeasePayloadLen])) + require.Equal(t, tsoLeaseAllocationFloor(ceilingMs), binary.BigEndian.Uint64(payload.Bytes()[hlcLeasePayloadLen:])) + }) + } +} + +func TestTSOStateMachineClassifiesOnlyFullLeaseEntriesAsVolatile(t *testing.T) { + t.Parallel() + + fsm := NewTSOStateMachine(NewHLC()) + require.True(t, fsm.IsVolatileOnlyPayload(marshalHLCLeaseRenew(1_700_000_123_456))) + require.True(t, fsm.IsVolatileOnlyPayload(marshalTSOAllocationFloor(1))) + require.False(t, fsm.IsVolatileOnlyPayload([]byte{raftEncodeHLCLease})) + require.False(t, fsm.IsVolatileOnlyPayload([]byte(tsoAllocationFloorEnvelope))) + require.False(t, fsm.IsVolatileOnlyPayload([]byte{raftEncodeSingle})) +} + +func requireTSOHaltError(t *testing.T, result any) error { + t.Helper() + + if _, ok := result.(error); ok { + t.Fatalf("expected HaltApply response, got plain error %T", result) + } + halt, ok := result.(interface{ HaltApply() error }) + require.Truef(t, ok, "expected HaltApply response, got %T", result) + err := halt.HaltApply() + require.Error(t, err) + return err +} diff --git a/main.go b/main.go index 479bcc69b..5addfc099 100644 --- a/main.go +++ b/main.go @@ -1390,106 +1390,149 @@ func buildShardGroups( // breaking the shared-cache invariant 6D-6c-1 relies on. encWiring = encWiring.withDefaultedCache() multi = effectiveMultiDataDirs(groups, multi) + builder := shardGroupBuilder{ + raftID: raftID, + raftDir: raftDir, + multi: multi, + bootstrap: bootstrap, + bootstrapCfg: bootstrapCfg, + factory: factory, + proposalObserverForGroup: proposalObserverForGroup, + clock: clock, + kekWrapper: kekWrapper, + keystore: keystore, + sidecarPath: sidecarPath, + encWiring: encWiring, + routeEngine: routeEngine, + applyObservers: applyObservers, + } runtimes := make([]*raftGroupRuntime, 0, len(groups)) shardGroups := make(map[uint64]*kv.ShardGroup, len(groups)) for _, g := range groups { - dir := groupDataDir(raftDir, raftID, g.id, multi) - if err := os.MkdirAll(dir, dirPerm); err != nil { - return nil, nil, errors.Wrapf(err, "failed to create fsm store dir for group %d", g.id) - } - st, err := store.NewPebbleStore(filepath.Join(dir, "fsm.db"), encWiring.pebbleOptions()...) + runtime, sg, err := builder.build(g) if err != nil { - return nil, nil, errors.Wrapf(err, "failed to open pebble fsm store for group %d", g.id) - } - // Each shard FSM shares the same HLC so any shard's lease renewal advances - // the global physicalCeiling. The logical counter remains in-memory only. - // - // §6.3 EncryptionApplier wiring (Stage 6A). Constructs an - // Applier backed by the FSM's Pebble store and threads it - // into the FSM dispatch via kv.WithEncryption. The applier's - // ApplyBootstrap / ApplyRotation paths return ErrKEKNotConfigured - // until Stage 6B threads in the KEK plumbing; only - // ApplyRegistration is fully functional in 6A, but the gRPC - // mutator gate (registerEncryptionAdminServer, Stage 5D - // posture) keeps the operator surface inert until 6B re-enables - // it gated on (--encryption-enabled AND KEKConfigured()). - reg, err := store.WriterRegistryFor(st) - if err != nil { - for _, rt := range runtimes { - rt.Close() - } - _ = st.Close() - return nil, nil, errors.Wrapf(err, "failed to construct writer registry for group %d", g.id) - } - // Stage 6B-2: thread WithKEK + WithKeystore + WithSidecarPath - // into the Applier when all three are supplied. Without - // any of them the applier stays in the Stage 6A posture - // (ApplyBootstrap / ApplyRotation return ErrKEKNotConfigured) - // which is exactly the desired apply-time fail-closed - // behaviour when an operator has not opted in to encryption. - // - // Stage 6D-6c-1: WithStateCache threads the process-shared - // StateCache so an encryption apply landing on this shard's - // FSM updates the atomics every shard's storage layer reads. - // Stage 7b' §3.1: WithLocalEpoch threads this process load's - // pinned storage write-path w.epoch into the Applier so - // applyRotateDEK (PurposeStorage) can write each node's own - // highest-emitted local_epoch into Keys[newDEK].LocalEpoch - // rather than the proposer's 0. Stage 6E does the same for - // raft rotations through WithRaftLocalEpoch using the raft - // DEK's separate per-process epoch. - applierOpts := encryptionApplierOptionsFor(kekWrapper, keystore, sidecarPath, encWiring) - applier, err := encryption.NewApplier(reg, applierOpts...) - if err != nil { - for _, rt := range runtimes { - rt.Close() - } - _ = st.Close() - return nil, nil, errors.Wrapf(err, "failed to construct encryption applier for group %d", g.id) - } - // Composed-1 M2 plumbing: wire the shared route catalog - // engine and this shard's owning group ID so M3's - // verifyComposed1 apply-time gate can resolve the - // observed-version owner-of-key without further plumbing - // work. At M2 the FSM stores both but does not consult them; - // see docs/design/2026_05_29_implemented_composed1_cross_group_commit_guard.md - // §M2. - sm := kv.NewKvFSMWithHLC(st, clock, fsmOptionsForGroup(applier, routeEngine, g.id, encWiring, applyObservers...)...) - groupBootstrap, groupBootstrapServers, groupBootstrapSeed := bootstrapSettingsForGroup(bootstrapCfg, g.id, bootstrap) - runtime, err := buildRuntimeForGroup( - raftID, g, raftDir, multi, groupBootstrap, - groupBootstrapServers, groupBootstrapSeed, - st, sm, factory, *raftJoinAsLearner) - if err != nil { - for _, rt := range runtimes { - rt.Close() - } - _ = st.Close() - return nil, nil, errors.Wrapf(err, "failed to start raft group %d", g.id) + closeRaftGroupRuntimes(runtimes) + return nil, nil, err } runtimes = append(runtimes, runtime) - // Stage 6E-2c: route every shard group's TransactionManager - // through NewLeaderProxyForShardGroup so the proposer chain - // consults sg.raftPayloadWrap on every Propose / ProposeAdmin. - // The cell is nil at startup (the Stage 3 default — payloads - // pass through cleartext); Stage 6E-2d's EnableRaftEnvelope - // handler will publish the active wrap closure the instant - // the cutover entry commits. Going through this constructor - // for every group avoids a future "I forgot to wire wrap on - // this code path" regression — if a new shard group bypasses - // this helper, the wrap install would silently no-op on - // proposals for that group. - sg := &kv.ShardGroup{ - Engine: runtime.engine, - Store: st, - } - sg.Txn = kv.NewLeaderProxyForShardGroup(sg, kv.WithProposalObserver(observerForGroup(proposalObserverForGroup, g.id))) shardGroups[g.id] = sg encWiring.attachRaftEnvelopeGroup(g.id, sg) } return runtimes, shardGroups, nil } +type shardGroupBuilder struct { + raftID string + raftDir string + multi bool + bootstrap bool + bootstrapCfg raftBootstrapConfig + factory raftengine.Factory + proposalObserverForGroup func(uint64) kv.ProposalObserver + clock *kv.HLC + kekWrapper kek.Wrapper + keystore *encryption.Keystore + sidecarPath string + encWiring encryptionWriteWiring + routeEngine *distribution.Engine + applyObservers []kv.ApplyObserver +} + +func (b shardGroupBuilder) build(group groupSpec) (*raftGroupRuntime, *kv.ShardGroup, error) { + groupBootstrap, groupBootstrapServers, groupBootstrapSeed := bootstrapSettingsForGroup(b.bootstrapCfg, group.id, b.bootstrap) + observer := observerForGroup(b.proposalObserverForGroup, group.id) + if group.id == dedicatedTSORaftGroupID { + runtime, sg, err := buildDedicatedTSOGroup( + b.raftID, group, b.raftDir, b.multi, groupBootstrap, + groupBootstrapServers, groupBootstrapSeed, + b.factory, b.clock, observer, + ) + return runtime, sg, errors.Wrap(err, "failed to start dedicated TSO group") + } + return b.buildDataGroup(group, groupBootstrap, groupBootstrapServers, groupBootstrapSeed, observer) +} + +func (b shardGroupBuilder) buildDataGroup( + group groupSpec, + bootstrap bool, + bootstrapServers []raftengine.Server, + bootstrapSeed []raftengine.Server, + proposalObserver kv.ProposalObserver, +) (*raftGroupRuntime, *kv.ShardGroup, error) { + dir := groupDataDir(b.raftDir, b.raftID, group.id, b.multi) + if err := os.MkdirAll(dir, dirPerm); err != nil { + return nil, nil, errors.Wrapf(err, "failed to create fsm store dir for group %d", group.id) + } + st, err := store.NewPebbleStore(filepath.Join(dir, "fsm.db"), b.encWiring.pebbleOptions()...) + if err != nil { + return nil, nil, errors.Wrapf(err, "failed to open pebble fsm store for group %d", group.id) + } + reg, err := store.WriterRegistryFor(st) + if err != nil { + _ = st.Close() + return nil, nil, errors.Wrapf(err, "failed to construct writer registry for group %d", group.id) + } + applierOpts := encryptionApplierOptionsFor(b.kekWrapper, b.keystore, b.sidecarPath, b.encWiring) + applier, err := encryption.NewApplier(reg, applierOpts...) + if err != nil { + _ = st.Close() + return nil, nil, errors.Wrapf(err, "failed to construct encryption applier for group %d", group.id) + } + sm := kv.NewKvFSMWithHLC(st, b.clock, fsmOptionsForGroup(applier, b.routeEngine, group.id, b.encWiring, b.applyObservers...)...) + runtime, err := buildRuntimeForGroup( + b.raftID, group, b.raftDir, b.multi, bootstrap, + bootstrapServers, bootstrapSeed, + st, sm, b.factory, *raftJoinAsLearner, + ) + if err != nil { + _ = st.Close() + return nil, nil, errors.Wrapf(err, "failed to start raft group %d", group.id) + } + // Every group goes through the wrap-aware proposer so HLC renewals and + // transactional proposals observe raft-envelope cutovers uniformly. + sg := &kv.ShardGroup{Engine: runtime.engine, Store: st} + sg.Txn = kv.NewLeaderProxyForShardGroup(sg, kv.WithProposalObserver(proposalObserver)) + return runtime, sg, nil +} + +func closeRaftGroupRuntimes(runtimes []*raftGroupRuntime) { + for _, runtime := range runtimes { + runtime.Close() + } +} + +// buildDedicatedTSOGroup constructs group 0 with the minimal TSO state machine. +// It intentionally does not open an MVCC store: shard routing validation keeps +// user data out of group 0, while the state machine persists only its ceiling +// and allocation floor in Raft snapshots. The ShardGroup still receives the +// normal proposer wrapper so lease renewals participate in raft-envelope +// cutovers exactly like data-group proposals. +func buildDedicatedTSOGroup( + raftID string, + group groupSpec, + baseDir string, + multi bool, + bootstrap bool, + bootstrapServers []raftengine.Server, + bootstrapSeed []raftengine.Server, + factory raftengine.Factory, + clock *kv.HLC, + proposalObserver kv.ProposalObserver, +) (*raftGroupRuntime, *kv.ShardGroup, error) { + sm := kv.NewTSOStateMachine(clock) + runtime, err := buildRuntimeForGroup( + raftID, group, baseDir, multi, bootstrap, + bootstrapServers, bootstrapSeed, + nil, sm, factory, *raftJoinAsLearner, + ) + if err != nil { + return nil, nil, err + } + sg := &kv.ShardGroup{Engine: runtime.engine} + sg.Txn = kv.NewLeaderProxyForShardGroup(sg, kv.WithProposalObserver(proposalObserver)) + return runtime, sg, nil +} + func bootstrapSettingsForGroup( cfg raftBootstrapConfig, groupID uint64, @@ -2557,12 +2600,17 @@ func startRaftServers( // Proposer + LeaderView for the cutover marker, while // ShardGroup.Proposer() supplies the wrap-aware post-cutover // path for normal admin entries. + runtimeMutators, runtimeEncryptionEngine := encryptionAdminWiringForGroup( + rt.spec.id, + enableMutators, + rt.engine, + ) registerEncryptionAdminServer( gs, etcdraftengine.DeriveNodeID(*raftId), *encryptionSidecarPath, - enableMutators, - rt.engine, + runtimeMutators, + runtimeEncryptionEngine, encryptionCapabilityFanout, adapter.WithEncryptionAdminLatestAppliedIndex(appliedIndexForEngine(rt.engine)), adapter.WithEncryptionAdminPostCutoverProposer(proposerForGroup(rt, shardGroups)), diff --git a/main_encryption_admin.go b/main_encryption_admin.go index 7a92976d6..80866d4f6 100644 --- a/main_encryption_admin.go +++ b/main_encryption_admin.go @@ -76,6 +76,20 @@ type encryptionAdminEngine interface { raftengine.LeaderView } +// encryptionAdminWiringForGroup keeps the read-only capability service on the +// dedicated TSO listener while withholding every mutating proposal path. The +// TSO FSM has no encryption applier or writer registry, so allowing a mutator +// to propose there would commit an entry that cannot update dedicated TSO +// state. The engine remains available as a LeaderView so ResyncSidecar keeps +// its quorum-confirmed leader-only freshness contract. Data groups retain the +// normal flag-driven mutator wiring. +func encryptionAdminWiringForGroup(groupID uint64, enableMutators bool, engine encryptionAdminEngine) (bool, encryptionAdminEngine) { + if groupID == dedicatedTSORaftGroupID { + return false, engine + } + return enableMutators, engine +} + // registerEncryptionAdminServer constructs and registers an // EncryptionAdminServer on the supplied gRPC server. The function // is intentionally per-shard: the §7.1 Phase-0 GetCapability @@ -92,15 +106,14 @@ type encryptionAdminEngine interface { // number — Codex r1 P1 on PR #760 caught the original wiring // passing the shard id by mistake. // -// Stage 6B-2: the mutator wiring (Proposer + LeaderView) is now -// gated on the supplied enableMutators boolean, which the caller -// in main.go computes as (--encryption-enabled AND --kekFile -// non-empty AND engine non-nil). When enableMutators is false, -// Proposer + LeaderView stay unwired and EncryptionAdminServer's -// BootstrapEncryption / RotateDEK / RegisterEncryptionWriter -// short-circuit at the gRPC boundary with FailedPrecondition — -// identical to the Stage 5D posture. When enableMutators is -// true, both options are wired and the mutators reach the §6.3 +// Stage 6B-2 gates the mutating Proposer on enableMutators, which +// the caller in main.go computes as (--encryption-enabled AND +// --kekFile non-empty AND engine non-nil). The LeaderView remains +// wired whenever an engine is present because read-only recovery +// RPCs also require quorum-confirmed leader freshness. When +// enableMutators is false, BootstrapEncryption / RotateDEK / +// RegisterEncryptionWriter short-circuit at the gRPC boundary with +// FailedPrecondition. When true, the proposer reaches the §6.3 // applier through the supplied engine. Callers may append a // wrap-aware post-cutover proposer option so non-cutover admin // entries wrap after EnableRaftEnvelope while the cutover marker @@ -122,9 +135,9 @@ type encryptionAdminEngine interface { // practice. // // Validate() panics on a misconfiguration that wires a proposer -// without a LeaderView. The wiring below pairs both options -// together (both wired iff enableMutators) so the invariant -// holds by construction. +// without a LeaderView. The wiring below always pairs a mutating +// proposer with the engine's LeaderView, while permitting a +// LeaderView-only read surface. // // sidecarPath controls only the read-only capability surface: // when empty, GetCapability reports encryption_capable=false @@ -138,14 +151,16 @@ func registerEncryptionAdminServer(gs *grpc.Server, fullNodeID uint64, sidecarPa if sidecarPath != "" { opts = append(opts, adapter.WithEncryptionAdminSidecarPath(sidecarPath)) } + if engine != nil { + opts = append(opts, adapter.WithEncryptionAdminLeaderView(engine)) + } if enableMutators && engine != nil { opts = append(opts, adapter.WithEncryptionAdminProposer(engine), - adapter.WithEncryptionAdminLeaderView(engine), ) // The §4 capability fan-out is only meaningful when the // cutover mutator is reachable (same enableMutators gate as - // the Proposer/LeaderView). Without it the cutover RPC + // the Proposer). Without it the cutover RPC // refuses with the §4 "fan-out not wired" FailedPrecondition. if capabilityFanout != nil { opts = append(opts, adapter.WithEncryptionAdminCapabilityFanout(capabilityFanout)) diff --git a/main_encryption_admin_test.go b/main_encryption_admin_test.go index df98b0d1c..3b4b77954 100644 --- a/main_encryption_admin_test.go +++ b/main_encryption_admin_test.go @@ -41,6 +41,10 @@ func (stubEncryptionAdminEngine) LinearizableRead(context.Context) (uint64, erro return 0, nil } +type followerEncryptionAdminEngine struct{ stubEncryptionAdminEngine } + +func (followerEncryptionAdminEngine) State() raftengine.State { return raftengine.StateFollower } + // TestEncryptionAdminFullNodeID_DistinctPerRaftId pins the // PR760 r1 Codex P1 regression: full_node_id must be derived from // --raftId (per-node-stable) and NOT from the Raft group id @@ -134,6 +138,42 @@ func TestEncryptionAdmin_MutatingRPCEnabledWhenGateOn(t *testing.T) { } } +func TestEncryptionAdmin_DedicatedTSOGroupAlwaysRefusesMutators(t *testing.T) { + want := stubEncryptionAdminEngine{} + enabled, engine := encryptionAdminWiringForGroup( + dedicatedTSORaftGroupID, + true, + want, + ) + if enabled || engine == nil { + t.Fatalf("dedicated TSO mutator wiring = (%v, %T), want (false, leader view)", enabled, engine) + } + err := callBootstrapAgainstNewServer(t, enabled, engine) + if got := status.Code(err); got != codes.FailedPrecondition { + t.Fatalf("BootstrapEncryption status=%v, want FailedPrecondition; err=%v", got, err) + } +} + +func TestEncryptionAdmin_DedicatedTSOGroupResyncRejectsFollower(t *testing.T) { + enabled, engine := encryptionAdminWiringForGroup( + dedicatedTSORaftGroupID, + true, + followerEncryptionAdminEngine{}, + ) + err := callResyncAgainstNewServer(t, enabled, engine) + if got := status.Code(err); got != codes.FailedPrecondition { + t.Fatalf("ResyncSidecar status=%v, want FailedPrecondition; err=%v", got, err) + } +} + +func TestEncryptionAdmin_DataGroupPreservesMutatorGate(t *testing.T) { + want := stubEncryptionAdminEngine{} + enabled, engine := encryptionAdminWiringForGroup(1, true, want) + if !enabled || engine == nil { + t.Fatalf("data-group mutator wiring = (%v, %T), want (true, non-nil)", enabled, engine) + } +} + func callBootstrapAgainstNewServer(t *testing.T, enableMutators bool, engine encryptionAdminEngine) error { t.Helper() listener := bufconn.Listen(1024 * 1024) @@ -162,6 +202,28 @@ func callBootstrapAgainstNewServer(t *testing.T, enableMutators bool, engine enc return err } +func callResyncAgainstNewServer(t *testing.T, enableMutators bool, engine encryptionAdminEngine) error { + t.Helper() + listener := bufconn.Listen(1024 * 1024) + gs := grpc.NewServer() + registerEncryptionAdminServer(gs, 1, "", enableMutators, engine, nil) + go func() { _ = gs.Serve(listener) }() + t.Cleanup(gs.Stop) + conn, err := grpc.NewClient( + "passthrough:///bufnet", + grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { + return listener.Dial() + }), + grpc.WithTransportCredentials(insecure.NewCredentials()), + ) + if err != nil { + t.Fatalf("grpc.NewClient: %v", err) + } + t.Cleanup(func() { _ = conn.Close() }) + _, err = pb.NewEncryptionAdminClient(conn).ResyncSidecar(context.Background(), &pb.ResyncSidecarRequest{}) + return err +} + func TestRegisterEncryptionAdminServer_Registers(t *testing.T) { gs := grpc.NewServer() registerEncryptionAdminServer(gs, 1, "", false, nil, nil) diff --git a/multiraft_runtime_test.go b/multiraft_runtime_test.go index 11670c9e9..c7b7da1e9 100644 --- a/multiraft_runtime_test.go +++ b/multiraft_runtime_test.go @@ -7,6 +7,7 @@ import ( "testing" "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/internal/encryption/fsmwire" "github.com/bootjp/elastickv/kv" "github.com/bootjp/elastickv/store" "github.com/stretchr/testify/require" @@ -74,14 +75,138 @@ func TestBuildShardGroupsWithDedicatedTSOPreservesSingleDataGroupDir(t *testing. }) require.Contains(t, shardGroups, uint64(dedicatedTSORaftGroupID)) require.Contains(t, shardGroups, uint64(1)) + require.Nil(t, shardGroups[dedicatedTSORaftGroupID].Store) + require.IsType(t, &kv.TSOStateMachine{}, runtimes[0].stateMachine) + require.Nil(t, runtimes[0].store) require.DirExists(t, filepath.Join(legacyDir, "fsm.db")) - require.DirExists(t, filepath.Join(legacyDir, "group-0", "fsm.db")) + require.DirExists(t, filepath.Join(legacyDir, "group-0")) + require.NoDirExists(t, filepath.Join(legacyDir, "group-0", "fsm.db")) require.NoDirExists(t, filepath.Join(legacyDir, "group-1"), "group 1 should stay in legacy dir") got, err := os.ReadFile(legacyMarker) require.NoError(t, err) require.Equal(t, []byte("keep"), got) } +func TestBuildShardGroupsWithDedicatedTSOPreservesLegacyGroupZeroStore(t *testing.T) { + baseDir := t.TempDir() + legacyStoreDir := filepath.Join(baseDir, "n1", "group-0", "fsm.db") + require.NoError(t, os.MkdirAll(legacyStoreDir, raftDirPerm)) + legacyMarker := filepath.Join(legacyStoreDir, "legacy.marker") + require.NoError(t, os.WriteFile(legacyMarker, []byte("preserve"), 0o600)) + + factory, err := newRaftFactory(raftEngineEtcd, nil) + require.NoError(t, err) + runtimes, shardGroups, err := buildShardGroups( + "n1", + baseDir, + []groupSpec{{id: dedicatedTSORaftGroupID, address: "127.0.0.1:17002"}}, + true, + true, + raftBootstrapConfig{}, + factory, + nil, + kv.NewHLC(), + nil, + nil, + "", + encryptionWriteWiring{}, + nil, + ) + require.NoError(t, err) + t.Cleanup(func() { closeRaftGroupRuntimes(runtimes) }) + require.Nil(t, shardGroups[dedicatedTSORaftGroupID].Store) + require.Nil(t, runtimes[0].store) + + got, err := os.ReadFile(legacyMarker) + require.NoError(t, err) + require.Equal(t, []byte("preserve"), got) +} + +func TestBuildShardGroupsClosesEarlierStoresWhenLaterGroupFails(t *testing.T) { + baseDir := t.TempDir() + badGroupDir := groupDataDir(baseDir, "n1", 2, true) + require.NoError(t, os.MkdirAll(badGroupDir, raftDirPerm)) + require.NoError(t, os.WriteFile( + filepath.Join(badGroupDir, raftEngineMarkerFile), + []byte("unsupported\n"), + raftEngineMarkerPerm, + )) + + factory, err := newRaftFactory(raftEngineEtcd, nil) + require.NoError(t, err) + runtimes, shardGroups, err := buildShardGroups( + "n1", + baseDir, + []groupSpec{ + {id: 1, address: "127.0.0.1:17011"}, + {id: 2, address: "127.0.0.1:17012"}, + }, + true, + true, + raftBootstrapConfig{}, + factory, + nil, + kv.NewHLC(), + nil, + nil, + "", + encryptionWriteWiring{}, + nil, + ) + require.ErrorIs(t, err, ErrUnsupportedRaftEngine) + require.Nil(t, runtimes) + require.Nil(t, shardGroups) + + // Reopening the first group's Pebble directory proves the error path ran + // raftGroupRuntime.Close, which owns both engine and store cleanup. + reopened, err := store.NewPebbleStore(filepath.Join( + groupDataDir(baseDir, "n1", 1, true), + "fsm.db", + )) + require.NoError(t, err) + require.NoError(t, reopened.Close()) +} + +func TestDedicatedTSOGroupContinuesAfterLegacyEncryptionEntry(t *testing.T) { + baseDir := t.TempDir() + factory, err := newRaftFactory(raftEngineEtcd, nil) + require.NoError(t, err) + clock := kv.NewHLC() + runtimes, _, err := buildShardGroups( + "n1", + baseDir, + []groupSpec{{id: dedicatedTSORaftGroupID, address: "127.0.0.1:17020"}}, + true, + true, + raftBootstrapConfig{}, + factory, + nil, + clock, + nil, + nil, + "", + encryptionWriteWiring{}, + nil, + ) + require.NoError(t, err) + t.Cleanup(func() { closeRaftGroupRuntimes(runtimes) }) + + legacyEntry := append([]byte{fsmwire.OpRegistration}, fsmwire.EncodeRegistration( + fsmwire.RegistrationPayload{DEKID: 1, FullNodeID: 2, LocalEpoch: 3}, + )...) + result, err := runtimes[0].engine.ProposeAdmin(context.Background(), legacyEntry) + require.NoError(t, err) + legacyErr, ok := result.Response.(error) + require.Truef(t, ok, "legacy response type = %T, want error", result.Response) + require.ErrorIs(t, legacyErr, kv.ErrTSOLegacyEncryptionEntryRejected) + + const ceilingMs = int64(1_700_000_123_456) + coordinator := kv.NewCoordinatorWithEngine(nil, runtimes[0].engine) + t.Cleanup(func() { require.NoError(t, coordinator.Close()) }) + require.NoError(t, coordinator.ProposeHLCLease(context.Background(), ceilingMs)) + require.Equal(t, ceilingMs, clock.PhysicalCeiling()) +} + func TestParseRaftEngineType(t *testing.T) { t.Run("default", func(t *testing.T) { engineType, err := parseRaftEngineType("") From 9182f378f75401617f38cfae9d80b92305c10c1b Mon Sep 17 00:00:00 2001 From: bootjp Date: Sun, 19 Jul 2026 21:20:45 +0900 Subject: [PATCH 02/14] tso: route durable timestamps through group leader --- adapter/distribution_server.go | 141 +++- adapter/distribution_server_test.go | 255 +++++++ adapter/grpc.go | 23 +- adapter/grpc_test.go | 58 ++ .../2026_04_16_partial_centralized_tso.md | 151 ++-- kv/leader_routed_store_test.go | 9 +- kv/shard_store.go | 147 ++++ kv/sharded_coordinator.go | 6 +- kv/tso.go | 22 +- kv/tso_floor_test.go | 237 ++++++ kv/tso_fsm.go | 125 +++- kv/tso_fsm_test.go | 40 +- kv/tso_raft.go | 686 ++++++++++++++++++ kv/tso_raft_test.go | 596 +++++++++++++++ main.go | 158 +++- main_tso_routing_test.go | 245 +++++++ proto/distribution.pb.go | 94 ++- proto/distribution.proto | 25 +- proto/service.pb.go | 33 +- proto/service.proto | 3 + 20 files changed, 2938 insertions(+), 116 deletions(-) create mode 100644 kv/tso_floor_test.go create mode 100644 kv/tso_raft.go create mode 100644 kv/tso_raft_test.go create mode 100644 main_tso_routing_test.go diff --git a/adapter/distribution_server.go b/adapter/distribution_server.go index 0e92f78de..f3f8bd851 100644 --- a/adapter/distribution_server.go +++ b/adapter/distribution_server.go @@ -20,13 +20,14 @@ import ( // DistributionServer serves distribution related gRPC APIs. type DistributionServer struct { - mu sync.Mutex - engine *distribution.Engine - catalog *distribution.CatalogStore - coordinator kv.Coordinator - readTracker *kv.ActiveTimestampTracker - fsObserver DistributionFilesystemObserver - reloadRetry struct { + mu sync.Mutex + engine *distribution.Engine + catalog *distribution.CatalogStore + coordinator kv.Coordinator + timestampAllocator kv.TSOAllocator + readTracker *kv.ActiveTimestampTracker + fsObserver DistributionFilesystemObserver + reloadRetry struct { attempts int interval time.Duration } @@ -50,6 +51,15 @@ func WithDistributionCoordinator(coordinator kv.Coordinator) DistributionServerO } } +// WithDistributionTimestampAllocator exposes the local dedicated TSO +// allocator through GetTimestamp. The allocator itself rejects followers, so +// clients can re-resolve the group-0 leader without a forwarding loop. +func WithDistributionTimestampAllocator(allocator kv.TSOAllocator) DistributionServerOption { + return func(s *DistributionServer) { + s.timestampAllocator = allocator + } +} + func WithDistributionActiveTimestampTracker(tracker *kv.ActiveTimestampTracker) DistributionServerOption { return func(s *DistributionServer) { s.readTracker = tracker @@ -129,10 +139,121 @@ func (s *DistributionServer) GetRoute(ctx context.Context, req *pb.GetRouteReque }, nil } -// GetTimestamp returns monotonically increasing timestamp. +// GetTimestamp returns the base of a consecutive timestamp window. When a +// dedicated allocator is configured, only the local group-0 leader can serve +// the request and the returned window is already durable in that Raft group. func (s *DistributionServer) GetTimestamp(ctx context.Context, req *pb.GetTimestampRequest) (*pb.GetTimestampResponse, error) { - ts := s.engine.NextTimestamp() - return &pb.GetTimestampResponse{Timestamp: ts}, nil + count, minTimestamp, err := timestampRequestValues(req) + if err != nil { + return nil, err + } + activateCutover := req != nil && req.GetActivateCutover() + if s.timestampAllocator == nil { + if count != 1 || minTimestamp != 0 || activateCutover { + return nil, errors.WithStack(status.Error(codes.FailedPrecondition, "dedicated TSO allocator is not configured")) + } + if s.engine == nil { + return nil, errors.WithStack(status.Error(codes.Unavailable, "distribution engine is not configured")) + } + return &pb.GetTimestampResponse{Timestamp: s.engine.NextTimestamp()}, nil + } + + reservation, err := s.allocateTimestampReservation(ctx, count, minTimestamp, activateCutover) + if err != nil { + return nil, timestampRPCError(err) + } + if err := validateTimestampReservation(reservation, count, minTimestamp); err != nil { + return nil, errors.WithStack(status.Error(codes.Internal, err.Error())) + } + return &pb.GetTimestampResponse{ + Timestamp: reservation.Base, + CommittedByDedicatedTso: true, + Count: uint32(count), //nolint:gosec // count is bounded by MaxTSOBatchSize. + PreviousAllocationFloor: reservation.PreviousAllocationFloor, + CutoverActive: reservation.CutoverActive, + }, nil +} + +func (s *DistributionServer) allocateTimestampReservation( + ctx context.Context, + count int, + minTimestamp uint64, + activateCutover bool, +) (kv.TSOReservation, error) { + if allocator, ok := s.timestampAllocator.(kv.TSOReservationAllocator); ok { + reservation, err := allocator.ReserveBatchAfter(ctx, count, minTimestamp, activateCutover) + return reservation, errors.Wrap(err, "reserve dedicated TSO window") + } + reservation := kv.TSOReservation{Count: count} + switch { + case activateCutover: + return reservation, errors.WithStack(status.Error(codes.FailedPrecondition, + "TSO allocator does not support durable cutover")) + case minTimestamp > 0: + return reservation, errors.WithStack(status.Error(codes.FailedPrecondition, + "TSO allocator does not support durable reservation metadata")) + default: + base, err := s.timestampAllocator.NextBatch(ctx, count) + reservation.Base = base + return reservation, errors.Wrap(err, "reserve TSO window") + } +} + +func validateTimestampWindow(base uint64, count int, minTimestamp uint64) error { + if base == 0 || base <= minTimestamp { + return errors.Errorf("dedicated TSO returned base %d at or below minimum %d", base, minTimestamp) + } + size := uint64(count) //nolint:gosec // count is positive and bounded by timestampRequestValues. + if base > math.MaxUint64-(size-1) { + return errors.Errorf("dedicated TSO returned overflowing window base=%d count=%d", base, count) + } + return nil +} + +func validateTimestampReservation(reservation kv.TSOReservation, count int, minTimestamp uint64) error { + if err := validateTimestampWindow(reservation.Base, count, minTimestamp); err != nil { + return err + } + if reservation.Count != count { + return errors.Errorf("dedicated TSO returned count %d, want %d", reservation.Count, count) + } + if reservation.PreviousAllocationFloor >= reservation.Base { + return errors.Errorf("dedicated TSO returned base %d at or below previous floor %d", + reservation.Base, reservation.PreviousAllocationFloor) + } + return nil +} + +func timestampRequestValues(req *pb.GetTimestampRequest) (int, uint64, error) { + if req == nil { + return 1, 0, nil + } + count := req.GetCount() + if count == 0 { + count = 1 + } + if count > uint32(kv.MaxTSOBatchSize) { + return 0, 0, status.Errorf(codes.InvalidArgument, "timestamp count %d exceeds maximum %d", count, kv.MaxTSOBatchSize) + } + return int(count), req.GetMinTimestamp(), nil +} + +func timestampRPCError(err error) error { + if code := status.Code(err); code != codes.Unknown { + return err + } + switch { + case errors.Is(err, context.Canceled): + return errors.WithStack(status.Error(codes.Canceled, err.Error())) + case errors.Is(err, context.DeadlineExceeded): + return errors.WithStack(status.Error(codes.DeadlineExceeded, err.Error())) + case errors.Is(err, kv.ErrInvalidTSOBatchSize), errors.Is(err, kv.ErrTxnCommitTSRequired): + return errors.WithStack(status.Error(codes.InvalidArgument, err.Error())) + case errors.Is(err, kv.ErrTSONotLeader), errors.Is(err, kv.ErrLeaderNotFound): + return errors.WithStack(status.Error(codes.FailedPrecondition, err.Error())) + default: + return errors.WithStack(status.Error(codes.Unavailable, err.Error())) + } } // ListRoutes returns all durable routes from catalog storage. diff --git a/adapter/distribution_server_test.go b/adapter/distribution_server_test.go index e645d4d96..2ab5323ac 100644 --- a/adapter/distribution_server_test.go +++ b/adapter/distribution_server_test.go @@ -3,15 +3,18 @@ package adapter import ( "context" "encoding/binary" + "net" "testing" "time" "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/internal/fskeys" + "github.com/bootjp/elastickv/internal/raftengine" "github.com/bootjp/elastickv/kv" pb "github.com/bootjp/elastickv/proto" "github.com/bootjp/elastickv/store" "github.com/stretchr/testify/require" + "google.golang.org/grpc" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" ) @@ -73,6 +76,159 @@ func TestDistributionServerGetTimestamp_IsMonotonic(t *testing.T) { require.Greater(t, second.Timestamp, first.Timestamp) } +func TestDistributionServerGetTimestamp_UsesDedicatedBatchAllocator(t *testing.T) { + t.Parallel() + + alloc := &distributionTSOAllocator{base: 101, leader: true} + s := NewDistributionServer( + distribution.NewEngine(), + nil, + WithDistributionTimestampAllocator(alloc), + ) + resp, err := s.GetTimestamp(context.Background(), &pb.GetTimestampRequest{ + Count: 8, + MinTimestamp: 100, + }) + require.NoError(t, err) + require.Equal(t, uint64(101), resp.GetTimestamp()) + require.True(t, resp.GetCommittedByDedicatedTso()) + require.Equal(t, 8, alloc.count) + require.Equal(t, uint64(100), alloc.min) +} + +func TestDistributionServerGetTimestamp_CommitsCutoverAndReturnsPriorFloor(t *testing.T) { + t.Parallel() + + alloc := &distributionTSOAllocator{base: 501, previousFloor: 499, leader: true} + s := NewDistributionServer( + distribution.NewEngine(), + nil, + WithDistributionTimestampAllocator(alloc), + ) + resp, err := s.GetTimestamp(context.Background(), &pb.GetTimestampRequest{ + Count: 1, + MinTimestamp: 500, + ActivateCutover: true, + }) + require.NoError(t, err) + require.True(t, alloc.activate) + require.True(t, resp.GetCutoverActive()) + require.Equal(t, uint64(499), resp.GetPreviousAllocationFloor()) +} + +func TestDistributionServerGetTimestamp_RejectsFollower(t *testing.T) { + t.Parallel() + + alloc := &distributionTSOAllocator{err: kv.ErrTSONotLeader} + s := NewDistributionServer( + distribution.NewEngine(), + nil, + WithDistributionTimestampAllocator(alloc), + ) + _, err := s.GetTimestamp(context.Background(), &pb.GetTimestampRequest{Count: 1}) + require.Error(t, err) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) +} + +func TestDistributionServerGetTimestamp_RejectsUnsupportedBatch(t *testing.T) { + t.Parallel() + + s := NewDistributionServer(distribution.NewEngine(), nil) + _, err := s.GetTimestamp(context.Background(), &pb.GetTimestampRequest{Count: 2}) + require.Error(t, err) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) + + _, err = s.GetTimestamp(context.Background(), &pb.GetTimestampRequest{Count: uint32(kv.MaxTSOBatchSize + 1)}) + require.Error(t, err) + require.Equal(t, codes.InvalidArgument, status.Code(err)) +} + +func TestDistributionServerGetTimestamp_RejectsMinimumWithoutReservationMetadata(t *testing.T) { + t.Parallel() + + allocator := &distributionTSOAllocator{base: 101, leader: true} + server := NewDistributionServer(distribution.NewEngine(), nil, + WithDistributionTimestampAllocator(&batchOnlyDistributionTSOAllocator{delegate: allocator})) + + _, err := server.GetTimestamp(context.Background(), &pb.GetTimestampRequest{MinTimestamp: 100}) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) +} + +func TestDistributionServerGetTimestamp_LeaderRoutedRPC(t *testing.T) { + serverAlloc := &distributionTSOAllocator{base: 501, leader: true} + addr := serveDistributionTestServer(t, NewDistributionServer( + distribution.NewEngine(), + nil, + WithDistributionTimestampAllocator(serverAlloc), + )) + + local := &distributionTSOAllocator{leader: false} + routed, err := kv.NewLeaderRoutedTSOAllocator(local, distributionLeaderView{addr: addr}) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, routed.Close()) }) + + base, err := routed.NextBatchAfter(context.Background(), 4, 500) + require.NoError(t, err) + require.Equal(t, uint64(501), base) + require.Equal(t, 4, serverAlloc.count) + require.Equal(t, uint64(500), serverAlloc.min) +} + +func TestDistributionServerGetTimestamp_LeaderRoutedRejectsLegacyServer(t *testing.T) { + addr := serveDistributionTestServer(t, NewDistributionServer(distribution.NewEngine(), nil)) + local := &distributionTSOAllocator{leader: false} + routed, err := kv.NewLeaderRoutedTSOAllocator(local, distributionLeaderView{addr: addr}) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, routed.Close()) }) + + ctx, cancel := context.WithTimeout(context.Background(), 75*time.Millisecond) + defer cancel() + _, err = routed.NextBatch(ctx, 4) + require.ErrorIs(t, err, context.DeadlineExceeded) +} + +func TestDistributionServerGetTimestamp_LeaderRoutedActivatesCutover(t *testing.T) { + serverAlloc := &distributionTSOAllocator{base: 701, previousFloor: 699, leader: true} + addr := serveDistributionTestServer(t, NewDistributionServer( + distribution.NewEngine(), + nil, + WithDistributionTimestampAllocator(serverAlloc), + )) + + local := &distributionTSOAllocator{leader: false} + routed, err := kv.NewLeaderRoutedTSOAllocator( + local, + distributionLeaderView{addr: addr}, + kv.WithTSOCutoverActivation(), + ) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, routed.Close()) }) + + base, err := routed.Next(context.Background()) + require.NoError(t, err) + require.Equal(t, uint64(701), base) + require.True(t, serverAlloc.activate) +} + +func serveDistributionTestServer(t *testing.T, server *DistributionServer) string { + t.Helper() + listener, err := new(net.ListenConfig).Listen(context.Background(), "tcp", "127.0.0.1:0") + require.NoError(t, err) + grpcServer := grpc.NewServer() + pb.RegisterDistributionServer(grpcServer, server) + serveErr := make(chan error, 1) + go func() { serveErr <- grpcServer.Serve(listener) }() + t.Cleanup(func() { + grpcServer.Stop() + _ = listener.Close() + err := <-serveErr + if err != nil { + require.ErrorIs(t, err, grpc.ErrServerStopped) + } + }) + return listener.Addr().String() +} + func TestNewDistributionServer_DefaultCatalogReloadRetryPolicy(t *testing.T) { t.Parallel() @@ -995,6 +1151,105 @@ func (s *distributionCoordinatorStub) LeaseReadForKey(ctx context.Context, _ []b return s.LinearizableRead(ctx) } +type distributionTSOAllocator struct { + base uint64 + previousFloor uint64 + count int + min uint64 + err error + leader bool + activate bool + cutover bool +} + +type batchOnlyDistributionTSOAllocator struct { + delegate *distributionTSOAllocator +} + +func (a *batchOnlyDistributionTSOAllocator) Next(ctx context.Context) (uint64, error) { + return a.delegate.Next(ctx) +} + +func (a *batchOnlyDistributionTSOAllocator) NextBatch(ctx context.Context, n int) (uint64, error) { + return a.delegate.NextBatch(ctx, n) +} + +func (a *batchOnlyDistributionTSOAllocator) IsLeader() bool { + return a.delegate.IsLeader() +} + +func (a *batchOnlyDistributionTSOAllocator) RunLeaseRenewal(ctx context.Context) { + a.delegate.RunLeaseRenewal(ctx) +} + +func (a *distributionTSOAllocator) Next(ctx context.Context) (uint64, error) { + return a.NextBatch(ctx, 1) +} + +func (a *distributionTSOAllocator) NextBatch(_ context.Context, n int) (uint64, error) { + a.count = n + if a.err != nil { + return 0, a.err + } + return a.base, nil +} + +func (a *distributionTSOAllocator) NextBatchAfter(_ context.Context, n int, min uint64) (uint64, error) { + a.count = n + a.min = min + if a.err != nil { + return 0, a.err + } + return a.base, nil +} + +func (a *distributionTSOAllocator) ReserveBatchAfter( + _ context.Context, + n int, + min uint64, + activate bool, +) (kv.TSOReservation, error) { + a.count = n + a.min = min + a.activate = activate + if a.err != nil { + return kv.TSOReservation{}, a.err + } + if activate { + a.cutover = true + } + return kv.TSOReservation{ + Base: a.base, + Count: n, + PreviousAllocationFloor: a.previousFloor, + CutoverActive: a.cutover, + }, nil +} + +func (a *distributionTSOAllocator) IsLeader() bool { return a.leader } + +func (a *distributionTSOAllocator) RunLeaseRenewal(ctx context.Context) { + <-ctx.Done() +} + +type distributionLeaderView struct { + addr string +} + +func (distributionLeaderView) State() raftengine.State { return raftengine.StateFollower } + +func (v distributionLeaderView) Leader() raftengine.LeaderInfo { + return raftengine.LeaderInfo{ID: "leader", Address: v.addr} +} + +func (distributionLeaderView) VerifyLeader(context.Context) error { + return raftengine.ErrNotLeader +} + +func (distributionLeaderView) LinearizableRead(context.Context) (uint64, error) { + return 0, raftengine.ErrNotLeader +} + type recordingDistributionFilesystemObserver struct { reasons []string } diff --git a/adapter/grpc.go b/adapter/grpc.go index 48e384200..6db612660 100644 --- a/adapter/grpc.go +++ b/adapter/grpc.go @@ -52,6 +52,10 @@ type rawGroupKeyScanner interface { ScanGroupKeysAt(ctx context.Context, groupID uint64, start []byte, end []byte, limit int, ts uint64) ([][]byte, error) } +type rawGroupCommitFloorReader interface { + GroupCommittedTimestampFloor(ctx context.Context, groupID uint64) (uint64, error) +} + func WithCloseStore() GRPCServerOption { return func(s *GRPCServer) { s.closeStore = true @@ -133,6 +137,23 @@ 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 groupID := req.GetGroupId(); groupID != 0 { + reader, ok := r.store.(rawGroupCommitFloorReader) + if !ok { + return nil, errors.WithStack(status.Error(codes.FailedPrecondition, + "group watermark requires a group-aware store")) + } + ts, err := reader.GroupCommittedTimestampFloor(ctx, groupID) + if err != nil { + return nil, errors.WithStack(status.Error(codes.FailedPrecondition, err.Error())) + } + return &pb.RawLatestCommitTSResponse{ + Ts: ts, + Exists: ts > 0, + GroupId: groupID, + LeaderFenced: true, + }, nil + } key := req.GetKey() if len(key) == 0 { // No key: return the store's global last-committed watermark. @@ -166,7 +187,6 @@ func (r *GRPCServer) RawScanAt(ctx context.Context, req *pb.RawScanAtRequest) (* if readTS == 0 { readTS = globalSnapshotTS(ctx, r.clock(), r.store) } - if req.GetKeysOnly() { keys, err := r.rawScanKeysAt(ctx, req, limit, readTS) if err != nil { @@ -179,7 +199,6 @@ func (r *GRPCServer) RawScanAt(ctx context.Context, req *pb.RawScanAtRequest) (* if err != nil { return rawScanErrorResponse(err) } - return &pb.RawScanAtResponse{Kv: rawKvPairs(res)}, nil } diff --git a/adapter/grpc_test.go b/adapter/grpc_test.go index 7d936c913..b48bee020 100644 --- a/adapter/grpc_test.go +++ b/adapter/grpc_test.go @@ -12,8 +12,10 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "google.golang.org/grpc" + "google.golang.org/grpc/codes" "google.golang.org/grpc/credentials/insecure" _ "google.golang.org/grpc/health" + "google.golang.org/grpc/status" ) func TestRawKeyPairsPreservesNilAndEmptyKeys(t *testing.T) { @@ -139,6 +141,8 @@ func TestGRPCServer_RawLatestCommitTS_EmptyKeyReturnsGlobalWatermark(t *testing. assert.NoError(t, err) assert.Equal(t, uint64(77), resp.GetTs()) assert.True(t, resp.GetExists()) + assert.Zero(t, resp.GetGroupId()) + assert.False(t, resp.GetLeaderFenced()) // Non-empty key should still work as before. resp, err = s.RawLatestCommitTS(ctx, &pb.RawLatestCommitTSRequest{Key: []byte("k")}) @@ -146,6 +150,32 @@ func TestGRPCServer_RawLatestCommitTS_EmptyKeyReturnsGlobalWatermark(t *testing. assert.Equal(t, uint64(77), resp.GetTs()) } +func TestGRPCServer_RawLatestCommitTS_ExplicitGroupUsesLeaderFencedReader(t *testing.T) { + t.Parallel() + + st := &recordingRawGroupStore{ + MVCCStore: store.NewMVCCStore(), + floorTS: 88, + } + s := NewGRPCServer(st, nil) + + resp, err := s.RawLatestCommitTS(context.Background(), &pb.RawLatestCommitTSRequest{GroupId: 7}) + require.NoError(t, err) + require.Equal(t, uint64(88), resp.GetTs()) + require.True(t, resp.GetExists()) + require.Equal(t, uint64(7), resp.GetGroupId()) + require.True(t, resp.GetLeaderFenced()) + require.Equal(t, uint64(7), st.floorGroupID) +} + +func TestGRPCServer_RawLatestCommitTS_ExplicitGroupRequiresAwareStore(t *testing.T) { + t.Parallel() + + s := NewGRPCServer(store.NewMVCCStore(), nil) + _, err := s.RawLatestCommitTS(context.Background(), &pb.RawLatestCommitTSRequest{GroupId: 1}) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) +} + func TestGRPCServer_RawScanAt_RejectsOversizedLimit(t *testing.T) { t.Parallel() @@ -169,9 +199,16 @@ type recordingRawGroupStore struct { keyScanGroup bool fallbackGet bool fallbackScan bool + floorGroupID uint64 + floorTS uint64 reverseScan bool } +func (s *recordingRawGroupStore) GroupCommittedTimestampFloor(_ context.Context, groupID uint64) (uint64, error) { + s.floorGroupID = groupID + return s.floorTS, nil +} + func (s *recordingRawGroupStore) GetAt(ctx context.Context, key []byte, ts uint64) ([]byte, error) { s.fallbackGet = true return s.MVCCStore.GetAt(ctx, key, ts) @@ -306,6 +343,27 @@ func TestGRPCServer_RawScanAt_KeysOnlyUsesExplicitGroup(t *testing.T) { require.Equal(t, []byte("z"), st.scanEnd) } +func TestGRPCServer_RawScanAt_KeysOnlyFallbackOmitsValues(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("large-value"), 9, 0)) + s := NewGRPCServer(st, nil) + + resp, err := s.RawScanAt(ctx, &pb.RawScanAtRequest{ + StartKey: []byte("a"), + EndKey: []byte("z"), + Limit: 10, + Ts: 9, + KeysOnly: true, + }) + require.NoError(t, err) + require.Len(t, resp.GetKv(), 1) + require.Equal(t, []byte("a"), resp.GetKv()[0].GetKey()) + require.Empty(t, resp.GetKv()[0].GetValue()) +} + func TestGRPCServer_RawScanAt_ReverseKeysOnlyUsesExplicitGroup(t *testing.T) { t.Parallel() diff --git a/docs/design/2026_04_16_partial_centralized_tso.md b/docs/design/2026_04_16_partial_centralized_tso.md index d3d6732d8..f4d38328c 100644 --- a/docs/design/2026_04_16_partial_centralized_tso.md +++ b/docs/design/2026_04_16_partial_centralized_tso.md @@ -1,9 +1,10 @@ # Centralized Timestamp Oracle (TSO) Design -- Status: Partial — M1 all-led-group HLC renewal, the M2 reserved-group - bootstrap bridge, and the M3-M5 TSO allocator/batch cutover are implemented; - the minimal TSO-only FSM and runtime group-0 wiring are implemented, while - follower redirect/admin exposure and shadow validation remain open +- Status: Partial — M1-M6 are implemented, including the dedicated group-0 + FSM, leader-routed durable windows, strict term bootstrap, serialized shadow + migration, and the one-way rolling cutover marker. M7 legacy cleanup and + cross-shard SSI timestamp validation remain open; runtime config reload and + production latency/alerting work also remain open. - Author: bootjp - Date: 2026-04-16 - Updated: 2026-07-18 @@ -41,10 +42,9 @@ Implemented: bridge, and the all-led-group HLC renewal loop keeps its ceiling warm alongside data groups. Adding only group 0 to an existing single-data-group deployment keeps the data group on its legacy `base/raftID` directory while - isolating group 0 under `base/raftID/group-0`. The current `--tsoEnabled` - bridge still issues timestamps from the locally led data shard through - `LocalTSOAllocator`; pinning timestamp issuance to group 0 remains deferred - until the TSO-leader redirect path exists. + isolating group 0 under `base/raftID/group-0`. When group 0 is present, + `--tsoEnabled` now routes issuance to its current leader; without group 0, + the flag preserves the earlier local/default-group compatibility bridge. 9. `TSOStateMachine` implements the dedicated-group FSM contract: it accepts HLC lease entries and explicit allocation-floor entries, halt-fails invalid entries, snapshots/restores TSO-owned ceiling/floor state, and classifies @@ -60,13 +60,52 @@ Implemented: unwired and return `FailedPrecondition`. During upgrade replay, valid encryption control entries committed by the earlier compatibility FSM are decoded and deterministically rejected without halting the TSO apply loop; - malformed control entries still halt fail-closed. + malformed control entries still halt fail-closed. +11. `RaftTSOAllocator` verifies group-0 leadership and commits every returned + window's inclusive end before exposing it. On the first request of every + leader term it obtains a strict, leader-fenced maximum `LastCommitTS` from + every data group; failure to reach any authoritative group leader blocks + issuance rather than falling back to a stale replica watermark. Remote + watermark requests carry the explicit data-group ID, and the receiving + node revalidates local leadership with a linearizable ReadIndex before + returning that group's store watermark, so dialing a recently-demoted + leader cannot satisfy the fence with stale local state. The response echoes + the group ID and carries an explicit leader-fenced marker; a legacy server + that ignores the request field is rejected fail-closed. The allocator also + revalidates the same group-0 term after the remote floor/cutover work and + after committing the allocation floor. A term change may leak a committed + window, but that window is never returned and its floor prevents reuse. +12. `LeaderRoutedTSOAllocator` serves the local group-0 leader directly and + redirects followers through `Distribution.GetTimestamp`. The response + carries an explicit durable-TSO marker and echoed count, so a new client + rejects a rolling-upgrade response from a legacy server that ignored the + batch/minimum request fields. +13. `--tsoShadowEnabled` serializes each legacy candidate through group 0 + before returning it. The response includes the allocation floor that + preceded the reservation. A candidate at or below that floor is discarded, + the local HLC observes the newer reservation, and allocation retries. TSO + unavailability fails the write closed because an unmirrored legacy value + cannot establish cutover readiness. +14. The group-0 FSM persists a one-way production cutover marker. The first + `--tsoEnabled` refill commits that marker before its window. Nodes still in + shadow mode observe the marker in their next reservation response and + return the TSO value instead of a legacy candidate, preserving safety + during a rolling flag transition. Routed responses are also observed into + the local legacy HLC so the fallback clock never moves backward. On + restart, the restored marker takes precedence over startup flags: a node + started with neither mode flag, or with only `--tsoShadowEnabled`, installs + production group-0 routing instead of re-entering legacy or shadow + issuance. The marker uses a versioned TSO envelope with the same legacy + fail-closed prefix as allocation-floor entries and a distinct magic, so an + old or misrouted data FSM halts while encryption bytes cannot activate it. Remaining: -1. Add follower redirect/admin exposure for the dedicated TSO leader. -2. Add Phase B shadow-read validation before making dedicated TSO the only - production timestamp path. +1. M7 Phase-D removal of per-shard renewal/legacy issuance after the migration + compatibility window closes. +2. Cross-shard SSI read-timestamp validation through the dedicated TSO. +3. Runtime config reload for the mode switch; current flags are startup-only. +4. Production benchmark, divergence metrics, and alert thresholds. ### 1.1 Original Limitation @@ -531,7 +570,10 @@ New TSO leader elected via Raft ├─ FSM.Restore() or Raft log replay │ → physicalCeiling restored to the last committed value │ - └─ HLC.Next() uses max(now, physicalCeiling) + ├─ strict read of every data-group leader LastCommitTS + │ → failure blocks issuance + │ + └─ HLC.Next() uses max(now, physicalCeiling, global commit floor) → new leader issues timestamps strictly above the old leader's window ✅ ``` @@ -632,7 +674,7 @@ Migrating from the current per-shard ceiling model to a centralized TSO must not interrupt writes or violate timestamp monotonicity. The following phased approach enables a live cutover. -### 7.1 Phase A — Dual-Write Bridge (no cutover risk) +### 7.1 Phase A — Group-0 Warm-up (no read cutover) ``` ┌──────────────────────────────────────────────────────────┐ @@ -641,32 +683,47 @@ approach enables a live cutover. │ startTS = legacyHLC.Next() (existing path, unchanged) │ │ │ │ RunHLCLeaseRenewal(): │ -│ ├─ propose to all shard groups (M1 fix, Section 6) │ -│ └─ also propose to TSO group (new, write-only) │ -│ ↳ TSO FSM advances its ceiling in parallel │ +│ └─ each local Raft leader proposes to the groups it leads│ +│ ↳ the group-0 leader warms the TSO FSM ceiling │ └──────────────────────────────────────────────────────────┘ ``` - The TSO group receives ceiling proposals but **no reads are served from it**. - This allows TSO FSM state to warm up and be validated in production before the cutover. -- Rollback: stop proposing to the TSO group; no state change on data path. - -### 7.2 Phase B — Shadow Read Validation - -- Both `legacyHLC.Next()` and `tso.Next()` are called per transaction. -- Results are compared in a shadow log; divergences are alerted but the legacy - value is used. -- This phase validates that the TSO ceiling is always ≥ the legacy ceiling. - -### 7.3 Phase C — TSO Cutover (feature flag) - -- A runtime feature flag (`tso.enabled`) switches `startTS` to `tso.Next()`. -- The flag can be toggled per-node via config reload (no process restart). -- Because the TSO ceiling was kept ≥ the legacy ceiling throughout Phase A/B, - there is no timestamp regression at the moment of cutover. -- Rollback: flip the flag back; the legacy HLC has been monotonically advancing - in parallel so it remains safe to resume. +- These independent proposals do **not** prove that the TSO ceiling is greater + than every data-group ceiling. Phase A is warm-up only; Phase B provides the + issuance-level serialization required for a safe transition. +- Rollback: remove group 0 before entering Phase B; no timestamp path changed. + +### 7.2 Phase B — Serialized Shadow Migration + +- Every node must run `--tsoShadowEnabled` before any node enables cutover. +- A legacy candidate is sent to the TSO leader as `min_timestamp`. Under the + group-0 allocator mutex, the leader samples the preceding durable allocation + floor, reserves and commits a timestamp above the candidate, and returns both. +- If the candidate is at or below the preceding floor, it may overlap an older + TSO/shadow reservation and is discarded. The caller observes the newer TSO + value into its HLC and retries; only a candidate proven above the preceding + floor is returned to the legacy write path. +- Shadow RPC/proposal failure fails timestamp issuance closed. A fail-open + shadow cannot prove the migration invariant and is therefore not eligible + for production cutover. + +### 7.3 Phase C — Durable TSO Cutover + +- The startup flag `--tsoEnabled` switches coordinator issuance to the + leader-routed allocator. Runtime config reload remains future work. +- The first production refill commits the group-0 cutover marker before + committing and returning its timestamp window. +- The marker is encoded in a versioned TSO-specific envelope whose leading + reserved byte remains fail-closed on pre-TSO and data-group FSMs. +- A node still running Phase B receives `cutover_active=true` on its next + shadow reservation and returns the reserved TSO timestamp. Therefore a + rolling restart does not keep issuing legacy values after the marker. +- The marker is intentionally one-way. Before it commits, rollback means + returning every node to Phase B. After it commits, rollback must preserve the + TSO path; clearing it requires a separate cluster-wide quiescence protocol. ### 7.4 Phase D — Legacy Cleanup @@ -676,15 +733,17 @@ approach enables a live cutover. ### 7.5 Monotonicity Invariant Across Phases -At every phase boundary, the following invariant must hold: +At every Phase-B/C issuance boundary, the following invariant must hold: ``` -tso_ceiling ≥ max(ceiling committed by any shard group leader) +returned_legacy_ts > previous_tso_allocation_floor +new_tso_allocation_floor >= returned_legacy_ts ``` -This is enforced by Phase A's dual-write: every ceiling update that reaches -a shard group also reaches the TSO group, so the TSO ceiling is always at -least as large as the maximum shard ceiling. +Both comparisons occur in one group-0-serialized reservation before the legacy +candidate is returned. Once the cutover marker applies, shadow callers stop +returning legacy candidates. Independently, every new TSO leader term fences +its first window above the strict maximum committed data timestamp. --- @@ -693,12 +752,12 @@ least as large as the maximum shard ceiling. | Phase | Scope | Priority | |-------|-------|----------| | M1 — shipped | Extend `RunHLCLeaseRenewal` to all shard groups with parallel proposals (Section 6) | High | -| M2 — shipped for reserved group 0 | Phase A dual-write bridge: when group 0 is configured, all-led HLC renewal also proposes ceiling updates to it while shard range validation prevents user data routes to group 0 (Section 7.1) | High | +| M2 — shipped | Reserve group 0, keep its ceiling warm as a pre-migration compatibility step, and prevent user-data routes from entering the control group (Section 7.1). This warm-up alone is not a cutover proof. | High | | M3 — shipped | Define `TSOAllocator` interface; implement backed by `defaultGroup` | Medium | | M4 — shipped | `BatchAllocator` with atomic counter for low-latency timestamp serving | Medium | -| M5 — shipped for default-group bridge | Coordinator feature-flag cutover via `--tsoEnabled`; shadow validation against a dedicated group remains deferred to M6 | Medium | -| M6 — partial | Dedicated TSO Raft group (`groupID = 0`) is reserved/bootstrap-capable, warmed by the HLC renewal bridge, and runs the minimal `TSOStateMachine` without an MVCC store; TSO-leader redirect/admin exposure and TSO-leader-only timestamp issuance remain open | Low | -| M7 | Phase D legacy cleanup + cross-shard SSI read-timestamp validation via TSO | Low | +| M5 — shipped | Preserve the default-group `LocalTSOAllocator` compatibility bridge when group 0 is absent; route coordinator-owned timestamp call sites through the allocator abstraction. | Medium | +| M6 — shipped | Run the dedicated group-0 FSM, fence each new TSO leader term above all authoritative data-group commit floors, redirect follower requests to the TSO leader over gRPC, synchronously serialize fail-closed shadow issuance, and commit the one-way rolling cutover marker before production windows. | Low | +| M7 — open | Phase D legacy cleanup + cross-shard SSI read-timestamp validation via TSO | Low | --- @@ -714,9 +773,9 @@ least as large as the maximum shard ceiling. stricter form (`ceiling + 1`) guarantees no overlap even if wall clocks drift, at the cost of one extra millisecond per renewal window. -4. **Non-leader TSO requests:** Should follower nodes redirect to the TSO - leader via gRPC, or support follower reads with a known-safe timestamp - bound? +4. **Non-leader TSO requests (resolved):** Followers redirect to the current + group-0 leader through `Distribution.GetTimestamp`. Timestamp allocation is + a consensus write, so follower reads cannot replace the leader proposal. 5. **Backward compatibility (resolved for group 0):** Existing group-0 Raft logs already carry compatible HLC lease entries. `TSOStateMachine.Restore` diff --git a/kv/leader_routed_store_test.go b/kv/leader_routed_store_test.go index 22ddc9535..eee9cd172 100644 --- a/kv/leader_routed_store_test.go +++ b/kv/leader_routed_store_test.go @@ -2,6 +2,7 @@ package kv import ( "context" + "fmt" "net" "sync" "testing" @@ -91,6 +92,9 @@ type fakeRawKVServer struct { getResp *pb.RawGetResponse scanResp *pb.RawScanAtResponse latestResp *pb.RawLatestCommitTSResponse + // wantLatestGroupID, when non-zero, makes the fake reject a watermark + // request that is not pinned to the expected Raft group. + wantLatestGroupID uint64 } func (f *fakeRawKVServer) RawGet(context.Context, *pb.RawGetRequest) (*pb.RawGetResponse, error) { @@ -115,10 +119,13 @@ func (f *fakeRawKVServer) RawScanAt(_ context.Context, req *pb.RawScanAtRequest) return &pb.RawScanAtResponse{}, nil } -func (f *fakeRawKVServer) RawLatestCommitTS(context.Context, *pb.RawLatestCommitTSRequest) (*pb.RawLatestCommitTSResponse, error) { +func (f *fakeRawKVServer) RawLatestCommitTS(_ context.Context, req *pb.RawLatestCommitTSRequest) (*pb.RawLatestCommitTSResponse, error) { f.mu.Lock() defer f.mu.Unlock() f.latestCalls++ + if f.wantLatestGroupID != 0 && req.GetGroupId() != f.wantLatestGroupID { + return nil, fmt.Errorf("watermark group_id=%d, want %d", req.GetGroupId(), f.wantLatestGroupID) + } if f.latestResp != nil { return f.latestResp, nil } diff --git a/kv/shard_store.go b/kv/shard_store.go index 14c72332f..c641dc44a 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -4,10 +4,12 @@ import ( "bytes" "context" "encoding/binary" + stderrors "errors" "io" "math" "slices" "sort" + "sync" "time" "github.com/bootjp/elastickv/distribution" @@ -17,10 +19,13 @@ import ( pb "github.com/bootjp/elastickv/proto" "github.com/bootjp/elastickv/store" "github.com/cockroachdb/errors" + "golang.org/x/sync/errgroup" ) const proxyForwardTimeout = 5 * time.Second +const maxTSOCommitFloorConcurrency = 16 + // ShardStore routes MVCC reads to shard-specific stores and proxies to leaders when needed. type ShardStore struct { engine *distribution.Engine @@ -34,6 +39,14 @@ var ( ErrFilesystemPlacementTargetNotFound = errors.New("filesystem placement target group has no routable home slot") ) +var ErrTSOCommitFloorUnavailable = errors.New("tso: authoritative data-group commit floor is unavailable") + +// TSOCutoverFloorProvider supplies the highest commit timestamp that the +// dedicated TSO must exceed before a leader term can issue a window. +type TSOCutoverFloorProvider interface { + GlobalCommittedTimestampFloor(context.Context) (uint64, error) +} + // NewShardStore creates a sharded MVCC store wrapper. func NewShardStore(engine *distribution.Engine, groups map[uint64]*ShardGroup) *ShardStore { return &ShardStore{ @@ -2255,6 +2268,140 @@ func (s *ShardStore) LastCommitTS() uint64 { return max } +// GlobalCommittedTimestampFloor returns a strict, leader-fenced maximum over +// every data group. Unlike GlobalLastCommitTS helpers used by best-effort read +// snapshots, this method never falls back to a stale local follower watermark: +// inability to reach any group's authoritative leader fails the TSO term +// initialization closed. +func (s *ShardStore) GlobalCommittedTimestampFloor(ctx context.Context) (uint64, error) { + if s == nil { + return 0, errors.WithStack(ErrTSOCommitFloorUnavailable) + } + ids, err := s.tsoCommitFloorGroupIDs() + if err != nil { + return 0, err + } + if len(ids) == 0 { + return 0, nil + } + if len(ids) == 1 { + return s.singleGroupCommittedTimestampFloor(ctx, ids[0]) + } + return s.parallelCommittedTimestampFloor(ctx, ids) +} + +func (s *ShardStore) tsoCommitFloorGroupIDs() ([]uint64, error) { + ids := make([]uint64, 0, len(s.groups)) + for groupID, group := range s.groups { + if groupID == 0 { + continue + } + if group == nil || group.Store == nil { + return nil, errors.Wrapf(ErrTSOCommitFloorUnavailable, + "data group %d has no local store", groupID) + } + ids = append(ids, groupID) + } + sort.Slice(ids, func(i, j int) bool { return ids[i] < ids[j] }) + return ids, nil +} + +func (s *ShardStore) singleGroupCommittedTimestampFloor(ctx context.Context, groupID uint64) (uint64, error) { + ts, err := s.authoritativeGroupLastCommitTS(ctx, groupID, s.groups[groupID]) + if err != nil { + return 0, errors.Wrapf(err, "tso commit floor: group %d", groupID) + } + return ts, nil +} + +// GroupCommittedTimestampFloor returns the local group's watermark only after +// a ReadIndex fence proves this node is still that group's leader. It is the +// server-side contract for remote TSO term initialization; callers must not +// substitute a node-global or follower-local watermark. +func (s *ShardStore) GroupCommittedTimestampFloor(ctx context.Context, groupID uint64) (uint64, error) { + if s == nil || groupID == 0 { + return 0, errors.WithStack(ErrTSOCommitFloorUnavailable) + } + group, ok := s.groups[groupID] + if !ok || group == nil || group.Store == nil { + return 0, errors.Wrapf(ErrTSOCommitFloorUnavailable, + "data group %d is not available on this node", groupID) + } + engine := engineForGroup(group) + if !isLeaderEngine(engine) { + return 0, errors.Wrapf(ErrTSOCommitFloorUnavailable, + "data group %d is not led by this node", groupID) + } + if _, err := linearizableReadEngineCtx(nonNilTSOContext(ctx), engine); err != nil { + return 0, errors.Wrapf(ErrTSOCommitFloorUnavailable, + "fence data group %d leader: %v", groupID, err) + } + return group.Store.LastCommitTS(), nil +} + +func (s *ShardStore) parallelCommittedTimestampFloor(ctx context.Context, ids []uint64) (uint64, error) { + eg, egctx := errgroup.WithContext(nonNilTSOContext(ctx)) + eg.SetLimit(min(len(ids), maxTSOCommitFloorConcurrency)) + var mu sync.Mutex + var maxTS uint64 + for _, groupID := range ids { + eg.Go(func() error { + ts, err := s.authoritativeGroupLastCommitTS(egctx, groupID, s.groups[groupID]) + if err != nil { + return errors.Wrapf(err, "tso commit floor: group %d", groupID) + } + mu.Lock() + maxTS = max(maxTS, ts) + mu.Unlock() + return nil + }) + } + if err := eg.Wait(); err != nil { + return 0, errors.WithStack(err) + } + return maxTS, nil +} + +func (s *ShardStore) authoritativeGroupLastCommitTS(ctx context.Context, groupID uint64, group *ShardGroup) (uint64, error) { + engine := engineForGroup(group) + if engine == nil { + return 0, errors.Wrapf(ErrTSOCommitFloorUnavailable, + "data group %d has no raft engine", groupID) + } + if isLeaderEngine(engine) { + if ts, err := s.GroupCommittedTimestampFloor(ctx, groupID); err == nil { + return ts, nil + } + // Leadership may have changed between State and ReadIndex. Resolve the + // newly published leader below instead of trusting the local watermark. + } + addr := leaderAddrFromEngine(engine) + if addr == "" { + return 0, errors.Wrapf(ErrTSOCommitFloorUnavailable, + "data group %d has no known leader", groupID) + } + conn, err := s.connCache.ConnFor(addr) + if err != nil { + return 0, errors.Wrapf(ErrTSOCommitFloorUnavailable, + "dial data group %d leader %s: %v", groupID, addr, err) + } + requestCtx, cancel := context.WithTimeout(nonNilTSOContext(ctx), proxyForwardTimeout) + defer cancel() + resp, err := pb.NewRawKVClient(conn).RawLatestCommitTS(requestCtx, &pb.RawLatestCommitTSRequest{GroupId: groupID}) + if err != nil { + return 0, errors.Wrapf(ErrTSOCommitFloorUnavailable, + "read data group %d leader %s watermark: %v", groupID, addr, err) + } + if !resp.GetLeaderFenced() || resp.GetGroupId() != groupID { + return 0, errors.Wrapf( + stderrors.Join(ErrTSOCommitFloorUnavailable, ErrTSOProtocolUnsupported), + "data group %d leader %s returned unfenced watermark for group %d", + groupID, addr, resp.GetGroupId(), + ) + } + return resp.GetTs(), nil +} + // LastAppliedIndex aggregates the durable applied-index across every // shard group, returning the MIN over all groups that report one. // diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index 37b963c90..60427c0ab 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -24,7 +24,11 @@ type ShardGroup struct { Engine raftengine.Engine Store store.MVCCStore Txn Transactional - lease leaseState + // TSOState is set only for reserved group 0. Keeping the applied floor and + // cutover marker next to the group's engine lets the leader allocator read + // consensus-owned state without treating the shared HLC mirror as durable. + TSOState *TSOStateMachine + lease leaseState // lp caches the Engine's optional LeaseProvider capability so the // groupLeaseRead / maybeRefresh hot paths test a single field for // nil instead of performing an interface type assertion per call. diff --git a/kv/tso.go b/kv/tso.go index e07cb34f9..f9646951c 100644 --- a/kv/tso.go +++ b/kv/tso.go @@ -11,6 +11,10 @@ import ( const defaultTSOLeaderPollInterval = 25 * time.Millisecond +// MaxTSOBatchSize is the largest consecutive timestamp window accepted by a +// TSO allocator or the Distribution.GetTimestamp RPC. +const MaxTSOBatchSize = maxHLCBatchSize + var ( ErrTSOAllocatorRequired = errors.New("tso: allocator is required") ErrTSOCoordinatorNil = errors.New("tso: coordinator is required") @@ -34,6 +38,12 @@ type TimestampAllocator interface { Next(ctx context.Context) (uint64, error) } +// TimestampAllocatorProvider lets coordinator decorators preserve access to +// the configured allocator without making the decorator itself an allocator. +type TimestampAllocatorProvider interface { + TimestampAllocator() TimestampAllocator +} + type TimestampAfterAllocator interface { NextAfter(ctx context.Context, min uint64) (uint64, error) } @@ -91,7 +101,13 @@ func NextTimestampAfterThrough(ctx context.Context, coord Coordinator, startTS u return nextTimestampAfterObserved(ctx, coord, clock, startTS, label) } -func coordinatorTimestampAllocator(coord Coordinator) (TimestampAllocator, bool) { +// TimestampAllocatorThrough returns the allocator configured behind a +// coordinator or coordinator decorator. +func TimestampAllocatorThrough(coord Coordinator) (TimestampAllocator, bool) { + if provider, ok := coord.(TimestampAllocatorProvider); ok { + alloc := provider.TimestampAllocator() + return alloc, alloc != nil + } switch c := coord.(type) { case *Coordinate: if c != nil && c.tsAllocator != nil { @@ -109,6 +125,10 @@ func coordinatorTimestampAllocator(coord Coordinator) (TimestampAllocator, bool) } } +func coordinatorTimestampAllocator(coord Coordinator) (TimestampAllocator, bool) { + return TimestampAllocatorThrough(coord) +} + func nextTimestampFromAllocator(ctx context.Context, alloc TimestampAllocator, label string) (uint64, error) { ts, err := alloc.Next(ctx) if err != nil { diff --git a/kv/tso_floor_test.go b/kv/tso_floor_test.go new file mode 100644 index 000000000..15cf47b31 --- /dev/null +++ b/kv/tso_floor_test.go @@ -0,0 +1,237 @@ +package kv + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/bootjp/elastickv/internal/raftengine" + pb "github.com/bootjp/elastickv/proto" + "github.com/bootjp/elastickv/store" + "github.com/stretchr/testify/require" +) + +type blockingFloorRawKVServer struct { + pb.UnimplementedRawKVServer + + ts uint64 + started chan struct{} + release <-chan struct{} + once sync.Once +} + +func (s *blockingFloorRawKVServer) RawLatestCommitTS(ctx context.Context, req *pb.RawLatestCommitTSRequest) (*pb.RawLatestCommitTSResponse, error) { + s.once.Do(func() { close(s.started) }) + select { + case <-s.release: + return &pb.RawLatestCommitTSResponse{ + Ts: s.ts, + Exists: true, + GroupId: req.GetGroupId(), + LeaderFenced: true, + }, nil + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +func TestShardStoreGlobalCommittedTimestampFloorUsesEveryGroupLeader(t *testing.T) { + t.Parallel() + + local := store.NewMVCCStore() + require.NoError(t, local.PutAt(context.Background(), []byte("local"), []byte("v"), 42, 0)) + remote := &fakeRawKVServer{ + latestResp: &pb.RawLatestCommitTSResponse{ + Ts: 99, + Exists: true, + GroupId: 2, + LeaderFenced: true, + }, + wantLatestGroupID: 2, + } + addr, stop := startRawKVServer(t, remote) + t.Cleanup(stop) + + groups := map[uint64]*ShardGroup{ + 0: {Engine: &recordingTSOEngine{state: raftengine.StateFollower}}, + 1: { + Engine: &recordingTSOEngine{state: raftengine.StateLeader}, + Store: local, + }, + 2: { + Engine: &recordingTSOEngine{ + state: raftengine.StateFollower, + leader: raftengine.LeaderInfo{Address: addr}, + }, + Store: store.NewMVCCStore(), + }, + } + shards := NewShardStore(nil, groups) + t.Cleanup(func() { require.NoError(t, shards.Close()) }) + + floor, err := shards.GlobalCommittedTimestampFloor(context.Background()) + require.NoError(t, err) + require.Equal(t, uint64(99), floor) +} + +func TestShardStoreGlobalCommittedTimestampFloorRejectsUnfencedRemoteResponse(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + resp *pb.RawLatestCommitTSResponse + }{ + { + name: "legacy response", + resp: &pb.RawLatestCommitTSResponse{Ts: 99, Exists: true}, + }, + { + name: "wrong group echo", + resp: &pb.RawLatestCommitTSResponse{ + Ts: 99, + Exists: true, + GroupId: 3, + LeaderFenced: true, + }, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + remote := &fakeRawKVServer{ + latestResp: tc.resp, + wantLatestGroupID: 2, + } + addr, stop := startRawKVServer(t, remote) + t.Cleanup(stop) + + shards := NewShardStore(nil, map[uint64]*ShardGroup{ + 2: { + Engine: &recordingTSOEngine{ + state: raftengine.StateFollower, + leader: raftengine.LeaderInfo{Address: addr}, + }, + Store: store.NewMVCCStore(), + }, + }) + t.Cleanup(func() { require.NoError(t, shards.Close()) }) + + _, err := shards.GlobalCommittedTimestampFloor(context.Background()) + require.ErrorIs(t, err, ErrTSOCommitFloorUnavailable) + require.ErrorIs(t, err, ErrTSOProtocolUnsupported) + }) + } +} + +func TestShardStoreGroupCommittedTimestampFloorRequiresLocalLeadership(t *testing.T) { + t.Parallel() + + local := store.NewMVCCStore() + require.NoError(t, local.PutAt(context.Background(), []byte("stale"), []byte("v"), 77, 0)) + engine := &recordingTSOEngine{state: raftengine.StateFollower} + shards := NewShardStore(nil, map[uint64]*ShardGroup{ + 1: {Engine: engine, Store: local}, + }) + t.Cleanup(func() { require.NoError(t, shards.Close()) }) + + _, err := shards.GroupCommittedTimestampFloor(context.Background(), 1) + require.ErrorIs(t, err, ErrTSOCommitFloorUnavailable) + + engine.mu.Lock() + engine.state = raftengine.StateLeader + engine.mu.Unlock() + floor, err := shards.GroupCommittedTimestampFloor(context.Background(), 1) + require.NoError(t, err) + require.Equal(t, uint64(77), floor) +} + +func TestShardStoreGlobalCommittedTimestampFloorQueriesGroupsConcurrently(t *testing.T) { + t.Parallel() + + release := make(chan struct{}) + started1 := make(chan struct{}) + started2 := make(chan struct{}) + addr1, stop1 := startRawKVServer(t, &blockingFloorRawKVServer{ + ts: 11, started: started1, release: release, + }) + addr2, stop2 := startRawKVServer(t, &blockingFloorRawKVServer{ + ts: 22, started: started2, release: release, + }) + t.Cleanup(stop1) + t.Cleanup(stop2) + + shards := NewShardStore(nil, map[uint64]*ShardGroup{ + 1: { + Engine: &recordingTSOEngine{state: raftengine.StateFollower, leader: raftengine.LeaderInfo{Address: addr1}}, + Store: store.NewMVCCStore(), + }, + 2: { + Engine: &recordingTSOEngine{state: raftengine.StateFollower, leader: raftengine.LeaderInfo{Address: addr2}}, + Store: store.NewMVCCStore(), + }, + }) + t.Cleanup(func() { require.NoError(t, shards.Close()) }) + + type result struct { + floor uint64 + err error + } + resultCh := make(chan result, 1) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + t.Cleanup(cancel) + go func() { + floor, err := shards.GlobalCommittedTimestampFloor(ctx) + resultCh <- result{floor: floor, err: err} + }() + + requireChannelClosed(t, started1) + requireChannelClosed(t, started2) + close(release) + + got := <-resultCh + require.NoError(t, got.err) + require.Equal(t, uint64(22), got.floor) +} + +func requireChannelClosed(t *testing.T, ch <-chan struct{}) { + t.Helper() + select { + case <-ch: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for concurrent TSO floor request") + } +} + +func TestShardStoreGlobalCommittedTimestampFloorFailsWithoutGroupLeader(t *testing.T) { + t.Parallel() + + local := store.NewMVCCStore() + require.NoError(t, local.PutAt(context.Background(), []byte("stale"), []byte("v"), 7, 0)) + shards := NewShardStore(nil, map[uint64]*ShardGroup{ + 1: { + Engine: &recordingTSOEngine{state: raftengine.StateFollower}, + Store: local, + }, + }) + t.Cleanup(func() { require.NoError(t, shards.Close()) }) + + _, err := shards.GlobalCommittedTimestampFloor(context.Background()) + require.ErrorIs(t, err, ErrTSOCommitFloorUnavailable) +} + +func TestShardStoreGlobalCommittedTimestampFloorFailsWithoutRaftEngine(t *testing.T) { + t.Parallel() + + local := store.NewMVCCStore() + require.NoError(t, local.PutAt(context.Background(), []byte("unfenced"), []byte("v"), 11, 0)) + shards := NewShardStore(nil, map[uint64]*ShardGroup{ + 1: {Store: local}, + }) + t.Cleanup(func() { require.NoError(t, shards.Close()) }) + + _, err := shards.GlobalCommittedTimestampFloor(context.Background()) + require.ErrorIs(t, err, ErrTSOCommitFloorUnavailable) + require.ErrorContains(t, err, "no raft engine") +} diff --git a/kv/tso_fsm.go b/kv/tso_fsm.go index fa18e3176..238efaa27 100644 --- a/kv/tso_fsm.go +++ b/kv/tso_fsm.go @@ -25,8 +25,12 @@ const ( // range so old data-group FSMs halt on a misrouted entry. The remaining // magic and version bytes keep it distinct from every encryption opcode. tsoAllocationFloorEnvelope = "\x07TSOF\x01" - tsoSnapshotV1Len = hlcLeasePayloadLen - tsoSnapshotV2Len = hlcLeasePayloadLen * 2 + // tsoCutoverEnvelope uses the same legacy fail-closed prefix but a distinct + // magic, so an encryption entry cannot become a one-way TSO state change. + tsoCutoverEnvelope = "\x07TSOC\x01" + tsoSnapshotV1Len = hlcLeasePayloadLen + tsoSnapshotV2Len = hlcLeasePayloadLen * 2 + tsoSnapshotV3Len = tsoSnapshotV2Len + 1 ) // TSOStateMachine is the minimal state machine for the dedicated timestamp @@ -38,6 +42,7 @@ type TSOStateMachine struct { hlc *HLC ceilingMs atomic.Int64 allocationFloor atomic.Uint64 + cutoverActive atomic.Bool } func NewTSOStateMachine(hlc *HLC) *TSOStateMachine { @@ -53,6 +58,8 @@ func (f *TSOStateMachine) Apply(data []byte) any { return f.applyLeaseEntry(data) case bytes.HasPrefix(data, []byte(tsoAllocationFloorEnvelope)): return f.applyAllocationFloorEntry(data) + case bytes.Equal(data, []byte(tsoCutoverEnvelope)): + return f.applyCutoverEntry(data) case data[0] >= fsmwire.OpEncryptionMin && data[0] <= fsmwire.OpEncryptionMax: return rejectLegacyTSOEncryptionEntry(data) default: @@ -116,14 +123,47 @@ func (f *TSOStateMachine) applyAllocationFloorEntry(data []byte) any { return nil } +func (f *TSOStateMachine) applyCutoverEntry(data []byte) any { + if !bytes.Equal(data, []byte(tsoCutoverEnvelope)) { + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "invalid TSO cutover envelope length %d", len(data))) + } + if f != nil { + f.cutoverActive.Store(true) + } + return nil +} + +// AllocationFloor returns the highest timestamp window end applied by the +// dedicated TSO group. It is consensus-owned state, unlike HLC.Current(). +func (f *TSOStateMachine) AllocationFloor() uint64 { + if f == nil { + return 0 + } + return f.allocationFloor.Load() +} + +// CutoverActive reports whether production issuance has durably crossed the +// one-way migration marker. The marker cannot be cleared without a separate +// cluster-wide rollback protocol. +func (f *TSOStateMachine) CutoverActive() bool { + return f != nil && f.cutoverActive.Load() +} + func (f *TSOStateMachine) Snapshot() (raftengine.Snapshot, error) { var ceilingMs int64 var allocationFloor uint64 + var cutoverActive bool if f != nil { ceilingMs = f.ceilingMs.Load() allocationFloor = f.allocationFloor.Load() + cutoverActive = f.cutoverActive.Load() } - return &tsoFSMSnapshot{ceilingMs: ceilingMs, allocationFloor: allocationFloor}, nil + return &tsoFSMSnapshot{ + ceilingMs: ceilingMs, + allocationFloor: allocationFloor, + cutoverActive: cutoverActive, + }, nil } func (f *TSOStateMachine) Restore(r io.Reader) error { @@ -134,12 +174,12 @@ func (f *TSOStateMachine) Restore(r io.Reader) error { if legacy, err := restoreLegacyKVFSMSnapshot(f, br); legacy || err != nil { return err } - ceilingMs, allocationFloor, err := readTSOSnapshotState(br) + ceilingMs, allocationFloor, cutoverActive, err := readTSOSnapshotState(br) if err != nil { return err } if f != nil { - f.restoreSnapshotState(ceilingMs, allocationFloor) + f.restoreSnapshotState(ceilingMs, allocationFloor, cutoverActive) } return nil } @@ -151,13 +191,28 @@ func tsoSnapshotReader(r io.Reader) *bufio.Reader { return bufio.NewReader(r) } -func readTSOSnapshotState(br *bufio.Reader) (int64, uint64, error) { - payload, err := io.ReadAll(io.LimitReader(br, tsoSnapshotV2Len+1)) +func readTSOSnapshotState(br *bufio.Reader) (int64, uint64, bool, error) { + payload, err := io.ReadAll(io.LimitReader(br, tsoSnapshotV3Len+1)) + if err != nil { + return 0, 0, false, errors.Wrap(err, "restore tso fsm snapshot") + } + ceilingMs, allocationFloor, cutoverActive, legacySnapshot, err := decodeTSOSnapshotPayload(payload) if err != nil { - return 0, 0, errors.Wrap(err, "restore tso fsm snapshot") + return 0, 0, false, err + } + if ceilingMs < 0 { + return 0, 0, false, errors.Wrapf(ErrTSOStateMachineInvalidEntry, "tso fsm snapshot: negative ceiling %d", ceilingMs) + } + if legacySnapshot && ceilingMs > 0 { + allocationFloor = tsoLeaseAllocationFloor(ceilingMs) } + return ceilingMs, allocationFloor, cutoverActive, nil +} + +func decodeTSOSnapshotPayload(payload []byte) (int64, uint64, bool, bool, error) { var ceilingMs int64 var allocationFloor uint64 + var cutoverActive bool var legacySnapshot bool switch len(payload) { case tsoSnapshotV1Len: @@ -166,16 +221,32 @@ func readTSOSnapshotState(br *bufio.Reader) (int64, uint64, error) { case tsoSnapshotV2Len: ceilingMs = int64(binary.BigEndian.Uint64(payload[:hlcLeasePayloadLen])) //nolint:gosec // snapshot value. allocationFloor = binary.BigEndian.Uint64(payload[hlcLeasePayloadLen:]) + case tsoSnapshotV3Len: + ceilingMs = int64(binary.BigEndian.Uint64(payload[:hlcLeasePayloadLen])) //nolint:gosec // snapshot value. + allocationFloor = binary.BigEndian.Uint64(payload[hlcLeasePayloadLen:tsoSnapshotV2Len]) + var err error + cutoverActive, err = decodeTSOCutoverByte(payload[tsoSnapshotV2Len]) + if err != nil { + return 0, 0, false, false, err + } default: - return 0, 0, errors.Wrapf(ErrTSOStateMachineInvalidEntry, "tso fsm snapshot: expected %d or %d bytes, got %d", tsoSnapshotV1Len, tsoSnapshotV2Len, len(payload)) + return 0, 0, false, false, errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "tso fsm snapshot: expected %d, %d, or %d bytes, got %d", + tsoSnapshotV1Len, tsoSnapshotV2Len, tsoSnapshotV3Len, len(payload)) } - if ceilingMs < 0 { - return 0, 0, errors.Wrapf(ErrTSOStateMachineInvalidEntry, "tso fsm snapshot: negative ceiling %d", ceilingMs) - } - if legacySnapshot && ceilingMs > 0 { - allocationFloor = tsoLeaseAllocationFloor(ceilingMs) + return ceilingMs, allocationFloor, cutoverActive, legacySnapshot, nil +} + +func decodeTSOCutoverByte(value byte) (bool, error) { + switch value { + case 0: + return false, nil + case 1: + return true, nil + default: + return false, errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "tso fsm snapshot: invalid cutover byte %d", value) } - return ceilingMs, allocationFloor, nil } // restoreLegacyKVFSMSnapshot migrates snapshots produced while reserved group @@ -201,7 +272,7 @@ func restoreLegacyKVFSMSnapshot(f *TSOStateMachine, br *bufio.Reader) (bool, err if f == nil || ceilingMs == 0 { return true, nil } - f.restoreSnapshotState(ceilingMs, tsoLeaseAllocationFloor(ceilingMs)) + f.restoreSnapshotState(ceilingMs, tsoLeaseAllocationFloor(ceilingMs), false) return true, nil } @@ -226,6 +297,9 @@ func hasLegacyKVFSMSnapshotHeader(br *bufio.Reader) (bool, error) { } func (f *TSOStateMachine) IsVolatileOnlyPayload(payload []byte) bool { + if bytes.Equal(payload, []byte(tsoCutoverEnvelope)) { + return true + } return len(payload) == hlcLeaseEntryLen && payload[0] == raftEncodeHLCLease || len(payload) == len(tsoAllocationFloorEnvelope)+hlcLeasePayloadLen && bytes.HasPrefix(payload, []byte(tsoAllocationFloorEnvelope)) @@ -251,7 +325,7 @@ func (f *TSOStateMachine) applyAllocationFloor(floor uint64) { } } -func (f *TSOStateMachine) restoreSnapshotState(ceilingMs int64, allocationFloor uint64) { +func (f *TSOStateMachine) restoreSnapshotState(ceilingMs int64, allocationFloor uint64, cutoverActive bool) { if f == nil { return } @@ -261,6 +335,9 @@ func (f *TSOStateMachine) restoreSnapshotState(ceilingMs int64, allocationFloor if allocationFloor > 0 { storeMaxUint64(&f.allocationFloor, allocationFloor) } + if cutoverActive { + f.cutoverActive.Store(true) + } if f.hlc != nil { if currentCeiling := f.ceilingMs.Load(); currentCeiling > 0 { f.hlc.SetPhysicalCeiling(currentCeiling) @@ -306,9 +383,14 @@ func marshalTSOAllocationFloor(floor uint64) []byte { return out } +func marshalTSOCutover() []byte { + return []byte(tsoCutoverEnvelope) +} + type tsoFSMSnapshot struct { ceilingMs int64 allocationFloor uint64 + cutoverActive bool } func (s *tsoFSMSnapshot) WriteTo(w io.Writer) (int64, error) { @@ -317,13 +399,18 @@ func (s *tsoFSMSnapshot) WriteTo(w io.Writer) (int64, error) { } var ceilingMs int64 var allocationFloor uint64 + var cutoverActive bool if s != nil { ceilingMs = s.ceilingMs allocationFloor = s.allocationFloor + cutoverActive = s.cutoverActive } - var buf [tsoSnapshotV2Len]byte + var buf [tsoSnapshotV3Len]byte binary.BigEndian.PutUint64(buf[:], uint64(ceilingMs)) //nolint:gosec // ceilingMs is a Unix ms timestamp. - binary.BigEndian.PutUint64(buf[hlcLeasePayloadLen:], allocationFloor) + binary.BigEndian.PutUint64(buf[hlcLeasePayloadLen:tsoSnapshotV2Len], allocationFloor) + if cutoverActive { + buf[tsoSnapshotV2Len] = 1 + } n, err := w.Write(buf[:]) if err != nil { return int64(n), errors.Wrap(err, "write tso fsm snapshot") diff --git a/kv/tso_fsm_test.go b/kv/tso_fsm_test.go index 4c60bbba5..3cd346558 100644 --- a/kv/tso_fsm_test.go +++ b/kv/tso_fsm_test.go @@ -157,6 +157,29 @@ func TestTSOStateMachineRejectsBareEncryptionReservedAllocationFloor(t *testing. require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) } +func TestTSOStateMachineAppliesOneWayCutoverMarker(t *testing.T) { + t.Parallel() + + fsm := NewTSOStateMachine(NewHLC()) + require.False(t, fsm.CutoverActive()) + require.Nil(t, fsm.Apply(marshalTSOCutover())) + require.True(t, fsm.CutoverActive()) + require.Nil(t, fsm.Apply(marshalTSOCutover()), "cutover replay must be idempotent") + + err := requireTSOHaltError(t, fsm.Apply(append([]byte(tsoCutoverEnvelope), 1))) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) +} + +func TestTSOStateMachineRejectsInvalidCutoverSnapshotByte(t *testing.T) { + t.Parallel() + + payload := make([]byte, tsoSnapshotV3Len) + payload[tsoSnapshotV2Len] = 2 + err := NewTSOStateMachine(NewHLC()).Restore(bytes.NewReader(payload)) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) + require.ErrorContains(t, err, "invalid cutover byte") +} + func TestTSOStateMachineNilHLCDoesNotPanic(t *testing.T) { t.Parallel() @@ -173,6 +196,7 @@ func TestTSOStateMachineSnapshotRestoreRoundTrip(t *testing.T) { require.Nil(t, source.Apply(marshalHLCLeaseRenew(ceilingMs))) floor := tsoLeaseAllocationFloor(ceilingMs) require.Nil(t, source.Apply(marshalTSOAllocationFloor(floor))) + require.Nil(t, source.Apply(marshalTSOCutover())) snap, err := source.Snapshot() require.NoError(t, err) @@ -181,14 +205,15 @@ func TestTSOStateMachineSnapshotRestoreRoundTrip(t *testing.T) { var buf bytes.Buffer n, err := snap.WriteTo(&buf) require.NoError(t, err) - require.EqualValues(t, tsoSnapshotV2Len, n) - require.Len(t, buf.Bytes(), tsoSnapshotV2Len) + require.EqualValues(t, tsoSnapshotV3Len, n) + require.Len(t, buf.Bytes(), tsoSnapshotV3Len) targetHLC := NewHLC() target := NewTSOStateMachine(targetHLC) require.NoError(t, target.Restore(bytes.NewReader(buf.Bytes()))) require.Equal(t, ceilingMs, targetHLC.PhysicalCeiling()) require.Equal(t, floor, targetHLC.Current()) + require.True(t, target.CutoverActive()) } func TestTSOStateMachineSnapshotUsesTSOOwnedCeiling(t *testing.T) { @@ -274,9 +299,9 @@ func TestTSOStateMachineRestoreKeepsMonotonicCeiling(t *testing.T) { var snapBuf bytes.Buffer n, err := snap.WriteTo(&snapBuf) require.NoError(t, err) - require.EqualValues(t, tsoSnapshotV2Len, n) + require.EqualValues(t, tsoSnapshotV3Len, n) require.Equal(t, uint64(higherCeiling), binary.BigEndian.Uint64(snapBuf.Bytes()[:hlcLeasePayloadLen])) - require.Equal(t, tsoLeaseAllocationFloor(higherCeiling), binary.BigEndian.Uint64(snapBuf.Bytes()[hlcLeasePayloadLen:])) + require.Equal(t, tsoLeaseAllocationFloor(higherCeiling), binary.BigEndian.Uint64(snapBuf.Bytes()[hlcLeasePayloadLen:tsoSnapshotV2Len])) } func TestTSOStateMachineRestoresLegacyKVFSMSnapshot(t *testing.T) { @@ -313,9 +338,10 @@ func TestTSOStateMachineRestoresLegacyKVFSMSnapshot(t *testing.T) { var payload bytes.Buffer _, err = got.WriteTo(&payload) require.NoError(t, err) - require.Len(t, payload.Bytes(), tsoSnapshotV2Len) + require.Len(t, payload.Bytes(), tsoSnapshotV3Len) require.Equal(t, uint64(ceilingMs), binary.BigEndian.Uint64(payload.Bytes()[:hlcLeasePayloadLen])) - require.Equal(t, tsoLeaseAllocationFloor(ceilingMs), binary.BigEndian.Uint64(payload.Bytes()[hlcLeasePayloadLen:])) + require.Equal(t, tsoLeaseAllocationFloor(ceilingMs), binary.BigEndian.Uint64(payload.Bytes()[hlcLeasePayloadLen:tsoSnapshotV2Len])) + require.Zero(t, payload.Bytes()[tsoSnapshotV2Len]) }) } } @@ -326,8 +352,10 @@ func TestTSOStateMachineClassifiesOnlyFullLeaseEntriesAsVolatile(t *testing.T) { fsm := NewTSOStateMachine(NewHLC()) require.True(t, fsm.IsVolatileOnlyPayload(marshalHLCLeaseRenew(1_700_000_123_456))) require.True(t, fsm.IsVolatileOnlyPayload(marshalTSOAllocationFloor(1))) + require.True(t, fsm.IsVolatileOnlyPayload(marshalTSOCutover())) require.False(t, fsm.IsVolatileOnlyPayload([]byte{raftEncodeHLCLease})) require.False(t, fsm.IsVolatileOnlyPayload([]byte(tsoAllocationFloorEnvelope))) + require.False(t, fsm.IsVolatileOnlyPayload(append([]byte(tsoCutoverEnvelope), 1))) require.False(t, fsm.IsVolatileOnlyPayload([]byte{raftEncodeSingle})) } diff --git a/kv/tso_raft.go b/kv/tso_raft.go new file mode 100644 index 000000000..e2810854d --- /dev/null +++ b/kv/tso_raft.go @@ -0,0 +1,686 @@ +package kv + +import ( + "context" + stderrors "errors" + "log/slog" + "sync" + "time" + + "github.com/bootjp/elastickv/internal/raftengine" + pb "github.com/bootjp/elastickv/proto" + "github.com/cockroachdb/errors" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +const ( + defaultTSORouteRetryBudget = 5 * time.Second + defaultTSORouteRetryInterval = 25 * time.Millisecond +) + +var ( + ErrTSOGroupRequired = errors.New("tso: dedicated raft group is required") + ErrTSOStateRequired = errors.New("tso: dedicated state machine is required") + ErrTSOFloorProviderNeeded = errors.New("tso: authoritative commit floor provider is required") + ErrTSONotLeader = errors.New("tso: not leader") + ErrTSOProtocolUnsupported = errors.New("tso: leader does not support durable timestamp windows") +) + +// TSOReservation describes one group-0-serialized timestamp window. +// PreviousAllocationFloor is sampled before this request's reservation and is +// used by shadow migration to reject overlapping legacy candidates. +type TSOReservation struct { + Base uint64 + Count int + PreviousAllocationFloor uint64 + CutoverActive bool +} + +// TSOReservationAllocator is the migration-aware extension exposed by the +// dedicated group leader. Ordinary callers continue to use TSOAllocator. +type TSOReservationAllocator interface { + ReserveBatchAfter(context.Context, int, uint64, bool) (TSOReservation, error) +} + +type TSOShadowReservationAllocator interface { + ValidateShadowTimestamp(context.Context, uint64) (TSOReservation, error) +} + +// RaftTSOAllocator reserves timestamp windows on the dedicated TSO leader. +// A window is returned only after its inclusive end has committed to Raft. +// Failed proposals may leak a local window, but they can never expose an +// uncommitted timestamp or make a later leader reuse a returned timestamp. +type RaftTSOAllocator struct { + group *ShardGroup + clock *HLC + state *TSOStateMachine + floorProvider TSOCutoverFloorProvider + initializedTerm uint64 + mu sync.Mutex +} + +type RaftTSOAllocatorOption func(*RaftTSOAllocator) + +func WithTSOCutoverFloorProvider(provider TSOCutoverFloorProvider) RaftTSOAllocatorOption { + return func(a *RaftTSOAllocator) { + a.floorProvider = provider + } +} + +func NewRaftTSOAllocator(group *ShardGroup, clock *HLC, opts ...RaftTSOAllocatorOption) (*RaftTSOAllocator, error) { + if group == nil || group.Engine == nil { + return nil, errors.WithStack(ErrTSOGroupRequired) + } + if clock == nil { + return nil, errors.WithStack(ErrTSOClockNil) + } + if group.TSOState == nil { + return nil, errors.WithStack(ErrTSOStateRequired) + } + a := &RaftTSOAllocator{group: group, clock: clock, state: group.TSOState} + for _, opt := range opts { + if opt != nil { + opt(a) + } + } + if a.floorProvider == nil { + return nil, errors.WithStack(ErrTSOFloorProviderNeeded) + } + return a, nil +} + +func (a *RaftTSOAllocator) Next(ctx context.Context) (uint64, error) { + return a.NextBatch(ctx, 1) +} + +func (a *RaftTSOAllocator) NextBatch(ctx context.Context, n int) (uint64, error) { + return a.nextBatchAfter(ctx, n, 0) +} + +func (a *RaftTSOAllocator) NextAfter(ctx context.Context, min uint64) (uint64, error) { + return a.NextBatchAfter(ctx, 1, min) +} + +func (a *RaftTSOAllocator) NextBatchAfter(ctx context.Context, n int, min uint64) (uint64, error) { + if min == ^uint64(0) { + return 0, errors.WithStack(ErrTxnCommitTSRequired) + } + return a.nextBatchAfter(ctx, n, min) +} + +func (a *RaftTSOAllocator) nextBatchAfter(ctx context.Context, n int, min uint64) (uint64, error) { + reservation, err := a.ReserveBatchAfter(ctx, n, min, false) + return reservation.Base, err +} + +// ReserveBatchAfter serializes floor discovery, the one-way cutover marker, +// and window reservation under the TSO leader. A returned window is always +// above both the caller's minimum and every authoritative data-group commit +// observed when this leader term first serves a request. +func (a *RaftTSOAllocator) ReserveBatchAfter( + ctx context.Context, + n int, + min uint64, + activateCutover bool, +) (TSOReservation, error) { + var empty TSOReservation + if err := validateTSOBatchSize(n); err != nil { + return empty, err + } + if min == ^uint64(0) { + return empty, errors.WithStack(ErrTxnCommitTSRequired) + } + ctx = nonNilTSOContext(ctx) + a.mu.Lock() + defer a.mu.Unlock() + return a.reserveBatchAfterLocked(ctx, n, min, activateCutover) +} + +func (a *RaftTSOAllocator) reserveBatchAfterLocked( + ctx context.Context, + n int, + min uint64, + activateCutover bool, +) (TSOReservation, error) { + var empty TSOReservation + engine, err := a.verifiedLeader(ctx) + if err != nil { + return empty, err + } + term := engine.Status().Term + if term == 0 { + return empty, errors.Wrap(ErrTSONotLeader, "tso leader has no active term") + } + previousFloor, err := a.prepareLeaderTermReservation(ctx, engine, term, activateCutover) + if err != nil { + return empty, err + } + minimum := max(min, previousFloor) + if minimum > 0 { + a.clock.Observe(minimum) + } + base, err := a.clock.NextBatchFenced(n) + if err != nil { + return empty, errors.Wrap(err, "tso reserve local batch") + } + end := base + uint64(n) - 1 //nolint:gosec // n is positive and bounded above. + if err := a.commitAllocationFloor(ctx, engine, end); err != nil { + return empty, err + } + if err := verifyTSOLeaderTerm(ctx, engine, term, false); err != nil { + return empty, err + } + a.initializedTerm = term + return TSOReservation{ + Base: base, + Count: n, + PreviousAllocationFloor: previousFloor, + CutoverActive: a.state.CutoverActive(), + }, nil +} + +func (a *RaftTSOAllocator) prepareLeaderTermReservation( + ctx context.Context, + engine raftengine.Engine, + term uint64, + activateCutover bool, +) (uint64, error) { + termFloor, err := a.termCommitFloor(ctx, term) + if err != nil { + return 0, err + } + previousFloor := max(a.state.AllocationFloor(), termFloor) + if activateCutover && !a.state.CutoverActive() { + if err := a.commitCutover(ctx, engine); err != nil { + return 0, err + } + } + if err := verifyTSOLeaderTerm(ctx, engine, term, true); err != nil { + return 0, err + } + return previousFloor, nil +} + +func (a *RaftTSOAllocator) termCommitFloor(ctx context.Context, term uint64) (uint64, error) { + if a.initializedTerm == term { + return a.state.AllocationFloor(), nil + } + floor, err := a.floorProvider.GlobalCommittedTimestampFloor(ctx) + if err != nil { + return 0, errors.Wrap(err, "tso initialize leader term commit floor") + } + return max(floor, a.state.AllocationFloor()), nil +} + +func (a *RaftTSOAllocator) verifiedLeader(ctx context.Context) (raftengine.Engine, error) { + if err := ctx.Err(); err != nil { + return nil, errors.Wrap(err, "tso reserve batch") + } + engine := engineForGroup(a.group) + if !isLeaderEngine(engine) { + return nil, errors.WithStack(ErrTSONotLeader) + } + if err := verifyLeaderEngineCtx(ctx, engine); err != nil { + return nil, errors.Wrap(stderrors.Join(ErrTSONotLeader, err), "tso verify leader") + } + return engine, nil +} + +func verifyTSOLeaderTerm(ctx context.Context, engine raftengine.Engine, expectedTerm uint64, fence bool) error { + if err := ctx.Err(); err != nil { + return errors.Wrap(err, "tso verify leader term") + } + if !isLeaderEngine(engine) { + return errors.WithStack(ErrTSONotLeader) + } + if fence { + if err := verifyLeaderEngineCtx(ctx, engine); err != nil { + return errors.Wrap(stderrors.Join(ErrTSONotLeader, err), "tso fence leader term") + } + } + status := engine.Status() + if status.State != raftengine.StateLeader || status.Term != expectedTerm { + return errors.Wrapf(ErrTSONotLeader, + "tso leader term changed: expected=%d current=%d state=%v", + expectedTerm, status.Term, status.State) + } + return nil +} + +func (a *RaftTSOAllocator) commitAllocationFloor(ctx context.Context, engine raftengine.Engine, end uint64) error { + if _, err := a.group.Proposer().Propose(ctx, marshalTSOAllocationFloor(end)); err != nil { + wrapped := errors.Wrap(err, "tso commit allocation floor") + if tsoProposalLostLeadership(engine, err) { + return stderrors.Join(ErrTSONotLeader, wrapped) + } + return wrapped + } + return nil +} + +func (a *RaftTSOAllocator) commitCutover(ctx context.Context, engine raftengine.Engine) error { + if _, err := a.group.Proposer().Propose(ctx, marshalTSOCutover()); err != nil { + wrapped := errors.Wrap(err, "tso commit cutover marker") + if tsoProposalLostLeadership(engine, err) { + return stderrors.Join(ErrTSONotLeader, wrapped) + } + return wrapped + } + if !a.state.CutoverActive() { + return errors.New("tso cutover proposal committed without applied marker") + } + return nil +} + +func tsoProposalLostLeadership(engine raftengine.Engine, err error) bool { + return !isLeaderEngine(engine) || errors.Is(err, raftengine.ErrNotLeader) || isTransientLeaderError(err) +} + +func (a *RaftTSOAllocator) IsLeader() bool { + return a != nil && isLeaderEngine(engineForGroup(a.group)) +} + +func (a *RaftTSOAllocator) RunLeaseRenewal(ctx context.Context) { + <-nonNilTSOContext(ctx).Done() +} + +type tsoRemoteRequest func(context.Context, string, int, uint64, bool) (TSOReservation, error) + +// LeaderRoutedTSOAllocator serves local requests on the TSO leader and sends +// follower requests to the leader address published by the group-0 engine. +// It re-resolves that address after transient errors so a leadership change +// does not pin a BatchAllocator refill to a stale endpoint. +type LeaderRoutedTSOAllocator struct { + local TSOAllocator + leader raftengine.LeaderView + connCache GRPCConnCache + remoteRequest tsoRemoteRequest + retryBudget time.Duration + retryInterval time.Duration + activate bool + clock *HLC +} + +type LeaderRoutedTSOAllocatorOption func(*LeaderRoutedTSOAllocator) + +func WithTSOCutoverActivation() LeaderRoutedTSOAllocatorOption { + return func(a *LeaderRoutedTSOAllocator) { a.activate = true } +} + +func WithTSORoutedClock(clock *HLC) LeaderRoutedTSOAllocatorOption { + return func(a *LeaderRoutedTSOAllocator) { a.clock = clock } +} + +func NewLeaderRoutedTSOAllocator( + local TSOAllocator, + leader raftengine.LeaderView, + opts ...LeaderRoutedTSOAllocatorOption, +) (*LeaderRoutedTSOAllocator, error) { + if local == nil { + return nil, errors.WithStack(ErrTSOAllocatorRequired) + } + if leader == nil { + return nil, errors.WithStack(ErrTSOGroupRequired) + } + a := &LeaderRoutedTSOAllocator{ + local: local, + leader: leader, + retryBudget: defaultTSORouteRetryBudget, + retryInterval: defaultTSORouteRetryInterval, + } + for _, opt := range opts { + if opt != nil { + opt(a) + } + } + a.remoteRequest = a.requestRemoteBatch + return a, nil +} + +func (a *LeaderRoutedTSOAllocator) Next(ctx context.Context) (uint64, error) { + return a.NextBatch(ctx, 1) +} + +func (a *LeaderRoutedTSOAllocator) NextBatch(ctx context.Context, n int) (uint64, error) { + return a.nextBatchAfter(ctx, n, 0) +} + +func (a *LeaderRoutedTSOAllocator) NextAfter(ctx context.Context, min uint64) (uint64, error) { + return a.NextBatchAfter(ctx, 1, min) +} + +func (a *LeaderRoutedTSOAllocator) NextBatchAfter(ctx context.Context, n int, min uint64) (uint64, error) { + if min == ^uint64(0) { + return 0, errors.WithStack(ErrTxnCommitTSRequired) + } + return a.nextBatchAfter(ctx, n, min) +} + +func (a *LeaderRoutedTSOAllocator) nextBatchAfter(ctx context.Context, n int, min uint64) (uint64, error) { + reservation, err := a.nextReservation(ctx, n, min, a.activate, a.activate) + if err != nil { + return 0, err + } + a.observeReservation(reservation) + return reservation.Base, nil +} + +func (a *LeaderRoutedTSOAllocator) ValidateShadowTimestamp(ctx context.Context, min uint64) (TSOReservation, error) { + reservation, err := a.nextReservation(ctx, 1, min, false, true) + if err == nil { + a.observeReservation(reservation) + } + return reservation, err +} + +func (a *LeaderRoutedTSOAllocator) nextReservation( + ctx context.Context, + n int, + min uint64, + activate bool, + requireReservation bool, +) (TSOReservation, error) { + var empty TSOReservation + if err := validateTSOBatchSize(n); err != nil { + return empty, err + } + ctx = nonNilTSOContext(ctx) + deadline := time.Now().Add(a.retryBudget) + ctx, cancel := context.WithDeadline(ctx, deadline) + defer cancel() + + var lastErr error + for { + reservation, err := a.tryBatch(ctx, n, min, activate, requireReservation) + if err == nil { + return reservation, nil + } + lastErr = err + if !isTransientTSORouteError(err) { + return empty, err + } + if err := waitTSORouteRetry(ctx, a.retryInterval); err != nil { + if ctxErr := ctx.Err(); ctxErr != nil { + if lastErr != nil { + return empty, errors.Wrap(stderrors.Join(ctxErr, lastErr), "tso route retry exhausted") + } + return empty, errors.Wrap(ctxErr, "tso route retry exhausted") + } + return empty, err + } + } +} + +func (a *LeaderRoutedTSOAllocator) tryBatch( + ctx context.Context, + n int, + min uint64, + activate bool, + requireReservation bool, +) (TSOReservation, error) { + var empty TSOReservation + if a.local.IsLeader() { + return a.tryLocalBatch(ctx, n, min, activate, requireReservation) + } + addr := leaderAddrFromEngine(a.leader) + if addr == "" { + return empty, errors.WithStack(ErrLeaderNotFound) + } + reservation, err := a.remoteRequest(ctx, addr, n, min, activate) + if err != nil { + return empty, err + } + if err := validateTSOReservation(reservation, n, min, requireReservation); err != nil { + return empty, err + } + return reservation, nil +} + +func (a *LeaderRoutedTSOAllocator) tryLocalBatch( + ctx context.Context, + n int, + min uint64, + activate bool, + requireReservation bool, +) (TSOReservation, error) { + var empty TSOReservation + reservationAllocator, ok := a.local.(TSOReservationAllocator) + if !ok { + if requireReservation { + return empty, errors.WithStack(ErrTSOProtocolUnsupported) + } + base, err := nextTSOBatchAfter(ctx, a.local, n, min) + if err != nil { + return empty, err + } + reservation := TSOReservation{Base: base, Count: n} + return reservation, validateTSOReservation(reservation, n, min, false) + } + reservation, err := reservationAllocator.ReserveBatchAfter(ctx, n, min, activate) + if err != nil { + return empty, errors.Wrap(err, "reserve local TSO window") + } + return reservation, validateTSOReservation(reservation, n, min, true) +} + +func (a *LeaderRoutedTSOAllocator) requestRemoteBatch( + ctx context.Context, + addr string, + n int, + min uint64, + activate bool, +) (TSOReservation, error) { + var empty TSOReservation + conn, err := a.connCache.ConnFor(addr) + if err != nil { + return empty, errors.Wrap(err, "tso dial leader") + } + resp, err := pb.NewDistributionClient(conn).GetTimestamp(ctx, &pb.GetTimestampRequest{ + Count: uint32(n), //nolint:gosec // n is bounded by maxHLCBatchSize. + MinTimestamp: min, + ActivateCutover: activate, + }) + if err != nil { + return empty, errors.Wrap(err, "tso request leader batch") + } + count := uint32(n) //nolint:gosec // n is validated before this request. + if !resp.GetCommittedByDedicatedTso() || resp.GetCount() != count { + return empty, errors.Wrapf(ErrTSOProtocolUnsupported, + "leader response committed=%t count=%d, want committed=true count=%d", + resp.GetCommittedByDedicatedTso(), resp.GetCount(), n) + } + if activate && !resp.GetCutoverActive() { + return empty, errors.Wrap(ErrTSOProtocolUnsupported, + "leader did not confirm durable TSO cutover") + } + return TSOReservation{ + Base: resp.GetTimestamp(), + Count: n, + PreviousAllocationFloor: resp.GetPreviousAllocationFloor(), + CutoverActive: resp.GetCutoverActive(), + }, nil +} + +func (a *LeaderRoutedTSOAllocator) observeReservation(reservation TSOReservation) { + if a == nil || a.clock == nil || reservation.Base == 0 || reservation.Count <= 0 { + return + } + end := reservation.Base + uint64(reservation.Count) - 1 //nolint:gosec // validated reservation. + a.clock.Observe(end) +} + +func validateRoutedTSOWindow(base uint64, n int, min uint64) error { + if base == 0 || base <= min { + return errors.Wrapf(ErrTxnCommitTSRequired, + "tso leader returned base %d at or below minimum %d", base, min) + } + size := uint64(n) //nolint:gosec // n was validated as positive and bounded. + if base > ^uint64(0)-(size-1) { + return errors.Wrapf(ErrTxnCommitTSRequired, + "tso leader window base %d count %d overflows", base, n) + } + return nil +} + +func validateTSOReservation(reservation TSOReservation, n int, min uint64, requireMetadata bool) error { + if err := validateRoutedTSOWindow(reservation.Base, n, min); err != nil { + return err + } + if reservation.Count != n { + return errors.Wrapf(ErrTSOProtocolUnsupported, + "tso reservation count=%d, want %d", reservation.Count, n) + } + if requireMetadata && reservation.PreviousAllocationFloor >= reservation.Base { + return errors.Wrapf(ErrTSOProtocolUnsupported, + "tso reservation base=%d does not exceed previous floor=%d", + reservation.Base, reservation.PreviousAllocationFloor) + } + return nil +} + +func (a *LeaderRoutedTSOAllocator) IsLeader() bool { + return a != nil && a.local != nil && a.local.IsLeader() +} + +func (a *LeaderRoutedTSOAllocator) RunLeaseRenewal(ctx context.Context) { + if a == nil || a.local == nil { + <-nonNilTSOContext(ctx).Done() + return + } + a.local.RunLeaseRenewal(ctx) +} + +func (a *LeaderRoutedTSOAllocator) Close() error { + if a == nil { + return nil + } + return a.connCache.Close() +} + +func nextTSOBatchAfter(ctx context.Context, alloc TSOAllocator, n int, min uint64) (uint64, error) { + if after, ok := alloc.(tsoBatchAfterAllocator); ok && min > 0 { + base, err := after.NextBatchAfter(ctx, n, min) + return base, errors.Wrap(err, "tso allocate batch after minimum") + } + base, err := alloc.NextBatch(ctx, n) + if err != nil { + return 0, errors.Wrap(err, "tso allocate batch") + } + if base <= min { + return 0, errors.WithStack(ErrTxnCommitTSRequired) + } + return base, nil +} + +func isTransientTSORouteError(err error) bool { + if err == nil { + return false + } + if errors.Is(err, ErrTSONotLeader) || errors.Is(err, ErrTSOProtocolUnsupported) || isTransientLeaderError(err) { + return true + } + code := status.Code(err) + return code == codes.Aborted || code == codes.FailedPrecondition || code == codes.Unavailable +} + +func validateTSOBatchSize(n int) error { + if n <= 0 || n > maxHLCBatchSize { + return errors.WithStack(ErrInvalidTSOBatchSize) + } + return nil +} + +func waitTSORouteRetry(ctx context.Context, interval time.Duration) error { + timer := time.NewTimer(interval) + defer timer.Stop() + select { + case <-timer.C: + return nil + case <-ctx.Done(): + return errors.Wrap(ctx.Err(), "tso route retry") + } +} + +func nonNilTSOContext(ctx context.Context) context.Context { + if ctx == nil { + return context.Background() + } + return ctx +} + +// ShadowTimestampAllocator serializes each legacy candidate through group 0 +// before returning it. Candidates at or below a prior TSO floor are discarded +// and retried; once the durable cutover marker is active, the allocator returns +// the reserved TSO timestamp directly so rolling restarts cannot mix sources. +type ShadowTimestampAllocator struct { + legacy *HLC + shadow TSOShadowReservationAllocator + log *slog.Logger +} + +func NewShadowTimestampAllocator(legacy *HLC, shadow TSOShadowReservationAllocator, logger *slog.Logger) (*ShadowTimestampAllocator, error) { + if legacy == nil { + return nil, errors.WithStack(ErrTSOClockNil) + } + if shadow == nil { + return nil, errors.WithStack(ErrTSOAllocatorRequired) + } + if logger == nil { + logger = slog.Default() + } + return &ShadowTimestampAllocator{ + legacy: legacy, + shadow: shadow, + log: logger, + }, nil +} + +func (a *ShadowTimestampAllocator) Next(ctx context.Context) (uint64, error) { + return a.nextAfter(ctx, 0) +} + +func (a *ShadowTimestampAllocator) NextAfter(ctx context.Context, min uint64) (uint64, error) { + if min == ^uint64(0) { + return 0, errors.WithStack(ErrTxnCommitTSRequired) + } + return a.nextAfter(ctx, min) +} + +func (a *ShadowTimestampAllocator) nextAfter(ctx context.Context, min uint64) (uint64, error) { + ctx = nonNilTSOContext(ctx) + a.legacy.Observe(min) + for { + if err := ctx.Err(); err != nil { + return 0, errors.Wrap(err, "tso shadow migration") + } + legacyTS, err := a.legacy.NextFenced() + if err != nil { + return 0, errors.Wrap(err, "tso shadow allocate legacy timestamp") + } + reservation, err := a.shadow.ValidateShadowTimestamp(ctx, legacyTS) + if err != nil { + a.log.ErrorContext(ctx, "tso shadow allocation failed", + slog.Uint64("legacy_ts", legacyTS), + slog.Any("err", err), + ) + return 0, errors.Wrap(err, "tso shadow validation") + } + a.legacy.Observe(reservation.Base) + if reservation.CutoverActive { + return reservation.Base, nil + } + if legacyTS > reservation.PreviousAllocationFloor { + return legacyTS, nil + } + a.log.WarnContext(ctx, "tso shadow discarded overlapping legacy timestamp", + slog.Uint64("legacy_ts", legacyTS), + slog.Uint64("previous_tso_floor", reservation.PreviousAllocationFloor), + slog.Uint64("reserved_tso_ts", reservation.Base), + ) + } +} + +func (a *ShadowTimestampAllocator) Close() error { + return nil +} diff --git a/kv/tso_raft_test.go b/kv/tso_raft_test.go new file mode 100644 index 000000000..10f4925b0 --- /dev/null +++ b/kv/tso_raft_test.go @@ -0,0 +1,596 @@ +package kv + +import ( + "bytes" + "context" + "encoding/binary" + "io" + "log/slog" + "sync" + "testing" + "time" + + "github.com/bootjp/elastickv/internal/raftengine" + "github.com/cockroachdb/errors" + "github.com/stretchr/testify/require" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +func TestRaftTSOAllocatorCommitsWindowEndBeforeReturning(t *testing.T) { + clock := NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) + fsm := NewTSOStateMachine(clock) + engine := &recordingTSOEngine{state: raftengine.StateLeader, leader: raftengine.LeaderInfo{Address: "self"}} + engine.apply = func(payload []byte) error { + if result := fsm.Apply(payload); result != nil { + if err, ok := result.(error); ok { + return err + } + return errors.Newf("unexpected TSO FSM result %T", result) + } + return nil + } + alloc, err := newTestRaftTSOAllocator(&ShardGroup{Engine: engine, TSOState: fsm}, clock) + require.NoError(t, err) + + base, err := alloc.NextBatch(context.Background(), testTSOBatchSize) + require.NoError(t, err) + require.NotZero(t, base) + + payloads := engine.proposedPayloads() + require.Len(t, payloads, 1) + require.True(t, bytes.HasPrefix(payloads[0], []byte(tsoAllocationFloorEnvelope))) + wantEnd := base + uint64(testTSOBatchSize) - 1 + require.Equal(t, wantEnd, binary.BigEndian.Uint64(payloads[0][len(tsoAllocationFloorEnvelope):])) + + snapshot, err := fsm.Snapshot() + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, snapshot.Close()) }) + var encoded snapshotBuffer + _, err = snapshot.WriteTo(&encoded) + require.NoError(t, err) + require.Equal(t, wantEnd, binary.BigEndian.Uint64(encoded.Bytes()[hlcLeasePayloadLen:])) +} + +func TestRaftTSOAllocatorRequiresCommitFloorProvider(t *testing.T) { + clock := NewHLC() + fsm := NewTSOStateMachine(clock) + engine := &recordingTSOEngine{state: raftengine.StateLeader} + + _, err := NewRaftTSOAllocator(&ShardGroup{Engine: engine, TSOState: fsm}, clock) + require.ErrorIs(t, err, ErrTSOFloorProviderNeeded) +} + +func TestRaftTSOAllocatorDoesNotReturnFailedReservation(t *testing.T) { + clock := NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) + engine := &recordingTSOEngine{ + state: raftengine.StateLeader, + leader: raftengine.LeaderInfo{Address: "self"}, + proposeErr: raftengine.ErrNotLeader, + } + fsm := NewTSOStateMachine(clock) + engine.apply = applyTSOTestFSM(fsm) + alloc, err := newTestRaftTSOAllocator(&ShardGroup{Engine: engine, TSOState: fsm}, clock) + require.NoError(t, err) + + _, err = alloc.NextBatch(context.Background(), testTSOBatchSize) + require.ErrorIs(t, err, ErrTSONotLeader) + leakedEnd := clock.Current() + require.NotZero(t, leakedEnd) + + engine.setProposeError(nil) + base, err := alloc.NextBatch(context.Background(), testTSOBatchSize) + require.NoError(t, err) + require.Greater(t, base, leakedEnd) +} + +func TestRaftTSOAllocatorInitializesEveryLeaderTermAboveDataFloor(t *testing.T) { + clock := NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) + fsm := NewTSOStateMachine(clock) + engine := &recordingTSOEngine{ + state: raftengine.StateLeader, + leader: raftengine.LeaderInfo{Address: "self"}, + term: 7, + apply: applyTSOTestFSM(fsm), + } + floor := uint64(time.Now().Add(time.Minute).UnixMilli()) << hlcLogicalBits //nolint:gosec // future test HLC. + provider := &recordingTSOFloorProvider{floor: floor} + alloc, err := NewRaftTSOAllocator( + &ShardGroup{Engine: engine, TSOState: fsm}, + clock, + WithTSOCutoverFloorProvider(provider), + ) + require.NoError(t, err) + + first, err := alloc.Next(context.Background()) + require.NoError(t, err) + require.Greater(t, first, floor) + require.Equal(t, 1, provider.callCount()) + + _, err = alloc.Next(context.Background()) + require.NoError(t, err) + require.Equal(t, 1, provider.callCount(), "one leader term must read the data floor once") + + engine.setTerm(8) + _, err = alloc.Next(context.Background()) + require.NoError(t, err) + require.Equal(t, 2, provider.callCount(), "a new leader term must fence against data groups again") +} + +func TestRaftTSOAllocatorRejectsTermChangeDuringCommitFloorRead(t *testing.T) { + clock := NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) + fsm := NewTSOStateMachine(clock) + engine := &recordingTSOEngine{ + state: raftengine.StateLeader, + leader: raftengine.LeaderInfo{Address: "self"}, + term: 1, + apply: applyTSOTestFSM(fsm), + } + provider := &recordingTSOFloorProvider{afterRead: func() { engine.setTerm(2) }} + alloc, err := NewRaftTSOAllocator( + &ShardGroup{Engine: engine, TSOState: fsm}, + clock, + WithTSOCutoverFloorProvider(provider), + ) + require.NoError(t, err) + + _, err = alloc.Next(context.Background()) + require.ErrorIs(t, err, ErrTSONotLeader) + require.Empty(t, engine.proposedPayloads()) +} + +func TestRaftTSOAllocatorDropsCommittedWindowAfterTermChange(t *testing.T) { + clock := NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) + fsm := NewTSOStateMachine(clock) + engine := &recordingTSOEngine{ + state: raftengine.StateLeader, + leader: raftengine.LeaderInfo{Address: "self"}, + term: 1, + apply: applyTSOTestFSM(fsm), + } + var changeTerm sync.Once + engine.afterPropose = func(payload []byte) { + if bytes.HasPrefix(payload, []byte(tsoAllocationFloorEnvelope)) { + changeTerm.Do(func() { engine.setTerm(2) }) + } + } + alloc, err := newTestRaftTSOAllocator(&ShardGroup{Engine: engine, TSOState: fsm}, clock) + require.NoError(t, err) + + _, err = alloc.Next(context.Background()) + require.ErrorIs(t, err, ErrTSONotLeader) + leakedFloor := fsm.AllocationFloor() + require.NotZero(t, leakedFloor) + + next, err := alloc.Next(context.Background()) + require.NoError(t, err) + require.Greater(t, next, leakedFloor) +} + +func TestRaftTSOAllocatorCommitsCutoverBeforeProductionWindow(t *testing.T) { + clock := NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) + fsm := NewTSOStateMachine(clock) + engine := &recordingTSOEngine{ + state: raftengine.StateLeader, + leader: raftengine.LeaderInfo{Address: "self"}, + term: 1, + apply: applyTSOTestFSM(fsm), + } + alloc, err := newTestRaftTSOAllocator(&ShardGroup{Engine: engine, TSOState: fsm}, clock) + require.NoError(t, err) + + reservation, err := alloc.ReserveBatchAfter(context.Background(), testTSOBatchSize, 0, true) + require.NoError(t, err) + require.True(t, reservation.CutoverActive) + payloads := engine.proposedPayloads() + require.Len(t, payloads, 2) + require.Equal(t, []byte(tsoCutoverEnvelope), payloads[0]) + require.True(t, bytes.HasPrefix(payloads[1], []byte(tsoAllocationFloorEnvelope))) +} + +func TestLeaderRoutedTSOAllocatorReResolvesAfterStaleLeader(t *testing.T) { + local := &fakeTSOAllocator{leader: false, nextBase: testTSOInitialBase} + engine := &recordingTSOEngine{state: raftengine.StateFollower, leader: raftengine.LeaderInfo{Address: "old"}} + alloc, err := NewLeaderRoutedTSOAllocator(local, engine) + require.NoError(t, err) + alloc.retryBudget = time.Second + alloc.retryInterval = time.Millisecond + + var addresses []string + alloc.remoteRequest = func(_ context.Context, addr string, n int, min uint64, activate bool) (TSOReservation, error) { + addresses = append(addresses, addr) + require.Equal(t, testTSOBatchSize, n) + require.Equal(t, uint64(testTSOInitialBase), min) + require.False(t, activate) + if addr == "old" { + engine.setLeaderAddress("new") + return TSOReservation{}, status.Error(codes.FailedPrecondition, "tso: not leader") + } + return TSOReservation{Base: testTSOInitialBase + 1, Count: n}, nil + } + + base, err := alloc.NextBatchAfter(context.Background(), testTSOBatchSize, testTSOInitialBase) + require.NoError(t, err) + require.Equal(t, uint64(testTSOInitialBase+1), base) + require.Equal(t, []string{"old", "new"}, addresses) +} + +func TestLeaderRoutedTSOAllocatorUsesLocalLeader(t *testing.T) { + local := &fakeTSOAllocator{leader: true, nextBase: testTSOInitialBase} + engine := &recordingTSOEngine{state: raftengine.StateLeader, leader: raftengine.LeaderInfo{Address: "self"}} + alloc, err := NewLeaderRoutedTSOAllocator(local, engine) + require.NoError(t, err) + alloc.remoteRequest = func(context.Context, string, int, uint64, bool) (TSOReservation, error) { + t.Fatal("local TSO leader must not call remote RPC") + return TSOReservation{}, nil + } + + base, err := alloc.NextBatch(context.Background(), testTSOBatchSize) + require.NoError(t, err) + require.Equal(t, uint64(testTSOInitialBase), base) +} + +func TestLeaderRoutedTSOAllocatorRejectsLocalShadowWithoutReservationMetadata(t *testing.T) { + local := &fakeTSOAllocator{leader: true, nextBase: testTSOInitialBase} + engine := &recordingTSOEngine{state: raftengine.StateLeader, leader: raftengine.LeaderInfo{Address: "self"}} + alloc, err := NewLeaderRoutedTSOAllocator(local, engine) + require.NoError(t, err) + + _, err = alloc.ValidateShadowTimestamp(context.Background(), testTSOInitialBase-1) + require.ErrorIs(t, err, ErrTSOProtocolUnsupported) +} + +func TestLeaderRoutedTSOAllocatorRejectsRemoteWindowAtMinimum(t *testing.T) { + local := &fakeTSOAllocator{leader: false} + engine := &recordingTSOEngine{state: raftengine.StateFollower, leader: raftengine.LeaderInfo{Address: "leader"}} + alloc, err := NewLeaderRoutedTSOAllocator(local, engine) + require.NoError(t, err) + alloc.remoteRequest = func(context.Context, string, int, uint64, bool) (TSOReservation, error) { + return TSOReservation{Base: testTSOInitialBase, Count: testTSOBatchSize}, nil + } + + _, err = alloc.NextBatchAfter(context.Background(), testTSOBatchSize, testTSOInitialBase) + require.ErrorIs(t, err, ErrTxnCommitTSRequired) +} + +func TestLeaderRoutedTSOAllocatorPreservesDeadlineAfterTransientErrors(t *testing.T) { + local := &fakeTSOAllocator{leader: false} + engine := &recordingTSOEngine{state: raftengine.StateFollower, leader: raftengine.LeaderInfo{Address: "leader"}} + alloc, err := NewLeaderRoutedTSOAllocator(local, engine) + require.NoError(t, err) + alloc.retryBudget = time.Second + alloc.retryInterval = time.Millisecond + + var attempts int + alloc.remoteRequest = func(context.Context, string, int, uint64, bool) (TSOReservation, error) { + attempts++ + return TSOReservation{}, status.Error(codes.Unavailable, "leader restarting") + } + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + _, err = alloc.NextBatch(ctx, testTSOBatchSize) + require.ErrorIs(t, err, context.DeadlineExceeded) + require.Greater(t, attempts, 1) +} + +func TestShadowTimestampAllocatorReturnsLegacyAndAdvancesTSO(t *testing.T) { + legacy := NewHLC() + legacy.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) + shadow := &recordingShadowReservationAllocator{} + alloc, err := NewShadowTimestampAllocator(legacy, shadow, slog.New(slog.NewTextHandler(io.Discard, nil))) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, alloc.Close()) }) + + legacyTS, err := alloc.NextAfter(context.Background(), testTSOInitialBase) + require.NoError(t, err) + require.Greater(t, legacyTS, uint64(testTSOInitialBase)) + min, returned, calls := shadow.values() + require.Equal(t, legacyTS, min) + require.Equal(t, legacyTS+1, returned) + require.Equal(t, 1, calls) + require.Equal(t, returned, legacy.Current(), "shadow reservation must keep the rollback clock warm") +} + +func TestShadowTimestampAllocatorFailsClosedOnShadowError(t *testing.T) { + legacy := NewHLC() + legacy.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) + shadowErr := errors.New("shadow unavailable") + shadow := &recordingShadowReservationAllocator{err: shadowErr} + alloc, err := NewShadowTimestampAllocator(legacy, shadow, slog.New(slog.NewTextHandler(io.Discard, nil))) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, alloc.Close()) }) + + _, err = alloc.Next(context.Background()) + require.ErrorIs(t, err, shadowErr) +} + +func TestShadowTimestampAllocatorHonorsDeadlineWhileShadowBlocked(t *testing.T) { + legacy := NewHLC() + legacy.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) + shadow := &blockingShadowTimestampAllocator{started: make(chan struct{})} + alloc, err := NewShadowTimestampAllocator(legacy, shadow, slog.New(slog.NewTextHandler(io.Discard, nil))) + require.NoError(t, err) + + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + _, err = alloc.Next(ctx) + require.ErrorIs(t, err, context.DeadlineExceeded) + require.NoError(t, alloc.Close()) +} + +func TestShadowTimestampAllocatorDiscardsCandidateBelowPriorFloor(t *testing.T) { + legacy := NewHLC() + legacy.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) + shadow := &recordingShadowReservationAllocator{overlapFirst: true} + alloc, err := NewShadowTimestampAllocator(legacy, shadow, slog.New(slog.NewTextHandler(io.Discard, nil))) + require.NoError(t, err) + + legacyTS, err := alloc.Next(context.Background()) + require.NoError(t, err) + _, reserved, calls := shadow.values() + require.Equal(t, 2, calls) + require.Equal(t, legacyTS+1, reserved) +} + +func TestShadowTimestampAllocatorReturnsTSOAfterDurableCutover(t *testing.T) { + legacy := NewHLC() + legacy.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) + shadow := &recordingShadowReservationAllocator{cutover: true} + alloc, err := NewShadowTimestampAllocator(legacy, shadow, slog.New(slog.NewTextHandler(io.Discard, nil))) + require.NoError(t, err) + + issued, err := alloc.Next(context.Background()) + require.NoError(t, err) + legacyCandidate, reserved, calls := shadow.values() + require.Equal(t, 1, calls) + require.Equal(t, reserved, issued) + require.Greater(t, issued, legacyCandidate) +} + +func TestShadowAndCutoverAllocatorsSerializeMigrationOnGroupZero(t *testing.T) { + tsoClock := NewHLC() + tsoClock.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) + fsm := NewTSOStateMachine(tsoClock) + engine := &recordingTSOEngine{ + state: raftengine.StateLeader, + leader: raftengine.LeaderInfo{Address: "self"}, + term: 1, + apply: applyTSOTestFSM(fsm), + } + local, err := newTestRaftTSOAllocator(&ShardGroup{Engine: engine, TSOState: fsm}, tsoClock) + require.NoError(t, err) + + legacyClock := NewHLC() + legacyClock.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) + shadowRoute, err := NewLeaderRoutedTSOAllocator(local, engine, WithTSORoutedClock(legacyClock)) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, shadowRoute.Close()) }) + shadow, err := NewShadowTimestampAllocator(legacyClock, shadowRoute, slog.New(slog.NewTextHandler(io.Discard, nil))) + require.NoError(t, err) + + legacyIssued, err := shadow.Next(context.Background()) + require.NoError(t, err) + require.False(t, fsm.CutoverActive()) + + cutoverRoute, err := NewLeaderRoutedTSOAllocator(local, engine, WithTSOCutoverActivation()) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, cutoverRoute.Close()) }) + cutoverIssued, err := cutoverRoute.Next(context.Background()) + require.NoError(t, err) + require.True(t, fsm.CutoverActive()) + require.Greater(t, cutoverIssued, legacyIssued) + + shadowAfterCutover, err := shadow.Next(context.Background()) + require.NoError(t, err) + require.Greater(t, shadowAfterCutover, cutoverIssued) + require.Equal(t, fsm.AllocationFloor(), shadowAfterCutover) +} + +type recordingShadowReservationAllocator struct { + mu sync.Mutex + min uint64 + returned uint64 + calls int + err error + overlapFirst bool + cutover bool +} + +func (a *recordingShadowReservationAllocator) ValidateShadowTimestamp(_ context.Context, min uint64) (TSOReservation, error) { + a.mu.Lock() + defer a.mu.Unlock() + a.min = min + a.calls++ + if a.err != nil { + return TSOReservation{}, a.err + } + a.returned = min + 1 + previousFloor := min - 1 + if a.overlapFirst && a.calls == 1 { + previousFloor = min + } + return TSOReservation{ + Base: a.returned, + Count: 1, + PreviousAllocationFloor: previousFloor, + CutoverActive: a.cutover, + }, nil +} + +func (a *recordingShadowReservationAllocator) values() (uint64, uint64, int) { + a.mu.Lock() + defer a.mu.Unlock() + return a.min, a.returned, a.calls +} + +type blockingShadowTimestampAllocator struct { + once sync.Once + started chan struct{} +} + +func (a *blockingShadowTimestampAllocator) ValidateShadowTimestamp(ctx context.Context, _ uint64) (TSOReservation, error) { + a.once.Do(func() { close(a.started) }) + <-ctx.Done() + return TSOReservation{}, ctx.Err() +} + +type recordingTSOEngine struct { + mu sync.Mutex + state raftengine.State + leader raftengine.LeaderInfo + proposeErr error + proposals [][]byte + apply func([]byte) error + afterPropose func([]byte) + term uint64 +} + +func (e *recordingTSOEngine) Propose(_ context.Context, payload []byte) (*raftengine.ProposalResult, error) { + e.mu.Lock() + if e.proposeErr != nil { + e.mu.Unlock() + return nil, e.proposeErr + } + copyPayload := append([]byte(nil), payload...) + e.proposals = append(e.proposals, copyPayload) + if e.apply != nil { + if err := e.apply(copyPayload); err != nil { + e.mu.Unlock() + return nil, err + } + } + afterPropose := e.afterPropose + e.mu.Unlock() + if afterPropose != nil { + afterPropose(copyPayload) + } + return &raftengine.ProposalResult{}, nil +} + +func (e *recordingTSOEngine) ProposeAdmin(ctx context.Context, payload []byte) (*raftengine.ProposalResult, error) { + return e.Propose(ctx, payload) +} + +func (e *recordingTSOEngine) State() raftengine.State { + e.mu.Lock() + defer e.mu.Unlock() + return e.state +} + +func (e *recordingTSOEngine) Leader() raftengine.LeaderInfo { + e.mu.Lock() + defer e.mu.Unlock() + return e.leader +} + +func (e *recordingTSOEngine) VerifyLeader(context.Context) error { + if e.State() != raftengine.StateLeader { + return raftengine.ErrNotLeader + } + return nil +} + +func (e *recordingTSOEngine) LinearizableRead(context.Context) (uint64, error) { return 0, nil } + +func (e *recordingTSOEngine) Status() raftengine.Status { + e.mu.Lock() + defer e.mu.Unlock() + term := e.term + if term == 0 { + term = 1 + } + return raftengine.Status{State: e.state, Term: term} +} + +func (e *recordingTSOEngine) Configuration(context.Context) (raftengine.Configuration, error) { + return raftengine.Configuration{}, nil +} + +func (e *recordingTSOEngine) Close() error { return nil } + +func (e *recordingTSOEngine) setProposeError(err error) { + e.mu.Lock() + e.proposeErr = err + e.mu.Unlock() +} + +func (e *recordingTSOEngine) setLeaderAddress(addr string) { + e.mu.Lock() + e.leader.Address = addr + e.mu.Unlock() +} + +func (e *recordingTSOEngine) setTerm(term uint64) { + e.mu.Lock() + e.term = term + e.mu.Unlock() +} + +func (e *recordingTSOEngine) proposedPayloads() [][]byte { + e.mu.Lock() + defer e.mu.Unlock() + out := make([][]byte, len(e.proposals)) + for i := range e.proposals { + out[i] = append([]byte(nil), e.proposals[i]...) + } + return out +} + +type snapshotBuffer struct { + data []byte +} + +func (b *snapshotBuffer) Write(p []byte) (int, error) { + b.data = append(b.data, p...) + return len(p), nil +} + +func (b *snapshotBuffer) Bytes() []byte { return b.data } + +func applyTSOTestFSM(fsm *TSOStateMachine) func([]byte) error { + return func(payload []byte) error { + if result := fsm.Apply(payload); result != nil { + if err, ok := result.(error); ok { + return err + } + return errors.Newf("unexpected TSO FSM result %T", result) + } + return nil + } +} + +type recordingTSOFloorProvider struct { + mu sync.Mutex + floor uint64 + calls int + afterRead func() +} + +func newTestRaftTSOAllocator(group *ShardGroup, clock *HLC) (*RaftTSOAllocator, error) { + return NewRaftTSOAllocator(group, clock, + WithTSOCutoverFloorProvider(&recordingTSOFloorProvider{})) +} + +func (p *recordingTSOFloorProvider) GlobalCommittedTimestampFloor(context.Context) (uint64, error) { + p.mu.Lock() + p.calls++ + floor := p.floor + afterRead := p.afterRead + p.mu.Unlock() + if afterRead != nil { + afterRead() + } + return floor, nil +} + +func (p *recordingTSOFloorProvider) callCount() int { + p.mu.Lock() + defer p.mu.Unlock() + return p.calls +} diff --git a/main.go b/main.go index 5addfc099..003645269 100644 --- a/main.go +++ b/main.go @@ -124,8 +124,9 @@ var ( raftGroupPeers = flag.String("raftGroupPeers", "", "Semicolon-separated per-group bootstrap members (groupID=raftID@host:port,...)") raftJoinMembers = flag.String("raftJoinMembers", "", "Comma-separated raft members used only for transport discovery while this fresh node joins an existing single-group cluster (raftID=host:port,...); requires --raftJoinAsLearner") raftJoinAsLearner = flag.Bool("raftJoinAsLearner", false, "Local node expects to join an existing cluster as a learner; if a post-apply ConfState lists this node as a voter instead, an ERROR-level alarm fires (the node keeps running -- the flag is an operator alarm, not a consensus veto). See docs/design/2026_04_26_implemented_raft_learner.md §4.5.") - tsoEnabled = flag.Bool("tsoEnabled", false, "Issue coordinator-owned persistence timestamps through the local TSO batch allocator instead of direct HLC calls") - tsoBatchSize = flag.Int("tsoBatchSize", defaultTSOBatchSize, "Timestamp batch size used when --tsoEnabled is true") + tsoEnabled = flag.Bool("tsoEnabled", false, "Commit the one-way cutover marker and issue coordinator-owned persistence timestamps through the dedicated TSO leader when group 0 is configured") + tsoShadowEnabled = flag.Bool("tsoShadowEnabled", false, "Serialize legacy HLC issuance through the dedicated TSO; fail closed on TSO errors and switch to TSO values after cutover") + tsoBatchSize = flag.Int("tsoBatchSize", defaultTSOBatchSize, "Timestamp batch size used by TSO cutover and shadow validation") leaderBalance = flag.Bool("leaderBalance", false, "Enable automatic count-based Raft-group leader balancing on the default-group leader") leaderBalanceInterval = flag.Duration("leaderBalanceInterval", defaultLeaderBalanceInterval, "Interval between leader-balance scheduler evaluations") leaderBalanceGroupCooldown = flag.Duration("leaderBalanceGroupCooldown", defaultLeaderBalanceGroupCooldown, "Minimum time before the scheduler can move the same raft group again") @@ -513,9 +514,11 @@ func run() error { WithKeyVizLabelsEnabled(*keyvizLabelsEnabled). WithAllShardGroups(dataGroupIDs(cfg.groups)...). WithPartitionResolver(buildSQSPartitionResolver(cfg.sqsFifoPartitionMap)) - if err := configureCoordinatorTSO(coordinate); err != nil { + tsoWiring, err := configureCoordinatorTSO(coordinate, shardGroups, shardStore) + if err != nil { return err } + cleanup.Add(tsoWiring.CloseWithLog) // SQS HT-FIFO §8 leadership-refusal: install per-group // observers that step the local node down via @@ -562,6 +565,7 @@ func run() error { cfg.engine, distCatalog, adapter.WithDistributionCoordinator(coordinate), + adapter.WithDistributionTimestampAllocator(tsoWiring.serverAllocator), adapter.WithDistributionActiveTimestampTracker(readTracker), adapter.WithDistributionFilesystemObserver(metricsRegistry.FileSystemObserver()), ) @@ -1503,10 +1507,10 @@ func closeRaftGroupRuntimes(runtimes []*raftGroupRuntime) { // buildDedicatedTSOGroup constructs group 0 with the minimal TSO state machine. // It intentionally does not open an MVCC store: shard routing validation keeps -// user data out of group 0, while the state machine persists only its ceiling -// and allocation floor in Raft snapshots. The ShardGroup still receives the -// normal proposer wrapper so lease renewals participate in raft-envelope -// cutovers exactly like data-group proposals. +// user data out of group 0, while the state machine persists its ceiling, +// allocation floor, and one-way cutover marker in Raft snapshots. The +// ShardGroup still receives the normal proposer wrapper so lease renewals +// participate in raft-envelope cutovers exactly like data-group proposals. func buildDedicatedTSOGroup( raftID string, group groupSpec, @@ -1528,7 +1532,7 @@ func buildDedicatedTSOGroup( if err != nil { return nil, nil, err } - sg := &kv.ShardGroup{Engine: runtime.engine} + sg := &kv.ShardGroup{Engine: runtime.engine, TSOState: sm} sg.Txn = kv.NewLeaderProxyForShardGroup(sg, kv.WithProposalObserver(proposalObserver)) return runtime, sg, nil } @@ -2054,22 +2058,134 @@ func configurationContainsMember(configuration raftengine.Configuration, localID return false } -func configureCoordinatorTSO(coordinate *kv.ShardedCoordinator) error { - if !*tsoEnabled { +type coordinatorTSOWiring struct { + serverAllocator kv.TSOAllocator + routedAllocator *kv.LeaderRoutedTSOAllocator + shadowAllocator *kv.ShadowTimestampAllocator +} + +func (w coordinatorTSOWiring) Close() error { + if w.shadowAllocator != nil { + if err := w.shadowAllocator.Close(); err != nil { + return errors.Wrap(err, "close TSO shadow validator") + } + } + if w.routedAllocator == nil { return nil } - // Group 0 is reserved for TSO state, but data-shard leaders must keep - // issuing timestamps locally until a TSO-leader redirect path exists. - tso, err := kv.NewLocalTSOAllocator(coordinate) + return errors.Wrap(w.routedAllocator.Close(), "close TSO leader routing") +} + +func (w coordinatorTSOWiring) CloseWithLog() { + if err := w.Close(); err != nil { + slog.Warn("failed to close TSO routing connections", slog.Any("err", err)) + } +} + +func configureCoordinatorTSO( + coordinate *kv.ShardedCoordinator, + shardGroups map[uint64]*kv.ShardGroup, + floorProviders ...kv.TSOCutoverFloorProvider, +) (coordinatorTSOWiring, error) { + var wiring coordinatorTSOWiring + if coordinate == nil { + return wiring, errors.Wrap(kv.ErrTSOCoordinatorNil, "configure tso allocator") + } + if *tsoEnabled && *tsoShadowEnabled { + return wiring, errors.New("--tsoEnabled and --tsoShadowEnabled are mutually exclusive") + } + + tsoGroup, dedicated := shardGroups[dedicatedTSORaftGroupID] + if !dedicated { + return configureLegacyCoordinatorTSO(coordinate) + } + var floorProvider kv.TSOCutoverFloorProvider + if len(floorProviders) > 0 { + floorProvider = floorProviders[0] + } + return configureDedicatedCoordinatorTSO(coordinate, tsoGroup, floorProvider) +} + +func configureLegacyCoordinatorTSO(coordinate *kv.ShardedCoordinator) (coordinatorTSOWiring, error) { + var wiring coordinatorTSOWiring + if *tsoShadowEnabled { + return wiring, errors.Wrap(kv.ErrTSOGroupRequired, "configure tso shadow allocator") + } + if !*tsoEnabled { + return wiring, nil + } + // Preserve the pre-group-0 bridge for deployments that enable TSO + // batching before adding the dedicated group. + local, err := kv.NewLocalTSOAllocator(coordinate) if err != nil { - return errors.Wrap(err, "configure tso allocator") + return wiring, errors.Wrap(err, "configure local tso bridge") } - batch, err := kv.NewBatchAllocator(tso, *tsoBatchSize) + batch, err := kv.NewBatchAllocator(local, *tsoBatchSize) if err != nil { - return errors.Wrap(err, "configure tso batch allocator") + return wiring, errors.Wrap(err, "configure local tso bridge batch") } coordinate.WithTSOAllocator(batch) - return nil + return wiring, nil +} + +func configureDedicatedCoordinatorTSO( + coordinate *kv.ShardedCoordinator, + tsoGroup *kv.ShardGroup, + floorProvider kv.TSOCutoverFloorProvider, +) (coordinatorTSOWiring, error) { + var wiring coordinatorTSOWiring + local, err := kv.NewRaftTSOAllocator( + tsoGroup, + coordinate.Clock(), + kv.WithTSOCutoverFloorProvider(floorProvider), + ) + if err != nil { + return wiring, errors.Wrap(err, "configure dedicated tso allocator") + } + wiring.serverAllocator = local + cutoverActive := tsoGroup.TSOState.CutoverActive() + if !*tsoEnabled && !*tsoShadowEnabled && !cutoverActive { + return wiring, nil + } + + routedOpts := []kv.LeaderRoutedTSOAllocatorOption{ + kv.WithTSORoutedClock(coordinate.Clock()), + } + if *tsoEnabled { + routedOpts = append(routedOpts, kv.WithTSOCutoverActivation()) + } + routed, err := kv.NewLeaderRoutedTSOAllocator(local, tsoGroup.Engine, routedOpts...) + if err != nil { + return wiring, errors.Wrap(err, "configure tso leader routing") + } + wiring.routedAllocator = routed + coordinate.WithTimestampGroup(dedicatedTSORaftGroupID) + return installDedicatedCoordinatorTSO(coordinate, wiring, routed, cutoverActive) +} + +func installDedicatedCoordinatorTSO( + coordinate *kv.ShardedCoordinator, + wiring coordinatorTSOWiring, + routed *kv.LeaderRoutedTSOAllocator, + cutoverActive bool, +) (coordinatorTSOWiring, error) { + if *tsoShadowEnabled && !cutoverActive { + shadow, err := kv.NewShadowTimestampAllocator(coordinate.Clock(), routed, slog.Default()) + if err != nil { + _ = wiring.Close() + return coordinatorTSOWiring{}, errors.Wrap(err, "configure tso shadow allocator") + } + wiring.shadowAllocator = shadow + coordinate.WithTSOAllocator(shadow) + return wiring, nil + } + batch, err := kv.NewBatchAllocator(routed, *tsoBatchSize) + if err != nil { + _ = wiring.Close() + return coordinatorTSOWiring{}, errors.Wrap(err, "configure tso batch allocator") + } + coordinate.WithTSOAllocator(batch) + return wiring, nil } type hlcLeaseRenewalBlocker interface { @@ -2215,6 +2331,7 @@ var _ kv.Coordinator = (*startupGatedCoordinator)(nil) var _ kv.LeaseReadableCoordinator = (*startupGatedCoordinator)(nil) var _ kv.AllGroupsLeaseReadableCoordinator = (*startupGatedCoordinator)(nil) var _ kv.GroupRoutableCoordinator = (*startupGatedCoordinator)(nil) +var _ kv.TimestampAllocatorProvider = (*startupGatedCoordinator)(nil) func (c startupGatedCoordinator) Dispatch(ctx context.Context, reqs *kv.OperationGroup[kv.OP]) (*kv.CoordinateResponse, error) { if c.gate != nil && c.gate.blocked() { @@ -2255,6 +2372,11 @@ func (c startupGatedCoordinator) Clock() *kv.HLC { return c.inner.Clock() } +func (c startupGatedCoordinator) TimestampAllocator() kv.TimestampAllocator { + alloc, _ := kv.TimestampAllocatorThrough(c.inner) + return alloc +} + func (c startupGatedCoordinator) LeaseRead(ctx context.Context) (uint64, error) { return kv.LeaseReadThrough(c.inner, ctx) //nolint:wrapcheck // Pass through coordinator errors unchanged. } @@ -2659,7 +2781,7 @@ func startRaftServers( } func internalTimestampOptions(coordinate kv.Coordinator) []adapter.InternalOption { - if alloc, ok := coordinate.(kv.TimestampAllocator); ok { + if alloc, ok := kv.TimestampAllocatorThrough(coordinate); ok { return []adapter.InternalOption{adapter.WithInternalTimestampAllocator(alloc)} } return nil diff --git a/main_tso_routing_test.go b/main_tso_routing_test.go new file mode 100644 index 000000000..7c69c9340 --- /dev/null +++ b/main_tso_routing_test.go @@ -0,0 +1,245 @@ +package main + +import ( + "context" + "fmt" + "sync/atomic" + "testing" + "time" + + "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/internal/raftengine" + "github.com/bootjp/elastickv/kv" + "github.com/stretchr/testify/require" +) + +func TestConfigureCoordinatorTSORejectsConflictingModes(t *testing.T) { + setTSOModeFlags(t, true, true) + coord := newMainTSOCoordinator(kv.NewHLC(), nil) + + _, err := configureCoordinatorTSO(coord, nil) + require.ErrorContains(t, err, "mutually exclusive") +} + +func TestConfigureCoordinatorTSOShadowRequiresDedicatedGroup(t *testing.T) { + setTSOModeFlags(t, false, true) + coord := newMainTSOCoordinator(kv.NewHLC(), nil) + + _, err := configureCoordinatorTSO(coord, nil) + require.ErrorIs(t, err, kv.ErrTSOGroupRequired) +} + +func TestConfigureCoordinatorTSOCutoverRoutesThroughDedicatedGroup(t *testing.T) { + setTSOModeFlags(t, true, false) + clock := kv.NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(time.Minute).UnixMilli()) + fsm := kv.NewTSOStateMachine(clock) + engine := &mainTSOEngine{state: raftengine.StateLeader, tsoState: fsm} + groups := map[uint64]*kv.ShardGroup{ + dedicatedTSORaftGroupID: {Engine: engine, TSOState: fsm}, + } + coord := newMainTSOCoordinator(clock, groups) + + wiring, err := configureCoordinatorTSO(coord, groups, mainTSOFloorProvider{}) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, wiring.Close()) }) + require.NotNil(t, wiring.serverAllocator) + require.NotNil(t, wiring.routedAllocator) + require.True(t, coord.IsTimestampLeader()) + + ts, err := kv.NextTimestampThrough(context.Background(), coord, "test dedicated tso") + require.NoError(t, err) + require.NotZero(t, ts) + require.Equal(t, uint64(2), engine.proposals.Load(), "cutover marker must commit before the first window") + require.True(t, fsm.CutoverActive()) +} + +func TestConfigureCoordinatorTSOShadowReturnsLegacyTimestamp(t *testing.T) { + setTSOModeFlags(t, false, true) + clock := kv.NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(time.Minute).UnixMilli()) + fsm := kv.NewTSOStateMachine(clock) + engine := &mainTSOEngine{state: raftengine.StateLeader, tsoState: fsm} + groups := map[uint64]*kv.ShardGroup{ + dedicatedTSORaftGroupID: {Engine: engine, TSOState: fsm}, + } + coord := newMainTSOCoordinator(clock, groups) + + wiring, err := configureCoordinatorTSO(coord, groups, mainTSOFloorProvider{}) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, wiring.Close()) }) + + legacyTS, err := kv.NextTimestampThrough(context.Background(), coord, "test shadow tso") + require.NoError(t, err) + require.NotZero(t, legacyTS) + require.Greater(t, clock.Current(), legacyTS) + require.Equal(t, uint64(1), engine.proposals.Load()) +} + +func TestConfigureCoordinatorTSOExposesDedicatedServerWithoutCutover(t *testing.T) { + setTSOModeFlags(t, false, false) + clock := kv.NewHLC() + fsm := kv.NewTSOStateMachine(clock) + engine := &mainTSOEngine{state: raftengine.StateFollower, tsoState: fsm} + groups := map[uint64]*kv.ShardGroup{ + dedicatedTSORaftGroupID: {Engine: engine, TSOState: fsm}, + } + coord := newMainTSOCoordinator(clock, groups) + + wiring, err := configureCoordinatorTSO(coord, groups, mainTSOFloorProvider{}) + require.NoError(t, err) + require.NotNil(t, wiring.serverAllocator) + require.Nil(t, wiring.routedAllocator) + require.False(t, coord.IsTimestampLeader(), "timestamp leadership stays on the compatibility bridge until a mode is enabled") +} + +func TestConfigureCoordinatorTSORestoresDurableCutoverWithoutFlags(t *testing.T) { + setTSOModeFlags(t, false, false) + clock, fsm, engine, groups := newActiveMainTSOCutover(t) + coord := newMainTSOCoordinator(clock, groups) + + wiring, err := configureCoordinatorTSO(coord, groups, mainTSOFloorProvider{}) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, wiring.Close()) }) + require.NotNil(t, wiring.serverAllocator) + require.NotNil(t, wiring.routedAllocator) + require.Nil(t, wiring.shadowAllocator) + require.True(t, coord.IsTimestampLeader()) + require.True(t, fsm.CutoverActive()) + + ts, err := kv.NextTimestampThrough(context.Background(), coord, "test restored tso cutover") + require.NoError(t, err) + require.NotZero(t, ts) + require.Equal(t, uint64(1), engine.proposals.Load(), "restored cutover must only commit the allocation floor") +} + +func TestConfigureCoordinatorTSODurableCutoverOverridesShadowFlag(t *testing.T) { + setTSOModeFlags(t, false, true) + clock, fsm, engine, groups := newActiveMainTSOCutover(t) + coord := newMainTSOCoordinator(clock, groups) + + wiring, err := configureCoordinatorTSO(coord, groups, mainTSOFloorProvider{}) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, wiring.Close()) }) + require.NotNil(t, wiring.routedAllocator) + require.Nil(t, wiring.shadowAllocator, "durable cutover cannot return to legacy shadow issuance") + require.True(t, coord.IsTimestampLeader()) + require.True(t, fsm.CutoverActive()) + + ts, err := kv.NextTimestampThrough(context.Background(), coord, "test shadow flag after cutover") + require.NoError(t, err) + require.NotZero(t, ts) + require.Equal(t, uint64(1), engine.proposals.Load(), "restored cutover must only commit the allocation floor") +} + +func TestInternalTimestampOptionsPreservesTSOThroughStartupGate(t *testing.T) { + t.Parallel() + allocator := &mainTimestampAllocator{next: 123} + coord := newMainTSOCoordinator(kv.NewHLC(), nil).WithTSOAllocator(allocator) + gated := startupGatedCoordinator{inner: coord, gate: &startupPublicKVGate{}} + + got, ok := kv.TimestampAllocatorThrough(gated) + require.True(t, ok) + require.Same(t, allocator, got) + require.Len(t, internalTimestampOptions(gated), 1) + + legacy := startupGatedCoordinator{inner: newMainTSOCoordinator(kv.NewHLC(), nil)} + require.Empty(t, internalTimestampOptions(legacy)) +} + +func newActiveMainTSOCutover(t *testing.T) (*kv.HLC, *kv.TSOStateMachine, *mainTSOEngine, map[uint64]*kv.ShardGroup) { + t.Helper() + clock := kv.NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(time.Minute).UnixMilli()) + fsm := kv.NewTSOStateMachine(clock) + engine := &mainTSOEngine{state: raftengine.StateLeader, tsoState: fsm} + group := &kv.ShardGroup{Engine: engine, TSOState: fsm} + allocator, err := kv.NewRaftTSOAllocator(group, clock, kv.WithTSOCutoverFloorProvider(mainTSOFloorProvider{})) + require.NoError(t, err) + _, err = allocator.ReserveBatchAfter(context.Background(), 1, 0, true) + require.NoError(t, err) + require.True(t, fsm.CutoverActive()) + engine.proposals.Store(0) + return clock, fsm, engine, map[uint64]*kv.ShardGroup{dedicatedTSORaftGroupID: group} +} + +func setTSOModeFlags(t *testing.T, enabled, shadow bool) { + t.Helper() + oldEnabled := *tsoEnabled + oldShadow := *tsoShadowEnabled + oldBatchSize := *tsoBatchSize + *tsoEnabled = enabled + *tsoShadowEnabled = shadow + *tsoBatchSize = 8 + t.Cleanup(func() { + *tsoEnabled = oldEnabled + *tsoShadowEnabled = oldShadow + *tsoBatchSize = oldBatchSize + }) +} + +func newMainTSOCoordinator(clock *kv.HLC, groups map[uint64]*kv.ShardGroup) *kv.ShardedCoordinator { + return kv.NewShardedCoordinator(distribution.NewEngine(), groups, 1, clock, nil) +} + +type mainTSOEngine struct { + state raftengine.State + proposals atomic.Uint64 + tsoState *kv.TSOStateMachine +} + +type mainTSOFloorProvider struct{} + +type mainTimestampAllocator struct { + next uint64 +} + +func (a *mainTimestampAllocator) Next(context.Context) (uint64, error) { + return a.next, nil +} + +func (mainTSOFloorProvider) GlobalCommittedTimestampFloor(context.Context) (uint64, error) { + return 0, nil +} + +func (e *mainTSOEngine) Propose(_ context.Context, payload []byte) (*raftengine.ProposalResult, error) { + e.proposals.Add(1) + if e.tsoState != nil { + if result := e.tsoState.Apply(payload); result != nil { + if err, ok := result.(error); ok { + return nil, err + } + return nil, fmt.Errorf("unexpected TSO apply result %T", result) + } + } + return &raftengine.ProposalResult{}, nil +} + +func (e *mainTSOEngine) ProposeAdmin(ctx context.Context, payload []byte) (*raftengine.ProposalResult, error) { + return e.Propose(ctx, payload) +} + +func (e *mainTSOEngine) State() raftengine.State { return e.state } + +func (e *mainTSOEngine) Leader() raftengine.LeaderInfo { + return raftengine.LeaderInfo{ID: "self", Address: "127.0.0.1:50051"} +} + +func (e *mainTSOEngine) VerifyLeader(context.Context) error { + if e.state != raftengine.StateLeader { + return raftengine.ErrNotLeader + } + return nil +} + +func (e *mainTSOEngine) LinearizableRead(context.Context) (uint64, error) { return 0, nil } + +func (e *mainTSOEngine) Status() raftengine.Status { + return raftengine.Status{State: e.state, Term: 1} +} + +func (e *mainTSOEngine) Configuration(context.Context) (raftengine.Configuration, error) { + return raftengine.Configuration{}, nil +} + +func (e *mainTSOEngine) Close() error { return nil } diff --git a/proto/distribution.pb.go b/proto/distribution.pb.go index b6847386e..99aaa98b2 100644 --- a/proto/distribution.pb.go +++ b/proto/distribution.pb.go @@ -357,9 +357,19 @@ func (x *GetRouteResponse) GetRaftGroupId() uint64 { } type GetTimestampRequest struct { - state protoimpl.MessageState `protogen:"open.v1"` - unknownFields protoimpl.UnknownFields - sizeCache protoimpl.SizeCache + state protoimpl.MessageState `protogen:"open.v1"` + // Count requests a consecutive timestamp window and defaults to one. + Count uint32 `protobuf:"varint,1,opt,name=count,proto3" json:"count,omitempty"` + // MinTimestamp requires the returned window base to be strictly greater + // than this value. It lets commit timestamp callers preserve startTS < + // commitTS across a follower-to-TSO-leader RPC. + MinTimestamp uint64 `protobuf:"varint,2,opt,name=min_timestamp,json=minTimestamp,proto3" json:"min_timestamp,omitempty"` + // ActivateCutover durably switches the dedicated TSO group into production + // issuance. Once committed, shadow clients return TSO timestamps instead of + // legacy HLC candidates, so a rolling cutover cannot reintroduce them. + ActivateCutover bool `protobuf:"varint,3,opt,name=activate_cutover,json=activateCutover,proto3" json:"activate_cutover,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *GetTimestampRequest) Reset() { @@ -392,9 +402,42 @@ func (*GetTimestampRequest) Descriptor() ([]byte, []int) { return file_distribution_proto_rawDescGZIP(), []int{2} } +func (x *GetTimestampRequest) GetCount() uint32 { + if x != nil { + return x.Count + } + return 0 +} + +func (x *GetTimestampRequest) GetMinTimestamp() uint64 { + if x != nil { + return x.MinTimestamp + } + return 0 +} + +func (x *GetTimestampRequest) GetActivateCutover() bool { + if x != nil { + return x.ActivateCutover + } + return false +} + type GetTimestampResponse struct { - state protoimpl.MessageState `protogen:"open.v1"` - Timestamp uint64 `protobuf:"varint,1,opt,name=timestamp,proto3" json:"timestamp,omitempty"` + state protoimpl.MessageState `protogen:"open.v1"` + Timestamp uint64 `protobuf:"varint,1,opt,name=timestamp,proto3" json:"timestamp,omitempty"` + // CommittedByDedicatedTSO distinguishes a durable TSO window from the + // legacy catalog timestamp response during rolling upgrades. + CommittedByDedicatedTso bool `protobuf:"varint,2,opt,name=committed_by_dedicated_tso,json=committedByDedicatedTso,proto3" json:"committed_by_dedicated_tso,omitempty"` + // Count echoes the number of consecutive timestamps reserved at timestamp. + Count uint32 `protobuf:"varint,3,opt,name=count,proto3" json:"count,omitempty"` + // PreviousAllocationFloor is the highest timestamp already reserved before + // this request. Shadow validation uses it to reject legacy candidates that + // overlap any earlier dedicated or shadow reservation. + PreviousAllocationFloor uint64 `protobuf:"varint,4,opt,name=previous_allocation_floor,json=previousAllocationFloor,proto3" json:"previous_allocation_floor,omitempty"` + // CutoverActive reports the durable group-0 cutover marker after applying + // this request. Shadow nodes switch to the returned TSO timestamp when true. + CutoverActive bool `protobuf:"varint,5,opt,name=cutover_active,json=cutoverActive,proto3" json:"cutover_active,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -436,6 +479,34 @@ func (x *GetTimestampResponse) GetTimestamp() uint64 { return 0 } +func (x *GetTimestampResponse) GetCommittedByDedicatedTso() bool { + if x != nil { + return x.CommittedByDedicatedTso + } + return false +} + +func (x *GetTimestampResponse) GetCount() uint32 { + if x != nil { + return x.Count + } + return 0 +} + +func (x *GetTimestampResponse) GetPreviousAllocationFloor() uint64 { + if x != nil { + return x.PreviousAllocationFloor + } + return 0 +} + +func (x *GetTimestampResponse) GetCutoverActive() bool { + if x != nil { + return x.CutoverActive + } + return false +} + type RouteDescriptor struct { state protoimpl.MessageState `protogen:"open.v1"` RouteId uint64 `protobuf:"varint,1,opt,name=route_id,json=routeId,proto3" json:"route_id,omitempty"` @@ -1146,10 +1217,17 @@ const file_distribution_proto_rawDesc = "" + "\x10GetRouteResponse\x12\x14\n" + "\x05start\x18\x01 \x01(\fR\x05start\x12\x10\n" + "\x03end\x18\x02 \x01(\fR\x03end\x12\"\n" + - "\rraft_group_id\x18\x03 \x01(\x04R\vraftGroupId\"\x15\n" + - "\x13GetTimestampRequest\"4\n" + + "\rraft_group_id\x18\x03 \x01(\x04R\vraftGroupId\"{\n" + + "\x13GetTimestampRequest\x12\x14\n" + + "\x05count\x18\x01 \x01(\rR\x05count\x12#\n" + + "\rmin_timestamp\x18\x02 \x01(\x04R\fminTimestamp\x12)\n" + + "\x10activate_cutover\x18\x03 \x01(\bR\x0factivateCutover\"\xea\x01\n" + "\x14GetTimestampResponse\x12\x1c\n" + - "\ttimestamp\x18\x01 \x01(\x04R\ttimestamp\"\xe5\x01\n" + + "\ttimestamp\x18\x01 \x01(\x04R\ttimestamp\x12;\n" + + "\x1acommitted_by_dedicated_tso\x18\x02 \x01(\bR\x17committedByDedicatedTso\x12\x14\n" + + "\x05count\x18\x03 \x01(\rR\x05count\x12:\n" + + "\x19previous_allocation_floor\x18\x04 \x01(\x04R\x17previousAllocationFloor\x12%\n" + + "\x0ecutover_active\x18\x05 \x01(\bR\rcutoverActive\"\xe5\x01\n" + "\x0fRouteDescriptor\x12\x19\n" + "\broute_id\x18\x01 \x01(\x04R\arouteId\x12\x14\n" + "\x05start\x18\x02 \x01(\fR\x05start\x12\x10\n" + diff --git a/proto/distribution.proto b/proto/distribution.proto index f2199d909..493bfa274 100644 --- a/proto/distribution.proto +++ b/proto/distribution.proto @@ -21,10 +21,33 @@ message GetRouteResponse { uint64 raft_group_id = 3; } -message GetTimestampRequest {} +message GetTimestampRequest { + // Count requests a consecutive timestamp window and defaults to one. + uint32 count = 1; + // MinTimestamp requires the returned window base to be strictly greater + // than this value. It lets commit timestamp callers preserve startTS < + // commitTS across a follower-to-TSO-leader RPC. + uint64 min_timestamp = 2; + // ActivateCutover durably switches the dedicated TSO group into production + // issuance. Once committed, shadow clients return TSO timestamps instead of + // legacy HLC candidates, so a rolling cutover cannot reintroduce them. + bool activate_cutover = 3; +} message GetTimestampResponse { uint64 timestamp = 1; + // CommittedByDedicatedTSO distinguishes a durable TSO window from the + // legacy catalog timestamp response during rolling upgrades. + bool committed_by_dedicated_tso = 2; + // Count echoes the number of consecutive timestamps reserved at timestamp. + uint32 count = 3; + // PreviousAllocationFloor is the highest timestamp already reserved before + // this request. Shadow validation uses it to reject legacy candidates that + // overlap any earlier dedicated or shadow reservation. + uint64 previous_allocation_floor = 4; + // CutoverActive reports the durable group-0 cutover marker after applying + // this request. Shadow nodes switch to the returned TSO timestamp when true. + bool cutover_active = 5; } enum RouteState { diff --git a/proto/service.pb.go b/proto/service.pb.go index 00000ae2f..a70a3434b 100644 --- a/proto/service.pb.go +++ b/proto/service.pb.go @@ -399,6 +399,7 @@ func (x *RawDeleteResponse) GetSuccess() bool { type RawLatestCommitTSRequest struct { state protoimpl.MessageState `protogen:"open.v1"` Key []byte `protobuf:"bytes,1,opt,name=key,proto3" json:"key,omitempty"` + GroupId uint64 `protobuf:"varint,2,opt,name=group_id,json=groupId,proto3" json:"group_id,omitempty"` // optional explicit group for leader-fenced watermark reads unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -440,10 +441,19 @@ func (x *RawLatestCommitTSRequest) GetKey() []byte { return nil } +func (x *RawLatestCommitTSRequest) GetGroupId() uint64 { + if x != nil { + return x.GroupId + } + return 0 +} + type RawLatestCommitTSResponse struct { state protoimpl.MessageState `protogen:"open.v1"` Ts uint64 `protobuf:"varint,1,opt,name=ts,proto3" json:"ts,omitempty"` Exists bool `protobuf:"varint,2,opt,name=exists,proto3" json:"exists,omitempty"` + GroupId uint64 `protobuf:"varint,3,opt,name=group_id,json=groupId,proto3" json:"group_id,omitempty"` // echoes an explicit group watermark request + LeaderFenced bool `protobuf:"varint,4,opt,name=leader_fenced,json=leaderFenced,proto3" json:"leader_fenced,omitempty"` // true only after that group's leader ReadIndex fence unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -492,6 +502,20 @@ func (x *RawLatestCommitTSResponse) GetExists() bool { return false } +func (x *RawLatestCommitTSResponse) GetGroupId() uint64 { + if x != nil { + return x.GroupId + } + return 0 +} + +func (x *RawLatestCommitTSResponse) GetLeaderFenced() bool { + if x != nil { + return x.LeaderFenced + } + return false +} + type RawScanAtRequest struct { state protoimpl.MessageState `protogen:"open.v1"` StartKey []byte `protobuf:"bytes,1,opt,name=start_key,json=startKey,proto3" json:"start_key,omitempty"` @@ -2282,12 +2306,15 @@ const file_service_proto_rawDesc = "" + "\x03key\x18\x01 \x01(\fR\x03key\"P\n" + "\x11RawDeleteResponse\x12!\n" + "\fcommit_index\x18\x01 \x01(\x04R\vcommitIndex\x12\x18\n" + - "\asuccess\x18\x02 \x01(\bR\asuccess\",\n" + + "\asuccess\x18\x02 \x01(\bR\asuccess\"G\n" + "\x18RawLatestCommitTSRequest\x12\x10\n" + - "\x03key\x18\x01 \x01(\fR\x03key\"C\n" + + "\x03key\x18\x01 \x01(\fR\x03key\x12\x19\n" + + "\bgroup_id\x18\x02 \x01(\x04R\agroupId\"\x83\x01\n" + "\x19RawLatestCommitTSResponse\x12\x0e\n" + "\x02ts\x18\x01 \x01(\x04R\x02ts\x12\x16\n" + - "\x06exists\x18\x02 \x01(\bR\x06exists\"\xc0\x01\n" + + "\x06exists\x18\x02 \x01(\bR\x06exists\x12\x19\n" + + "\bgroup_id\x18\x03 \x01(\x04R\agroupId\x12#\n" + + "\rleader_fenced\x18\x04 \x01(\bR\fleaderFenced\"\xc0\x01\n" + "\x10RawScanAtRequest\x12\x1b\n" + "\tstart_key\x18\x01 \x01(\fR\bstartKey\x12\x17\n" + "\aend_key\x18\x02 \x01(\fR\x06endKey\x12\x14\n" + diff --git a/proto/service.proto b/proto/service.proto index 68d12d5f7..6b01e6012 100644 --- a/proto/service.proto +++ b/proto/service.proto @@ -64,11 +64,14 @@ message RawDeleteResponse { message RawLatestCommitTSRequest { bytes key = 1; + uint64 group_id = 2; // optional explicit group for leader-fenced watermark reads } message RawLatestCommitTSResponse { uint64 ts = 1; bool exists = 2; + uint64 group_id = 3; // echoes an explicit group watermark request + bool leader_fenced = 4; // true only after that group's leader ReadIndex fence } message RawScanAtRequest { From 07460ec650066207d917cc6bc4b780af7690e025 Mon Sep 17 00:00:00 2001 From: bootjp Date: Sun, 19 Jul 2026 21:37:30 +0900 Subject: [PATCH 03/14] tso: complete centralized runtime semantics --- adapter/distribution_server.go | 129 ++++++-- adapter/distribution_server_test.go | 180 ++++++++++- adapter/dynamodb_item_write.go | 12 +- adapter/dynamodb_locks.go | 5 + adapter/dynamodb_migration.go | 37 ++- adapter/dynamodb_schema.go | 15 +- adapter/dynamodb_transact.go | 6 +- adapter/redis_delta_compactor.go | 20 +- adapter/redis_expire_cmds.go | 50 +-- adapter/redis_hash_cmds.go | 20 +- adapter/redis_lists.go | 68 +++-- adapter/redis_lua_context.go | 6 +- adapter/redis_set_cmds.go | 97 +++--- adapter/redis_stream_cmds.go | 25 +- adapter/redis_strings.go | 10 +- adapter/redis_txn.go | 60 ++-- adapter/redis_zset_cmds.go | 65 ++-- adapter/s3.go | 49 ++- adapter/s3_admin.go | 12 +- adapter/s3_admin_objects.go | 8 +- adapter/s3_hlc_fence_test.go | 77 +++++ adapter/s3_multipart_complete.go | 9 +- adapter/s3_put_object.go | 4 +- adapter/s3_upload_part.go | 19 +- adapter/sqs_catalog.go | 36 ++- adapter/sqs_messages.go | 29 +- adapter/sqs_messages_batch.go | 12 +- adapter/sqs_purge.go | 6 +- adapter/sqs_reaper.go | 30 +- adapter/sqs_redrive.go | 2 +- adapter/sqs_tags.go | 6 +- distribution/catalog.go | 10 + distribution/catalog_test.go | 28 ++ .../2026_04_16_partial_centralized_tso.md | 103 +++++-- kv/coordinator.go | 16 + kv/keyviz_label.go | 13 + kv/lease_warmup_test.go | 28 ++ kv/sharded_coordinator.go | 165 +++++++++- kv/sharded_coordinator_txn_test.go | 213 +++++++++++++ kv/tso.go | 247 ++++++++++++++- kv/tso_fsm.go | 172 +++++++++-- kv/tso_fsm_test.go | 91 +++++- kv/tso_raft.go | 287 ++++++++++++++++-- kv/tso_raft_test.go | 170 ++++++++++- kv/tso_test.go | 142 +++++++++ main.go | 59 +++- main_tso_routing_test.go | 41 ++- proto/distribution.pb.go | 269 ++++++++++++---- proto/distribution.proto | 19 ++ proto/distribution_grpc.pb.go | 46 ++- 50 files changed, 2833 insertions(+), 390 deletions(-) diff --git a/adapter/distribution_server.go b/adapter/distribution_server.go index f3f8bd851..2e8ed35e6 100644 --- a/adapter/distribution_server.go +++ b/adapter/distribution_server.go @@ -147,18 +147,15 @@ func (s *DistributionServer) GetTimestamp(ctx context.Context, req *pb.GetTimest if err != nil { return nil, err } - activateCutover := req != nil && req.GetActivateCutover() + activateCutover, activatePhaseD, err := timestampActivationValues(req) + if err != nil { + return nil, err + } if s.timestampAllocator == nil { - if count != 1 || minTimestamp != 0 || activateCutover { - return nil, errors.WithStack(status.Error(codes.FailedPrecondition, "dedicated TSO allocator is not configured")) - } - if s.engine == nil { - return nil, errors.WithStack(status.Error(codes.Unavailable, "distribution engine is not configured")) - } - return &pb.GetTimestampResponse{Timestamp: s.engine.NextTimestamp()}, nil + return s.legacyTimestampResponse(count, minTimestamp, activateCutover, activatePhaseD) } - reservation, err := s.allocateTimestampReservation(ctx, count, minTimestamp, activateCutover) + reservation, err := s.allocateTimestampReservation(ctx, count, minTimestamp, activateCutover, activatePhaseD) if err != nil { return nil, timestampRPCError(err) } @@ -171,22 +168,79 @@ func (s *DistributionServer) GetTimestamp(ctx context.Context, req *pb.GetTimest Count: uint32(count), //nolint:gosec // count is bounded by MaxTSOBatchSize. PreviousAllocationFloor: reservation.PreviousAllocationFloor, CutoverActive: reservation.CutoverActive, + PhaseDActive: reservation.PhaseDActive, + PhaseDFloor: reservation.PhaseDFloor, }, nil } +func timestampActivationValues(req *pb.GetTimestampRequest) (bool, bool, error) { + activateCutover := req != nil && req.GetActivateCutover() + activatePhaseD := req != nil && req.GetActivatePhaseD() + if activatePhaseD && !activateCutover { + return false, false, errors.WithStack(status.Error(codes.InvalidArgument, + "phase D activation requires cutover activation")) + } + return activateCutover, activatePhaseD, nil +} + +func (s *DistributionServer) legacyTimestampResponse( + count int, + minTimestamp uint64, + activateCutover bool, + activatePhaseD bool, +) (*pb.GetTimestampResponse, error) { + if count != 1 || minTimestamp != 0 || activateCutover || activatePhaseD { + return nil, errors.WithStack(status.Error(codes.FailedPrecondition, "dedicated TSO allocator is not configured")) + } + if s.engine == nil { + return nil, errors.WithStack(status.Error(codes.Unavailable, "distribution engine is not configured")) + } + return &pb.GetTimestampResponse{Timestamp: s.engine.NextTimestamp()}, nil +} + +// ValidateTimestamp verifies a read/start timestamp against the durable M7 +// allocation range. The local allocator performs the group-0 leader fence; +// followers reject so clients re-resolve rather than validating stale state. +func (s *DistributionServer) ValidateTimestamp( + ctx context.Context, + req *pb.ValidateTimestampRequest, +) (*pb.ValidateTimestampResponse, error) { + if req == nil || req.GetTimestamp() == 0 { + return nil, errors.WithStack(status.Error(codes.InvalidArgument, "timestamp is required")) + } + validator, ok := s.timestampAllocator.(kv.DurableTimestampValidator) + if !ok { + return nil, errors.WithStack(status.Error(codes.FailedPrecondition, + "dedicated TSO timestamp validation is not configured")) + } + if err := validator.ValidateDurableTimestamp(ctx, req.GetTimestamp()); err != nil { + return nil, timestampRPCError(err) + } + resp := &pb.ValidateTimestampResponse{Valid: true, PhaseDActive: true} + if state, ok := s.timestampAllocator.(interface { + PhaseDFloor() uint64 + AllocationFloor() uint64 + }); ok { + resp.PhaseDFloor = state.PhaseDFloor() + resp.AllocationFloor = state.AllocationFloor() + } + return resp, nil +} + func (s *DistributionServer) allocateTimestampReservation( ctx context.Context, count int, minTimestamp uint64, activateCutover bool, + activatePhaseD bool, ) (kv.TSOReservation, error) { if allocator, ok := s.timestampAllocator.(kv.TSOReservationAllocator); ok { - reservation, err := allocator.ReserveBatchAfter(ctx, count, minTimestamp, activateCutover) + reservation, err := allocator.ReserveBatchAfter(ctx, count, minTimestamp, activateCutover, activatePhaseD) return reservation, errors.Wrap(err, "reserve dedicated TSO window") } reservation := kv.TSOReservation{Count: count} switch { - case activateCutover: + case activateCutover || activatePhaseD: return reservation, errors.WithStack(status.Error(codes.FailedPrecondition, "TSO allocator does not support durable cutover")) case minTimestamp > 0: @@ -251,6 +305,12 @@ func timestampRPCError(err error) error { return errors.WithStack(status.Error(codes.InvalidArgument, err.Error())) case errors.Is(err, kv.ErrTSONotLeader), errors.Is(err, kv.ErrLeaderNotFound): return errors.WithStack(status.Error(codes.FailedPrecondition, err.Error())) + case errors.Is(err, kv.ErrTSOTimestampPrePhaseD): + return errors.WithStack(status.Error(codes.OutOfRange, err.Error())) + case errors.Is(err, kv.ErrTSOTimestampInvalid): + return errors.WithStack(status.Error(codes.InvalidArgument, err.Error())) + case errors.Is(err, kv.ErrTSOPhaseDInactive): + return errors.WithStack(status.Error(codes.FailedPrecondition, err.Error())) default: return errors.WithStack(status.Error(codes.Unavailable, err.Error())) } @@ -280,7 +340,16 @@ func (s *DistributionServer) SplitRange(ctx context.Context, req *pb.SplitRangeR return nil, err } - snapshot, err := s.loadCatalogSnapshot(ctx) + readTimestamp, err := kv.BeginReadTimestampThrough( + ctx, + s.coordinator, + s.catalog.LatestCommitTS(), + "distribution split range: begin read timestamp", + ) + if err != nil { + return nil, grpcStatusErrorf(codes.Internal, "begin catalog split snapshot: %v", err) + } + snapshot, err := s.loadCatalogSnapshotAt(ctx, readTimestamp.Timestamp()) if err != nil { return nil, err } @@ -290,15 +359,8 @@ func (s *DistributionServer) SplitRange(ctx context.Context, req *pb.SplitRangeR return nil, err } - parent, found := findRouteByID(snapshot.Routes, req.GetRouteId()) - if !found { - return nil, grpcStatusError(codes.NotFound, errDistributionUnknownRoute.Error()) - } - - rawSplitKey := req.GetSplitKey() - splitKey := distribution.CloneBytes(fskeys.NormalizeSplitBoundary(kv.RouteKey(rawSplitKey))) - if err := validateSplitKey(parent, splitKey); err != nil { - s.observeFilePinnedHotspotIfNeeded(rawSplitKey, splitKey, err) + parent, splitKey, err := s.prepareSplitRange(snapshot.Routes, req) + if err != nil { return nil, err } @@ -327,6 +389,20 @@ func (s *DistributionServer) SplitRange(ctx context.Context, req *pb.SplitRangeR }, nil } +func (s *DistributionServer) prepareSplitRange(routes []distribution.RouteDescriptor, req *pb.SplitRangeRequest) (distribution.RouteDescriptor, []byte, error) { + parent, found := findRouteByID(routes, req.GetRouteId()) + if !found { + return distribution.RouteDescriptor{}, nil, grpcStatusError(codes.NotFound, errDistributionUnknownRoute.Error()) + } + rawSplitKey := req.GetSplitKey() + splitKey := distribution.CloneBytes(fskeys.NormalizeSplitBoundary(kv.RouteKey(rawSplitKey))) + if err := validateSplitKey(parent, splitKey); err != nil { + s.observeFilePinnedHotspotIfNeeded(rawSplitKey, splitKey, err) + return distribution.RouteDescriptor{}, nil, err + } + return parent, splitKey, nil +} + func (s *DistributionServer) pinReadTS(ts uint64) *kv.ActiveTimestampToken { if s == nil || s.readTracker == nil { return nil @@ -479,6 +555,17 @@ func (s *DistributionServer) loadCatalogSnapshot(ctx context.Context) (distribut return snapshot, nil } +func (s *DistributionServer) loadCatalogSnapshotAt(ctx context.Context, readTS uint64) (distribution.CatalogSnapshot, error) { + if s.catalog == nil { + return distribution.CatalogSnapshot{}, grpcStatusError(codes.FailedPrecondition, errDistributionCatalogNotConfigured.Error()) + } + snapshot, err := s.catalog.SnapshotAt(ctx, readTS) + if err != nil { + return distribution.CatalogSnapshot{}, grpcStatusErrorf(codes.Internal, "load route catalog: %v", err) + } + return snapshot, nil +} + func (s *DistributionServer) loadCatalogSnapshotAtVersion( ctx context.Context, readTS uint64, diff --git a/adapter/distribution_server_test.go b/adapter/distribution_server_test.go index 2ab5323ac..3aec7ea42 100644 --- a/adapter/distribution_server_test.go +++ b/adapter/distribution_server_test.go @@ -116,6 +116,89 @@ func TestDistributionServerGetTimestamp_CommitsCutoverAndReturnsPriorFloor(t *te require.Equal(t, uint64(499), resp.GetPreviousAllocationFloor()) } +func TestDistributionServerGetTimestamp_CommitsPhaseDAndReturnsFloor(t *testing.T) { + t.Parallel() + + alloc := &distributionTSOAllocator{base: 501, previousFloor: 499, leader: true} + s := NewDistributionServer( + distribution.NewEngine(), + nil, + WithDistributionTimestampAllocator(alloc), + ) + resp, err := s.GetTimestamp(context.Background(), &pb.GetTimestampRequest{ + Count: 1, + MinTimestamp: 500, + ActivateCutover: true, + ActivatePhaseD: true, + }) + require.NoError(t, err) + require.True(t, alloc.activate) + require.True(t, alloc.activatePhaseD) + require.True(t, resp.GetCutoverActive()) + require.True(t, resp.GetPhaseDActive()) + require.Equal(t, uint64(499), resp.GetPhaseDFloor()) +} + +func TestDistributionServerGetTimestamp_RejectsPhaseDWithoutCutover(t *testing.T) { + t.Parallel() + + s := NewDistributionServer( + distribution.NewEngine(), + nil, + WithDistributionTimestampAllocator(&distributionTSOAllocator{leader: true}), + ) + _, err := s.GetTimestamp(context.Background(), &pb.GetTimestampRequest{ActivatePhaseD: true}) + require.Equal(t, codes.InvalidArgument, status.Code(err)) +} + +func TestDistributionServerValidateTimestamp(t *testing.T) { + t.Parallel() + + alloc := &distributionTSOAllocator{ + base: 501, + leader: true, + phaseD: true, + phaseDFloor: 499, + } + s := NewDistributionServer( + distribution.NewEngine(), + nil, + WithDistributionTimestampAllocator(alloc), + ) + resp, err := s.ValidateTimestamp(context.Background(), &pb.ValidateTimestampRequest{Timestamp: 500}) + require.NoError(t, err) + require.True(t, resp.GetValid()) + require.True(t, resp.GetPhaseDActive()) + require.Equal(t, uint64(499), resp.GetPhaseDFloor()) + require.Equal(t, uint64(501), resp.GetAllocationFloor()) + + _, err = s.ValidateTimestamp(context.Background(), &pb.ValidateTimestampRequest{Timestamp: 499}) + require.Equal(t, codes.InvalidArgument, status.Code(err)) +} + +func TestDistributionServerValidateTimestamp_LeaderRoutedRPC(t *testing.T) { + serverAlloc := &distributionTSOAllocator{ + base: 701, + leader: true, + phaseD: true, + phaseDFloor: 699, + } + addr := serveDistributionTestServer(t, NewDistributionServer( + distribution.NewEngine(), + nil, + WithDistributionTimestampAllocator(serverAlloc), + )) + routed, err := kv.NewLeaderRoutedTSOAllocator( + &distributionTSOAllocator{leader: false}, + distributionLeaderView{addr: addr}, + ) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, routed.Close()) }) + + require.NoError(t, routed.ValidateDurableTimestamp(context.Background(), 700)) + require.ErrorIs(t, routed.ValidateDurableTimestamp(context.Background(), 699), kv.ErrTSOTimestampInvalid) +} + func TestDistributionServerGetTimestamp_RejectsFollower(t *testing.T) { t.Parallel() @@ -671,6 +754,53 @@ func TestDistributionServerSplitRange_UsesCoordinatorForCatalogWrites(t *testing require.Equal(t, coordinator.lastCommitTS, resp.Right.SplitAtHlc) } +func TestDistributionServerSplitRange_PhaseDReadsAtValidatedAppliedWatermark(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(""), + End: []byte("m"), + GroupID: 1, + State: distribution.RouteStateActive, + }, + { + RouteID: 2, + Start: []byte("m"), + End: nil, + GroupID: 2, + State: distribution.RouteStateActive, + }, + }) + require.NoError(t, err) + legacyFloor := catalog.LatestCommitTS() + allocator := &distributionTSOAllocator{ + base: legacyFloor + 100, + leader: true, + phaseD: true, + phaseDFloor: legacyFloor - 1, + } + coordinator := newDistributionCoordinatorStub(baseStore, true) + coordinator.allocator = allocator + s := NewDistributionServer( + distribution.NewEngine(), + catalog, + WithDistributionCoordinator(coordinator), + ) + + _, err = s.SplitRange(ctx, &pb.SplitRangeRequest{ + ExpectedCatalogVersion: saved.Version, + RouteId: 1, + SplitKey: []byte("g"), + }) + require.NoError(t, err) + require.Equal(t, legacyFloor, coordinator.lastStartTS) +} + func TestDistributionServerSplitRange_UsesPersistentNextRouteID(t *testing.T) { t.Parallel() @@ -966,6 +1096,7 @@ func seededDistributionServerWithoutCoordinator(t *testing.T) (*DistributionServ type distributionCoordinatorStub struct { store store.MVCCStore + allocator kv.TimestampAllocator leader bool clock *kv.HLC nextTS uint64 @@ -978,6 +1109,10 @@ type distributionCoordinatorStub struct { dispatchCalls int } +func (s *distributionCoordinatorStub) TimestampAllocator() kv.TimestampAllocator { + return s.allocator +} + func newDistributionCoordinatorStub(st store.MVCCStore, leader bool) *distributionCoordinatorStub { return &distributionCoordinatorStub{ store: st, @@ -1152,14 +1287,17 @@ func (s *distributionCoordinatorStub) LeaseReadForKey(ctx context.Context, _ []b } type distributionTSOAllocator struct { - base uint64 - previousFloor uint64 - count int - min uint64 - err error - leader bool - activate bool - cutover bool + base uint64 + previousFloor uint64 + count int + min uint64 + err error + leader bool + activate bool + activatePhaseD bool + cutover bool + phaseD bool + phaseDFloor uint64 } type batchOnlyDistributionTSOAllocator struct { @@ -1208,24 +1346,50 @@ func (a *distributionTSOAllocator) ReserveBatchAfter( n int, min uint64, activate bool, + activatePhaseD bool, ) (kv.TSOReservation, error) { a.count = n a.min = min a.activate = activate + a.activatePhaseD = activatePhaseD if a.err != nil { return kv.TSOReservation{}, a.err } if activate { a.cutover = true } + if activatePhaseD { + a.phaseD = true + a.phaseDFloor = a.previousFloor + } return kv.TSOReservation{ Base: a.base, Count: n, PreviousAllocationFloor: a.previousFloor, CutoverActive: a.cutover, + PhaseDActive: a.phaseD, + PhaseDFloor: a.phaseDFloor, }, nil } +func (a *distributionTSOAllocator) ValidateDurableTimestamp(_ context.Context, timestamp uint64) error { + if a.err != nil { + return a.err + } + if !a.phaseD || timestamp <= a.phaseDFloor || timestamp > a.base { + return kv.ErrTSOTimestampInvalid + } + return nil +} + +func (a *distributionTSOAllocator) PhaseDFloor() uint64 { return a.phaseDFloor } + +func (a *distributionTSOAllocator) AllocationFloor() uint64 { return a.base } + +func (a *distributionTSOAllocator) PhaseDActive() bool { return a.phaseD } + +func (a *distributionTSOAllocator) PhaseDRequired() bool { return a.phaseD } + func (a *distributionTSOAllocator) IsLeader() bool { return a.leader } func (a *distributionTSOAllocator) RunLeaseRenewal(ctx context.Context) { diff --git a/adapter/dynamodb_item_write.go b/adapter/dynamodb_item_write.go index a5aecda19..bee55d2fb 100644 --- a/adapter/dynamodb_item_write.go +++ b/adapter/dynamodb_item_write.go @@ -155,7 +155,11 @@ func (d *DynamoDBServer) retryItemWriteWithGenerationLegacy( backoff := transactRetryInitialBackoff deadline := time.Now().Add(transactRetryMaxDuration) for range transactRetryMaxAttempts { - readTS := d.nextTxnReadTS() + readTimestamp, err := d.beginTxnReadTimestamp(ctx, "dynamodb item-write legacy: begin read timestamp") + if err != nil { + return nil, errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() plan, err := prepare(readTS) if err != nil { return nil, err @@ -266,7 +270,11 @@ func (d *DynamoDBServer) itemWriteFirstAttempt( tableName string, prepare func(readTS uint64) (*itemWritePlan, error), ) (*itemWritePlan, *reusableItemWrite, error) { - readTS := d.nextTxnReadTS() + readTimestamp, err := d.beginTxnReadTimestamp(ctx, "dynamodb item-write: begin read timestamp") + if err != nil { + return nil, nil, errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() plan, err := prepare(readTS) if err != nil { return nil, nil, err diff --git a/adapter/dynamodb_locks.go b/adapter/dynamodb_locks.go index 8bbb3698f..b68db8000 100644 --- a/adapter/dynamodb_locks.go +++ b/adapter/dynamodb_locks.go @@ -115,6 +115,11 @@ func (d *DynamoDBServer) nextTxnReadTS() uint64 { return maxTS } +func (d *DynamoDBServer) beginTxnReadTimestamp(ctx context.Context, label string) (kv.ReadTimestamp, error) { + readTimestamp, err := kv.BeginReadTimestampThrough(ctx, d.coordinator, d.nextTxnReadTS(), label) + return readTimestamp, errors.WithStack(err) +} + func (d *DynamoDBServer) pinReadTS(ts uint64) *kv.ActiveTimestampToken { if d == nil || d.readTracker == nil { return &kv.ActiveTimestampToken{} diff --git a/adapter/dynamodb_migration.go b/adapter/dynamodb_migration.go index 6fcaff1d9..86b715de8 100644 --- a/adapter/dynamodb_migration.go +++ b/adapter/dynamodb_migration.go @@ -21,12 +21,11 @@ func (d *DynamoDBServer) ensureLegacyTableMigrationLocked(ctx context.Context, t backoff := transactRetryInitialBackoff deadline := time.Now().Add(transactRetryMaxDuration) for range transactRetryMaxAttempts { - readTS := d.nextTxnReadTS() - schema, exists, err := d.loadTableSchemaAt(ctx, tableName, readTS) + readTimestamp, schema, migrationRequired, err := d.legacyMigrationSnapshot(ctx, tableName) if err != nil { return errors.WithStack(err) } - if !exists || !schema.needsLegacyKeyMigration() { + if !migrationRequired { return nil } // Admin read-only callers (AdminScanTable) must not trigger @@ -39,7 +38,7 @@ func (d *DynamoDBServer) ensureLegacyTableMigrationLocked(ctx context.Context, t "table requires a one-time legacy-key migration before admin read endpoints are available; migrate via the SigV4 surface first") } if !schema.usesOrderedKeyEncoding() { - err = d.startLegacyTableKeyMigration(ctx, schema, readTS) + err = d.startLegacyTableKeyMigration(ctx, schema, readTimestamp) } else { err = d.migrateLegacyTableGeneration(ctx, schema) } @@ -57,14 +56,30 @@ func (d *DynamoDBServer) ensureLegacyTableMigrationLocked(ctx context.Context, t return newDynamoAPIError(http.StatusInternalServerError, dynamoErrInternal, "legacy table migration retry attempts exhausted") } +func (d *DynamoDBServer) legacyMigrationSnapshot( + ctx context.Context, + tableName string, +) (kv.ReadTimestamp, *dynamoTableSchema, bool, error) { + readTimestamp, err := d.beginTxnReadTimestamp(ctx, "dynamodb legacy-table migration: begin read timestamp") + if err != nil { + return kv.ReadTimestamp{}, nil, false, errors.WithStack(err) + } + schema, exists, err := d.loadTableSchemaAt(ctx, tableName, readTimestamp.Timestamp()) + if err != nil { + return kv.ReadTimestamp{}, nil, false, errors.WithStack(err) + } + return readTimestamp, schema, exists && schema.needsLegacyKeyMigration(), nil +} + func (d *DynamoDBServer) startLegacyTableKeyMigration( ctx context.Context, schema *dynamoTableSchema, - readTS uint64, + readTimestamp kv.ReadTimestamp, ) error { if schema == nil || schema.usesOrderedKeyEncoding() { return nil } + readTS := readTimestamp.Timestamp() nextGeneration, err := d.nextTableGenerationAt(ctx, schema.TableName, readTS) if err != nil { return err @@ -151,7 +166,11 @@ func (d *DynamoDBServer) migrateLegacyItem( backoff := transactRetryInitialBackoff deadline := time.Now().Add(transactRetryMaxDuration) for range transactRetryMaxAttempts { - readTS := d.nextTxnReadTS() + readTimestamp, err := d.beginTxnReadTimestamp(ctx, "dynamodb legacy-item migration: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() req, done, err := d.buildLegacyMigrationRequest(ctx, targetSchema, sourceSchema, targetKey, sourceKey, readTS) if err != nil { return err @@ -292,7 +311,11 @@ func (d *DynamoDBServer) finalizeLegacyTableMigration(ctx context.Context, schem if err != nil { return errors.WithStack(err) } - readTS := d.nextTxnReadTS() + readTimestamp, err := d.beginTxnReadTimestamp(ctx, "dynamodb finalize migration: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() req := &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: readTS, diff --git a/adapter/dynamodb_schema.go b/adapter/dynamodb_schema.go index 0ffa0f917..2100a4768 100644 --- a/adapter/dynamodb_schema.go +++ b/adapter/dynamodb_schema.go @@ -183,7 +183,11 @@ func (d *DynamoDBServer) createTableWithRetry(ctx context.Context, tableName str backoff := transactRetryInitialBackoff deadline := time.Now().Add(transactRetryMaxDuration) for range transactRetryMaxAttempts { - readTS := d.nextTxnReadTS() + readTimestamp, err := d.beginTxnReadTimestamp(ctx, "dynamodb create table: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() exists, err := d.tableExistsAt(ctx, tableName, readTS) if err != nil { return err @@ -199,6 +203,7 @@ func (d *DynamoDBServer) createTableWithRetry(ctx context.Context, tableName str if err != nil { return err } + req.StartTS = readTS if _, err := d.coordinator.Dispatch(ctx, req); err == nil { return nil } @@ -293,7 +298,11 @@ func (d *DynamoDBServer) deleteTableWithRetry(ctx context.Context, tableName str backoff := transactRetryInitialBackoff deadline := time.Now().Add(transactRetryMaxDuration) for range transactRetryMaxAttempts { - readTS := d.nextTxnReadTS() + readTimestamp, err := d.beginTxnReadTimestamp(ctx, "dynamodb delete table: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() schema, exists, err := d.loadTableSchemaAt(ctx, tableName, readTS) if err != nil { return errors.WithStack(err) @@ -304,7 +313,7 @@ func (d *DynamoDBServer) deleteTableWithRetry(ctx context.Context, tableName str req := &kv.OperationGroup[kv.OP]{ IsTxn: true, - StartTS: 0, + StartTS: readTS, Elems: []*kv.Elem[kv.OP]{ {Op: kv.Del, Key: dynamoTableMetaKey(tableName)}, }, diff --git a/adapter/dynamodb_transact.go b/adapter/dynamodb_transact.go index 3f58a688a..3d5cd82af 100644 --- a/adapter/dynamodb_transact.go +++ b/adapter/dynamodb_transact.go @@ -733,7 +733,11 @@ func (d *DynamoDBServer) buildTransactWriteItemsRequest(ctx context.Context, in return nil, nil, nil, err } } - readTS := d.nextTxnReadTS() + readTimestamp, err := d.beginTxnReadTimestamp(ctx, "dynamodb transact-write: begin read timestamp") + if err != nil { + return nil, nil, nil, errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() reqs := &kv.OperationGroup[kv.OP]{ IsTxn: true, // Keep transaction start aligned with the snapshot used to evaluate diff --git a/adapter/redis_delta_compactor.go b/adapter/redis_delta_compactor.go index 424b71a7c..fa1bebc07 100644 --- a/adapter/redis_delta_compactor.go +++ b/adapter/redis_delta_compactor.go @@ -249,7 +249,15 @@ func (c *DeltaCompactor) compactUrgentKey(ctx context.Context, req urgentCompact func (c *DeltaCompactor) compactUrgentKeyBatch(ctx context.Context, req urgentCompactionRequest, h *collectionDeltaHandler, prefix, end []byte) (int, bool) { // Use a fresh readTS each iteration so we observe the committed state from // the previous compaction pass and do not re-scan already-deleted delta keys. - readTS := snapshotTS(c.coord.Clock(), c.st) + readTimestamp, err := kv.BeginReadTimestampThrough(ctx, c.coord, snapshotTS(c.coord.Clock(), c.st), + "redis delta compactor urgent: begin read timestamp") + if err != nil { + c.logger.WarnContext(ctx, "delta compactor urgent: timestamp allocation failed", + "type", req.typeName, "key", string(req.userKey), "error", err) + return 0, true + } + ctx = readTimestamp.WithDispatchVoucher(ctx) + readTS := readTimestamp.Timestamp() // Scan one extra beyond MaxDeltaScanLimit to detect whether more remain. kvs, err := c.st.ScanAt(ctx, prefix, end, store.MaxDeltaScanLimit+1, readTS) @@ -305,9 +313,15 @@ func (c *DeltaCompactor) SyncOnce(ctx context.Context) error { "backoff_until", until) return nil } - readTS := snapshotTS(c.coord.Clock(), c.st) tickCtx, cancel := context.WithTimeout(ctx, c.timeout) defer cancel() + readTimestamp, err := kv.BeginReadTimestampThrough(tickCtx, c.coord, snapshotTS(c.coord.Clock(), c.st), + "redis delta compactor: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } + tickCtx = readTimestamp.WithDispatchVoucher(tickCtx) + readTS := readTimestamp.Timestamp() combined := c.compactBackgroundHandlers(tickCtx, readTS) if tickCtx.Err() == nil { @@ -604,7 +618,7 @@ func (c *DeltaCompactor) dispatchCompaction(ctx context.Context, readTS uint64, if err != nil { return errors.WithStack(err) } - _, err = c.coord.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + _, err = kv.DispatchWithReadTimestamp(ctx, c.coord, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: normalizeStartTS(readTS), CommitTS: commitTS, diff --git a/adapter/redis_expire_cmds.go b/adapter/redis_expire_cmds.go index a40c8adc0..e55a52dc3 100644 --- a/adapter/redis_expire_cmds.go +++ b/adapter/redis_expire_cmds.go @@ -47,24 +47,18 @@ func (r *RedisServer) getdel(conn redcon.Conn, cmd redcon.Command) { defer cancel() var v []byte err := r.retryRedisWrite(ctx, func() error { - readTS := r.readTS() - typ, err := r.keyTypeAt(ctx, key, readTS) + readTS, err := r.beginTxnStartTS(ctx, "redis getdel: begin read timestamp") + if err != nil { + return cockerrors.WithStack(err) + } + raw, exists, err := r.getdelValueAt(ctx, key, readTS) if err != nil { return err } - if typ == redisTypeNone { + if !exists { v = nil return nil } - if typ != redisTypeString { - return wrongTypeError() - } - raw, _, err := r.readRedisStringAt(key, readTS) - if err != nil { - // Key may have expired or been deleted between type check and read. - v = nil - return nil //nolint:nilerr // treat not-found/expired as nil value - } elems, _, err := r.deleteLogicalKeyElems(ctx, key, readTS) if err != nil { return err @@ -86,6 +80,25 @@ func (r *RedisServer) getdel(conn redcon.Conn, cmd redcon.Command) { conn.WriteBulk(v) } +func (r *RedisServer) getdelValueAt(ctx context.Context, key []byte, readTS uint64) ([]byte, bool, error) { + typ, err := r.keyTypeAt(ctx, key, readTS) + if err != nil { + return nil, false, err + } + if typ == redisTypeNone { + return nil, false, nil + } + if typ != redisTypeString { + return nil, false, wrongTypeError() + } + raw, _, err := r.readRedisStringAt(key, readTS) + if err != nil { + // Key may have expired or been deleted between type check and read. + return nil, false, nil //nolint:nilerr // treat not-found/expired as nil value + } + return raw, true, nil +} + // SETNX key value — set if not exists, returns 1 on success, 0 on failure func (r *RedisServer) setnx(conn redcon.Conn, cmd redcon.Command) { if r.proxyToLeader(conn, cmd, cmd.Args[1]) { @@ -185,9 +198,12 @@ func parseExpireTTL(raw []byte) (int64, error) { return ttl, nil } -func (r *RedisServer) prepareExpire(key []byte, nxOnly bool) (uint64, bool, error) { - readTS := r.readTS() - exists, err := r.logicalExistsAt(context.Background(), key, readTS) +func (r *RedisServer) prepareExpire(ctx context.Context, key []byte, nxOnly bool) (uint64, bool, error) { + readTS, err := r.beginTxnStartTS(ctx, "redis expire: begin read timestamp") + if err != nil { + return 0, false, cockerrors.WithStack(err) + } + exists, err := r.logicalExistsAt(ctx, key, readTS) if err != nil { return 0, false, err } @@ -199,7 +215,7 @@ func (r *RedisServer) prepareExpire(key []byte, nxOnly bool) (uint64, bool, erro return readTS, true, nil } - currentTTL, err := r.ttlAt(context.Background(), key, readTS) + currentTTL, err := r.ttlAt(ctx, key, readTS) if err != nil { return 0, false, err } @@ -253,7 +269,7 @@ func (r *RedisServer) setExpire(conn redcon.Conn, cmd redcon.Command, unit time. // then re-invokes doSetExpire with a fresh readTS, providing OCC safety without // an explicit mutex. Leadership is verified by coordinator.Dispatch itself. func (r *RedisServer) doSetExpire(ctx context.Context, key []byte, ttl int64, expireAt time.Time, nxOnly bool) (int, error) { - readTS, eligible, err := r.prepareExpire(key, nxOnly) + readTS, eligible, err := r.prepareExpire(ctx, key, nxOnly) if err != nil { return 0, err } diff --git a/adapter/redis_hash_cmds.go b/adapter/redis_hash_cmds.go index 9943d6ef5..a537d85f8 100644 --- a/adapter/redis_hash_cmds.go +++ b/adapter/redis_hash_cmds.go @@ -162,7 +162,10 @@ func (r *RedisServer) applyHashFieldPairs(key []byte, args [][]byte) (int, error defer cancel() var added int err := r.retryRedisWrite(ctx, func() error { - readTS := r.readTS() + readTS, err := r.beginTxnStartTS(ctx, "redis hash field write: begin read timestamp") + if err != nil { + return cockerrors.WithStack(err) + } typ, err := r.keyTypeOrEmptyAt(ctx, key, readTS, redisTypeHash) if err != nil { return err @@ -472,7 +475,10 @@ func (r *RedisServer) resolveHashFieldDelElems(ctx context.Context, key []byte, } func (r *RedisServer) hdelTxn(ctx context.Context, key []byte, fields [][]byte) (int, error) { - readTS := r.readTS() + readTS, err := r.beginTxnStartTS(ctx, "redis hash field delete: begin read timestamp") + if err != nil { + return 0, cockerrors.WithStack(err) + } typ, err := r.keyTypeAtExpect(ctx, key, readTS, redisTypeHash) if err != nil { return 0, err @@ -762,7 +768,10 @@ func (r *RedisServer) hincrbyWithMigration(ctx context.Context, key, fieldKey [] } func (r *RedisServer) hincrbyTxn(ctx context.Context, key, field []byte, increment int64) (int64, error) { - readTS := r.readTS() + readTS, err := r.beginTxnStartTS(ctx, "redis hash increment: begin read timestamp") + if err != nil { + return 0, cockerrors.WithStack(err) + } typ, err := r.keyTypeOrEmptyAt(ctx, key, readTS, redisTypeHash) if err != nil { return 0, err @@ -831,7 +840,10 @@ func (r *RedisServer) incrLegacy(conn redcon.Conn, cmd redcon.Command) { defer cancel() var current int64 if err := r.retryRedisWrite(ctx, func() error { - readTS := r.readTS() + readTS, err := r.beginTxnStartTS(ctx, "redis incr: begin read timestamp") + if err != nil { + return cockerrors.WithStack(err) + } typ, err := r.keyTypeAt(ctx, cmd.Args[1], readTS) if err != nil { return err diff --git a/adapter/redis_lists.go b/adapter/redis_lists.go index 1f7a6fcf2..7d6860c2a 100644 --- a/adapter/redis_lists.go +++ b/adapter/redis_lists.go @@ -99,7 +99,10 @@ func (r *RedisServer) listPushCore(ctx context.Context, key []byte, values [][]b var newLen int64 err := r.retryRedisWrite(ctx, func() error { - readTS := r.readTS() + readTS, err := r.beginTxnStartTS(ctx, "redis list push: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } meta, metaExists, typ, cleanupElems, err := r.listPushSnapshot(ctx, key, readTS) if err != nil { return err @@ -380,19 +383,18 @@ func (r *RedisServer) listPushCoreWithDedup(ctx context.Context, key []byte, val var newLen int64 var pending *reusableListPush err := r.retryRedisWrite(ctx, func() error { - if pending != nil { - length, drop, dispErr := r.dispatchListPushReuse(ctx, key, pending) - if drop { - pending = nil - } - if dispErr != nil { - return dispErr + length, handled, err := r.tryPendingListPush(ctx, key, &pending) + if handled { + if err != nil { + return err } newLen = length return nil } - - readTS := r.readTS() + readTS, err := r.beginTxnStartTS(ctx, "redis list push: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } meta, metaExists, typ, cleanupElems, err := r.listPushSnapshot(ctx, key, readTS) if err != nil { return err @@ -454,6 +456,21 @@ func (r *RedisServer) listPushCoreWithDedup(ctx context.Context, key []byte, val return newLen, err } +func (r *RedisServer) tryPendingListPush( + ctx context.Context, + key []byte, + pending **reusableListPush, +) (int64, bool, error) { + if *pending == nil { + return 0, false, nil + } + length, drop, err := r.dispatchListPushReuse(ctx, key, *pending) + if drop { + *pending = nil + } + return length, true, err +} + func (r *RedisServer) listRPush(ctx context.Context, key []byte, values [][]byte) (int64, error) { return r.listPushCore(ctx, key, values, r.buildRPushOps) } @@ -675,7 +692,11 @@ func (r *RedisServer) listPopClaim(ctx context.Context, key []byte, count int, l var popped []string err := r.retryRedisWrite(ctx, func() error { - result, popErr := r.listPopClaimOnce(ctx, key, count, left, r.readTS()) + readTS, readErr := r.beginTxnStartTS(ctx, "redis list pop: begin read timestamp") + if readErr != nil { + return errors.WithStack(readErr) + } + result, popErr := r.listPopClaimOnce(ctx, key, count, left, readTS) if popErr != nil { return popErr } @@ -824,9 +845,13 @@ func (r *RedisServer) ltrim(conn redcon.Conn, cmd redcon.Command) { } ctx, cancel := context.WithTimeout(context.Background(), redisDispatchTimeout) defer cancel() + key := cmd.Args[1] if err := r.retryRedisWrite(ctx, func() error { - readTS := r.readTS() - typ, err := r.keyTypeAt(ctx, cmd.Args[1], readTS) + readTS, err := r.beginTxnStartTS(ctx, "redis ltrim: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } + typ, err := r.keyTypeAt(ctx, key, readTS) if err != nil { return err } @@ -836,16 +861,11 @@ func (r *RedisServer) ltrim(conn redcon.Conn, cmd redcon.Command) { if typ != redisTypeList { return wrongTypeError() } - current, err := r.listValuesAt(ctx, cmd.Args[1], readTS) + current, err := r.listValuesAt(ctx, key, readTS) if err != nil { return err } - s, e := normalizeRankRange(start, stop, len(current)) - trimmed := []string{} - if e >= s { - trimmed = append(trimmed, current[s:e+1]...) - } - return r.rewriteListTxn(ctx, cmd.Args[1], readTS, trimmed) + return r.rewriteListTxn(ctx, key, readTS, trimListValues(current, start, stop)) }); err != nil { writeRedisError(conn, err) return @@ -853,6 +873,14 @@ func (r *RedisServer) ltrim(conn redcon.Conn, cmd redcon.Command) { conn.WriteString("OK") } +func trimListValues(current []string, start, stop int) []string { + s, e := normalizeRankRange(start, stop, len(current)) + if e < s { + return []string{} + } + return append([]string{}, current[s:e+1]...) +} + func (r *RedisServer) lindex(conn redcon.Conn, cmd redcon.Command) { if r.proxyToLeader(conn, cmd, cmd.Args[1]) { return diff --git a/adapter/redis_lua_context.go b/adapter/redis_lua_context.go index 3b10a46ac..e906e9b10 100644 --- a/adapter/redis_lua_context.go +++ b/adapter/redis_lua_context.go @@ -259,7 +259,11 @@ func newLuaScriptContext(ctx context.Context, server *RedisServer) (*luaScriptCo if _, err := kv.LeaseReadThrough(server.coordinator, ctx); err != nil { return nil, errors.WithStack(err) } - startTS := server.readTS() + readTimestamp, err := kv.BeginReadTimestampThrough(ctx, server.coordinator, server.readTS(), "redis lua: begin read timestamp") + if err != nil { + return nil, errors.WithStack(err) + } + startTS := readTimestamp.Timestamp() return &luaScriptContext{ server: server, startTS: startTS, diff --git a/adapter/redis_set_cmds.go b/adapter/redis_set_cmds.go index d29aba6e1..d8836bbe6 100644 --- a/adapter/redis_set_cmds.go +++ b/adapter/redis_set_cmds.go @@ -278,7 +278,10 @@ func applySetMemberMutation(elems []*kv.Elem[kv.OP], memberKey []byte, exists, a func (r *RedisServer) mutateExactSetLegacy(conn redcon.Conn, ctx context.Context, kind string, key []byte, members [][]byte, add bool) { var changed int if err := r.retryRedisWrite(ctx, func() error { - readTS := r.readTS() + readTS, err := r.beginTxnStartTS(ctx, "redis set mutation: begin read timestamp") + if err != nil { + return cockerrors.WithStack(err) + } if err := r.validateExactSetKind(kind, key, readTS); err != nil { return err } @@ -303,49 +306,24 @@ func (r *RedisServer) mutateExactSetLegacy(conn redcon.Conn, ctx context.Context func (r *RedisServer) mutateExactSetWide(conn redcon.Conn, ctx context.Context, key []byte, members [][]byte, add bool) { var changed int if err := r.retryRedisWrite(ctx, func() error { - readTS := r.readTS() - typ, err := r.keyTypeOrEmptyAt(ctx, key, readTS, redisTypeSet) + readTS, err := r.beginTxnStartTS(ctx, "redis set mutation: begin read timestamp") if err != nil { - return err + return cockerrors.WithStack(err) } - cleanupElems, migrationElems, legacyMemberBase, expiredRecreate, err := r.setWideMutationBase(ctx, key, readTS, typ) + prepared, err := r.prepareExactSetWideMutation(ctx, key, members, add, readTS) if err != nil { return err } - - startTS := normalizeStartTS(readTS) - commitTS, err := r.nextCommitTSAfter(ctx, startTS, "mutateExactSetWide: allocate commitTS") - if err != nil { - return cockerrors.WithStack(err) - } - - elems := make([]*kv.Elem[kv.OP], 0, len(cleanupElems)+len(migrationElems)+len(members)+setWideColOverhead) - elems = append(elems, cleanupElems...) - elems = append(elems, migrationElems...) - - var lenDelta int64 - var mutErr error - elems, changed, lenDelta, mutErr = r.applySetMemberMutations(ctx, key, members, add, readTS, elems, legacyMemberBase, expiredRecreate) - if mutErr != nil { - return mutErr - } - - if changed == 0 && len(migrationElems) == 0 && len(cleanupElems) == 0 { + changed = prepared.changed + if prepared.skip { return nil } - - elems = appendSetDeltaElems(elems, key, lenDelta, commitTS) - - if len(elems) == 0 { - return nil - } - _, dispatchErr := r.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ IsTxn: true, - StartTS: startTS, - CommitTS: commitTS, - ReadKeys: redisTxnWideCreateReadKeys(key, typ, redisTxnWideSetFenceKey), - Elems: elems, + StartTS: prepared.startTS, + CommitTS: prepared.commitTS, + ReadKeys: redisTxnWideCreateReadKeys(key, prepared.typ, redisTxnWideSetFenceKey), + Elems: prepared.elems, }) return cockerrors.WithStack(dispatchErr) }); err != nil { @@ -355,6 +333,50 @@ func (r *RedisServer) mutateExactSetWide(conn redcon.Conn, ctx context.Context, conn.WriteInt(changed) } +type exactSetWideMutation struct { + typ redisValueType + startTS uint64 + commitTS uint64 + changed int + elems []*kv.Elem[kv.OP] + skip bool +} + +func (r *RedisServer) prepareExactSetWideMutation( + ctx context.Context, + key []byte, + members [][]byte, + add bool, + readTS uint64, +) (exactSetWideMutation, error) { + var prepared exactSetWideMutation + typ, err := r.keyTypeOrEmptyAt(ctx, key, readTS, redisTypeSet) + if err != nil { + return prepared, err + } + cleanupElems, migrationElems, legacyMemberBase, expiredRecreate, err := r.setWideMutationBase(ctx, key, readTS, typ) + if err != nil { + return prepared, err + } + startTS := normalizeStartTS(readTS) + commitTS, err := r.nextCommitTSAfter(ctx, startTS, "mutateExactSetWide: allocate commitTS") + if err != nil { + return prepared, cockerrors.WithStack(err) + } + elems := make([]*kv.Elem[kv.OP], 0, len(cleanupElems)+len(migrationElems)+len(members)+setWideColOverhead) + elems = append(elems, cleanupElems...) + elems = append(elems, migrationElems...) + elems, changed, lenDelta, err := r.applySetMemberMutations(ctx, key, members, add, readTS, elems, legacyMemberBase, expiredRecreate) + if err != nil { + return prepared, err + } + elems = appendSetDeltaElems(elems, key, lenDelta, commitTS) + return exactSetWideMutation{ + typ: typ, startTS: startTS, commitTS: commitTS, changed: changed, elems: elems, + skip: (changed == 0 && len(migrationElems) == 0 && len(cleanupElems) == 0) || len(elems) == 0, + }, nil +} + func (r *RedisServer) setWideMutationBase( ctx context.Context, key []byte, @@ -718,7 +740,10 @@ func (r *RedisServer) pfadd(conn redcon.Conn, cmd redcon.Command) { defer cancel() var changed int if err := r.retryRedisWrite(ctx, func() error { - readTS := r.readTS() + readTS, err := r.beginTxnStartTS(ctx, "redis pfadd: begin read timestamp") + if err != nil { + return cockerrors.WithStack(err) + } if err := r.validateExactSetKind(hllKind, cmd.Args[1], readTS); err != nil { return err } diff --git a/adapter/redis_stream_cmds.go b/adapter/redis_stream_cmds.go index 1e27e0c96..0cfea772a 100644 --- a/adapter/redis_stream_cmds.go +++ b/adapter/redis_stream_cmds.go @@ -247,7 +247,10 @@ func (r *RedisServer) prepareXAdd( key []byte, req xaddRequest, ) (string, uint64, []*kv.Elem[kv.OP], error) { - readTS := r.readTS() + readTS, err := r.xaddReadTimestamp(ctx, key) + if err != nil { + return "", 0, nil, err + } typ, err := r.streamTypeForXAdd(ctx, key, readTS) if err != nil { return "", 0, nil, err @@ -472,6 +475,21 @@ func (r *RedisServer) streamCleanupForExpiredRecreate( return cleanup, store.StreamMeta{}, false, nil } +func (r *RedisServer) xaddReadTimestamp(ctx context.Context, key []byte) (uint64, error) { + readTS, err := r.beginTxnStartTS(ctx, "redis xadd: begin read timestamp") + if err != nil { + return 0, cockerrors.WithStack(err) + } + typ, err := r.keyTypeAtExpect(ctx, key, readTS, redisTypeStream) + if err != nil { + return 0, err + } + if typ != redisTypeNone && typ != redisTypeStream { + return 0, wrongTypeError() + } + return readTS, nil +} + // dispatchAndSignalStream dispatches the elems through the coordinator // and, on success, wakes any XREAD BLOCK waiter on the same node. // dispatchElems blocks until the FSM applies locally, so by the time @@ -811,7 +829,10 @@ func (r *RedisServer) flushLegacyCleanupOnTrimNoOp( } func (r *RedisServer) xtrimTxn(ctx context.Context, key []byte, maxLen int) (int, error) { - readTS := r.readTS() + readTS, err := r.beginTxnStartTS(ctx, "redis xtrim: begin read timestamp") + if err != nil { + return 0, cockerrors.WithStack(err) + } proceed, err := r.streamTypeForWrite(ctx, key, readTS) if err != nil || !proceed { return 0, err diff --git a/adapter/redis_strings.go b/adapter/redis_strings.go index a01fa41a1..a61ce4066 100644 --- a/adapter/redis_strings.go +++ b/adapter/redis_strings.go @@ -154,7 +154,10 @@ func (r *RedisServer) replaceWithStringTxn(ctx context.Context, key, value []byt func (r *RedisServer) executeSet(ctx context.Context, key, value []byte, opts redisSetOptions) (redisSetExecution, error) { var result redisSetExecution err := r.retryRedisWrite(ctx, func() error { - readTS := r.readTS() + readTS, err := r.beginTxnStartTS(ctx, "redis set string: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } state, err := r.loadRedisSetState(ctx, key, readTS, opts.returnOld) if err != nil { return err @@ -519,7 +522,10 @@ func (r *RedisServer) delLocal(keys [][]byte) (int, error) { err := r.retryRedisWrite(ctx, func() error { elems := []*kv.Elem[kv.OP]{} nextRemoved := 0 - readTS := r.readTS() + readTS, err := r.beginTxnStartTS(ctx, "redis del: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } for _, key := range keys { keyElems, existed, err := r.deleteLogicalKeyElems(ctx, key, readTS) if err != nil { diff --git a/adapter/redis_txn.go b/adapter/redis_txn.go index a38001517..99f17894b 100644 --- a/adapter/redis_txn.go +++ b/adapter/redis_txn.go @@ -1577,7 +1577,7 @@ func (t *txnContext) commit() error { CommitTS: prepared.commitTS, ReadKeys: prepared.readKeys, } - if _, err := t.server.coordinator.Dispatch(prepared.ctx, group); err != nil { + if _, err := kv.DispatchWithReadTimestamp(prepared.ctx, t.server.coordinator, group); err != nil { return errors.WithStack(err) } return nil @@ -2379,13 +2379,18 @@ func (r *RedisServer) runTransactionDirect(queue []redcon.Command) ([]redisResul var results []redisResult err := r.retryRedisWrite(dispatchCtx, func() error { - startTS := r.txnStartTS() + readTimestamp, err := r.beginTxnReadTimestamp(dispatchCtx, "redis exec: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } + attemptCtx := readTimestamp.WithDispatchVoucher(dispatchCtx) + startTS := readTimestamp.Timestamp() readPin := r.pinReadTS(startTS) defer readPin.Release() txn := &txnContext{ server: r, - ctx: dispatchCtx, + ctx: attemptCtx, working: map[string]*txnValue{}, replacers: map[string]*stringReplacement{}, listStates: map[string]*listTxnState{}, @@ -2412,7 +2417,7 @@ func (r *RedisServer) runTransactionDirect(queue []redcon.Command) ([]redisResul nextResults = append(nextResults, res) } - if err := txn.validateReadSet(dispatchCtx); err != nil { + if err := txn.validateReadSet(attemptCtx); err != nil { return err } if err := txn.commit(); err != nil { @@ -2446,11 +2451,12 @@ func (r *RedisServer) runTransactionDirect(queue []redcon.Command) ([]redisResul // results are only returned when reuse actually represents the // outcome of attempt 1's intent. type reusableExecTxn struct { - elems []*kv.Elem[kv.OP] - startTS uint64 - commitTS uint64 - readKeys [][]byte - results []redisResult + elems []*kv.Elem[kv.OP] + startTS uint64 + commitTS uint64 + readKeys [][]byte + results []redisResult + readTimestamp kv.ReadTimestamp } // dispatchExecReuse runs one iteration of the option-2 reuse path for @@ -2468,6 +2474,7 @@ type reusableExecTxn struct { // is the current length" question; the client-visible result IS the // cached results array. func (r *RedisServer) dispatchExecReuse(ctx context.Context, pending *reusableExecTxn) (results []redisResult, drop bool, err error) { + ctx = pending.readTimestamp.WithDispatchVoucher(ctx) // gemini PR-A HIGH: persistence-grade commit_ts allocation must honor the // HLC-4 physical-ceiling fence (see kv/hlc.go NextFenced + the TLA proof // at tla/hlc/MCHLC_gap.cfg). Clock().Next() bypasses the ceiling and @@ -2478,7 +2485,7 @@ func (r *RedisServer) dispatchExecReuse(ctx context.Context, pending *reusableEx if allocErr != nil { return nil, false, errors.WithStack(allocErr) } - _, dispErr := r.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + _, dispErr := kv.DispatchWithReadTimestamp(ctx, r.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: pending.startTS, CommitTS: commitTS, @@ -2593,7 +2600,12 @@ func (r *RedisServer) runTransactionWithDedup(queue []redcon.Command) ([]redisRe // from runTransactionWithDedup to keep that loop under the cyclop // budget; the dedup rationale lives there. func (r *RedisServer) firstExecAttempt(dispatchCtx context.Context, queue []redcon.Command) ([]redisResult, *reusableExecTxn, error) { - startTS := r.txnStartTS() + readTimestamp, err := r.beginTxnReadTimestamp(dispatchCtx, "redis exec: begin read timestamp") + if err != nil { + return nil, nil, errors.WithStack(err) + } + dispatchCtx = readTimestamp.WithDispatchVoucher(dispatchCtx) + startTS := readTimestamp.Timestamp() readPin := r.pinReadTS(startTS) defer readPin.Release() @@ -2647,7 +2659,7 @@ func (r *RedisServer) firstExecAttempt(dispatchCtx context.Context, queue []redc CommitTS: prepared.commitTS, ReadKeys: prepared.readKeys, } - if _, dispErr := r.coordinator.Dispatch(prepared.ctx, group); dispErr != nil { + if _, dispErr := kv.DispatchWithReadTimestamp(prepared.ctx, r.coordinator, group); dispErr != nil { // Preserve the exact attempt for a forwarded conflict only after // restoring its typed form. runTransactionWithDedup can then reuse this // write set instead of replaying the EXEC body from a new snapshot. @@ -2660,11 +2672,12 @@ func (r *RedisServer) firstExecAttempt(dispatchCtx context.Context, queue []redc // escape to the client are out of scope for this loop. if isRetryableRedisTxnErr(dispErr) { return nil, &reusableExecTxn{ - elems: prepared.elems, - startTS: txn.startTS, - commitTS: prepared.commitTS, - readKeys: prepared.readKeys, - results: nextResults, + elems: prepared.elems, + startTS: txn.startTS, + commitTS: prepared.commitTS, + readKeys: prepared.readKeys, + results: nextResults, + readTimestamp: readTimestamp, }, errors.WithStack(dispErr) } return nil, nil, errors.WithStack(dispErr) @@ -2710,6 +2723,19 @@ func (r *RedisServer) txnStartTS() uint64 { return maxTS } +func (r *RedisServer) beginTxnReadTimestamp(ctx context.Context, label string) (kv.ReadTimestamp, error) { + readTimestamp, err := kv.BeginReadTimestampThrough(ctx, r.coordinator, r.txnStartTS(), label) + return readTimestamp, errors.WithStack(err) +} + +func (r *RedisServer) beginTxnStartTS(ctx context.Context, label string) (uint64, error) { + readTimestamp, err := r.beginTxnReadTimestamp(ctx, label) + if err != nil { + return 0, err + } + return readTimestamp.Timestamp(), nil +} + func (r *RedisServer) writeResults(conn redcon.Conn, results []redisResult) { conn.WriteArray(len(results)) for _, res := range results { diff --git a/adapter/redis_zset_cmds.go b/adapter/redis_zset_cmds.go index 193d1f90e..be6f5cd39 100644 --- a/adapter/redis_zset_cmds.go +++ b/adapter/redis_zset_cmds.go @@ -600,7 +600,10 @@ func (r *RedisServer) applyZAddPair(ctx context.Context, key []byte, p zaddPair, } func (r *RedisServer) zaddTxn(ctx context.Context, key []byte, flags zaddFlags, pairs []zaddPair) (int, error) { - readTS := r.readTS() + readTS, err := r.beginTxnStartTS(ctx, "redis zadd: begin read timestamp") + if err != nil { + return 0, cockerrors.WithStack(err) + } base, err := r.prepareZSetWriteBase(ctx, key, readTS, len(pairs)) if err != nil { return 0, err @@ -722,7 +725,10 @@ func (r *RedisServer) dispatchAndSignalZSet( // zincrbyTxn performs one attempt of ZINCRBY in wide-column format. // Returns the new score after applying increment. func (r *RedisServer) zincrbyTxn(ctx context.Context, key []byte, member string, increment float64) (float64, error) { - readTS := r.readTS() + readTS, err := r.beginTxnStartTS(ctx, "redis zincrby: begin read timestamp") + if err != nil { + return 0, cockerrors.WithStack(err) + } base, err := r.prepareZSetWriteBase(ctx, key, readTS, 0) if err != nil { return 0, err @@ -1016,7 +1022,10 @@ func (r *RedisServer) zrem(conn redcon.Conn, cmd redcon.Command) { defer cancel() var removed int if err := r.retryRedisWrite(ctx, func() error { - readTS := r.readTS() + readTS, err := r.beginTxnStartTS(ctx, "redis zrem: begin read timestamp") + if err != nil { + return cockerrors.WithStack(err) + } typ, err := r.keyTypeAtExpect(ctx, cmd.Args[1], readTS, redisTypeZSet) if err != nil { return err @@ -1065,22 +1074,18 @@ func (r *RedisServer) zremrangebyrank(conn redcon.Conn, cmd redcon.Command) { defer cancel() var removed int if err := r.retryRedisWrite(ctx, func() error { - readTS := r.readTS() - typ, err := r.keyTypeAtExpect(ctx, cmd.Args[1], readTS, redisTypeZSet) + readTS, err := r.beginTxnStartTS(ctx, "redis zremrangebyrank: begin read timestamp") + if err != nil { + return cockerrors.WithStack(err) + } + value, exists, err := r.zsetMutationValueAt(ctx, cmd.Args[1], readTS) if err != nil { return err } - if typ == redisTypeNone { + if !exists { removed = 0 return nil } - if typ != redisTypeZSet { - return wrongTypeError() - } - value, _, err := r.loadZSetAt(ctx, cmd.Args[1], readTS) - if err != nil { - return err - } s, e := normalizeRankRange(start, stop, len(value.Entries)) if e < s { removed = 0 @@ -1098,6 +1103,25 @@ func (r *RedisServer) zremrangebyrank(conn redcon.Conn, cmd redcon.Command) { conn.WriteInt(removed) } +func (r *RedisServer) zsetMutationValueAt( + ctx context.Context, + key []byte, + readTS uint64, +) (redisZSetValue, bool, error) { + typ, err := r.keyTypeAtExpect(ctx, key, readTS, redisTypeZSet) + if err != nil { + return redisZSetValue{}, false, err + } + if typ == redisTypeNone { + return redisZSetValue{}, false, nil + } + if typ != redisTypeZSet { + return redisZSetValue{}, false, wrongTypeError() + } + value, _, err := r.loadZSetAt(ctx, key, readTS) + return value, true, err +} + // tryBZPopMinWithMode runs one BZPOPMIN attempt against key. The // fast flag selects keyTypeAtExpectFast (no slow-path fallback, no // wrongType detection) when true; the caller MUST guarantee that the @@ -1110,16 +1134,19 @@ func (r *RedisServer) tryBZPopMinWithMode(key []byte, fast bool) (*bzpopminResul defer cancel() var result *bzpopminResult err := r.retryRedisWrite(ctx, func() error { - readTS := r.readTS() + readTS, err := r.beginTxnStartTS(ctx, "redis bzpopmin: begin read timestamp") + if err != nil { + return cockerrors.WithStack(err) + } var typ redisValueType - var err error + var typeErr error if fast { - typ, err = r.keyTypeAtExpectFast(ctx, key, readTS, redisTypeZSet) + typ, typeErr = r.keyTypeAtExpectFast(ctx, key, readTS, redisTypeZSet) } else { - typ, err = r.keyTypeAtExpect(ctx, key, readTS, redisTypeZSet) + typ, typeErr = r.keyTypeAtExpect(ctx, key, readTS, redisTypeZSet) } - if err != nil { - return err + if typeErr != nil { + return typeErr } if typ == redisTypeNone { result = nil diff --git a/adapter/s3.go b/adapter/s3.go index 6c4874157..bcb5b0355 100644 --- a/adapter/s3.go +++ b/adapter/s3.go @@ -596,10 +596,12 @@ func (s *S3Server) createBucket(w http.ResponseWriter, r *http.Request, bucket s } err := s.retryS3Mutation(r.Context(), func() error { readTS := s.readTS() - startTS, err := s.txnStartTS(r.Context(), readTS) + readTimestamp, err := s.beginTxnReadTimestamp(r.Context(), readTS, "s3 create bucket: begin read timestamp") if err != nil { return errors.Wrap(err, "s3: allocate startTS for mutation") } + readTS = readTimestamp.Timestamp() + startTS := readTS readPin := s.pinReadTS(readTS) defer readPin.Release() @@ -675,10 +677,12 @@ func (s *S3Server) deleteBucket(w http.ResponseWriter, r *http.Request, bucket s var deletedGeneration uint64 err := s.retryS3Mutation(r.Context(), func() error { readTS := s.readTS() - startTS, err := s.txnStartTS(r.Context(), readTS) + readTimestamp, err := s.beginTxnReadTimestamp(r.Context(), readTS, "s3 delete bucket: begin read timestamp") if err != nil { return errors.Wrap(err, "s3: allocate startTS for mutation") } + readTS = readTimestamp.Timestamp() + startTS := readTS readPin := s.pinReadTS(readTS) defer readPin.Release() @@ -823,10 +827,12 @@ func (s *S3Server) putBucketAcl(w http.ResponseWriter, r *http.Request, bucket s err := s.retryS3Mutation(r.Context(), func() error { readTS := s.readTS() - startTS, err := s.txnStartTS(r.Context(), readTS) + readTimestamp, err := s.beginTxnReadTimestamp(r.Context(), readTS, "s3 put bucket acl: begin read timestamp") if err != nil { return errors.Wrap(err, "s3: allocate startTS for mutation") } + readTS = readTimestamp.Timestamp() + startTS := readTS readPin := s.pinReadTS(readTS) defer readPin.Release() @@ -1095,10 +1101,12 @@ func (s *S3Server) deleteObject(w http.ResponseWriter, r *http.Request, bucket s var generation uint64 err := s.retryS3Mutation(r.Context(), func() error { readTS := s.readTS() - startTS, err := s.txnStartTS(r.Context(), readTS) + readTimestamp, err := s.beginTxnReadTimestamp(r.Context(), readTS, "s3 delete object: begin read timestamp") if err != nil { return errors.Wrap(err, "s3: allocate startTS for mutation") } + readTS = readTimestamp.Timestamp() + startTS := readTS readPin := s.pinReadTS(readTS) defer readPin.Release() @@ -1151,6 +1159,12 @@ func (s *S3Server) deleteObject(w http.ResponseWriter, r *http.Request, bucket s func (s *S3Server) createMultipartUpload(w http.ResponseWriter, r *http.Request, bucket string, objectKey string) { readTS := s.readTS() + readTimestamp, err := s.beginTxnReadTimestamp(r.Context(), readTS, "s3 create multipart upload: begin read timestamp") + if err != nil { + writeS3InternalError(w, err) + return + } + readTS = readTimestamp.Timestamp() readPin := s.pinReadTS(readTS) defer readPin.Release() @@ -1165,11 +1179,7 @@ func (s *S3Server) createMultipartUpload(w http.ResponseWriter, r *http.Request, } uploadID := newS3UploadID(s.clock()) - startTS, err := s.txnStartTS(r.Context(), readTS) - if err != nil { - writeS3InternalError(w, err) - return - } + startTS := readTS commitTS, err := s.nextTxnCommitTS(r.Context(), startTS) if err != nil { writeS3InternalError(w, err) @@ -1285,10 +1295,12 @@ func (s *S3Server) abortMultipartUpload(w http.ResponseWriter, r *http.Request, var generation uint64 err := s.retryS3Mutation(r.Context(), func() error { readTS := s.readTS() - startTS, err := s.txnStartTS(r.Context(), readTS) + readTimestamp, err := s.beginTxnReadTimestamp(r.Context(), readTS, "s3 abort multipart upload: begin read timestamp") if err != nil { return errors.Wrap(err, "s3: allocate startTS for mutation") } + readTS = readTimestamp.Timestamp() + startTS := readTS readPin := s.pinReadTS(readTS) defer readPin.Release() @@ -2431,6 +2443,23 @@ func (s *S3Server) txnStartTS(ctx context.Context, readTS uint64) (uint64, error return ts, nil } +func (s *S3Server) beginTxnReadTimestamp(ctx context.Context, readTS uint64, label string) (kv.ReadTimestamp, error) { + if readTS == ^uint64(0) { + if alloc, ok := kv.TimestampAllocatorThrough(s.coordinator); ok { + if phaseD, phaseDOK := alloc.(kv.TSOPhaseDState); phaseDOK && (phaseD.PhaseDRequired() || phaseD.PhaseDActive()) { + readTimestamp, err := kv.BeginReadTimestampThrough(ctx, s.coordinator, readTS, label) + return readTimestamp, errors.WithStack(err) + } + } + } + startTS, err := s.txnStartTS(ctx, readTS) + if err != nil { + return kv.ReadTimestamp{}, err + } + readTimestamp, err := kv.BeginReadTimestampThrough(ctx, s.coordinator, startTS, label) + return readTimestamp, errors.WithStack(err) +} + // newS3UploadID generates an upload identifier for multipart uploads. // This is NOT a persistence timestamp — it is an opaque identifier // returned to the client and is only used as a lookup key for diff --git a/adapter/s3_admin.go b/adapter/s3_admin.go index 0d8c58d28..fe72e8b3a 100644 --- a/adapter/s3_admin.go +++ b/adapter/s3_admin.go @@ -259,10 +259,12 @@ func (s *S3Server) AdminCreateBucket(ctx context.Context, principal AdminPrincip // error path that the wrapping retry harness needs to see. func (s *S3Server) adminCreateBucketTxn(ctx context.Context, principal AdminPrincipal, name, acl string) (*AdminBucketSummary, error) { readTS := s.readTS() - startTS, err := s.txnStartTS(ctx, readTS) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, readTS, "s3 admin create bucket: begin read timestamp") if err != nil { return nil, errors.Wrap(err, "s3 admin: allocate startTS for adminCreateBucketTxn") } + readTS = readTimestamp.Timestamp() + startTS := readTS readPin := s.pinReadTS(readTS) defer readPin.Release() @@ -327,10 +329,12 @@ func (s *S3Server) AdminPutBucketAcl(ctx context.Context, principal AdminPrincip err := s.retryS3Mutation(ctx, func() error { readTS := s.readTS() - startTS, err := s.txnStartTS(ctx, readTS) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, readTS, "s3 admin put bucket acl: begin read timestamp") if err != nil { return errors.Wrap(err, "s3 admin: allocate startTS for mutation") } + readTS = readTimestamp.Timestamp() + startTS := readTS readPin := s.pinReadTS(readTS) defer readPin.Release() @@ -438,10 +442,12 @@ func (s *S3Server) AdminDeleteBucket(ctx context.Context, principal AdminPrincip // allocation (PR #867 Phase 2a). func (s *S3Server) adminDeleteBucketTxnBody(ctx context.Context, name string, deletedGeneration *uint64) error { readTS := s.readTS() - startTS, err := s.txnStartTS(ctx, readTS) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, readTS, "s3 admin delete bucket: begin read timestamp") if err != nil { return errors.Wrap(err, "s3 admin: allocate startTS for adminDeleteBucket") } + readTS = readTimestamp.Timestamp() + startTS := readTS readPin := s.pinReadTS(readTS) defer readPin.Release() diff --git a/adapter/s3_admin_objects.go b/adapter/s3_admin_objects.go index ecac803be..118e0092a 100644 --- a/adapter/s3_admin_objects.go +++ b/adapter/s3_admin_objects.go @@ -116,10 +116,12 @@ func (s *S3Server) AdminDeleteObject(ctx context.Context, principal AdminPrincip // silent-no-op semantics and AWS S3. func (s *S3Server) adminDeleteObjectTxn(ctx context.Context, bucket, key string) (*s3ObjectManifest, uint64, error) { readTS := s.readTS() - startTS, err := s.txnStartTS(ctx, readTS) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, readTS, "s3 admin delete object: begin read timestamp") if err != nil { return nil, 0, errors.Wrap(err, "s3 admin: allocate startTS for adminDeleteObjectTxn") } + readTS = readTimestamp.Timestamp() + startTS := readTS readPin := s.pinReadTS(readTS) defer readPin.Release() @@ -215,10 +217,12 @@ func (s *S3Server) AdminPutObject(ctx context.Context, principal AdminPrincipal, //nolint:cyclop,gocognit,nestif // see comment above func (s *S3Server) adminPutObjectStream(ctx context.Context, bucket, key string, body io.Reader, contentType string) (*s3ObjectManifest, uint64, error) { readTS := s.readTS() - startTS, err := s.txnStartTS(ctx, readTS) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, readTS, "s3 admin put object: begin read timestamp") if err != nil { return nil, 0, errors.Wrap(err, "s3 admin: allocate startTS for adminPutObjectStream") } + readTS = readTimestamp.Timestamp() + startTS := readTS readPin := s.pinReadTS(readTS) defer readPin.Release() diff --git a/adapter/s3_hlc_fence_test.go b/adapter/s3_hlc_fence_test.go index 30f6b3b4b..348a48bb7 100644 --- a/adapter/s3_hlc_fence_test.go +++ b/adapter/s3_hlc_fence_test.go @@ -5,7 +5,9 @@ import ( "testing" "time" + "github.com/bootjp/elastickv/internal/s3keys" "github.com/bootjp/elastickv/kv" + "github.com/bootjp/elastickv/store" "github.com/stretchr/testify/require" ) @@ -54,6 +56,33 @@ func TestS3TxnStartTSPassesThroughExplicitReadTS(t *testing.T) { require.Equal(t, uint64(42), ts) } +func TestS3BeginTxnReadTimestampPhaseDPreservesAppliedWatermark(t *testing.T) { + t.Parallel() + + allocator := &distributionTSOAllocator{base: 100, phaseD: true, phaseDFloor: 10} + coord := newDistributionCoordinatorStub(nil, true) + coord.allocator = allocator + srv := &S3Server{coordinator: coord} + + readTimestamp, err := srv.beginTxnReadTimestamp(context.Background(), 42, "test") + require.NoError(t, err) + require.Equal(t, uint64(42), readTimestamp.Timestamp()) + require.Zero(t, allocator.count, "an applied Phase-D read watermark must not allocate ahead of Raft apply") +} + +func TestS3BeginTxnReadTimestampPhaseDRejectsLatestSentinel(t *testing.T) { + t.Parallel() + + allocator := &distributionTSOAllocator{base: 100, phaseD: true, phaseDFloor: 10} + coord := newDistributionCoordinatorStub(nil, true) + coord.allocator = allocator + srv := &S3Server{coordinator: coord} + + _, err := srv.beginTxnReadTimestamp(context.Background(), ^uint64(0), "test") + require.ErrorIs(t, err, kv.ErrTSOTimestampInvalid) + require.Zero(t, allocator.count, "the latest sentinel must fail closed instead of allocating an unapplied read timestamp") +} + // TestS3NextTxnCommitTSFailsClosedOnExpiredCeiling verifies that // nextTxnCommitTS surfaces ErrCeilingExpired through the // NextFenced() it calls after Observe(startTS). This is the @@ -81,3 +110,51 @@ func TestS3NextTxnCommitTSFailsClosedOnExpiredCeiling(t *testing.T) { require.ErrorIs(t, err, kv.ErrCeilingExpired, "s3 nextTxnCommitTS must propagate ErrCeilingExpired from NextFenced") } + +func TestS3CommitUploadPartRechecksUploadAtLatestAppliedWatermark(t *testing.T) { + st := store.NewMVCCStore() + const generation = uint64(1) + uploadMetaKey := s3keys.UploadMetaKey("bucket", generation, "object", "upload") + require.NoError(t, st.PutAt(context.Background(), uploadMetaKey, []byte("meta"), 10, 0)) + require.NoError(t, st.DeleteAt(context.Background(), uploadMetaKey, 20)) + coord := &recordingS3DispatchCoordinator{} + server := NewS3Server(nil, "", st, coord, nil) + + _, err := server.commitS3UploadPart(context.Background(), &s3UploadPartState{ + partNo: 1, + readTS: 10, + meta: &s3BucketMeta{Generation: generation}, + uploadMetaKey: uploadMetaKey, + }, s3ChunkUploadResult{}, "bucket", "object", "upload", 10, 30) + require.ErrorContains(t, err, "upload not found") + require.Nil(t, coord.request, "an upload removed after startTS must not dispatch an orphan part") +} + +func TestS3CommitUploadPartIncludesUploadMetaInReadSet(t *testing.T) { + st := store.NewMVCCStore() + const generation = uint64(1) + uploadMetaKey := s3keys.UploadMetaKey("bucket", generation, "object", "upload") + require.NoError(t, st.PutAt(context.Background(), uploadMetaKey, []byte("meta"), 10, 0)) + coord := &recordingS3DispatchCoordinator{} + server := NewS3Server(nil, "", st, coord, nil) + + _, err := server.commitS3UploadPart(context.Background(), &s3UploadPartState{ + partNo: 1, + readTS: 10, + meta: &s3BucketMeta{Generation: generation}, + uploadMetaKey: uploadMetaKey, + }, s3ChunkUploadResult{}, "bucket", "object", "upload", 10, 30) + require.NoError(t, err) + require.NotNil(t, coord.request) + require.Equal(t, [][]byte{uploadMetaKey}, coord.request.ReadKeys) +} + +type recordingS3DispatchCoordinator struct { + stubAdapterCoordinator + request *kv.OperationGroup[kv.OP] +} + +func (c *recordingS3DispatchCoordinator) Dispatch(_ context.Context, request *kv.OperationGroup[kv.OP]) (*kv.CoordinateResponse, error) { + c.request = request + return &kv.CoordinateResponse{}, nil +} diff --git a/adapter/s3_multipart_complete.go b/adapter/s3_multipart_complete.go index 7cded1e9a..acc6fe67a 100644 --- a/adapter/s3_multipart_complete.go +++ b/adapter/s3_multipart_complete.go @@ -67,6 +67,11 @@ func validateS3MultipartCompletionParts(parts []s3CompleteMultipartUploadPart, b func (s *S3Server) loadS3MultipartCompletion(ctx context.Context, bucket, objectKey, uploadID string, request s3CompleteMultipartUploadRequest) (s3MultipartCompletion, error) { completion := s3MultipartCompletion{bucket: bucket, objectKey: objectKey, uploadID: uploadID} readTS := s.readTS() + readTimestamp, err := s.beginTxnReadTimestamp(ctx, readTS, "s3 complete multipart upload preparation: begin read timestamp") + if err != nil { + return completion, errors.WithStack(err) + } + readTS = readTimestamp.Timestamp() readPin := s.pinReadTS(readTS) defer readPin.Release() @@ -164,10 +169,12 @@ func (s *S3Server) commitS3MultipartCompletion(ctx context.Context, completion s func (s *S3Server) commitS3MultipartCompletionAttempt(ctx context.Context, completion s3MultipartCompletion) (*s3ObjectManifest, error) { readTS := s.readTS() - startTS, err := s.txnStartTS(ctx, readTS) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, readTS, "s3 complete multipart upload: begin read timestamp") if err != nil { return nil, errors.Wrap(err, "s3: allocate startTS for completeMultipartUpload retry") } + readTS = readTimestamp.Timestamp() + startTS := readTS readPin := s.pinReadTS(readTS) defer readPin.Release() diff --git a/adapter/s3_put_object.go b/adapter/s3_put_object.go index bb74e8db6..8b883d99f 100644 --- a/adapter/s3_put_object.go +++ b/adapter/s3_put_object.go @@ -20,10 +20,12 @@ type s3PutObjectState struct { func (s *S3Server) prepareS3PutObject(ctx context.Context, request *http.Request, bucket, objectKey string) (*s3PutObjectState, error) { readTS := s.readTS() - startTS, err := s.txnStartTS(ctx, readTS) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, readTS, "s3 put object: begin read timestamp") if err != nil { return nil, errors.WithStack(err) } + readTS = readTimestamp.Timestamp() + startTS := readTS state := &s3PutObjectState{startTS: startTS, readPin: s.pinReadTS(readTS)} prepared := false defer func() { diff --git a/adapter/s3_upload_part.go b/adapter/s3_upload_part.go index 1f491faa4..80250e48f 100644 --- a/adapter/s3_upload_part.go +++ b/adapter/s3_upload_part.go @@ -25,7 +25,11 @@ func (s *S3Server) prepareS3UploadPart(ctx context.Context, bucket, objectKey, u if err != nil { return nil, err } - state := &s3UploadPartState{partNo: partNo, readTS: s.readTS()} + readTimestamp, err := s.beginTxnReadTimestamp(ctx, s.readTS(), "s3 upload part preparation: begin read timestamp") + if err != nil { + return nil, errors.WithStack(err) + } + state := &s3UploadPartState{partNo: partNo, readTS: readTimestamp.Timestamp()} state.readPin = s.pinReadTS(state.readTS) prepared := false defer func() { @@ -94,10 +98,12 @@ func (s *S3Server) storeS3UploadPart(ctx context.Context, request *http.Request, func (s *S3Server) allocateS3UploadPartVersion(ctx context.Context) (uint64, uint64, error) { readTS := s.readTS() - startTS, err := s.txnStartTS(ctx, readTS) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, readTS, "s3 upload part: begin read timestamp") if err != nil { return 0, 0, errors.WithStack(err) } + readTS = readTimestamp.Timestamp() + startTS := readTS commitTS, err := s.nextTxnCommitTS(ctx, startTS) if err != nil { return 0, 0, errors.WithStack(err) @@ -116,12 +122,13 @@ func (s *S3Server) commitS3UploadPart(ctx context.Context, state *s3UploadPartSt } partKey := s3keys.UploadPartKey(bucket, state.meta.Generation, objectKey, uploadID, state.partNo) previous := s.loadPreviousS3PartDescriptor(ctx, partKey, state.readTS) - if err := s.verifyS3UploadStillExists(ctx, state.uploadMetaKey, bucket, objectKey); err != nil { + if err := s.verifyS3UploadStillExists(ctx, state.uploadMetaKey, bucket, objectKey, s.readTS()); err != nil { return nil, err } _, err = s.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, CommitTS: commitTS, - Elems: []*kv.Elem[kv.OP]{{Op: kv.Put, Key: partKey, Value: body}}, + Elems: []*kv.Elem[kv.OP]{{Op: kv.Put, Key: partKey, Value: body}}, + ReadKeys: [][]byte{state.uploadMetaKey}, }) if err != nil { return nil, errors.WithStack(err) @@ -141,8 +148,8 @@ func (s *S3Server) loadPreviousS3PartDescriptor(ctx context.Context, partKey []b return &descriptor } -func (s *S3Server) verifyS3UploadStillExists(ctx context.Context, uploadMetaKey []byte, bucket, objectKey string) error { - if _, err := s.store.GetAt(ctx, uploadMetaKey, s.readTS()); err != nil { +func (s *S3Server) verifyS3UploadStillExists(ctx context.Context, uploadMetaKey []byte, bucket, objectKey string, readTS uint64) error { + if _, err := s.store.GetAt(ctx, uploadMetaKey, readTS); err != nil { if errors.Is(err, store.ErrKeyNotFound) { return newS3ResponseError(http.StatusNotFound, "NoSuchUpload", "upload not found", bucket, objectKey) } diff --git a/adapter/sqs_catalog.go b/adapter/sqs_catalog.go index 7356d6b5e..8ba22f8ca 100644 --- a/adapter/sqs_catalog.go +++ b/adapter/sqs_catalog.go @@ -874,6 +874,11 @@ func (s *SQSServer) nextTxnReadTS(ctx context.Context) uint64 { return maxTS } +func (s *SQSServer) beginTxnReadTimestamp(ctx context.Context, label string) (kv.ReadTimestamp, error) { + readTimestamp, err := kv.BeginReadTimestampThrough(ctx, s.coordinator, s.nextTxnReadTS(ctx), label) + return readTimestamp, errors.WithStack(err) +} + func (s *SQSServer) loadQueueMetaAt(ctx context.Context, queueName string, ts uint64) (*sqsQueueMeta, bool, error) { b, err := s.store.GetAt(ctx, sqsQueueMetaKey(queueName), ts) if err != nil { @@ -979,11 +984,11 @@ func (s *SQSServer) createQueueWithRetry(ctx context.Context, requested *sqsQueu // with the requested attributes, false means the dispatch hit a retryable // conflict and should be retried after backoff. func (s *SQSServer) tryCreateQueueOnce(ctx context.Context, requested *sqsQueueMeta) (bool, error) { - readTS := s.nextTxnReadTS(ctx) - existing, exists, err := s.loadQueueMetaAt(ctx, requested.Name, readTS) + readTimestamp, existing, exists, err := s.createQueueSnapshot(ctx, requested.Name) if err != nil { return false, errors.WithStack(err) } + readTS := readTimestamp.Timestamp() if exists { if attributesEqual(existing, requested) { return true, nil @@ -1075,6 +1080,21 @@ func (s *SQSServer) tryCreateQueueOnce(ctx context.Context, requested *sqsQueueM return true, nil } +func (s *SQSServer) createQueueSnapshot( + ctx context.Context, + queueName string, +) (kv.ReadTimestamp, *sqsQueueMeta, bool, error) { + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs create queue: begin read timestamp") + if err != nil { + return kv.ReadTimestamp{}, nil, false, errors.WithStack(err) + } + existing, exists, err := s.loadQueueMetaAt(ctx, queueName, readTimestamp.Timestamp()) + if err != nil { + return kv.ReadTimestamp{}, nil, false, errors.WithStack(err) + } + return readTimestamp, existing, exists, nil +} + func (s *SQSServer) deleteQueue(w http.ResponseWriter, r *http.Request) { var in sqsDeleteQueueInput if err := decodeSQSJSONInput(r, &in); err != nil { @@ -1120,7 +1140,11 @@ func (s *SQSServer) deleteQueueWithRetry(ctx context.Context, queueName string) backoff := transactRetryInitialBackoff deadline := time.Now().Add(transactRetryMaxDuration) for range transactRetryMaxAttempts { - readTS := s.nextTxnReadTS(ctx) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs delete queue: begin read timestamp") + if err != nil { + return 0, errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() existing, exists, err := s.loadQueueMetaAt(ctx, queueName, readTS) if err != nil { return 0, errors.WithStack(err) @@ -1652,7 +1676,11 @@ func applyAndValidateSetAttributes(meta *sqsQueueMeta, attrs map[string]string) // successful Dispatch so post-commit request-path gauges are not removed by // the caller's cleanup. func (s *SQSServer) trySetQueueAttributesOnce(ctx context.Context, queueName string, attrs map[string]string) (bool, *sqsQueueThrottle, []string, uint64, bool, error) { - readTS := s.nextTxnReadTS(ctx) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs set queue attributes: begin read timestamp") + if err != nil { + return false, nil, nil, 0, false, errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() meta, exists, err := s.loadQueueMetaAt(ctx, queueName, readTS) if err != nil { return false, nil, nil, 0, false, errors.WithStack(err) diff --git a/adapter/sqs_messages.go b/adapter/sqs_messages.go index 516d53d87..89ad6dc40 100644 --- a/adapter/sqs_messages.go +++ b/adapter/sqs_messages.go @@ -631,7 +631,11 @@ func (s *SQSServer) sendMessageFifoLoop(w http.ResponseWriter, r *http.Request, // concurrent DeleteQueue / PurgeQueue could slip in between our read // and the write, storing a message under a dead generation. func (s *SQSServer) loadQueueMetaForSend(ctx context.Context, queueName string, body []byte) (*sqsQueueMeta, uint64, error) { - readTS := s.nextTxnReadTS(ctx) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs send message: begin read timestamp") + if err != nil { + return nil, 0, errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() meta, exists, err := s.loadQueueMetaAt(ctx, queueName, readTS) if err != nil { return nil, readTS, errors.WithStack(err) @@ -991,7 +995,12 @@ func (s *SQSServer) longPollReceive(ctx context.Context, queueName string, opts // elapses prevents false-empty returns under poison-message backlogs // or hot-FIFO-group fan-in. func (s *SQSServer) scanAndDeliverOnce(ctx context.Context, queueName string, opts sqsReceiveOptions) ([]map[string]any, error) { - readTS := s.nextTxnReadTS(ctx) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs receive message: begin read timestamp") + if err != nil { + return nil, errors.WithStack(err) + } + ctx = readTimestamp.WithDispatchVoucher(ctx) + readTS := readTimestamp.Timestamp() meta, exists, err := s.loadQueueMetaAt(ctx, queueName, readTS) if err != nil { return nil, errors.WithStack(err) @@ -1291,7 +1300,7 @@ func (s *SQSServer) expireMessage(ctx context.Context, queueName string, meta *s ReadKeys: readKeys, Elems: elems, } - if _, err := s.coordinator.Dispatch(ctx, req); err != nil { + if _, err := kv.DispatchWithReadTimestamp(ctx, s.coordinator, req); err != nil { if isRetryableTransactWriteError(err) { return nil } @@ -1440,7 +1449,7 @@ func (s *SQSServer) commitReceiveRotation(ctx context.Context, queueName string, if err != nil { return nil, false, err } - if _, err := s.coordinator.Dispatch(ctx, req); err != nil { + if _, err := kv.DispatchWithReadTimestamp(ctx, s.coordinator, req); err != nil { if isRetryableTransactWriteError(err) { return nil, true, nil } @@ -1699,7 +1708,11 @@ const ( // error — silently succeeding would let misrouted deletes ack messages // that cannot possibly be deleted on this queue. func (s *SQSServer) loadMessageForDelete(ctx context.Context, queueName string, handle *decodedReceiptHandle) (*sqsQueueMeta, *sqsMessageRecord, []byte, uint64, sqsDeleteOutcome, error) { - readTS := s.nextTxnReadTS(ctx) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs delete message: begin read timestamp") + if err != nil { + return nil, nil, nil, 0, sqsDeleteProceed, errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() meta, exists, err := s.loadQueueMetaAt(ctx, queueName, readTS) if err != nil { return nil, nil, nil, readTS, sqsDeleteProceed, errors.WithStack(err) @@ -1853,7 +1866,11 @@ func (s *SQSServer) parseQueueAndReceipt(queueUrl, receiptHandle string) (string // be rejected with ReceiptHandleIsInvalid instead of silently // mutating the orphan record. func (s *SQSServer) loadAndVerifyMessage(ctx context.Context, queueName string, handle *decodedReceiptHandle) (*sqsQueueMeta, *sqsMessageRecord, []byte, uint64, error) { - readTS := s.nextTxnReadTS(ctx) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs change message visibility: begin read timestamp") + if err != nil { + return nil, nil, nil, 0, errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() meta, exists, err := s.loadQueueMetaAt(ctx, queueName, readTS) if err != nil { return nil, nil, nil, readTS, errors.WithStack(err) diff --git a/adapter/sqs_messages_batch.go b/adapter/sqs_messages_batch.go index ea2de434b..1afbe51e6 100644 --- a/adapter/sqs_messages_batch.go +++ b/adapter/sqs_messages_batch.go @@ -181,7 +181,11 @@ func (s *SQSServer) trySendMessageBatchOnce( entries []sqsSendMessageBatchEntryInput, identities []sqsSendIdentity, ) ([]sqsSendMessageBatchResultEntry, []sqsBatchResultErrorEntry, bool, error) { - readTS := s.nextTxnReadTS(ctx) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs send message batch: begin read timestamp") + if err != nil { + return nil, nil, false, errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() meta, exists, err := s.loadQueueMetaAt(ctx, queueName, readTS) if err != nil { return nil, nil, false, errors.WithStack(err) @@ -367,7 +371,11 @@ func (s *SQSServer) runFifoSendWithRetry( backoff := transactRetryInitialBackoff deadline := time.Now().Add(transactRetryMaxDuration) for range transactRetryMaxAttempts { - readTS := s.nextTxnReadTS(ctx) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs fifo send: begin read timestamp") + if err != nil { + return nil, errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() meta, dedupID, delay, err := s.resolveFreshFifoSnapshot(ctx, queueName, in, readTS) if err != nil { return nil, err diff --git a/adapter/sqs_purge.go b/adapter/sqs_purge.go index 10dc0a06d..f2e993ab8 100644 --- a/adapter/sqs_purge.go +++ b/adapter/sqs_purge.go @@ -91,7 +91,11 @@ func (s *SQSServer) purgeQueueWithRetry(ctx context.Context, queueName string) ( // pre-bump and post-bump generations so the caller can audit-log the // committed value without a second meta read. func (s *SQSServer) tryPurgeQueueOnce(ctx context.Context, queueName string) (bool, uint64, uint64, error) { - readTS := s.nextTxnReadTS(ctx) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs purge queue: begin read timestamp") + if err != nil { + return false, 0, 0, errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() meta, exists, err := s.loadQueueMetaAt(ctx, queueName, readTS) if err != nil { return false, 0, 0, errors.WithStack(err) diff --git a/adapter/sqs_reaper.go b/adapter/sqs_reaper.go index ae7f6c20d..5c5d43d06 100644 --- a/adapter/sqs_reaper.go +++ b/adapter/sqs_reaper.go @@ -73,8 +73,13 @@ func (s *SQSServer) reapAllQueues(ctx context.Context) error { if err := ctx.Err(); err != nil { return errors.WithStack(err) } - readTS := s.nextTxnReadTS(ctx) - meta, exists, err := s.loadQueueMetaAt(ctx, name, readTS) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs reaper queue: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } + queueCtx := readTimestamp.WithDispatchVoucher(ctx) + readTS := readTimestamp.Timestamp() + meta, exists, err := s.loadQueueMetaAt(queueCtx, name, readTS) if err != nil || !exists { // Even when meta is gone (DeleteQueue), prior-generation // orphans need reaping; reapTombstonedQueues (called @@ -82,10 +87,10 @@ func (s *SQSServer) reapAllQueues(ctx context.Context) error { // the queue if loading itself failed (transient). continue } - if err := s.reapQueue(ctx, name, meta, readTS); err != nil { + if err := s.reapQueue(queueCtx, name, meta, readTS); err != nil { slog.Warn("sqs reaper queue pass failed", "queue", name, "err", err) } - if err := s.reapExpiredDedup(ctx, name, meta, readTS); err != nil { + if err := s.reapExpiredDedup(queueCtx, name, meta, readTS); err != nil { slog.Warn("sqs dedup reaper pass failed", "queue", name, "err", err) } } @@ -108,8 +113,13 @@ func (s *SQSServer) reapTombstonedQueues(ctx context.Context) error { upper := prefixScanEnd(prefix) start := bytes.Clone(prefix) for { - readTS := s.nextTxnReadTS(ctx) - page, err := s.store.ScanAt(ctx, start, upper, sqsReaperPageLimit, readTS) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs tombstone reaper: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } + pageCtx := readTimestamp.WithDispatchVoucher(ctx) + readTS := readTimestamp.Timestamp() + page, err := s.store.ScanAt(pageCtx, start, upper, sqsReaperPageLimit, readTS) if err != nil { return errors.WithStack(err) } @@ -129,7 +139,7 @@ func (s *SQSServer) reapTombstonedQueues(ctx context.Context) error { // non-canonical values to 1 so pre-PR-6a tombstones // retain their byte-identical legacy reaper path. partitionCount := decodeQueueTombstoneValue(kvp.Value) - s.reapTombstonedGeneration(ctx, queueName, gen, partitionCount, kvp.Key, readTS) + s.reapTombstonedGeneration(pageCtx, queueName, gen, partitionCount, kvp.Key, readTS) } if len(page) < sqsReaperPageLimit { return nil @@ -663,7 +673,7 @@ func (s *SQSServer) reapOneRecord(ctx context.Context, queueName string, meta *s if err != nil { return err } - if _, err := s.coordinator.Dispatch(ctx, req); err != nil { + if _, err := kv.DispatchWithReadTimestamp(ctx, s.coordinator, req); err != nil { if isRetryableTransactWriteError(err) { return nil } @@ -717,7 +727,7 @@ func (s *SQSServer) dispatchOrphanByAgeDrop(ctx context.Context, byAgeKey []byte {Op: kv.Del, Key: byAgeKey}, }, } - _, _ = s.coordinator.Dispatch(ctx, req) + _, _ = kv.DispatchWithReadTimestamp(ctx, s.coordinator, req) } // reapExpiredDedup walks every FIFO dedup record under the given @@ -872,7 +882,7 @@ func (s *SQSServer) dispatchDedupDelete(ctx context.Context, key []byte, readTS {Op: kv.Del, Key: key}, }, } - if _, err := s.coordinator.Dispatch(ctx, req); err != nil { + if _, err := kv.DispatchWithReadTimestamp(ctx, s.coordinator, req); err != nil { if isRetryableTransactWriteError(err) { return nil } diff --git a/adapter/sqs_redrive.go b/adapter/sqs_redrive.go index 3fb13b8d3..c87547a55 100644 --- a/adapter/sqs_redrive.go +++ b/adapter/sqs_redrive.go @@ -236,7 +236,7 @@ func (s *SQSServer) redriveCandidateToDLQ( if err != nil { return false, err } - if _, err := s.coordinator.Dispatch(ctx, req); err != nil { + if _, err := kv.DispatchWithReadTimestamp(ctx, s.coordinator, req); err != nil { if isRetryableTransactWriteError(err) { return true, nil } diff --git a/adapter/sqs_tags.go b/adapter/sqs_tags.go index 7635b89f0..eec8b74f2 100644 --- a/adapter/sqs_tags.go +++ b/adapter/sqs_tags.go @@ -164,7 +164,11 @@ func (s *SQSServer) tryMutateQueueTagsOnce( queueName string, mutate func(*sqsQueueMeta) error, ) (bool, error) { - readTS := s.nextTxnReadTS(ctx) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs mutate queue tags: begin read timestamp") + if err != nil { + return false, errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() meta, exists, err := s.loadQueueMetaAt(ctx, queueName, readTS) if err != nil { return false, errors.WithStack(err) diff --git a/distribution/catalog.go b/distribution/catalog.go index b92f85f23..5d4836003 100644 --- a/distribution/catalog.go +++ b/distribution/catalog.go @@ -240,6 +240,16 @@ func (s *CatalogStore) Snapshot(ctx context.Context) (CatalogSnapshot, error) { return s.SnapshotAt(ctx, s.store.LastCommitTS()) } +// LatestCommitTS returns the local catalog store watermark without reading +// catalog data. Callers use it as the pre-Phase-D legacy timestamp input before +// selecting the transaction snapshot through the dedicated TSO. +func (s *CatalogStore) LatestCommitTS() uint64 { + if s == nil || s.store == nil { + return 0 + } + return s.store.LastCommitTS() +} + // SnapshotAt reads a consistent route catalog snapshot at a specific MVCC // timestamp. func (s *CatalogStore) SnapshotAt(ctx context.Context, ts uint64) (CatalogSnapshot, error) { diff --git a/distribution/catalog_test.go b/distribution/catalog_test.go index a3ffe852f..33b59e8c3 100644 --- a/distribution/catalog_test.go +++ b/distribution/catalog_test.go @@ -8,6 +8,7 @@ import ( "github.com/bootjp/elastickv/store" "github.com/cockroachdb/errors" + "github.com/stretchr/testify/require" ) func TestCatalogVersionCodecRoundTrip(t *testing.T) { @@ -345,6 +346,33 @@ func TestCatalogStoreSaveAndSnapshot(t *testing.T) { assertRouteEqual(t, saved.Routes[1], snapshot.Routes[1]) } +func TestCatalogStoreSnapshotAtPreservesRequestedVersion(t *testing.T) { + cs := NewCatalogStore(store.NewMVCCStore()) + ctx := context.Background() + first, err := cs.Save(ctx, 0, []RouteDescriptor{{ + RouteID: 1, + Start: []byte(""), + GroupID: 1, + State: RouteStateActive, + }}) + require.NoError(t, err) + firstTS := cs.LatestCommitTS() + _, err = cs.Save(ctx, first.Version, []RouteDescriptor{{ + RouteID: 2, + Start: []byte(""), + GroupID: 2, + State: RouteStateActive, + }}) + require.NoError(t, err) + + snapshot, err := cs.SnapshotAt(ctx, firstTS) + require.NoError(t, err) + require.Equal(t, first.Version, snapshot.Version) + require.Equal(t, firstTS, snapshot.ReadTS) + require.Equal(t, uint64(1), snapshot.Routes[0].RouteID) + require.GreaterOrEqual(t, cs.LatestCommitTS(), firstTS) +} + func TestCatalogStoreSaveAndSnapshotSortsRoutesByStart(t *testing.T) { cs := NewCatalogStore(store.NewMVCCStore()) ctx := context.Background() diff --git a/docs/design/2026_04_16_partial_centralized_tso.md b/docs/design/2026_04_16_partial_centralized_tso.md index f4d38328c..979f4dec5 100644 --- a/docs/design/2026_04_16_partial_centralized_tso.md +++ b/docs/design/2026_04_16_partial_centralized_tso.md @@ -1,13 +1,13 @@ # Centralized Timestamp Oracle (TSO) Design -- Status: Partial — M1-M6 are implemented, including the dedicated group-0 +- Status: Partial — M1-M7 are implemented, including the dedicated group-0 FSM, leader-routed durable windows, strict term bootstrap, serialized shadow - migration, and the one-way rolling cutover marker. M7 legacy cleanup and - cross-shard SSI timestamp validation remain open; runtime config reload and - production latency/alerting work also remain open. + migration, one-way rolling cutover, durable Phase-D retirement, and + cross-shard SSI timestamp validation. Runtime config reload and production + latency/alerting work remain open. - Author: bootjp - Date: 2026-04-16 -- Updated: 2026-07-18 +- Updated: 2026-07-19 --- @@ -98,14 +98,50 @@ Implemented: issuance. The marker uses a versioned TSO envelope with the same legacy fail-closed prefix as allocation-floor entries and a distinct magic, so an old or misrouted data FSM halts while encryption bytes cannot activate it. +15. `--tsoPhaseDEnabled` commits a second one-way group-0 marker after cutover. + The marker captures the maximum pre-Phase-D timestamp floor. Its state and + floor are persisted in the V4 TSO snapshot. Before the marker applies, the + FSM continues writing V3 snapshots so older followers can install them + during the rolling upgrade. Malformed ordering or a changed replay floor + halts the TSO FSM fail-closed. The marker has its own versioned TSO envelope + with the same legacy fail-closed prefix, so old or misrouted data FSMs halt + and no reserved encryption byte can activate Phase D. +16. Once Phase D is durable, data groups no longer receive HLC ceiling renewal, + coordinator legacy HLC issuance is rejected, and shadow mode bypasses + legacy candidate generation/comparison. Group 0 remains renewed while it + is needed for the dedicated allocator's physical-ceiling fence. +17. Cross-shard transactions with a caller-supplied `StartTS` are accepted only + when the group-0 leader verifies `phase_d_floor < StartTS <= + allocation_floor`. The upper bound is a committed TSO reservation floor and + the lower bound excludes every legacy or pre-Phase-D value. Follower-local + state is never authoritative for this check. +18. Adapter read-modify-write paths call `BeginReadTimestampThrough` before the + first read and pass an applied store/catalog watermark, never a freshly + allocated clock value. If Phase D is required but not yet locally active, + the read boundary first forces the marker and its post-marker allocation + window to commit, then discards that allocation. The adapter still uses the + exact applied watermark for every read and `StartTS`. When that watermark + predates the Phase-D floor, the read boundary registers a bounded, + process-local capability that identifies the audited applied snapshot. The + adapter binds that capability to the logical operation and reserves exactly + one one-use coordinator voucher immediately before each dispatch that + shares the snapshot. Unvouched external `StartTS` values remain subject to + group-0 validation and values at/below the floor fail closed. A zero, + latest-sentinel, unavailable activation, or otherwise unprovable watermark + also fails closed. Retries reload and revalidate the applied watermark; + retries that intentionally reuse an exact OCC write set reserve a fresh use + without widening the capability to another timestamp. Coordinator + decorators forward both the allocator provider and voucher operation so + startup and keyviz wrappers cannot silently fall back to an unvalidated + value. +19. `BatchAllocator` validates cached candidates once Phase D is required. A + candidate at/below the Phase-D floor invalidates the entire local window and + forces a refill above the marker before any timestamp is returned. Remaining: -1. M7 Phase-D removal of per-shard renewal/legacy issuance after the migration - compatibility window closes. -2. Cross-shard SSI read-timestamp validation through the dedicated TSO. -3. Runtime config reload for the mode switch; current flags are startup-only. -4. Production benchmark, divergence metrics, and alert thresholds. +1. Runtime config reload for the mode switch; current flags are startup-only. +2. Production benchmark, divergence metrics, and alert thresholds. ### 1.1 Original Limitation @@ -123,10 +159,10 @@ Node B: not in defaultGroup → ceiling never updated ❌ → may collide with the previous leader's committed window ``` -M1 fixes that per-node renewal gap by proposing to every group this node -currently leads. **Global timestamp monotonicity is still not guaranteed when -different coordinators on different nodes allocate timestamps for cross-group -work**; that is the dedicated TSO / single-oracle work left open by M2-M7. +M1 fixed that per-node renewal gap by proposing to every group this node +currently leads. It did not by itself guarantee global timestamp monotonicity +when different coordinators allocated timestamps for cross-group work; M2-M7 +add the dedicated TSO and retire that distributed issuance path. ### 1.2 Near-Term Workaround @@ -727,9 +763,40 @@ approach enables a live cutover. ### 7.4 Phase D — Legacy Cleanup -- Remove per-shard ceiling proposals. -- Remove `legacyHLC` from `ShardedCoordinator`. -- Remove the shadow comparison code. +- Roll every member to a binary that understands the Phase-D entry and + `ValidateTimestamp` RPC. Run all members on the dedicated TSO path before + enabling `--tsoPhaseDEnabled`; an older member would correctly halt on the + unknown control entry, so mixed-version activation is prohibited. +- The first allocation with `--tsoPhaseDEnabled` commits cutover (if needed), + then the Phase-D marker and its pre-Phase-D floor, then a new allocation + window strictly above that floor. The switch is one-way and survives restart + through the TSO V4 snapshot. +- `ShardedCoordinator` dynamically stops data-group ceiling proposals and + fails closed instead of issuing from its legacy HLC. It continues group-0 + renewal only. Shadow mode observes durable cutover and directly uses the + dedicated allocator without generating or comparing a legacy candidate. +- Every adapter read-modify-write attempt obtains its transaction snapshot + from an applied store/catalog watermark before reading. If Phase D is required + but not yet active locally, the read boundary first commits the marker and a + post-marker allocation window, then discards the allocated value; it is never + substituted for data-group apply progress. The same applied watermark is used + for direct MVCC reads, active-snapshot pinning, and `OperationGroup.StartTS`; + a retry reloads and revalidates the watermark and repeats all reads. An applied + watermark at/below the Phase-D floor is admitted only through a bounded, + process-local capability registered by that audited read boundary. Every + dispatch sharing the snapshot reserves and consumes its own one-use voucher; + the capability is timestamp-bound and cannot authorize another value. + Arbitrary or repeated unvouched caller values at/below the floor still fail + closed at group 0. Pre-Phase-D deployments retain the prior + committed-watermark behavior. +- After Phase D, the coordinator asks the current group-0 leader to validate a + caller-supplied cross-shard `StartTS` before allocating `CommitTS` or proposing + to any data group. The allocator's configured `PhaseDRequired` state activates + this check before the local group-0 replica has applied the marker. Values + at/below the marker floor, beyond the committed TSO allocation floor, zero + values, inactive Phase-D state, missing validation support, and unavailable + leadership all fail closed. Existing single-shard caller-supplied `StartTS` + remains compatible because it does not establish a cross-shard SSI snapshot. ### 7.5 Monotonicity Invariant Across Phases @@ -757,7 +824,7 @@ its first window above the strict maximum committed data timestamp. | M4 — shipped | `BatchAllocator` with atomic counter for low-latency timestamp serving | Medium | | M5 — shipped | Preserve the default-group `LocalTSOAllocator` compatibility bridge when group 0 is absent; route coordinator-owned timestamp call sites through the allocator abstraction. | Medium | | M6 — shipped | Run the dedicated group-0 FSM, fence each new TSO leader term above all authoritative data-group commit floors, redirect follower requests to the TSO leader over gRPC, synchronously serialize fail-closed shadow issuance, and commit the one-way rolling cutover marker before production windows. | Low | -| M7 — open | Phase D legacy cleanup + cross-shard SSI read-timestamp validation via TSO | Low | +| M7 — shipped | Commit the durable Phase-D floor marker, preserve V3 snapshots until activation, retire data-shard HLC renewal and legacy/shadow issuance after cutover, activate before read validation while preserving exact applied snapshots through timestamp-bound per-dispatch vouchers, invalidate pre-Phase-D batch windows, and validate unvouched caller-supplied cross-shard SSI timestamps at the group-0 leader from activation onward. | Low | --- diff --git a/kv/coordinator.go b/kv/coordinator.go index 5ee5c4f46..af91a09fd 100644 --- a/kv/coordinator.go +++ b/kv/coordinator.go @@ -68,6 +68,15 @@ func WithTSOAllocator(alloc TimestampAllocator) CoordinatorOption { } } +// TimestampAllocator exposes the configured allocator to coordinator +// decorators without widening the Coordinator interface. +func (c *Coordinate) TimestampAllocator() TimestampAllocator { + if c == nil { + return nil + } + return c.tsAllocator +} + // LeaseReadObserver records lease-read fast-path vs slow-path outcomes // without coupling kv to a concrete monitoring backend. It is called once // per LeaseRead invocation that actually evaluates the lease (the initial @@ -230,6 +239,13 @@ type Coordinate struct { } var _ Coordinator = (*Coordinate)(nil) +var _ AppliedReadTimestampVoucher = (*Coordinate)(nil) + +// VouchAppliedReadTimestamp is a no-op for the single-group coordinator. The +// sharded coordinator consumes vouchers before cross-group StartTS validation. +func (c *Coordinate) VouchAppliedReadTimestamp(uint64) error { + return nil +} type Coordinator interface { Dispatch(ctx context.Context, reqs *OperationGroup[OP]) (*CoordinateResponse, error) diff --git a/kv/keyviz_label.go b/kv/keyviz_label.go index eb1949076..ba40e30d5 100644 --- a/kv/keyviz_label.go +++ b/kv/keyviz_label.go @@ -90,6 +90,19 @@ func (c keyVizLabeledCoordinator) RaftLeaderForKey(key []byte) string { func (c keyVizLabeledCoordinator) Clock() *HLC { return c.inner.Clock() } +func (c keyVizLabeledCoordinator) TimestampAllocator() TimestampAllocator { + alloc, _ := TimestampAllocatorThrough(c.inner) + return alloc +} + +func (c keyVizLabeledCoordinator) VouchAppliedReadTimestamp(timestamp uint64) error { + voucher, ok := c.inner.(AppliedReadTimestampVoucher) + if !ok { + return errors.WithStack(ErrTSOProtocolUnsupported) + } + return errors.WithStack(voucher.VouchAppliedReadTimestamp(timestamp)) +} + func (c keyVizLabeledCoordinator) LeaseRead(ctx context.Context) (uint64, error) { if lr, ok := c.inner.(LeaseReadableCoordinator); ok { idx, err := lr.LeaseRead(ctx) diff --git a/kv/lease_warmup_test.go b/kv/lease_warmup_test.go index 5558a913c..779266336 100644 --- a/kv/lease_warmup_test.go +++ b/kv/lease_warmup_test.go @@ -8,6 +8,7 @@ import ( "testing" "time" + "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/internal/monoclock" "github.com/bootjp/elastickv/internal/raftengine" "github.com/stretchr/testify/require" @@ -294,6 +295,33 @@ func TestShardedCoordinator_RenewHLCLeases_ProposesToEveryLedGroup(t *testing.T) "the non-default group lease must be warmed by all-group renewal") } +func TestShardedCoordinator_RenewHLCLeases_PhaseDOnlyRenewsTimestampGroup(t *testing.T) { + t.Parallel() + eng0 := newShardedLeaseEngine(50) + eng1 := newShardedLeaseEngine(100) + eng2 := newShardedLeaseEngine(200) + distEngine := distribution.NewEngine() + distEngine.UpdateRoute([]byte("a"), []byte("m"), 1) + distEngine.UpdateRoute([]byte("m"), nil, 2) + phaseD := NewTSOStateMachine(NewHLC()) + require.Nil(t, phaseD.Apply(marshalTSOCutover())) + require.Nil(t, phaseD.Apply(marshalTSOPhaseD(0))) + coord := NewShardedCoordinator(distEngine, map[uint64]*ShardGroup{ + 0: {Engine: eng0}, + 1: {Engine: eng1}, + 2: {Engine: eng2}, + }, 1, NewHLC(), nil). + WithTimestampGroup(0). + WithTSOCutoverState(phaseD) + + done := coord.renewHLCLeases(context.Background()) + requireRenewalDone(t, done) + + require.Equal(t, int32(1), eng0.proposeCalls.Load()) + require.Equal(t, int32(0), eng1.proposeCalls.Load()) + require.Equal(t, int32(0), eng2.proposeCalls.Load()) +} + 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 60427c0ab..3705fcf4c 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -364,7 +364,8 @@ const ( // surfaces unchanged so the client (or a wrapping retry harness in // the adapter) sees the failure rather than the coordinator spinning // on a persistent route shift. - composed1RetryAttempts = 1 + composed1RetryAttempts = 1 + maxAppliedReadTimestampVouches = 4096 ) // ShardedCoordinator routes operations to shard-specific raft groups. @@ -384,6 +385,14 @@ type ShardedCoordinator struct { // behavior where any locally-led shard group can issue TSO timestamps. timestampGroup uint64 timestampGroupConfigured bool + // tsoCutoverState is consensus-owned group-0 migration state. Phase D uses + // it to retire data-shard HLC renewal and reject legacy cross-shard startTS. + tsoCutoverState interface { + CutoverActive() bool + PhaseDActive() bool + } + appliedReadVoucherMu sync.Mutex + appliedReadVouchers map[uint64]uint64 // allShardGroupIDs, when configured, is the explicit set of data groups // that whole-keyspace operations must visit. It lets callers keep // non-data groups (for example a reserved timestamp group) in c.groups @@ -467,6 +476,55 @@ func (c *ShardedCoordinator) WithTSOAllocator(alloc TimestampAllocator) *Sharded return c } +// TimestampAllocator exposes the configured allocator to coordinator +// decorators without widening the Coordinator interface. +func (c *ShardedCoordinator) TimestampAllocator() TimestampAllocator { + if c == nil { + return nil + } + return c.tsAllocator +} + +// VouchAppliedReadTimestamp records one use of an audited adapter watermark. +// The bounded map prevents abandoned requests from growing process memory. +func (c *ShardedCoordinator) VouchAppliedReadTimestamp(timestamp uint64) error { + if c == nil || timestamp == 0 || timestamp == ^uint64(0) { + return errors.WithStack(ErrTSOTimestampInvalid) + } + c.appliedReadVoucherMu.Lock() + defer c.appliedReadVoucherMu.Unlock() + if c.appliedReadVouchers == nil { + c.appliedReadVouchers = make(map[uint64]uint64) + } + if uses, ok := c.appliedReadVouchers[timestamp]; ok { + c.appliedReadVouchers[timestamp] = uses + 1 + return nil + } + if len(c.appliedReadVouchers) >= maxAppliedReadTimestampVouches { + return errors.WithStack(ErrTSOReadVoucherLimit) + } + c.appliedReadVouchers[timestamp] = 1 + return nil +} + +func (c *ShardedCoordinator) consumeAppliedReadTimestampVoucher(timestamp uint64) bool { + if c == nil { + return false + } + c.appliedReadVoucherMu.Lock() + defer c.appliedReadVoucherMu.Unlock() + uses, ok := c.appliedReadVouchers[timestamp] + if !ok { + return false + } + if uses <= 1 { + delete(c.appliedReadVouchers, timestamp) + } else { + c.appliedReadVouchers[timestamp] = uses - 1 + } + return true +} + // WithTimestampGroup pins timestamp issuance leadership to one Raft group. // Callers should only enable this once a data-shard leader can redirect // timestamp allocation to that group; otherwise data leaders would stop being @@ -477,6 +535,17 @@ func (c *ShardedCoordinator) WithTimestampGroup(groupID uint64) *ShardedCoordina return c } +// WithTSOCutoverState wires the durable group-0 migration state. The state is +// read dynamically so a marker applied after startup changes renewal and +// validation behavior without a process-local mode race. +func (c *ShardedCoordinator) WithTSOCutoverState(state interface { + CutoverActive() bool + PhaseDActive() bool +}) *ShardedCoordinator { + c.tsoCutoverState = state + return c +} + // WithAllShardGroups restricts whole-keyspace operations to the supplied data // groups. When unset, the coordinator preserves the legacy behaviour and uses // every group it owns. @@ -888,7 +957,17 @@ func (c *ShardedCoordinator) dispatchTxnWithComposed1Retry(ctx context.Context, c.maybeAutoPinObservedRouteVersion(reqs, callerSuppliedStartTS) for attempt := 0; attempt <= composed1RetryAttempts; attempt++ { - resp, err := c.dispatchTxn(ctx, reqs.StartTS, reqs.CommitTS, reqs.PrevCommitTS, reqs.Elems, reqs.ReadKeys, reqs.ObservedRouteVersion, reqs.KeyVizLabel) + resp, err := c.dispatchTxn( + ctx, + reqs.StartTS, + reqs.CommitTS, + reqs.PrevCommitTS, + reqs.Elems, + reqs.ReadKeys, + reqs.ObservedRouteVersion, + reqs.KeyVizLabel, + callerSuppliedStartTS, + ) if err == nil { return resp, nil } @@ -1117,7 +1196,17 @@ func (c *ShardedCoordinator) broadcastToAllGroups(ctx context.Context, requests return &CoordinateResponse{CommitIndex: maxIndex.Load()}, nil } -func (c *ShardedCoordinator) dispatchTxn(ctx context.Context, startTS uint64, commitTS uint64, prevCommitTS uint64, elems []*Elem[OP], readKeys [][]byte, observedRouteVersion uint64, label keyviz.Label) (*CoordinateResponse, error) { +func (c *ShardedCoordinator) dispatchTxn( + ctx context.Context, + startTS uint64, + commitTS uint64, + prevCommitTS uint64, + elems []*Elem[OP], + readKeys [][]byte, + observedRouteVersion uint64, + label keyviz.Label, + callerSuppliedStartTS bool, +) (*CoordinateResponse, error) { if len(readKeys) > maxReadKeys { return nil, errors.WithStack(ErrInvalidRequest) } @@ -1129,16 +1218,17 @@ func (c *ShardedCoordinator) dispatchTxn(ctx context.Context, startTS uint64, co if len(primaryKey) == 0 { return nil, errors.WithStack(ErrTxnPrimaryKeyRequired) } - - commitTS, err = c.resolveTxnCommitTS(ctx, startTS, commitTS) - if err != nil { + singleShard := len(gids) == 1 && c.allReadKeysInShard(readKeys, gids[0]) + if err := c.validateCallerSuppliedTxnStart(ctx, startTS, singleShard, callerSuppliedStartTS); err != nil { return nil, err } - if err := ValidateElemCommitTSPatches(elems, commitTS); err != nil { + + commitTS, err = c.prepareTxnCommitTimestamp(ctx, startTS, commitTS, elems) + if err != nil { return nil, err } - if len(gids) == 1 && c.allReadKeysInShard(readKeys, gids[0]) { + if singleShard { // Fast path: all mutations and read keys are in a single shard. // Use the one-phase path without allocating a grouped-read-keys map. // If any read key belongs to a different shard the 2PC path is required @@ -1152,9 +1242,47 @@ func (c *ShardedCoordinator) dispatchTxn(ctx context.Context, startTS uint64, co return c.dispatchMultiShardTxn(ctx, startTS, commitTS, prevCommitTS, primaryKey, grouped, gids, readKeys, observedRouteVersion) } +func (c *ShardedCoordinator) validateCallerSuppliedTxnStart(ctx context.Context, startTS uint64, singleShard, callerSupplied bool) error { + if !callerSupplied { + return nil + } + if c.consumeAppliedReadTimestampVoucher(startTS) || singleShard { + return nil + } + return c.validateCrossShardReadTimestamp(ctx, startTS) +} + +func (c *ShardedCoordinator) prepareTxnCommitTimestamp(ctx context.Context, startTS, commitTS uint64, elems []*Elem[OP]) (uint64, error) { + resolved, err := c.resolveTxnCommitTS(ctx, startTS, commitTS) + if err != nil { + return 0, err + } + if err := ValidateElemCommitTSPatches(elems, resolved); err != nil { + return 0, err + } + return resolved, nil +} + +func (c *ShardedCoordinator) validateCrossShardReadTimestamp( + ctx context.Context, + startTS uint64, +) error { + validator, _, required, err := phaseDTimestampValidator(c.tsAllocator) + if err != nil { + return errors.Wrap(err, "cross-shard read timestamp validator is unavailable") + } + if !required { + return nil + } + if err := validator.ValidateDurableTimestamp(ctx, startTS); err != nil { + return errors.Wrap(err, "validate cross-shard read/start timestamp") + } + return nil +} + // dispatchMultiShardTxn runs the 2PC path. Extracted from dispatchTxn to keep -// that function under the cyclop budget after the prevCommitTS reject (codex -// P2 round-10) was added; the multi-shard branch already carries five linear +// that function under the cyclop budget after the prevCommitTS reject was +// added; the multi-shard branch already carries five linear // error checks (groupReadKeys, prewrite, commitPrimary, abortCleanup, // commitSecondaries) that pushed the parent over the 10-edge limit. func (c *ShardedCoordinator) dispatchMultiShardTxn(ctx context.Context, startTS, commitTS, prevCommitTS uint64, primaryKey []byte, grouped map[uint64][]*pb.Mutation, gids []uint64, readKeys [][]byte, observedRouteVersion uint64) (*CoordinateResponse, error) { @@ -1507,6 +1635,10 @@ func (c *ShardedCoordinator) allocateTimestampAfter(ctx context.Context, label s } return nextTimestampFromAllocator(ctx, c.tsAllocator, label) } + if c.tsoCutoverState != nil && c.tsoCutoverState.PhaseDActive() { + return 0, errors.Wrap(ErrTSOAllocatorRequired, + "legacy HLC issuance is disabled by durable TSO phase D") + } if c.clock == nil { return 0, errors.Wrap(ErrTSOClockNil, label) } @@ -2276,10 +2408,7 @@ func (c *ShardedCoordinator) renewHLCLeases(ctx context.Context) <-chan struct{} done := make(chan struct{}) var wg sync.WaitGroup for gid, group := range c.groups { - if group == nil || group.Engine == nil { - continue - } - if group.Engine.State() != raftengine.StateLeader { + if !c.shouldRenewHLCGroup(gid, group) { continue } if !c.startHLCLeaseRenewal(gid) { @@ -2301,6 +2430,14 @@ func (c *ShardedCoordinator) renewHLCLeases(ctx context.Context) <-chan struct{} return done } +func (c *ShardedCoordinator) shouldRenewHLCGroup(gid uint64, group *ShardGroup) bool { + if c.tsoCutoverState != nil && c.tsoCutoverState.PhaseDActive() && + (!c.timestampGroupConfigured || gid != c.timestampGroup) { + return false + } + return group != nil && group.Engine != nil && group.Engine.State() == raftengine.StateLeader +} + func (c *ShardedCoordinator) startHLCLeaseRenewal(gid uint64) bool { c.hlcRenewalMu.Lock() defer c.hlcRenewalMu.Unlock() diff --git a/kv/sharded_coordinator_txn_test.go b/kv/sharded_coordinator_txn_test.go index 42b65f328..82826272a 100644 --- a/kv/sharded_coordinator_txn_test.go +++ b/kv/sharded_coordinator_txn_test.go @@ -203,6 +203,219 @@ func TestShardedCoordinatorDispatchTxn_CrossShardPhasesAndCommitIndex(t *testing require.Zero(t, binary.BigEndian.Uint64(value2[4:12])) } +func TestShardedCoordinatorDispatchTxn_PhaseDRejectsInvalidCallerStartTSBeforeProposal(t *testing.T) { + t.Parallel() + coord, g1Txn, g2Txn, alloc := newPhaseDCrossShardCoordinator(t, ErrTSOTimestampInvalid) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + IsTxn: true, + StartTS: 10, + Elems: []*Elem[OP]{ + {Op: Put, Key: []byte("b"), Value: []byte("v1")}, + {Op: Put, Key: []byte("x"), Value: []byte("v2")}, + }, + }) + require.ErrorIs(t, err, ErrTSOTimestampInvalid) + require.Equal(t, uint64(1), alloc.validateCalls.Load()) + require.Empty(t, g1Txn.requests) + require.Empty(t, g2Txn.requests) +} + +func TestShardedCoordinatorDispatchTxn_PhaseDAcceptsValidatedCallerStartTS(t *testing.T) { + t.Parallel() + coord, g1Txn, g2Txn, alloc := newPhaseDCrossShardCoordinator(t, nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + IsTxn: true, + StartTS: 100, + Elems: []*Elem[OP]{ + {Op: Put, Key: []byte("b"), Value: []byte("v1")}, + {Op: Put, Key: []byte("x"), Value: []byte("v2")}, + }, + }) + require.NoError(t, err) + require.Equal(t, uint64(1), alloc.validateCalls.Load()) + require.Len(t, g1Txn.requests, 2) + require.Len(t, g2Txn.requests, 2) + require.Equal(t, uint64(100), g1Txn.requests[0].Ts) + require.Equal(t, uint64(100), g2Txn.requests[0].Ts) +} + +func TestShardedCoordinatorDispatchTxn_PhaseDAcceptsVouchedAppliedWatermarkOnce(t *testing.T) { + t.Parallel() + prePhaseDErr := errors.Join(ErrTSOTimestampInvalid, ErrTSOTimestampPrePhaseD) + coord, g1Txn, g2Txn, alloc := newPhaseDCrossShardCoordinator(t, prePhaseDErr) + + readTS, err := BeginReadTimestampThrough(context.Background(), coord, 10, "vouch applied watermark") + require.NoError(t, err) + require.Equal(t, uint64(10), readTS.Timestamp()) + require.Equal(t, uint64(1), alloc.validateCalls.Load()) + + request := func() *OperationGroup[OP] { + return &OperationGroup[OP]{ + IsTxn: true, + StartTS: readTS.Timestamp(), + Elems: []*Elem[OP]{ + {Op: Put, Key: []byte("b"), Value: []byte("v1")}, + {Op: Put, Key: []byte("x"), Value: []byte("v2")}, + }, + } + } + _, err = coord.Dispatch(context.Background(), request()) + require.NoError(t, err) + require.Equal(t, uint64(1), alloc.validateCalls.Load(), "voucher must bypass numeric Phase-D validation") + require.Len(t, g1Txn.requests, 2) + require.Len(t, g2Txn.requests, 2) + + _, err = coord.Dispatch(context.Background(), request()) + require.ErrorIs(t, err, ErrTSOTimestampPrePhaseD) + require.Equal(t, uint64(2), alloc.validateCalls.Load(), "voucher must be single-use") +} + +func TestDispatchWithReadTimestampVouchesEveryBoundDispatch(t *testing.T) { + t.Parallel() + prePhaseDErr := errors.Join(ErrTSOTimestampInvalid, ErrTSOTimestampPrePhaseD) + coord, g1Txn, g2Txn, alloc := newPhaseDCrossShardCoordinator(t, prePhaseDErr) + + readTimestamp, err := BeginReadTimestampThrough(context.Background(), coord, 10, "vouch reused applied watermark") + require.NoError(t, err) + ctx := readTimestamp.WithDispatchVoucher(context.Background()) + request := func(startTS uint64) *OperationGroup[OP] { + return &OperationGroup[OP]{ + IsTxn: true, + StartTS: startTS, + Elems: []*Elem[OP]{ + {Op: Put, Key: []byte("b"), Value: []byte("v1")}, + {Op: Put, Key: []byte("x"), Value: []byte("v2")}, + }, + } + } + + for range 2 { + _, err = DispatchWithReadTimestamp(ctx, coord, request(readTimestamp.Timestamp())) + require.NoError(t, err) + } + require.Equal(t, uint64(1), alloc.validateCalls.Load(), "each dispatch must consume a reserved voucher") + require.Len(t, g1Txn.requests, 4) + require.Len(t, g2Txn.requests, 4) + + _, err = DispatchWithReadTimestamp(ctx, coord, request(readTimestamp.Timestamp()+1)) + require.ErrorIs(t, err, ErrTSOTimestampInvalid, "the bound capability must not authorize another timestamp") + + _, err = coord.Dispatch(context.Background(), request(readTimestamp.Timestamp())) + require.ErrorIs(t, err, ErrTSOTimestampPrePhaseD, "no unused voucher may remain after the bound dispatches") + require.Equal(t, uint64(2), alloc.validateCalls.Load()) +} + +func TestReadTimestampVoucherBindingShadowsParentCapability(t *testing.T) { + prePhaseDErr := errors.Join(ErrTSOTimestampInvalid, ErrTSOTimestampPrePhaseD) + coord, _, _, alloc := newPhaseDCrossShardCoordinator(t, prePhaseDErr) + + oldRead, err := BeginReadTimestampThrough(context.Background(), coord, 10, "vouch old applied watermark") + require.NoError(t, err) + ctx := oldRead.WithDispatchVoucher(context.Background()) + + alloc.validateErr = nil + currentRead, err := BeginReadTimestampThrough(ctx, coord, 100, "validate current durable watermark") + require.NoError(t, err) + ctx = currentRead.WithDispatchVoucher(ctx) + + _, err = DispatchWithReadTimestamp(ctx, coord, &OperationGroup[OP]{ + IsTxn: true, + StartTS: currentRead.Timestamp(), + Elems: []*Elem[OP]{ + {Op: Put, Key: []byte("b"), Value: []byte("v1")}, + {Op: Put, Key: []byte("x"), Value: []byte("v2")}, + }, + }) + require.NoError(t, err, "the current timestamp must shadow the parent capability") +} + +func TestShardedCoordinatorDispatchTxn_PhaseDPreservesSingleShardCallerStartTS(t *testing.T) { + t.Parallel() + coord, g1Txn, _, alloc := newPhaseDCrossShardCoordinator(t, ErrTSOTimestampInvalid) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + IsTxn: true, + StartTS: 10, + Elems: []*Elem[OP]{ + {Op: Put, Key: []byte("b"), Value: []byte("v1")}, + }, + }) + require.NoError(t, err) + require.Zero(t, alloc.validateCalls.Load()) + require.Len(t, g1Txn.requests, 1) +} + +func TestShardedCoordinatorDispatchTxn_PrePhaseDPreservesCrossShardCallerStartTS(t *testing.T) { + t.Parallel() + coord, g1Txn, g2Txn, alloc := newPhaseDCrossShardCoordinator(t, ErrTSOTimestampInvalid) + coord.WithTSOCutoverState(NewTSOStateMachine(NewHLC())) + alloc.phaseDActive = false + alloc.phaseDRequired = false + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + IsTxn: true, + StartTS: 10, + Elems: []*Elem[OP]{ + {Op: Put, Key: []byte("b"), Value: []byte("v1")}, + {Op: Put, Key: []byte("x"), Value: []byte("v2")}, + }, + }) + require.NoError(t, err) + require.Zero(t, alloc.validateCalls.Load()) + require.Len(t, g1Txn.requests, 2) + require.Len(t, g2Txn.requests, 2) +} + +func TestShardedCoordinatorDispatchTxn_PhaseDActivationValidatesBeforeLocalMarker(t *testing.T) { + t.Parallel() + coord, g1Txn, g2Txn, alloc := newPhaseDCrossShardCoordinator(t, ErrTSOTimestampInvalid) + coord.WithTSOCutoverState(NewTSOStateMachine(NewHLC())) + alloc.phaseDActive = false + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + IsTxn: true, + StartTS: 10, + Elems: []*Elem[OP]{ + {Op: Put, Key: []byte("b"), Value: []byte("v1")}, + {Op: Put, Key: []byte("x"), Value: []byte("v2")}, + }, + }) + require.ErrorIs(t, err, ErrTSOTimestampInvalid) + require.Equal(t, uint64(1), alloc.validateCalls.Load()) + require.Empty(t, g1Txn.requests) + require.Empty(t, g2Txn.requests) +} + +func newPhaseDCrossShardCoordinator( + t *testing.T, + validateErr error, +) (*ShardedCoordinator, *recordingTransactional, *recordingTransactional, *phaseDTestAllocator) { + t.Helper() + engine := distribution.NewEngine() + engine.UpdateRoute([]byte("a"), []byte("m"), 1) + engine.UpdateRoute([]byte("m"), nil, 2) + g1Txn := &recordingTransactional{} + g2Txn := &recordingTransactional{} + alloc := &phaseDTestAllocator{ + next: 200, + phaseDActive: true, + phaseDRequired: true, + validateErr: validateErr, + } + state := NewTSOStateMachine(NewHLC()) + require.Nil(t, state.Apply(marshalTSOCutover())) + require.Nil(t, state.Apply(marshalTSOPhaseD(0))) + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + }, 1, NewHLC(), nil). + WithTSOAllocator(alloc). + WithTSOCutoverState(state) + return coord, g1Txn, g2Txn, alloc +} + func TestShardedCoordinatorDispatchTxn_SingleShardUsesOnePhase(t *testing.T) { t.Parallel() diff --git a/kv/tso.go b/kv/tso.go index f9646951c..c92840db1 100644 --- a/kv/tso.go +++ b/kv/tso.go @@ -16,10 +16,14 @@ const defaultTSOLeaderPollInterval = 25 * time.Millisecond const MaxTSOBatchSize = maxHLCBatchSize var ( - ErrTSOAllocatorRequired = errors.New("tso: allocator is required") - ErrTSOCoordinatorNil = errors.New("tso: coordinator is required") - ErrTSOClockNil = errors.New("tso: coordinator clock is nil") - ErrInvalidTSOBatchSize = errors.New("tso: invalid batch size") + ErrTSOAllocatorRequired = errors.New("tso: allocator is required") + ErrTSOCoordinatorNil = errors.New("tso: coordinator is required") + ErrTSOClockNil = errors.New("tso: coordinator clock is nil") + ErrInvalidTSOBatchSize = errors.New("tso: invalid batch size") + ErrTSOPhaseDInactive = errors.New("tso: phase D is not active") + ErrTSOTimestampInvalid = errors.New("tso: timestamp is not a durable phase-D allocation") + ErrTSOTimestampPrePhaseD = errors.New("tso: timestamp predates phase D") + ErrTSOReadVoucherLimit = errors.New("tso: applied read timestamp voucher limit reached") ) // TSOAllocator issues globally monotonic timestamps. NextBatch returns the @@ -48,6 +52,100 @@ type TimestampAfterAllocator interface { NextAfter(ctx context.Context, min uint64) (uint64, error) } +// DurableTimestampValidator verifies that a timestamp belongs to the durable +// post-Phase-D allocation range owned by the dedicated TSO group. +type DurableTimestampValidator interface { + ValidateDurableTimestamp(context.Context, uint64) error +} + +// TSOPhaseDState exposes the one-way Phase-D state to coordinators and adapter +// migration helpers without coupling them to TSOStateMachine. +type TSOPhaseDState interface { + PhaseDActive() bool + PhaseDRequired() bool +} + +// AppliedReadTimestampVoucher records an adapter-provided applied watermark. +// It is a process-local capability used only to distinguish audited adapter +// snapshots from arbitrary caller-supplied StartTS values during Phase D. +type AppliedReadTimestampVoucher interface { + VouchAppliedReadTimestamp(uint64) error +} + +// ReadTimestamp is the adapter-side result of beginning a transaction snapshot. +// When it represents an applied pre-Phase-D watermark, it also carries a +// process-local capability that can reserve exactly one coordinator voucher per +// DispatchWithReadTimestamp call. The capability cannot be constructed outside +// this package because both the timestamp and voucher state are private. +type ReadTimestamp struct { + timestamp uint64 + voucher *appliedReadDispatchVoucher +} + +func (t ReadTimestamp) Timestamp() uint64 { + return t.timestamp +} + +type appliedReadDispatchVoucher struct { + mu sync.Mutex + prepared uint64 +} + +type appliedReadDispatchVoucherContextKey struct{} + +// WithDispatchVoucher binds this read timestamp's process-local capability to +// ctx. A timestamp without a voucher is still bound so it shadows any parent +// capability instead of accidentally inheriting authority for an older read. +func (t ReadTimestamp) WithDispatchVoucher(ctx context.Context) context.Context { + return context.WithValue(nonNilTSOContext(ctx), appliedReadDispatchVoucherContextKey{}, t) +} + +// DispatchWithReadTimestamp dispatches an OCC operation under the applied-read +// capability bound by ReadTimestamp.WithDispatchVoucher. The first dispatch +// consumes the voucher reserved by BeginReadTimestampThrough; every subsequent +// dispatch reserves one additional use immediately before dispatching. +func DispatchWithReadTimestamp( + ctx context.Context, + coord Coordinator, + reqs *OperationGroup[OP], +) (*CoordinateResponse, error) { + if ctx == nil { + resp, err := coord.Dispatch(ctx, reqs) + return resp, errors.WithStack(err) + } + readTimestamp, ok := ctx.Value(appliedReadDispatchVoucherContextKey{}).(ReadTimestamp) + if !ok || readTimestamp.voucher == nil { + resp, err := coord.Dispatch(ctx, reqs) + return resp, errors.WithStack(err) + } + if reqs == nil || reqs.StartTS != readTimestamp.timestamp { + return nil, errors.WithStack(ErrTSOTimestampInvalid) + } + if err := readTimestamp.voucher.prepare(coord, readTimestamp.timestamp); err != nil { + return nil, err + } + resp, err := coord.Dispatch(ctx, reqs) + return resp, errors.WithStack(err) +} + +func (v *appliedReadDispatchVoucher) prepare(coord Coordinator, timestamp uint64) error { + v.mu.Lock() + defer v.mu.Unlock() + if v.prepared == 0 { + v.prepared = 1 + return nil + } + voucher, ok := coord.(AppliedReadTimestampVoucher) + if !ok { + return errors.WithStack(ErrTSOProtocolUnsupported) + } + if err := voucher.VouchAppliedReadTimestamp(timestamp); err != nil { + return errors.WithStack(err) + } + v.prepared++ + return nil +} + type tsoBatchAfterAllocator interface { NextBatchAfter(ctx context.Context, n int, min uint64) (uint64, error) } @@ -129,6 +227,97 @@ func coordinatorTimestampAllocator(coord Coordinator) (TimestampAllocator, bool) return TimestampAllocatorThrough(coord) } +// BeginReadTimestampThrough preserves the caller's applied-snapshot watermark. +// Once Phase D is requested, this boundary activates it before validation. An +// applied pre-D watermark receives a bounded one-use coordinator voucher; +// arbitrary caller timestamps remain subject to group-0 numeric validation. +// The returned timestamp must be used for every read and OperationGroup.StartTS. +func BeginReadTimestampThrough( + ctx context.Context, + coord Coordinator, + legacyTimestamp uint64, + label string, +) (ReadTimestamp, error) { + alloc, ok := coordinatorTimestampAllocator(coord) + if !ok { + return ReadTimestamp{timestamp: legacyTimestamp}, nil + } + validator, phaseD, phaseDRequired, err := phaseDTimestampValidator(alloc) + if err != nil { + return ReadTimestamp{}, errors.Wrap(err, label) + } + if !phaseDRequired { + return ReadTimestamp{timestamp: legacyTimestamp}, nil + } + if legacyTimestamp == 0 || legacyTimestamp == ^uint64(0) { + return ReadTimestamp{}, errors.Wrap(ErrTSOTimestampInvalid, label) + } + if err := activatePhaseDForRead(ctx, alloc, phaseD, label); err != nil { + return ReadTimestamp{}, err + } + vouched, err := validateAppliedReadTimestamp(ctx, coord, validator, legacyTimestamp, label) + if err != nil { + return ReadTimestamp{}, err + } + readTimestamp := ReadTimestamp{timestamp: legacyTimestamp} + if vouched { + readTimestamp.voucher = &appliedReadDispatchVoucher{} + } + return readTimestamp, nil +} + +func activatePhaseDForRead( + ctx context.Context, + alloc TimestampAllocator, + phaseD TSOPhaseDState, + label string, +) error { + if phaseD.PhaseDActive() { + return nil + } + // Reserve and discard one post-D timestamp to commit the marker. The + // caller's applied watermark remains the read snapshot; using this fresh + // allocation for reads could run ahead of data-group apply. + _, err := nextTimestampFromAllocator(nonNilTSOContext(ctx), alloc, label+": activate phase D") + return err +} + +func validateAppliedReadTimestamp( + ctx context.Context, + coord Coordinator, + validator DurableTimestampValidator, + timestamp uint64, + label string, +) (bool, error) { + err := validator.ValidateDurableTimestamp(nonNilTSOContext(ctx), timestamp) + if err == nil { + return false, nil + } + if !errors.Is(err, ErrTSOTimestampPrePhaseD) { + return false, errors.Wrap(err, label) + } + voucher, ok := coord.(AppliedReadTimestampVoucher) + if !ok { + return false, errors.Wrap(ErrTSOProtocolUnsupported, label+": applied read voucher unavailable") + } + if err := voucher.VouchAppliedReadTimestamp(timestamp); err != nil { + return false, errors.Wrap(err, label) + } + return true, nil +} + +func phaseDTimestampValidator(alloc TimestampAllocator) (DurableTimestampValidator, TSOPhaseDState, bool, error) { + phaseD, ok := alloc.(TSOPhaseDState) + if !ok || (!phaseD.PhaseDRequired() && !phaseD.PhaseDActive()) { + return nil, nil, false, nil + } + validator, ok := alloc.(DurableTimestampValidator) + if !ok { + return nil, phaseD, false, ErrTSOProtocolUnsupported + } + return validator, phaseD, true, nil +} + func nextTimestampFromAllocator(ctx context.Context, alloc TimestampAllocator, label string) (uint64, error) { ts, err := alloc.Next(ctx) if err != nil { @@ -315,10 +504,11 @@ type windowSnapshot struct { // TSOAllocator. The hot path is lock-free: callers claim a slot with atomic Add // on the currently published window. type BatchAllocator struct { - tso TSOAllocator - batchSize int - win atomic.Pointer[windowSnapshot] - epoch atomic.Uint64 + tso TSOAllocator + batchSize int + win atomic.Pointer[windowSnapshot] + epoch atomic.Uint64 + phaseDSeen atomic.Bool mu sync.Mutex refillDone chan struct{} @@ -353,12 +543,38 @@ func (b *BatchAllocator) NextAfter(ctx context.Context, min uint64) (uint64, err return b.nextAfter(ctx, min) } +func (b *BatchAllocator) ValidateDurableTimestamp(ctx context.Context, timestamp uint64) error { + validator, ok := b.tso.(DurableTimestampValidator) + if !ok { + return errors.WithStack(ErrTSOProtocolUnsupported) + } + return errors.WithStack(validator.ValidateDurableTimestamp(ctx, timestamp)) +} + +func (b *BatchAllocator) PhaseDActive() bool { + state, ok := b.tso.(TSOPhaseDState) + return ok && state.PhaseDActive() +} + +func (b *BatchAllocator) PhaseDRequired() bool { + state, ok := b.tso.(TSOPhaseDState) + return ok && state.PhaseDRequired() +} + func (b *BatchAllocator) nextAfter(ctx context.Context, min uint64) (uint64, error) { for { if err := ctxErr(ctx); err != nil { return 0, err } + phaseDAtStart, invalidated := b.ensurePhaseDTransition() + if invalidated { + continue + } if ts, ok := b.tryWindowAfter(min); ok { + phaseDAfterClaim, invalidated := b.ensurePhaseDTransition() + if invalidated || phaseDAfterClaim != phaseDAtStart { + continue + } return ts, nil } if err := b.refill(ctx, min); err != nil { @@ -367,6 +583,21 @@ func (b *BatchAllocator) nextAfter(ctx context.Context, min uint64) (uint64, err } } +func (b *BatchAllocator) ensurePhaseDTransition() (required, invalidated bool) { + if !b.PhaseDRequired() { + return false, false + } + if b.phaseDSeen.Load() { + return true, false + } + // Publish phaseDSeen only after the old epoch is invalidated. Concurrent + // callers may invalidate redundantly, but none can observe the transition as + // complete while a pre-Phase-D window is still current. + b.Invalidate() + b.phaseDSeen.CompareAndSwap(false, true) + return true, true +} + func (b *BatchAllocator) tryWindowAfter(min uint64) (uint64, bool) { w := b.win.Load() if w == nil { diff --git a/kv/tso_fsm.go b/kv/tso_fsm.go index 238efaa27..75a32a766 100644 --- a/kv/tso_fsm.go +++ b/kv/tso_fsm.go @@ -28,9 +28,13 @@ const ( // tsoCutoverEnvelope uses the same legacy fail-closed prefix but a distinct // magic, so an encryption entry cannot become a one-way TSO state change. tsoCutoverEnvelope = "\x07TSOC\x01" - tsoSnapshotV1Len = hlcLeasePayloadLen - tsoSnapshotV2Len = hlcLeasePayloadLen * 2 - tsoSnapshotV3Len = tsoSnapshotV2Len + 1 + // tsoPhaseDEnvelope retains the rolling-upgrade halt prefix and gives the + // irreversible Phase-D marker its own exact wire identity. + tsoPhaseDEnvelope = "\x07TSOD\x01" + tsoSnapshotV1Len = hlcLeasePayloadLen + tsoSnapshotV2Len = hlcLeasePayloadLen * 2 + tsoSnapshotV3Len = tsoSnapshotV2Len + 1 + tsoSnapshotV4Len = tsoSnapshotV3Len + 1 + hlcLeasePayloadLen ) // TSOStateMachine is the minimal state machine for the dedicated timestamp @@ -43,6 +47,8 @@ type TSOStateMachine struct { ceilingMs atomic.Int64 allocationFloor atomic.Uint64 cutoverActive atomic.Bool + phaseDActive atomic.Bool + phaseDFloor atomic.Uint64 } func NewTSOStateMachine(hlc *HLC) *TSOStateMachine { @@ -60,6 +66,8 @@ func (f *TSOStateMachine) Apply(data []byte) any { return f.applyAllocationFloorEntry(data) case bytes.Equal(data, []byte(tsoCutoverEnvelope)): return f.applyCutoverEntry(data) + case bytes.HasPrefix(data, []byte(tsoPhaseDEnvelope)): + return f.applyPhaseDEntry(data) case data[0] >= fsmwire.OpEncryptionMin && data[0] <= fsmwire.OpEncryptionMax: return rejectLegacyTSOEncryptionEntry(data) default: @@ -134,6 +142,33 @@ func (f *TSOStateMachine) applyCutoverEntry(data []byte) any { return nil } +func (f *TSOStateMachine) applyPhaseDEntry(data []byte) any { + expectedLen := len(tsoPhaseDEnvelope) + hlcLeasePayloadLen + if len(data) != expectedLen { + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "expected TSO phase-D entry length %d, got %d", expectedLen, len(data))) + } + if f == nil { + return nil + } + if !f.cutoverActive.Load() { + return haltErr(errors.Wrap(ErrTSOStateMachineInvalidEntry, + "TSO phase-D marker requires the durable cutover marker")) + } + floor := binary.BigEndian.Uint64(data[len(tsoPhaseDEnvelope):]) + if f.phaseDActive.Load() && f.phaseDFloor.Load() != floor { + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "TSO phase-D floor changed from %d to %d", f.phaseDFloor.Load(), floor)) + } + if !f.phaseDActive.Load() && floor < f.allocationFloor.Load() { + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "TSO phase-D floor %d is below allocation floor %d", floor, f.allocationFloor.Load())) + } + f.phaseDFloor.Store(floor) + f.phaseDActive.Store(true) + return nil +} + // AllocationFloor returns the highest timestamp window end applied by the // dedicated TSO group. It is consensus-owned state, unlike HLC.Current(). func (f *TSOStateMachine) AllocationFloor() uint64 { @@ -150,19 +185,42 @@ func (f *TSOStateMachine) CutoverActive() bool { return f != nil && f.cutoverActive.Load() } +// PhaseDActive reports whether the compatibility window has been durably +// closed. Once active, data-shard HLC renewal and caller-supplied cross-shard +// timestamps may no longer use legacy issuance semantics. +func (f *TSOStateMachine) PhaseDActive() bool { + return f != nil && f.phaseDActive.Load() +} + +// PhaseDFloor is the highest allocation floor that existed when Phase D was +// activated. Only timestamps reserved strictly above it are valid M7 durable +// read/start allocations. +func (f *TSOStateMachine) PhaseDFloor() uint64 { + if f == nil { + return 0 + } + return f.phaseDFloor.Load() +} + func (f *TSOStateMachine) Snapshot() (raftengine.Snapshot, error) { var ceilingMs int64 var allocationFloor uint64 var cutoverActive bool + var phaseDActive bool + var phaseDFloor uint64 if f != nil { ceilingMs = f.ceilingMs.Load() allocationFloor = f.allocationFloor.Load() cutoverActive = f.cutoverActive.Load() + phaseDActive = f.phaseDActive.Load() + phaseDFloor = f.phaseDFloor.Load() } return &tsoFSMSnapshot{ ceilingMs: ceilingMs, allocationFloor: allocationFloor, cutoverActive: cutoverActive, + phaseDActive: phaseDActive, + phaseDFloor: phaseDFloor, }, nil } @@ -174,12 +232,12 @@ func (f *TSOStateMachine) Restore(r io.Reader) error { if legacy, err := restoreLegacyKVFSMSnapshot(f, br); legacy || err != nil { return err } - ceilingMs, allocationFloor, cutoverActive, err := readTSOSnapshotState(br) + ceilingMs, allocationFloor, cutoverActive, phaseDActive, phaseDFloor, err := readTSOSnapshotState(br) if err != nil { return err } if f != nil { - f.restoreSnapshotState(ceilingMs, allocationFloor, cutoverActive) + f.restoreSnapshotState(ceilingMs, allocationFloor, cutoverActive, phaseDActive, phaseDFloor) } return nil } @@ -191,28 +249,30 @@ func tsoSnapshotReader(r io.Reader) *bufio.Reader { return bufio.NewReader(r) } -func readTSOSnapshotState(br *bufio.Reader) (int64, uint64, bool, error) { - payload, err := io.ReadAll(io.LimitReader(br, tsoSnapshotV3Len+1)) +func readTSOSnapshotState(br *bufio.Reader) (int64, uint64, bool, bool, uint64, error) { + payload, err := io.ReadAll(io.LimitReader(br, tsoSnapshotV4Len+1)) if err != nil { - return 0, 0, false, errors.Wrap(err, "restore tso fsm snapshot") + return 0, 0, false, false, 0, errors.Wrap(err, "restore tso fsm snapshot") } - ceilingMs, allocationFloor, cutoverActive, legacySnapshot, err := decodeTSOSnapshotPayload(payload) + ceilingMs, allocationFloor, cutoverActive, phaseDActive, phaseDFloor, legacySnapshot, err := decodeTSOSnapshotPayload(payload) if err != nil { - return 0, 0, false, err + return 0, 0, false, false, 0, err } if ceilingMs < 0 { - return 0, 0, false, errors.Wrapf(ErrTSOStateMachineInvalidEntry, "tso fsm snapshot: negative ceiling %d", ceilingMs) + return 0, 0, false, false, 0, errors.Wrapf(ErrTSOStateMachineInvalidEntry, "tso fsm snapshot: negative ceiling %d", ceilingMs) } if legacySnapshot && ceilingMs > 0 { allocationFloor = tsoLeaseAllocationFloor(ceilingMs) } - return ceilingMs, allocationFloor, cutoverActive, nil + return ceilingMs, allocationFloor, cutoverActive, phaseDActive, phaseDFloor, nil } -func decodeTSOSnapshotPayload(payload []byte) (int64, uint64, bool, bool, error) { +func decodeTSOSnapshotPayload(payload []byte) (int64, uint64, bool, bool, uint64, bool, error) { var ceilingMs int64 var allocationFloor uint64 var cutoverActive bool + var phaseDActive bool + var phaseDFloor uint64 var legacySnapshot bool switch len(payload) { case tsoSnapshotV1Len: @@ -227,17 +287,46 @@ func decodeTSOSnapshotPayload(payload []byte) (int64, uint64, bool, bool, error) var err error cutoverActive, err = decodeTSOCutoverByte(payload[tsoSnapshotV2Len]) if err != nil { - return 0, 0, false, false, err + return 0, 0, false, false, 0, false, err } + case tsoSnapshotV4Len: + return decodeTSOSnapshotV4(payload) default: - return 0, 0, false, false, errors.Wrapf(ErrTSOStateMachineInvalidEntry, - "tso fsm snapshot: expected %d, %d, or %d bytes, got %d", - tsoSnapshotV1Len, tsoSnapshotV2Len, tsoSnapshotV3Len, len(payload)) + return 0, 0, false, false, 0, false, errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "tso fsm snapshot: expected %d, %d, %d, or %d bytes, got %d", + tsoSnapshotV1Len, tsoSnapshotV2Len, tsoSnapshotV3Len, tsoSnapshotV4Len, len(payload)) + } + return ceilingMs, allocationFloor, cutoverActive, phaseDActive, phaseDFloor, legacySnapshot, nil +} + +func decodeTSOSnapshotV4(payload []byte) (int64, uint64, bool, bool, uint64, bool, error) { + ceilingMs := int64(binary.BigEndian.Uint64(payload[:hlcLeasePayloadLen])) //nolint:gosec // snapshot value. + allocationFloor := binary.BigEndian.Uint64(payload[hlcLeasePayloadLen:tsoSnapshotV2Len]) + cutoverActive, err := decodeTSOCutoverByte(payload[tsoSnapshotV2Len]) + if err != nil { + return 0, 0, false, false, 0, false, err + } + phaseDActive, err := decodeTSOBooleanByte("phase-D", payload[tsoSnapshotV3Len]) + if err != nil { + return 0, 0, false, false, 0, false, err + } + phaseDFloor := binary.BigEndian.Uint64(payload[tsoSnapshotV3Len+1:]) + if phaseDActive && !cutoverActive { + return 0, 0, false, false, 0, false, errors.Wrap(ErrTSOStateMachineInvalidEntry, + "tso fsm snapshot: phase-D active without cutover") } - return ceilingMs, allocationFloor, cutoverActive, legacySnapshot, nil + if !phaseDActive && phaseDFloor != 0 { + return 0, 0, false, false, 0, false, errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "tso fsm snapshot: inactive phase-D has floor %d", phaseDFloor) + } + return ceilingMs, allocationFloor, cutoverActive, phaseDActive, phaseDFloor, false, nil } func decodeTSOCutoverByte(value byte) (bool, error) { + return decodeTSOBooleanByte("cutover", value) +} + +func decodeTSOBooleanByte(name string, value byte) (bool, error) { switch value { case 0: return false, nil @@ -245,7 +334,7 @@ func decodeTSOCutoverByte(value byte) (bool, error) { return true, nil default: return false, errors.Wrapf(ErrTSOStateMachineInvalidEntry, - "tso fsm snapshot: invalid cutover byte %d", value) + "tso fsm snapshot: invalid %s byte %d", name, value) } } @@ -272,7 +361,7 @@ func restoreLegacyKVFSMSnapshot(f *TSOStateMachine, br *bufio.Reader) (bool, err if f == nil || ceilingMs == 0 { return true, nil } - f.restoreSnapshotState(ceilingMs, tsoLeaseAllocationFloor(ceilingMs), false) + f.restoreSnapshotState(ceilingMs, tsoLeaseAllocationFloor(ceilingMs), false, false, 0) return true, nil } @@ -302,7 +391,9 @@ func (f *TSOStateMachine) IsVolatileOnlyPayload(payload []byte) bool { } return len(payload) == hlcLeaseEntryLen && payload[0] == raftEncodeHLCLease || len(payload) == len(tsoAllocationFloorEnvelope)+hlcLeasePayloadLen && - bytes.HasPrefix(payload, []byte(tsoAllocationFloorEnvelope)) + bytes.HasPrefix(payload, []byte(tsoAllocationFloorEnvelope)) || + len(payload) == len(tsoPhaseDEnvelope)+hlcLeasePayloadLen && + bytes.HasPrefix(payload, []byte(tsoPhaseDEnvelope)) } func (f *TSOStateMachine) applyLeaseCeiling(ceilingMs int64) { @@ -325,7 +416,13 @@ func (f *TSOStateMachine) applyAllocationFloor(floor uint64) { } } -func (f *TSOStateMachine) restoreSnapshotState(ceilingMs int64, allocationFloor uint64, cutoverActive bool) { +func (f *TSOStateMachine) restoreSnapshotState( + ceilingMs int64, + allocationFloor uint64, + cutoverActive bool, + phaseDActive bool, + phaseDFloor uint64, +) { if f == nil { return } @@ -338,6 +435,10 @@ func (f *TSOStateMachine) restoreSnapshotState(ceilingMs int64, allocationFloor if cutoverActive { f.cutoverActive.Store(true) } + if phaseDActive { + f.phaseDFloor.Store(phaseDFloor) + f.phaseDActive.Store(true) + } if f.hlc != nil { if currentCeiling := f.ceilingMs.Load(); currentCeiling > 0 { f.hlc.SetPhysicalCeiling(currentCeiling) @@ -387,10 +488,19 @@ func marshalTSOCutover() []byte { return []byte(tsoCutoverEnvelope) } +func marshalTSOPhaseD(floor uint64) []byte { + out := make([]byte, len(tsoPhaseDEnvelope)+hlcLeasePayloadLen) + copy(out, tsoPhaseDEnvelope) + binary.BigEndian.PutUint64(out[len(tsoPhaseDEnvelope):], floor) + return out +} + type tsoFSMSnapshot struct { ceilingMs int64 allocationFloor uint64 cutoverActive bool + phaseDActive bool + phaseDFloor uint64 } func (s *tsoFSMSnapshot) WriteTo(w io.Writer) (int64, error) { @@ -400,18 +510,30 @@ func (s *tsoFSMSnapshot) WriteTo(w io.Writer) (int64, error) { var ceilingMs int64 var allocationFloor uint64 var cutoverActive bool + var phaseDActive bool + var phaseDFloor uint64 if s != nil { ceilingMs = s.ceilingMs allocationFloor = s.allocationFloor cutoverActive = s.cutoverActive + phaseDActive = s.phaseDActive + phaseDFloor = s.phaseDFloor } - var buf [tsoSnapshotV3Len]byte - binary.BigEndian.PutUint64(buf[:], uint64(ceilingMs)) //nolint:gosec // ceilingMs is a Unix ms timestamp. + snapshotLen := tsoSnapshotV3Len + if phaseDActive { + snapshotLen = tsoSnapshotV4Len + } + buf := make([]byte, snapshotLen) + binary.BigEndian.PutUint64(buf, uint64(ceilingMs)) //nolint:gosec // ceilingMs is a Unix ms timestamp. binary.BigEndian.PutUint64(buf[hlcLeasePayloadLen:tsoSnapshotV2Len], allocationFloor) if cutoverActive { buf[tsoSnapshotV2Len] = 1 } - n, err := w.Write(buf[:]) + if phaseDActive { + buf[tsoSnapshotV3Len] = 1 + binary.BigEndian.PutUint64(buf[tsoSnapshotV3Len+1:], phaseDFloor) + } + n, err := w.Write(buf) if err != nil { return int64(n), errors.Wrap(err, "write tso fsm snapshot") } diff --git a/kv/tso_fsm_test.go b/kv/tso_fsm_test.go index 3cd346558..52666e144 100644 --- a/kv/tso_fsm_test.go +++ b/kv/tso_fsm_test.go @@ -170,6 +170,41 @@ func TestTSOStateMachineAppliesOneWayCutoverMarker(t *testing.T) { require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) } +func TestTSOStateMachineAppliesOneWayPhaseDMarker(t *testing.T) { + t.Parallel() + + const floor = uint64(1234) + fsm := NewTSOStateMachine(NewHLC()) + + err := requireTSOHaltError(t, fsm.Apply(marshalTSOPhaseD(floor))) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) + require.False(t, fsm.PhaseDActive()) + + require.Nil(t, fsm.Apply(marshalTSOCutover())) + require.Nil(t, fsm.Apply(marshalTSOPhaseD(floor))) + require.True(t, fsm.PhaseDActive()) + require.Equal(t, floor, fsm.PhaseDFloor()) + require.Nil(t, fsm.Apply(marshalTSOPhaseD(floor)), "phase-D replay must be idempotent") + + err = requireTSOHaltError(t, fsm.Apply(marshalTSOPhaseD(floor+1))) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) + require.ErrorContains(t, err, "floor changed") + err = requireTSOHaltError(t, fsm.Apply([]byte(tsoPhaseDEnvelope))) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) +} + +func TestTSOStateMachineRejectsPhaseDFloorBelowExistingAllocation(t *testing.T) { + t.Parallel() + + fsm := NewTSOStateMachine(NewHLC()) + require.Nil(t, fsm.Apply(marshalTSOAllocationFloor(100))) + require.Nil(t, fsm.Apply(marshalTSOCutover())) + err := requireTSOHaltError(t, fsm.Apply(marshalTSOPhaseD(99))) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) + require.ErrorContains(t, err, "below allocation floor") + require.False(t, fsm.PhaseDActive()) +} + func TestTSOStateMachineRejectsInvalidCutoverSnapshotByte(t *testing.T) { t.Parallel() @@ -180,6 +215,29 @@ func TestTSOStateMachineRejectsInvalidCutoverSnapshotByte(t *testing.T) { require.ErrorContains(t, err, "invalid cutover byte") } +func TestTSOStateMachineRejectsInvalidPhaseDSnapshot(t *testing.T) { + t.Parallel() + + payload := make([]byte, tsoSnapshotV4Len) + payload[tsoSnapshotV2Len] = 1 + payload[tsoSnapshotV3Len] = 2 + err := NewTSOStateMachine(NewHLC()).Restore(bytes.NewReader(payload)) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) + require.ErrorContains(t, err, "invalid phase-D byte") + + payload[tsoSnapshotV2Len] = 0 + payload[tsoSnapshotV3Len] = 1 + err = NewTSOStateMachine(NewHLC()).Restore(bytes.NewReader(payload)) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) + require.ErrorContains(t, err, "phase-D active without cutover") + + payload[tsoSnapshotV3Len] = 0 + binary.BigEndian.PutUint64(payload[tsoSnapshotV3Len+1:], 1) + err = NewTSOStateMachine(NewHLC()).Restore(bytes.NewReader(payload)) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) + require.ErrorContains(t, err, "inactive phase-D has floor") +} + func TestTSOStateMachineNilHLCDoesNotPanic(t *testing.T) { t.Parallel() @@ -197,6 +255,9 @@ func TestTSOStateMachineSnapshotRestoreRoundTrip(t *testing.T) { floor := tsoLeaseAllocationFloor(ceilingMs) require.Nil(t, source.Apply(marshalTSOAllocationFloor(floor))) require.Nil(t, source.Apply(marshalTSOCutover())) + require.Nil(t, source.Apply(marshalTSOPhaseD(floor))) + postPhaseDFloor := floor + 10 + require.Nil(t, source.Apply(marshalTSOAllocationFloor(postPhaseDFloor))) snap, err := source.Snapshot() require.NoError(t, err) @@ -205,15 +266,17 @@ func TestTSOStateMachineSnapshotRestoreRoundTrip(t *testing.T) { var buf bytes.Buffer n, err := snap.WriteTo(&buf) require.NoError(t, err) - require.EqualValues(t, tsoSnapshotV3Len, n) - require.Len(t, buf.Bytes(), tsoSnapshotV3Len) + require.EqualValues(t, tsoSnapshotV4Len, n) + require.Len(t, buf.Bytes(), tsoSnapshotV4Len) targetHLC := NewHLC() target := NewTSOStateMachine(targetHLC) require.NoError(t, target.Restore(bytes.NewReader(buf.Bytes()))) require.Equal(t, ceilingMs, targetHLC.PhysicalCeiling()) - require.Equal(t, floor, targetHLC.Current()) + require.Equal(t, postPhaseDFloor, targetHLC.Current()) require.True(t, target.CutoverActive()) + require.True(t, target.PhaseDActive()) + require.Equal(t, floor, target.PhaseDFloor()) } func TestTSOStateMachineSnapshotUsesTSOOwnedCeiling(t *testing.T) { @@ -273,6 +336,26 @@ func TestTSOStateMachineRestoreLegacySnapshotDerivesAllocationFloor(t *testing.T require.Equal(t, tsoLeaseAllocationFloor(ceilingMs), hlc.Current()) } +func TestTSOStateMachineRestoresV3SnapshotWithoutPhaseD(t *testing.T) { + t.Parallel() + + const ( + ceilingMs = int64(1_700_000_654_321) + floor = uint64(4567) + ) + payload := make([]byte, tsoSnapshotV3Len) + binary.BigEndian.PutUint64(payload[:hlcLeasePayloadLen], uint64(ceilingMs)) + binary.BigEndian.PutUint64(payload[hlcLeasePayloadLen:tsoSnapshotV2Len], floor) + payload[tsoSnapshotV2Len] = 1 + + fsm := NewTSOStateMachine(NewHLC()) + require.NoError(t, fsm.Restore(bytes.NewReader(payload))) + require.Equal(t, floor, fsm.AllocationFloor()) + require.True(t, fsm.CutoverActive()) + require.False(t, fsm.PhaseDActive()) + require.Zero(t, fsm.PhaseDFloor()) +} + func TestTSOStateMachineRestoreKeepsMonotonicCeiling(t *testing.T) { t.Parallel() @@ -353,9 +436,11 @@ func TestTSOStateMachineClassifiesOnlyFullLeaseEntriesAsVolatile(t *testing.T) { require.True(t, fsm.IsVolatileOnlyPayload(marshalHLCLeaseRenew(1_700_000_123_456))) require.True(t, fsm.IsVolatileOnlyPayload(marshalTSOAllocationFloor(1))) require.True(t, fsm.IsVolatileOnlyPayload(marshalTSOCutover())) + require.True(t, fsm.IsVolatileOnlyPayload(marshalTSOPhaseD(1))) require.False(t, fsm.IsVolatileOnlyPayload([]byte{raftEncodeHLCLease})) require.False(t, fsm.IsVolatileOnlyPayload([]byte(tsoAllocationFloorEnvelope))) require.False(t, fsm.IsVolatileOnlyPayload(append([]byte(tsoCutoverEnvelope), 1))) + require.False(t, fsm.IsVolatileOnlyPayload([]byte(tsoPhaseDEnvelope))) require.False(t, fsm.IsVolatileOnlyPayload([]byte{raftEncodeSingle})) } diff --git a/kv/tso_raft.go b/kv/tso_raft.go index e2810854d..3975f5733 100644 --- a/kv/tso_raft.go +++ b/kv/tso_raft.go @@ -35,18 +35,24 @@ type TSOReservation struct { Count int PreviousAllocationFloor uint64 CutoverActive bool + PhaseDActive bool + PhaseDFloor uint64 } // TSOReservationAllocator is the migration-aware extension exposed by the // dedicated group leader. Ordinary callers continue to use TSOAllocator. type TSOReservationAllocator interface { - ReserveBatchAfter(context.Context, int, uint64, bool) (TSOReservation, error) + ReserveBatchAfter(context.Context, int, uint64, bool, bool) (TSOReservation, error) } type TSOShadowReservationAllocator interface { ValidateShadowTimestamp(context.Context, uint64) (TSOReservation, error) } +type tsoCutoverState interface { + CutoverActive() bool +} + // RaftTSOAllocator reserves timestamp windows on the dedicated TSO leader. // A window is returned only after its inclusive end has committed to Raft. // Failed proposals may leak a local window, but they can never expose an @@ -110,7 +116,7 @@ func (a *RaftTSOAllocator) NextBatchAfter(ctx context.Context, n int, min uint64 } func (a *RaftTSOAllocator) nextBatchAfter(ctx context.Context, n int, min uint64) (uint64, error) { - reservation, err := a.ReserveBatchAfter(ctx, n, min, false) + reservation, err := a.ReserveBatchAfter(ctx, n, min, false, false) return reservation.Base, err } @@ -123,6 +129,7 @@ func (a *RaftTSOAllocator) ReserveBatchAfter( n int, min uint64, activateCutover bool, + activatePhaseD bool, ) (TSOReservation, error) { var empty TSOReservation if err := validateTSOBatchSize(n); err != nil { @@ -134,7 +141,7 @@ func (a *RaftTSOAllocator) ReserveBatchAfter( ctx = nonNilTSOContext(ctx) a.mu.Lock() defer a.mu.Unlock() - return a.reserveBatchAfterLocked(ctx, n, min, activateCutover) + return a.reserveBatchAfterLocked(ctx, n, min, activateCutover, activatePhaseD) } func (a *RaftTSOAllocator) reserveBatchAfterLocked( @@ -142,6 +149,7 @@ func (a *RaftTSOAllocator) reserveBatchAfterLocked( n int, min uint64, activateCutover bool, + activatePhaseD bool, ) (TSOReservation, error) { var empty TSOReservation engine, err := a.verifiedLeader(ctx) @@ -152,7 +160,7 @@ func (a *RaftTSOAllocator) reserveBatchAfterLocked( if term == 0 { return empty, errors.Wrap(ErrTSONotLeader, "tso leader has no active term") } - previousFloor, err := a.prepareLeaderTermReservation(ctx, engine, term, activateCutover) + previousFloor, err := a.prepareLeaderTermReservation(ctx, engine, term, activateCutover, activatePhaseD) if err != nil { return empty, err } @@ -177,6 +185,8 @@ func (a *RaftTSOAllocator) reserveBatchAfterLocked( Count: n, PreviousAllocationFloor: previousFloor, CutoverActive: a.state.CutoverActive(), + PhaseDActive: a.state.PhaseDActive(), + PhaseDFloor: a.state.PhaseDFloor(), }, nil } @@ -185,16 +195,15 @@ func (a *RaftTSOAllocator) prepareLeaderTermReservation( engine raftengine.Engine, term uint64, activateCutover bool, + activatePhaseD bool, ) (uint64, error) { termFloor, err := a.termCommitFloor(ctx, term) if err != nil { return 0, err } - previousFloor := max(a.state.AllocationFloor(), termFloor) - if activateCutover && !a.state.CutoverActive() { - if err := a.commitCutover(ctx, engine); err != nil { - return 0, err - } + previousFloor := max(a.state.AllocationFloor(), termFloor, a.state.PhaseDFloor()) + if err := a.activateDurableMarkers(ctx, engine, previousFloor, activateCutover, activatePhaseD); err != nil { + return 0, err } if err := verifyTSOLeaderTerm(ctx, engine, term, true); err != nil { return 0, err @@ -202,6 +211,27 @@ func (a *RaftTSOAllocator) prepareLeaderTermReservation( return previousFloor, nil } +func (a *RaftTSOAllocator) activateDurableMarkers( + ctx context.Context, + engine raftengine.Engine, + previousFloor uint64, + activateCutover bool, + activatePhaseD bool, +) error { + if activateCutover && !a.state.CutoverActive() { + if err := a.commitCutover(ctx, engine); err != nil { + return err + } + } + if !activatePhaseD || a.state.PhaseDActive() { + return nil + } + if !a.state.CutoverActive() { + return errors.Wrap(ErrTSOPhaseDInactive, "phase D activation requires durable cutover") + } + return a.commitPhaseD(ctx, engine, previousFloor) +} + func (a *RaftTSOAllocator) termCommitFloor(ctx context.Context, term uint64) (uint64, error) { if a.initializedTerm == term { return a.state.AllocationFloor(), nil @@ -273,6 +303,65 @@ func (a *RaftTSOAllocator) commitCutover(ctx context.Context, engine raftengine. return nil } +func (a *RaftTSOAllocator) commitPhaseD(ctx context.Context, engine raftengine.Engine, floor uint64) error { + if _, err := a.group.Proposer().Propose(ctx, marshalTSOPhaseD(floor)); err != nil { + wrapped := errors.Wrap(err, "tso commit phase-D marker") + if tsoProposalLostLeadership(engine, err) { + return stderrors.Join(ErrTSONotLeader, wrapped) + } + return wrapped + } + if !a.state.PhaseDActive() || a.state.PhaseDFloor() != floor { + return errors.New("tso phase-D proposal committed without applied marker") + } + return nil +} + +func (a *RaftTSOAllocator) ValidateDurableTimestamp(ctx context.Context, timestamp uint64) error { + ctx = nonNilTSOContext(ctx) + a.mu.Lock() + defer a.mu.Unlock() + if _, err := a.verifiedLeader(ctx); err != nil { + return err + } + if !a.state.PhaseDActive() { + return errors.WithStack(ErrTSOPhaseDInactive) + } + floor := a.state.PhaseDFloor() + end := a.state.AllocationFloor() + if timestamp == 0 || timestamp > end { + return errors.Wrapf(ErrTSOTimestampInvalid, + "timestamp=%d phase_d_floor=%d allocation_floor=%d", timestamp, floor, end) + } + if timestamp <= floor { + return errors.Wrapf(stderrors.Join(ErrTSOTimestampInvalid, ErrTSOTimestampPrePhaseD), + "timestamp=%d phase_d_floor=%d allocation_floor=%d", timestamp, floor, end) + } + return nil +} + +func (a *RaftTSOAllocator) PhaseDActive() bool { + return a != nil && a.state != nil && a.state.PhaseDActive() +} + +func (a *RaftTSOAllocator) PhaseDRequired() bool { + return a.PhaseDActive() +} + +func (a *RaftTSOAllocator) PhaseDFloor() uint64 { + if a == nil || a.state == nil { + return 0 + } + return a.state.PhaseDFloor() +} + +func (a *RaftTSOAllocator) AllocationFloor() uint64 { + if a == nil || a.state == nil { + return 0 + } + return a.state.AllocationFloor() +} + func tsoProposalLostLeadership(engine raftengine.Engine, err error) bool { return !isLeaderEngine(engine) || errors.Is(err, raftengine.ErrNotLeader) || isTransientLeaderError(err) } @@ -285,21 +374,25 @@ func (a *RaftTSOAllocator) RunLeaseRenewal(ctx context.Context) { <-nonNilTSOContext(ctx).Done() } -type tsoRemoteRequest func(context.Context, string, int, uint64, bool) (TSOReservation, error) +type tsoRemoteRequest func(context.Context, string, int, uint64, bool, bool) (TSOReservation, error) + +type tsoRemoteValidation func(context.Context, string, uint64) error // LeaderRoutedTSOAllocator serves local requests on the TSO leader and sends // follower requests to the leader address published by the group-0 engine. // It re-resolves that address after transient errors so a leadership change // does not pin a BatchAllocator refill to a stale endpoint. type LeaderRoutedTSOAllocator struct { - local TSOAllocator - leader raftengine.LeaderView - connCache GRPCConnCache - remoteRequest tsoRemoteRequest - retryBudget time.Duration - retryInterval time.Duration - activate bool - clock *HLC + local TSOAllocator + leader raftengine.LeaderView + connCache GRPCConnCache + remoteRequest tsoRemoteRequest + retryBudget time.Duration + retryInterval time.Duration + activate bool + activatePhaseD bool + clock *HLC + remoteValidate tsoRemoteValidation } type LeaderRoutedTSOAllocatorOption func(*LeaderRoutedTSOAllocator) @@ -308,6 +401,13 @@ func WithTSOCutoverActivation() LeaderRoutedTSOAllocatorOption { return func(a *LeaderRoutedTSOAllocator) { a.activate = true } } +func WithTSOPhaseDActivation() LeaderRoutedTSOAllocatorOption { + return func(a *LeaderRoutedTSOAllocator) { + a.activate = true + a.activatePhaseD = true + } +} + func WithTSORoutedClock(clock *HLC) LeaderRoutedTSOAllocatorOption { return func(a *LeaderRoutedTSOAllocator) { a.clock = clock } } @@ -335,6 +435,7 @@ func NewLeaderRoutedTSOAllocator( } } a.remoteRequest = a.requestRemoteBatch + a.remoteValidate = a.requestRemoteValidation return a, nil } @@ -358,7 +459,7 @@ func (a *LeaderRoutedTSOAllocator) NextBatchAfter(ctx context.Context, n int, mi } func (a *LeaderRoutedTSOAllocator) nextBatchAfter(ctx context.Context, n int, min uint64) (uint64, error) { - reservation, err := a.nextReservation(ctx, n, min, a.activate, a.activate) + reservation, err := a.nextReservation(ctx, n, min, a.activate, a.activatePhaseD, a.activate) if err != nil { return 0, err } @@ -367,7 +468,7 @@ func (a *LeaderRoutedTSOAllocator) nextBatchAfter(ctx context.Context, n int, mi } func (a *LeaderRoutedTSOAllocator) ValidateShadowTimestamp(ctx context.Context, min uint64) (TSOReservation, error) { - reservation, err := a.nextReservation(ctx, 1, min, false, true) + reservation, err := a.nextReservation(ctx, 1, min, false, false, true) if err == nil { a.observeReservation(reservation) } @@ -379,6 +480,7 @@ func (a *LeaderRoutedTSOAllocator) nextReservation( n int, min uint64, activate bool, + activatePhaseD bool, requireReservation bool, ) (TSOReservation, error) { var empty TSOReservation @@ -392,7 +494,7 @@ func (a *LeaderRoutedTSOAllocator) nextReservation( var lastErr error for { - reservation, err := a.tryBatch(ctx, n, min, activate, requireReservation) + reservation, err := a.tryBatch(ctx, n, min, activate, activatePhaseD, requireReservation) if err == nil { return reservation, nil } @@ -417,17 +519,18 @@ func (a *LeaderRoutedTSOAllocator) tryBatch( n int, min uint64, activate bool, + activatePhaseD bool, requireReservation bool, ) (TSOReservation, error) { var empty TSOReservation if a.local.IsLeader() { - return a.tryLocalBatch(ctx, n, min, activate, requireReservation) + return a.tryLocalBatch(ctx, n, min, activate, activatePhaseD, requireReservation) } addr := leaderAddrFromEngine(a.leader) if addr == "" { return empty, errors.WithStack(ErrLeaderNotFound) } - reservation, err := a.remoteRequest(ctx, addr, n, min, activate) + reservation, err := a.remoteRequest(ctx, addr, n, min, activate, activatePhaseD) if err != nil { return empty, err } @@ -442,6 +545,7 @@ func (a *LeaderRoutedTSOAllocator) tryLocalBatch( n int, min uint64, activate bool, + activatePhaseD bool, requireReservation bool, ) (TSOReservation, error) { var empty TSOReservation @@ -457,7 +561,7 @@ func (a *LeaderRoutedTSOAllocator) tryLocalBatch( reservation := TSOReservation{Base: base, Count: n} return reservation, validateTSOReservation(reservation, n, min, false) } - reservation, err := reservationAllocator.ReserveBatchAfter(ctx, n, min, activate) + reservation, err := reservationAllocator.ReserveBatchAfter(ctx, n, min, activate, activatePhaseD) if err != nil { return empty, errors.Wrap(err, "reserve local TSO window") } @@ -470,6 +574,7 @@ func (a *LeaderRoutedTSOAllocator) requestRemoteBatch( n int, min uint64, activate bool, + activatePhaseD bool, ) (TSOReservation, error) { var empty TSOReservation conn, err := a.connCache.ConnFor(addr) @@ -480,6 +585,7 @@ func (a *LeaderRoutedTSOAllocator) requestRemoteBatch( Count: uint32(n), //nolint:gosec // n is bounded by maxHLCBatchSize. MinTimestamp: min, ActivateCutover: activate, + ActivatePhaseD: activatePhaseD, }) if err != nil { return empty, errors.Wrap(err, "tso request leader batch") @@ -494,14 +600,88 @@ func (a *LeaderRoutedTSOAllocator) requestRemoteBatch( return empty, errors.Wrap(ErrTSOProtocolUnsupported, "leader did not confirm durable TSO cutover") } + if activatePhaseD && !resp.GetPhaseDActive() { + return empty, errors.Wrap(ErrTSOProtocolUnsupported, + "leader did not confirm durable TSO phase D") + } return TSOReservation{ Base: resp.GetTimestamp(), Count: n, PreviousAllocationFloor: resp.GetPreviousAllocationFloor(), CutoverActive: resp.GetCutoverActive(), + PhaseDActive: resp.GetPhaseDActive(), + PhaseDFloor: resp.GetPhaseDFloor(), }, nil } +func (a *LeaderRoutedTSOAllocator) ValidateDurableTimestamp(ctx context.Context, timestamp uint64) error { + if timestamp == 0 { + return errors.WithStack(ErrTSOTimestampInvalid) + } + ctx = nonNilTSOContext(ctx) + deadline := time.Now().Add(a.retryBudget) + ctx, cancel := context.WithDeadline(ctx, deadline) + defer cancel() + var lastErr error + for { + if a.local.IsLeader() { + validator, ok := a.local.(DurableTimestampValidator) + if !ok { + return errors.WithStack(ErrTSOProtocolUnsupported) + } + lastErr = errors.WithStack(validator.ValidateDurableTimestamp(ctx, timestamp)) + } else { + addr := leaderAddrFromEngine(a.leader) + if addr == "" { + lastErr = errors.WithStack(ErrLeaderNotFound) + } else { + lastErr = a.remoteValidate(ctx, addr, timestamp) + } + } + if lastErr == nil { + return nil + } + if !isTransientTSORouteError(lastErr) { + return lastErr + } + if err := waitTSORouteRetry(ctx, a.retryInterval); err != nil { + return errors.Wrap(stderrors.Join(ctx.Err(), lastErr), "tso durable timestamp validation") + } + } +} + +func (a *LeaderRoutedTSOAllocator) requestRemoteValidation(ctx context.Context, addr string, timestamp uint64) error { + conn, err := a.connCache.ConnFor(addr) + if err != nil { + return errors.Wrap(err, "tso dial validation leader") + } + resp, err := pb.NewDistributionClient(conn).ValidateTimestamp(ctx, &pb.ValidateTimestampRequest{Timestamp: timestamp}) + if err != nil { + code := status.Code(err) + if code == codes.OutOfRange { + return errors.Wrap(stderrors.Join(ErrTSOTimestampInvalid, ErrTSOTimestampPrePhaseD), err.Error()) + } + if code == codes.InvalidArgument { + return errors.Wrap(ErrTSOTimestampInvalid, err.Error()) + } + return errors.Wrap(err, "tso validate timestamp at leader") + } + if !resp.GetValid() || !resp.GetPhaseDActive() { + return errors.Wrapf(ErrTSOTimestampInvalid, + "leader validation valid=%t phase_d_active=%t", resp.GetValid(), resp.GetPhaseDActive()) + } + return nil +} + +func (a *LeaderRoutedTSOAllocator) PhaseDActive() bool { + state, ok := a.local.(TSOPhaseDState) + return ok && state.PhaseDActive() +} + +func (a *LeaderRoutedTSOAllocator) PhaseDRequired() bool { + return a != nil && (a.activatePhaseD || a.PhaseDActive()) +} + func (a *LeaderRoutedTSOAllocator) observeReservation(reservation TSOReservation) { if a == nil || a.clock == nil || reservation.Base == 0 || reservation.Count <= 0 { return @@ -614,12 +794,24 @@ func nonNilTSOContext(ctx context.Context) context.Context { // and retried; once the durable cutover marker is active, the allocator returns // the reserved TSO timestamp directly so rolling restarts cannot mix sources. type ShadowTimestampAllocator struct { - legacy *HLC - shadow TSOShadowReservationAllocator - log *slog.Logger + legacy *HLC + shadow TSOShadowReservationAllocator + cutover tsoCutoverState + log *slog.Logger } -func NewShadowTimestampAllocator(legacy *HLC, shadow TSOShadowReservationAllocator, logger *slog.Logger) (*ShadowTimestampAllocator, error) { +type ShadowTimestampAllocatorOption func(*ShadowTimestampAllocator) + +func WithTSOShadowCutoverState(state tsoCutoverState) ShadowTimestampAllocatorOption { + return func(a *ShadowTimestampAllocator) { a.cutover = state } +} + +func NewShadowTimestampAllocator( + legacy *HLC, + shadow TSOShadowReservationAllocator, + logger *slog.Logger, + opts ...ShadowTimestampAllocatorOption, +) (*ShadowTimestampAllocator, error) { if legacy == nil { return nil, errors.WithStack(ErrTSOClockNil) } @@ -629,11 +821,17 @@ func NewShadowTimestampAllocator(legacy *HLC, shadow TSOShadowReservationAllocat if logger == nil { logger = slog.Default() } - return &ShadowTimestampAllocator{ + a := &ShadowTimestampAllocator{ legacy: legacy, shadow: shadow, log: logger, - }, nil + } + for _, opt := range opts { + if opt != nil { + opt(a) + } + } + return a, nil } func (a *ShadowTimestampAllocator) Next(ctx context.Context) (uint64, error) { @@ -649,6 +847,9 @@ func (a *ShadowTimestampAllocator) NextAfter(ctx context.Context, min uint64) (u func (a *ShadowTimestampAllocator) nextAfter(ctx context.Context, min uint64) (uint64, error) { ctx = nonNilTSOContext(ctx) + if a.cutover != nil && a.cutover.CutoverActive() { + return a.nextDedicatedAfter(ctx, min) + } a.legacy.Observe(min) for { if err := ctx.Err(); err != nil { @@ -681,6 +882,32 @@ func (a *ShadowTimestampAllocator) nextAfter(ctx context.Context, min uint64) (u } } +func (a *ShadowTimestampAllocator) nextDedicatedAfter(ctx context.Context, min uint64) (uint64, error) { + alloc, ok := a.shadow.(TimestampAllocator) + if !ok { + return 0, errors.WithStack(ErrTSOProtocolUnsupported) + } + return nextTimestampAfterFromAllocator(ctx, alloc, min, "tso post-cutover shadow bypass") +} + +func (a *ShadowTimestampAllocator) ValidateDurableTimestamp(ctx context.Context, timestamp uint64) error { + validator, ok := a.shadow.(DurableTimestampValidator) + if !ok { + return errors.WithStack(ErrTSOProtocolUnsupported) + } + return errors.WithStack(validator.ValidateDurableTimestamp(ctx, timestamp)) +} + +func (a *ShadowTimestampAllocator) PhaseDActive() bool { + state, ok := a.shadow.(TSOPhaseDState) + return ok && state.PhaseDActive() +} + +func (a *ShadowTimestampAllocator) PhaseDRequired() bool { + state, ok := a.shadow.(TSOPhaseDState) + return ok && state.PhaseDRequired() +} + func (a *ShadowTimestampAllocator) Close() error { return nil } diff --git a/kv/tso_raft_test.go b/kv/tso_raft_test.go index 10f4925b0..56215a38f 100644 --- a/kv/tso_raft_test.go +++ b/kv/tso_raft_test.go @@ -7,6 +7,7 @@ import ( "io" "log/slog" "sync" + "sync/atomic" "testing" "time" @@ -172,6 +173,36 @@ func TestRaftTSOAllocatorDropsCommittedWindowAfterTermChange(t *testing.T) { require.Greater(t, next, leakedFloor) } +func TestRaftTSOAllocatorRejectsTermChangeAfterPhaseDMarker(t *testing.T) { + clock := NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) + fsm := NewTSOStateMachine(clock) + engine := &recordingTSOEngine{ + state: raftengine.StateLeader, + leader: raftengine.LeaderInfo{Address: "self"}, + term: 1, + apply: applyTSOTestFSM(fsm), + } + var changeTerm sync.Once + engine.afterPropose = func(payload []byte) { + if bytes.HasPrefix(payload, []byte(tsoPhaseDEnvelope)) { + changeTerm.Do(func() { engine.setTerm(2) }) + } + } + alloc, err := newTestRaftTSOAllocator(&ShardGroup{Engine: engine, TSOState: fsm}, clock) + require.NoError(t, err) + + _, err = alloc.ReserveBatchAfter(context.Background(), testTSOBatchSize, 0, true, true) + require.ErrorIs(t, err, ErrTSONotLeader) + require.True(t, fsm.PhaseDActive()) + require.Zero(t, fsm.AllocationFloor()) + + payloads := engine.proposedPayloads() + require.Len(t, payloads, 2) + require.Equal(t, []byte(tsoCutoverEnvelope), payloads[0]) + require.Equal(t, marshalTSOPhaseD(0), payloads[1]) +} + func TestRaftTSOAllocatorCommitsCutoverBeforeProductionWindow(t *testing.T) { clock := NewHLC() clock.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) @@ -185,7 +216,7 @@ func TestRaftTSOAllocatorCommitsCutoverBeforeProductionWindow(t *testing.T) { alloc, err := newTestRaftTSOAllocator(&ShardGroup{Engine: engine, TSOState: fsm}, clock) require.NoError(t, err) - reservation, err := alloc.ReserveBatchAfter(context.Background(), testTSOBatchSize, 0, true) + reservation, err := alloc.ReserveBatchAfter(context.Background(), testTSOBatchSize, 0, true, false) require.NoError(t, err) require.True(t, reservation.CutoverActive) payloads := engine.proposedPayloads() @@ -194,6 +225,66 @@ func TestRaftTSOAllocatorCommitsCutoverBeforeProductionWindow(t *testing.T) { require.True(t, bytes.HasPrefix(payloads[1], []byte(tsoAllocationFloorEnvelope))) } +func TestRaftTSOAllocatorCommitsPhaseDBeforeWindowAndValidatesRange(t *testing.T) { + clock := NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) + fsm := NewTSOStateMachine(clock) + engine := &recordingTSOEngine{ + state: raftengine.StateLeader, + leader: raftengine.LeaderInfo{Address: "self"}, + term: 1, + apply: applyTSOTestFSM(fsm), + } + alloc, err := newTestRaftTSOAllocator(&ShardGroup{Engine: engine, TSOState: fsm}, clock) + require.NoError(t, err) + + reservation, err := alloc.ReserveBatchAfter(context.Background(), testTSOBatchSize, 0, true, true) + require.NoError(t, err) + require.True(t, reservation.CutoverActive) + require.True(t, reservation.PhaseDActive) + require.Equal(t, reservation.PreviousAllocationFloor, reservation.PhaseDFloor) + require.Greater(t, reservation.Base, reservation.PhaseDFloor) + + payloads := engine.proposedPayloads() + require.Len(t, payloads, 3) + require.Equal(t, []byte(tsoCutoverEnvelope), payloads[0]) + require.Equal(t, marshalTSOPhaseD(reservation.PreviousAllocationFloor), payloads[1]) + require.True(t, bytes.HasPrefix(payloads[2], []byte(tsoAllocationFloorEnvelope))) + + end := reservation.Base + positiveIntToUint64(reservation.Count) - 1 + require.NoError(t, alloc.ValidateDurableTimestamp(context.Background(), reservation.Base)) + require.NoError(t, alloc.ValidateDurableTimestamp(context.Background(), end)) + err = alloc.ValidateDurableTimestamp(context.Background(), reservation.PhaseDFloor) + require.ErrorIs(t, err, ErrTSOTimestampInvalid) + require.ErrorIs(t, alloc.ValidateDurableTimestamp(context.Background(), end+1), ErrTSOTimestampInvalid) +} + +func TestRaftTSOAllocatorFencesAboveRestoredPhaseDFloor(t *testing.T) { + clock := NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) + fsm := NewTSOStateMachine(clock) + phaseDFloor := uint64(time.Now().Add(30*time.Minute).UnixMilli()) << hlcLogicalBits //nolint:gosec // future test HLC. + require.Nil(t, fsm.Apply(marshalTSOCutover())) + require.Nil(t, fsm.Apply(marshalTSOPhaseD(phaseDFloor))) + engine := &recordingTSOEngine{ + state: raftengine.StateLeader, + leader: raftengine.LeaderInfo{Address: "self"}, + term: 1, + apply: applyTSOTestFSM(fsm), + } + alloc, err := newTestRaftTSOAllocator(&ShardGroup{Engine: engine, TSOState: fsm}, clock) + require.NoError(t, err) + + reservation, err := alloc.ReserveBatchAfter(context.Background(), 1, 0, false, false) + require.NoError(t, err) + require.Equal(t, phaseDFloor, reservation.PreviousAllocationFloor) + require.Greater(t, reservation.Base, phaseDFloor) + require.NoError(t, alloc.ValidateDurableTimestamp(context.Background(), reservation.Base)) + err = alloc.ValidateDurableTimestamp(context.Background(), phaseDFloor) + require.ErrorIs(t, err, ErrTSOTimestampInvalid) + require.ErrorIs(t, err, ErrTSOTimestampPrePhaseD) +} + func TestLeaderRoutedTSOAllocatorReResolvesAfterStaleLeader(t *testing.T) { local := &fakeTSOAllocator{leader: false, nextBase: testTSOInitialBase} engine := &recordingTSOEngine{state: raftengine.StateFollower, leader: raftengine.LeaderInfo{Address: "old"}} @@ -203,11 +294,12 @@ func TestLeaderRoutedTSOAllocatorReResolvesAfterStaleLeader(t *testing.T) { alloc.retryInterval = time.Millisecond var addresses []string - alloc.remoteRequest = func(_ context.Context, addr string, n int, min uint64, activate bool) (TSOReservation, error) { + alloc.remoteRequest = func(_ context.Context, addr string, n int, min uint64, activate, activatePhaseD bool) (TSOReservation, error) { addresses = append(addresses, addr) require.Equal(t, testTSOBatchSize, n) require.Equal(t, uint64(testTSOInitialBase), min) require.False(t, activate) + require.False(t, activatePhaseD) if addr == "old" { engine.setLeaderAddress("new") return TSOReservation{}, status.Error(codes.FailedPrecondition, "tso: not leader") @@ -226,7 +318,7 @@ func TestLeaderRoutedTSOAllocatorUsesLocalLeader(t *testing.T) { engine := &recordingTSOEngine{state: raftengine.StateLeader, leader: raftengine.LeaderInfo{Address: "self"}} alloc, err := NewLeaderRoutedTSOAllocator(local, engine) require.NoError(t, err) - alloc.remoteRequest = func(context.Context, string, int, uint64, bool) (TSOReservation, error) { + alloc.remoteRequest = func(context.Context, string, int, uint64, bool, bool) (TSOReservation, error) { t.Fatal("local TSO leader must not call remote RPC") return TSOReservation{}, nil } @@ -251,7 +343,7 @@ func TestLeaderRoutedTSOAllocatorRejectsRemoteWindowAtMinimum(t *testing.T) { engine := &recordingTSOEngine{state: raftengine.StateFollower, leader: raftengine.LeaderInfo{Address: "leader"}} alloc, err := NewLeaderRoutedTSOAllocator(local, engine) require.NoError(t, err) - alloc.remoteRequest = func(context.Context, string, int, uint64, bool) (TSOReservation, error) { + alloc.remoteRequest = func(context.Context, string, int, uint64, bool, bool) (TSOReservation, error) { return TSOReservation{Base: testTSOInitialBase, Count: testTSOBatchSize}, nil } @@ -268,7 +360,7 @@ func TestLeaderRoutedTSOAllocatorPreservesDeadlineAfterTransientErrors(t *testin alloc.retryInterval = time.Millisecond var attempts int - alloc.remoteRequest = func(context.Context, string, int, uint64, bool) (TSOReservation, error) { + alloc.remoteRequest = func(context.Context, string, int, uint64, bool, bool) (TSOReservation, error) { attempts++ return TSOReservation{}, status.Error(codes.Unavailable, "leader restarting") } @@ -279,6 +371,29 @@ func TestLeaderRoutedTSOAllocatorPreservesDeadlineAfterTransientErrors(t *testin require.Greater(t, attempts, 1) } +func TestLeaderRoutedTSOAllocatorRetriesValidationAfterLocalLeadershipLoss(t *testing.T) { + local := &leadershipLosingValidationTSO{fakeTSOAllocator: &fakeTSOAllocator{leader: true}} + engine := &recordingTSOEngine{ + state: raftengine.StateFollower, + leader: raftengine.LeaderInfo{Address: "new-leader"}, + } + alloc, err := NewLeaderRoutedTSOAllocator(local, engine) + require.NoError(t, err) + alloc.retryBudget = time.Second + alloc.retryInterval = time.Millisecond + var remoteCalls atomic.Uint64 + alloc.remoteValidate = func(_ context.Context, addr string, timestamp uint64) error { + remoteCalls.Add(1) + require.Equal(t, "new-leader", addr) + require.Equal(t, uint64(42), timestamp) + return nil + } + + require.NoError(t, alloc.ValidateDurableTimestamp(context.Background(), 42)) + require.Equal(t, uint64(1), local.validateCalls.Load()) + require.Equal(t, uint64(1), remoteCalls.Load()) +} + func TestShadowTimestampAllocatorReturnsLegacyAndAdvancesTSO(t *testing.T) { legacy := NewHLC() legacy.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) @@ -353,6 +468,26 @@ func TestShadowTimestampAllocatorReturnsTSOAfterDurableCutover(t *testing.T) { require.Greater(t, issued, legacyCandidate) } +func TestShadowTimestampAllocatorBypassesLegacyAfterObservedCutover(t *testing.T) { + legacy := NewHLC() + fsm := NewTSOStateMachine(NewHLC()) + require.Nil(t, fsm.Apply(marshalTSOCutover())) + dedicated := &dedicatedShadowAllocator{next: testTSOInitialBase} + alloc, err := NewShadowTimestampAllocator( + legacy, + dedicated, + slog.New(slog.NewTextHandler(io.Discard, nil)), + WithTSOShadowCutoverState(fsm), + ) + require.NoError(t, err) + + issued, err := alloc.NextAfter(context.Background(), testTSOInitialBase-1) + require.NoError(t, err) + require.Equal(t, uint64(testTSOInitialBase), issued) + require.Zero(t, legacy.Current(), "post-cutover issuance must not sample legacy HLC") + require.Zero(t, dedicated.shadowCalls, "post-cutover issuance must not run shadow comparison") +} + func TestShadowAndCutoverAllocatorsSerializeMigrationOnGroupZero(t *testing.T) { tsoClock := NewHLC() tsoClock.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) @@ -402,6 +537,20 @@ type recordingShadowReservationAllocator struct { cutover bool } +type dedicatedShadowAllocator struct { + next uint64 + shadowCalls int +} + +func (a *dedicatedShadowAllocator) Next(context.Context) (uint64, error) { + return a.next, nil +} + +func (a *dedicatedShadowAllocator) ValidateShadowTimestamp(context.Context, uint64) (TSOReservation, error) { + a.shadowCalls++ + return TSOReservation{}, errors.New("shadow comparison must not run after cutover") +} + func (a *recordingShadowReservationAllocator) ValidateShadowTimestamp(_ context.Context, min uint64) (TSOReservation, error) { a.mu.Lock() defer a.mu.Unlock() @@ -451,6 +600,17 @@ type recordingTSOEngine struct { term uint64 } +type leadershipLosingValidationTSO struct { + *fakeTSOAllocator + validateCalls atomic.Uint64 +} + +func (a *leadershipLosingValidationTSO) ValidateDurableTimestamp(context.Context, uint64) error { + a.validateCalls.Add(1) + a.leader = false + return errors.WithStack(ErrTSONotLeader) +} + func (e *recordingTSOEngine) Propose(_ context.Context, payload []byte) (*raftengine.ProposalResult, error) { e.mu.Lock() if e.proposeErr != nil { diff --git a/kv/tso_test.go b/kv/tso_test.go index ef4e58510..78717844d 100644 --- a/kv/tso_test.go +++ b/kv/tso_test.go @@ -7,6 +7,7 @@ import ( "time" "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/keyviz" "github.com/cockroachdb/errors" "github.com/stretchr/testify/require" ) @@ -291,6 +292,82 @@ func TestNextTimestampAfterThroughUsesAllocatorFloor(t *testing.T) { require.EqualValues(t, testTSOInitialBase+6, got) } +func TestBeginReadTimestampThroughPreservesLegacyBeforePhaseD(t *testing.T) { + alloc := &phaseDTestAllocator{next: testTSOInitialBase} + coord := &Coordinate{tsAllocator: alloc} + + readTS, err := BeginReadTimestampThrough(context.Background(), coord, 42, "test read timestamp") + require.NoError(t, err) + require.Equal(t, uint64(42), readTS.Timestamp()) + require.Zero(t, alloc.nextCalls.Load()) + require.Zero(t, alloc.validateCalls.Load()) +} + +func TestBeginReadTimestampThroughValidatesAppliedWatermarkDuringPhaseD(t *testing.T) { + alloc := &phaseDTestAllocator{next: testTSOInitialBase, phaseDActive: true, phaseDRequired: true} + coord := &Coordinate{tsAllocator: alloc} + + readTS, err := BeginReadTimestampThrough(context.Background(), coord, 42, "test read timestamp") + require.NoError(t, err) + require.Equal(t, uint64(42), readTS.Timestamp()) + require.Zero(t, alloc.nextCalls.Load()) + require.Equal(t, uint64(1), alloc.validateCalls.Load()) + require.Equal(t, uint64(42), alloc.validated.Load()) +} + +func TestBeginReadTimestampThroughFailsClosedOnValidationError(t *testing.T) { + alloc := &phaseDTestAllocator{ + next: testTSOInitialBase, + phaseDActive: true, + phaseDRequired: true, + validateErr: ErrTSOTimestampInvalid, + } + coord := &Coordinate{tsAllocator: alloc} + + _, err := BeginReadTimestampThrough(context.Background(), coord, 42, "test read timestamp") + require.ErrorIs(t, err, ErrTSOTimestampInvalid) + require.Zero(t, alloc.nextCalls.Load()) + require.Equal(t, uint64(1), alloc.validateCalls.Load()) +} + +func TestBeginReadTimestampThroughActivatesPhaseDBeforeValidation(t *testing.T) { + alloc := &phaseDTestAllocator{next: testTSOInitialBase, phaseDRequired: true} + coord := &Coordinate{tsAllocator: alloc} + + readTS, err := BeginReadTimestampThrough(context.Background(), coord, 42, "test phase D activation") + require.NoError(t, err) + require.Equal(t, uint64(42), readTS.Timestamp()) + require.Equal(t, uint64(1), alloc.nextCalls.Load()) + require.Equal(t, uint64(1), alloc.validateCalls.Load()) +} + +func TestBatchAllocatorDropsPrePhaseDWindowAfterActivation(t *testing.T) { + raw := &phaseDWindowTSO{nextBase: testTSOInitialBase, leader: true} + alloc, err := NewBatchAllocator(raw, testTSOBatchSize) + require.NoError(t, err) + + first, err := alloc.Next(context.Background()) + require.NoError(t, err) + require.Equal(t, uint64(testTSOInitialBase), first) + raw.phaseD = true + raw.floor = testTSOInitialBase + testTSOBatchSize - 1 + + readTS, err := alloc.Next(context.Background()) + require.NoError(t, err) + require.Equal(t, uint64(testTSOInitialBase+testTSOBatchSize), readTS) + require.Equal(t, uint64(2), raw.calls.Load()) +} + +func TestBeginReadTimestampThroughFindsAllocatorBehindKeyVizDecorator(t *testing.T) { + alloc := &phaseDTestAllocator{next: testTSOInitialBase, phaseDActive: true, phaseDRequired: true} + coord := WithKeyVizLabel(&Coordinate{tsAllocator: alloc}, keyviz.LabelRedis) + + readTS, err := BeginReadTimestampThrough(context.Background(), coord, 42, "test decorated read timestamp") + require.NoError(t, err) + require.Equal(t, uint64(42), readTS.Timestamp()) + require.Equal(t, uint64(1), alloc.validateCalls.Load()) +} + func TestCoordinateUsesTSOAllocatorForIssuedTimestamps(t *testing.T) { t.Parallel() @@ -392,6 +469,71 @@ type fakeTSOAllocator struct { leader bool } +type phaseDTestAllocator struct { + next uint64 + phaseDActive bool + phaseDRequired bool + validateErr error + nextCalls atomic.Uint64 + validateCalls atomic.Uint64 + validated atomic.Uint64 +} + +func (a *phaseDTestAllocator) Next(context.Context) (uint64, error) { + a.nextCalls.Add(1) + return a.next, nil +} + +func (a *phaseDTestAllocator) NextAfter(_ context.Context, min uint64) (uint64, error) { + a.nextCalls.Add(1) + if a.next <= min { + return min + 1, nil + } + return a.next, nil +} + +func (a *phaseDTestAllocator) ValidateDurableTimestamp(_ context.Context, timestamp uint64) error { + a.validateCalls.Add(1) + a.validated.Store(timestamp) + return a.validateErr +} + +func (a *phaseDTestAllocator) PhaseDActive() bool { return a.phaseDActive } +func (a *phaseDTestAllocator) PhaseDRequired() bool { return a.phaseDRequired } + +type phaseDWindowTSO struct { + nextBase uint64 + floor uint64 + phaseD bool + leader bool + calls atomic.Uint64 +} + +func (a *phaseDWindowTSO) Next(ctx context.Context) (uint64, error) { + return a.NextBatch(ctx, 1) +} + +func (a *phaseDWindowTSO) NextBatch(_ context.Context, n int) (uint64, error) { + a.calls.Add(1) + base := a.nextBase + a.nextBase += uint64(n) //nolint:gosec // positive test batch size. + return base, nil +} + +func (a *phaseDWindowTSO) IsLeader() bool { return a.leader } + +func (a *phaseDWindowTSO) RunLeaseRenewal(ctx context.Context) { <-ctx.Done() } + +func (a *phaseDWindowTSO) ValidateDurableTimestamp(_ context.Context, timestamp uint64) error { + if timestamp <= a.floor { + return ErrTSOTimestampInvalid + } + return nil +} + +func (a *phaseDWindowTSO) PhaseDActive() bool { return a.phaseD } +func (a *phaseDWindowTSO) PhaseDRequired() bool { return a.phaseD } + func (f *fakeTSOAllocator) Next(ctx context.Context) (uint64, error) { return f.NextBatch(ctx, 1) } diff --git a/main.go b/main.go index 003645269..5f515a91e 100644 --- a/main.go +++ b/main.go @@ -126,6 +126,7 @@ var ( raftJoinAsLearner = flag.Bool("raftJoinAsLearner", false, "Local node expects to join an existing cluster as a learner; if a post-apply ConfState lists this node as a voter instead, an ERROR-level alarm fires (the node keeps running -- the flag is an operator alarm, not a consensus veto). See docs/design/2026_04_26_implemented_raft_learner.md §4.5.") tsoEnabled = flag.Bool("tsoEnabled", false, "Commit the one-way cutover marker and issue coordinator-owned persistence timestamps through the dedicated TSO leader when group 0 is configured") tsoShadowEnabled = flag.Bool("tsoShadowEnabled", false, "Serialize legacy HLC issuance through the dedicated TSO; fail closed on TSO errors and switch to TSO values after cutover") + tsoPhaseDEnabled = flag.Bool("tsoPhaseDEnabled", false, "Close the centralized-TSO compatibility window after every group-0 member understands the M7 marker") tsoBatchSize = flag.Int("tsoBatchSize", defaultTSOBatchSize, "Timestamp batch size used by TSO cutover and shadow validation") leaderBalance = flag.Bool("leaderBalance", false, "Enable automatic count-based Raft-group leader balancing on the default-group leader") leaderBalanceInterval = flag.Duration("leaderBalanceInterval", defaultLeaderBalanceInterval, "Interval between leader-balance scheduler evaluations") @@ -2094,6 +2095,9 @@ func configureCoordinatorTSO( if *tsoEnabled && *tsoShadowEnabled { return wiring, errors.New("--tsoEnabled and --tsoShadowEnabled are mutually exclusive") } + if *tsoPhaseDEnabled && !*tsoEnabled { + return wiring, errors.New("--tsoPhaseDEnabled requires --tsoEnabled") + } tsoGroup, dedicated := shardGroups[dedicatedTSORaftGroupID] if !dedicated { @@ -2108,6 +2112,9 @@ func configureCoordinatorTSO( func configureLegacyCoordinatorTSO(coordinate *kv.ShardedCoordinator) (coordinatorTSOWiring, error) { var wiring coordinatorTSOWiring + if *tsoPhaseDEnabled { + return wiring, errors.Wrap(kv.ErrTSOGroupRequired, "configure tso phase D") + } if *tsoShadowEnabled { return wiring, errors.Wrap(kv.ErrTSOGroupRequired, "configure tso shadow allocator") } @@ -2143,34 +2150,49 @@ func configureDedicatedCoordinatorTSO( return wiring, errors.Wrap(err, "configure dedicated tso allocator") } wiring.serverAllocator = local - cutoverActive := tsoGroup.TSOState.CutoverActive() - if !*tsoEnabled && !*tsoShadowEnabled && !cutoverActive { + coordinate.WithTimestampGroup(dedicatedTSORaftGroupID) + coordinate.WithTSOCutoverState(tsoGroup.TSOState) + if !dedicatedTSOAllocatorRequired(tsoGroup.TSOState) { return wiring, nil } - routedOpts := []kv.LeaderRoutedTSOAllocatorOption{ - kv.WithTSORoutedClock(coordinate.Clock()), - } - if *tsoEnabled { - routedOpts = append(routedOpts, kv.WithTSOCutoverActivation()) - } + routedOpts := dedicatedTSORoutingOptions(coordinate.Clock()) routed, err := kv.NewLeaderRoutedTSOAllocator(local, tsoGroup.Engine, routedOpts...) if err != nil { return wiring, errors.Wrap(err, "configure tso leader routing") } wiring.routedAllocator = routed - coordinate.WithTimestampGroup(dedicatedTSORaftGroupID) - return installDedicatedCoordinatorTSO(coordinate, wiring, routed, cutoverActive) + return installDedicatedCoordinatorAllocator(coordinate, tsoGroup.TSOState, wiring, routed) +} + +func dedicatedTSOAllocatorRequired(state *kv.TSOStateMachine) bool { + return *tsoEnabled || *tsoShadowEnabled || (state != nil && state.CutoverActive()) } -func installDedicatedCoordinatorTSO( +func dedicatedTSORoutingOptions(clock *kv.HLC) []kv.LeaderRoutedTSOAllocatorOption { + opts := []kv.LeaderRoutedTSOAllocatorOption{kv.WithTSORoutedClock(clock)} + if *tsoEnabled { + opts = append(opts, kv.WithTSOCutoverActivation()) + } + if *tsoPhaseDEnabled { + opts = append(opts, kv.WithTSOPhaseDActivation()) + } + return opts +} + +func installDedicatedCoordinatorAllocator( coordinate *kv.ShardedCoordinator, + state *kv.TSOStateMachine, wiring coordinatorTSOWiring, routed *kv.LeaderRoutedTSOAllocator, - cutoverActive bool, ) (coordinatorTSOWiring, error) { - if *tsoShadowEnabled && !cutoverActive { - shadow, err := kv.NewShadowTimestampAllocator(coordinate.Clock(), routed, slog.Default()) + if *tsoShadowEnabled && (state == nil || !state.CutoverActive()) { + shadow, err := kv.NewShadowTimestampAllocator( + coordinate.Clock(), + routed, + slog.Default(), + kv.WithTSOShadowCutoverState(state), + ) if err != nil { _ = wiring.Close() return coordinatorTSOWiring{}, errors.Wrap(err, "configure tso shadow allocator") @@ -2332,6 +2354,7 @@ var _ kv.LeaseReadableCoordinator = (*startupGatedCoordinator)(nil) var _ kv.AllGroupsLeaseReadableCoordinator = (*startupGatedCoordinator)(nil) var _ kv.GroupRoutableCoordinator = (*startupGatedCoordinator)(nil) var _ kv.TimestampAllocatorProvider = (*startupGatedCoordinator)(nil) +var _ kv.AppliedReadTimestampVoucher = (*startupGatedCoordinator)(nil) func (c startupGatedCoordinator) Dispatch(ctx context.Context, reqs *kv.OperationGroup[kv.OP]) (*kv.CoordinateResponse, error) { if c.gate != nil && c.gate.blocked() { @@ -2377,6 +2400,14 @@ func (c startupGatedCoordinator) TimestampAllocator() kv.TimestampAllocator { return alloc } +func (c startupGatedCoordinator) VouchAppliedReadTimestamp(timestamp uint64) error { + voucher, ok := c.inner.(kv.AppliedReadTimestampVoucher) + if !ok { + return errors.WithStack(kv.ErrTSOProtocolUnsupported) + } + return errors.WithStack(voucher.VouchAppliedReadTimestamp(timestamp)) +} + func (c startupGatedCoordinator) LeaseRead(ctx context.Context) (uint64, error) { return kv.LeaseReadThrough(c.inner, ctx) //nolint:wrapcheck // Pass through coordinator errors unchanged. } diff --git a/main_tso_routing_test.go b/main_tso_routing_test.go index 7c69c9340..a2527fa5f 100644 --- a/main_tso_routing_test.go +++ b/main_tso_routing_test.go @@ -9,6 +9,7 @@ import ( "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/internal/raftengine" + "github.com/bootjp/elastickv/keyviz" "github.com/bootjp/elastickv/kv" "github.com/stretchr/testify/require" ) @@ -54,6 +55,41 @@ func TestConfigureCoordinatorTSOCutoverRoutesThroughDedicatedGroup(t *testing.T) require.True(t, fsm.CutoverActive()) } +func TestConfigureCoordinatorTSOPhaseDWorksThroughAdapterDecorators(t *testing.T) { + setTSOModeFlags(t, true, false) + *tsoPhaseDEnabled = true + clock := kv.NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(time.Minute).UnixMilli()) + fsm := kv.NewTSOStateMachine(clock) + engine := &mainTSOEngine{state: raftengine.StateLeader, tsoState: fsm} + groups := map[uint64]*kv.ShardGroup{ + dedicatedTSORaftGroupID: {Engine: engine, TSOState: fsm}, + } + coord := newMainTSOCoordinator(clock, groups) + + wiring, err := configureCoordinatorTSO(coord, groups, mainTSOFloorProvider{}) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, wiring.Close()) }) + decorated := kv.WithKeyVizLabel(startupGatedCoordinator{inner: coord}, keyviz.LabelRedis) + + readTS, err := kv.BeginReadTimestampThrough(context.Background(), decorated, 42, "test decorated phase D") + require.NoError(t, err) + require.NotZero(t, readTS.Timestamp()) + require.True(t, fsm.CutoverActive()) + require.True(t, fsm.PhaseDActive()) + require.Equal(t, uint64(3), engine.proposals.Load(), + "cutover and phase-D markers must commit before the timestamp window") +} + +func TestConfigureCoordinatorTSOPhaseDRequiresCutoverMode(t *testing.T) { + setTSOModeFlags(t, false, false) + *tsoPhaseDEnabled = true + coord := newMainTSOCoordinator(kv.NewHLC(), nil) + + _, err := configureCoordinatorTSO(coord, nil) + require.ErrorContains(t, err, "requires --tsoEnabled") +} + func TestConfigureCoordinatorTSOShadowReturnsLegacyTimestamp(t *testing.T) { setTSOModeFlags(t, false, true) clock := kv.NewHLC() @@ -156,7 +192,7 @@ func newActiveMainTSOCutover(t *testing.T) (*kv.HLC, *kv.TSOStateMachine, *mainT group := &kv.ShardGroup{Engine: engine, TSOState: fsm} allocator, err := kv.NewRaftTSOAllocator(group, clock, kv.WithTSOCutoverFloorProvider(mainTSOFloorProvider{})) require.NoError(t, err) - _, err = allocator.ReserveBatchAfter(context.Background(), 1, 0, true) + _, err = allocator.ReserveBatchAfter(context.Background(), 1, 0, true, false) require.NoError(t, err) require.True(t, fsm.CutoverActive()) engine.proposals.Store(0) @@ -167,13 +203,16 @@ func setTSOModeFlags(t *testing.T, enabled, shadow bool) { t.Helper() oldEnabled := *tsoEnabled oldShadow := *tsoShadowEnabled + oldPhaseD := *tsoPhaseDEnabled oldBatchSize := *tsoBatchSize *tsoEnabled = enabled *tsoShadowEnabled = shadow + *tsoPhaseDEnabled = false *tsoBatchSize = 8 t.Cleanup(func() { *tsoEnabled = oldEnabled *tsoShadowEnabled = oldShadow + *tsoPhaseDEnabled = oldPhaseD *tsoBatchSize = oldBatchSize }) } diff --git a/proto/distribution.pb.go b/proto/distribution.pb.go index 99aaa98b2..bc0ccbfaa 100644 --- a/proto/distribution.pb.go +++ b/proto/distribution.pb.go @@ -368,8 +368,11 @@ type GetTimestampRequest struct { // issuance. Once committed, shadow clients return TSO timestamps instead of // legacy HLC candidates, so a rolling cutover cannot reintroduce them. ActivateCutover bool `protobuf:"varint,3,opt,name=activate_cutover,json=activateCutover,proto3" json:"activate_cutover,omitempty"` - unknownFields protoimpl.UnknownFields - sizeCache protoimpl.SizeCache + // ActivatePhaseD closes the legacy compatibility window after every member + // understands the M7 FSM entry. It requires ActivateCutover and is one-way. + ActivatePhaseD bool `protobuf:"varint,4,opt,name=activate_phase_d,json=activatePhaseD,proto3" json:"activate_phase_d,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *GetTimestampRequest) Reset() { @@ -423,6 +426,13 @@ func (x *GetTimestampRequest) GetActivateCutover() bool { return false } +func (x *GetTimestampRequest) GetActivatePhaseD() bool { + if x != nil { + return x.ActivatePhaseD + } + return false +} + type GetTimestampResponse struct { state protoimpl.MessageState `protogen:"open.v1"` Timestamp uint64 `protobuf:"varint,1,opt,name=timestamp,proto3" json:"timestamp,omitempty"` @@ -438,6 +448,10 @@ type GetTimestampResponse struct { // CutoverActive reports the durable group-0 cutover marker after applying // this request. Shadow nodes switch to the returned TSO timestamp when true. CutoverActive bool `protobuf:"varint,5,opt,name=cutover_active,json=cutoverActive,proto3" json:"cutover_active,omitempty"` + // PhaseDActive reports that legacy per-shard issuance has been retired. + PhaseDActive bool `protobuf:"varint,6,opt,name=phase_d_active,json=phaseDActive,proto3" json:"phase_d_active,omitempty"` + // PhaseDFloor is the allocation floor captured by the Phase-D marker. + PhaseDFloor uint64 `protobuf:"varint,7,opt,name=phase_d_floor,json=phaseDFloor,proto3" json:"phase_d_floor,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -507,6 +521,132 @@ func (x *GetTimestampResponse) GetCutoverActive() bool { return false } +func (x *GetTimestampResponse) GetPhaseDActive() bool { + if x != nil { + return x.PhaseDActive + } + return false +} + +func (x *GetTimestampResponse) GetPhaseDFloor() uint64 { + if x != nil { + return x.PhaseDFloor + } + return 0 +} + +type ValidateTimestampRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + Timestamp uint64 `protobuf:"varint,1,opt,name=timestamp,proto3" json:"timestamp,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ValidateTimestampRequest) Reset() { + *x = ValidateTimestampRequest{} + mi := &file_distribution_proto_msgTypes[4] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ValidateTimestampRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ValidateTimestampRequest) ProtoMessage() {} + +func (x *ValidateTimestampRequest) ProtoReflect() protoreflect.Message { + mi := &file_distribution_proto_msgTypes[4] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ValidateTimestampRequest.ProtoReflect.Descriptor instead. +func (*ValidateTimestampRequest) Descriptor() ([]byte, []int) { + return file_distribution_proto_rawDescGZIP(), []int{4} +} + +func (x *ValidateTimestampRequest) GetTimestamp() uint64 { + if x != nil { + return x.Timestamp + } + return 0 +} + +type ValidateTimestampResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Valid bool `protobuf:"varint,1,opt,name=valid,proto3" json:"valid,omitempty"` + PhaseDActive bool `protobuf:"varint,2,opt,name=phase_d_active,json=phaseDActive,proto3" json:"phase_d_active,omitempty"` + PhaseDFloor uint64 `protobuf:"varint,3,opt,name=phase_d_floor,json=phaseDFloor,proto3" json:"phase_d_floor,omitempty"` + AllocationFloor uint64 `protobuf:"varint,4,opt,name=allocation_floor,json=allocationFloor,proto3" json:"allocation_floor,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ValidateTimestampResponse) Reset() { + *x = ValidateTimestampResponse{} + mi := &file_distribution_proto_msgTypes[5] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ValidateTimestampResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ValidateTimestampResponse) ProtoMessage() {} + +func (x *ValidateTimestampResponse) ProtoReflect() protoreflect.Message { + mi := &file_distribution_proto_msgTypes[5] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ValidateTimestampResponse.ProtoReflect.Descriptor instead. +func (*ValidateTimestampResponse) Descriptor() ([]byte, []int) { + return file_distribution_proto_rawDescGZIP(), []int{5} +} + +func (x *ValidateTimestampResponse) GetValid() bool { + if x != nil { + return x.Valid + } + return false +} + +func (x *ValidateTimestampResponse) GetPhaseDActive() bool { + if x != nil { + return x.PhaseDActive + } + return false +} + +func (x *ValidateTimestampResponse) GetPhaseDFloor() uint64 { + if x != nil { + return x.PhaseDFloor + } + return 0 +} + +func (x *ValidateTimestampResponse) GetAllocationFloor() uint64 { + if x != nil { + return x.AllocationFloor + } + return 0 +} + type RouteDescriptor struct { state protoimpl.MessageState `protogen:"open.v1"` RouteId uint64 `protobuf:"varint,1,opt,name=route_id,json=routeId,proto3" json:"route_id,omitempty"` @@ -522,7 +662,7 @@ type RouteDescriptor struct { func (x *RouteDescriptor) Reset() { *x = RouteDescriptor{} - mi := &file_distribution_proto_msgTypes[4] + mi := &file_distribution_proto_msgTypes[6] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -534,7 +674,7 @@ func (x *RouteDescriptor) String() string { func (*RouteDescriptor) ProtoMessage() {} func (x *RouteDescriptor) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[4] + mi := &file_distribution_proto_msgTypes[6] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -547,7 +687,7 @@ func (x *RouteDescriptor) ProtoReflect() protoreflect.Message { // Deprecated: Use RouteDescriptor.ProtoReflect.Descriptor instead. func (*RouteDescriptor) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{4} + return file_distribution_proto_rawDescGZIP(), []int{6} } func (x *RouteDescriptor) GetRouteId() uint64 { @@ -615,7 +755,7 @@ type SplitJobBracketProgress struct { func (x *SplitJobBracketProgress) Reset() { *x = SplitJobBracketProgress{} - mi := &file_distribution_proto_msgTypes[5] + mi := &file_distribution_proto_msgTypes[7] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -627,7 +767,7 @@ func (x *SplitJobBracketProgress) String() string { func (*SplitJobBracketProgress) ProtoMessage() {} func (x *SplitJobBracketProgress) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[5] + mi := &file_distribution_proto_msgTypes[7] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -640,7 +780,7 @@ func (x *SplitJobBracketProgress) ProtoReflect() protoreflect.Message { // Deprecated: Use SplitJobBracketProgress.ProtoReflect.Descriptor instead. func (*SplitJobBracketProgress) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{5} + return file_distribution_proto_rawDescGZIP(), []int{7} } func (x *SplitJobBracketProgress) GetBracketId() uint64 { @@ -740,7 +880,7 @@ type SplitJob struct { func (x *SplitJob) Reset() { *x = SplitJob{} - mi := &file_distribution_proto_msgTypes[6] + mi := &file_distribution_proto_msgTypes[8] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -752,7 +892,7 @@ func (x *SplitJob) String() string { func (*SplitJob) ProtoMessage() {} func (x *SplitJob) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[6] + mi := &file_distribution_proto_msgTypes[8] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -765,7 +905,7 @@ func (x *SplitJob) ProtoReflect() protoreflect.Message { // Deprecated: Use SplitJob.ProtoReflect.Descriptor instead. func (*SplitJob) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{6} + return file_distribution_proto_rawDescGZIP(), []int{8} } func (x *SplitJob) GetJobId() uint64 { @@ -1007,7 +1147,7 @@ type ListRoutesRequest struct { func (x *ListRoutesRequest) Reset() { *x = ListRoutesRequest{} - mi := &file_distribution_proto_msgTypes[7] + mi := &file_distribution_proto_msgTypes[9] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1019,7 +1159,7 @@ func (x *ListRoutesRequest) String() string { func (*ListRoutesRequest) ProtoMessage() {} func (x *ListRoutesRequest) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[7] + mi := &file_distribution_proto_msgTypes[9] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1032,7 +1172,7 @@ func (x *ListRoutesRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use ListRoutesRequest.ProtoReflect.Descriptor instead. func (*ListRoutesRequest) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{7} + return file_distribution_proto_rawDescGZIP(), []int{9} } type ListRoutesResponse struct { @@ -1045,7 +1185,7 @@ type ListRoutesResponse struct { func (x *ListRoutesResponse) Reset() { *x = ListRoutesResponse{} - mi := &file_distribution_proto_msgTypes[8] + mi := &file_distribution_proto_msgTypes[10] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1057,7 +1197,7 @@ func (x *ListRoutesResponse) String() string { func (*ListRoutesResponse) ProtoMessage() {} func (x *ListRoutesResponse) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[8] + mi := &file_distribution_proto_msgTypes[10] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1070,7 +1210,7 @@ func (x *ListRoutesResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use ListRoutesResponse.ProtoReflect.Descriptor instead. func (*ListRoutesResponse) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{8} + return file_distribution_proto_rawDescGZIP(), []int{10} } func (x *ListRoutesResponse) GetCatalogVersion() uint64 { @@ -1098,7 +1238,7 @@ type SplitRangeRequest struct { func (x *SplitRangeRequest) Reset() { *x = SplitRangeRequest{} - mi := &file_distribution_proto_msgTypes[9] + mi := &file_distribution_proto_msgTypes[11] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1110,7 +1250,7 @@ func (x *SplitRangeRequest) String() string { func (*SplitRangeRequest) ProtoMessage() {} func (x *SplitRangeRequest) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[9] + mi := &file_distribution_proto_msgTypes[11] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1123,7 +1263,7 @@ func (x *SplitRangeRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use SplitRangeRequest.ProtoReflect.Descriptor instead. func (*SplitRangeRequest) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{9} + return file_distribution_proto_rawDescGZIP(), []int{11} } func (x *SplitRangeRequest) GetExpectedCatalogVersion() uint64 { @@ -1158,7 +1298,7 @@ type SplitRangeResponse struct { func (x *SplitRangeResponse) Reset() { *x = SplitRangeResponse{} - mi := &file_distribution_proto_msgTypes[10] + mi := &file_distribution_proto_msgTypes[12] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1170,7 +1310,7 @@ func (x *SplitRangeResponse) String() string { func (*SplitRangeResponse) ProtoMessage() {} func (x *SplitRangeResponse) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[10] + mi := &file_distribution_proto_msgTypes[12] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1183,7 +1323,7 @@ func (x *SplitRangeResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use SplitRangeResponse.ProtoReflect.Descriptor instead. func (*SplitRangeResponse) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{10} + return file_distribution_proto_rawDescGZIP(), []int{12} } func (x *SplitRangeResponse) GetCatalogVersion() uint64 { @@ -1217,17 +1357,27 @@ const file_distribution_proto_rawDesc = "" + "\x10GetRouteResponse\x12\x14\n" + "\x05start\x18\x01 \x01(\fR\x05start\x12\x10\n" + "\x03end\x18\x02 \x01(\fR\x03end\x12\"\n" + - "\rraft_group_id\x18\x03 \x01(\x04R\vraftGroupId\"{\n" + + "\rraft_group_id\x18\x03 \x01(\x04R\vraftGroupId\"\xa5\x01\n" + "\x13GetTimestampRequest\x12\x14\n" + "\x05count\x18\x01 \x01(\rR\x05count\x12#\n" + "\rmin_timestamp\x18\x02 \x01(\x04R\fminTimestamp\x12)\n" + - "\x10activate_cutover\x18\x03 \x01(\bR\x0factivateCutover\"\xea\x01\n" + + "\x10activate_cutover\x18\x03 \x01(\bR\x0factivateCutover\x12(\n" + + "\x10activate_phase_d\x18\x04 \x01(\bR\x0eactivatePhaseD\"\xb4\x02\n" + "\x14GetTimestampResponse\x12\x1c\n" + "\ttimestamp\x18\x01 \x01(\x04R\ttimestamp\x12;\n" + "\x1acommitted_by_dedicated_tso\x18\x02 \x01(\bR\x17committedByDedicatedTso\x12\x14\n" + "\x05count\x18\x03 \x01(\rR\x05count\x12:\n" + "\x19previous_allocation_floor\x18\x04 \x01(\x04R\x17previousAllocationFloor\x12%\n" + - "\x0ecutover_active\x18\x05 \x01(\bR\rcutoverActive\"\xe5\x01\n" + + "\x0ecutover_active\x18\x05 \x01(\bR\rcutoverActive\x12$\n" + + "\x0ephase_d_active\x18\x06 \x01(\bR\fphaseDActive\x12\"\n" + + "\rphase_d_floor\x18\a \x01(\x04R\vphaseDFloor\"8\n" + + "\x18ValidateTimestampRequest\x12\x1c\n" + + "\ttimestamp\x18\x01 \x01(\x04R\ttimestamp\"\xa6\x01\n" + + "\x19ValidateTimestampResponse\x12\x14\n" + + "\x05valid\x18\x01 \x01(\bR\x05valid\x12$\n" + + "\x0ephase_d_active\x18\x02 \x01(\bR\fphaseDActive\x12\"\n" + + "\rphase_d_floor\x18\x03 \x01(\x04R\vphaseDFloor\x12)\n" + + "\x10allocation_floor\x18\x04 \x01(\x04R\x0fallocationFloor\"\xe5\x01\n" + "\x0fRouteDescriptor\x12\x19\n" + "\broute_id\x18\x01 \x01(\x04R\arouteId\x12\x14\n" + "\x05start\x18\x02 \x01(\fR\x05start\x12\x10\n" + @@ -1326,10 +1476,11 @@ const file_distribution_proto_rawDesc = "" + "\x13SplitJobExportPhase\x12\x1f\n" + "\x1bSPLIT_JOB_EXPORT_PHASE_NONE\x10\x00\x12#\n" + "\x1fSPLIT_JOB_EXPORT_PHASE_BACKFILL\x10\x01\x12%\n" + - "!SPLIT_JOB_EXPORT_PHASE_DELTA_COPY\x10\x022\xf2\x01\n" + + "!SPLIT_JOB_EXPORT_PHASE_DELTA_COPY\x10\x022\xc0\x02\n" + "\fDistribution\x121\n" + "\bGetRoute\x12\x10.GetRouteRequest\x1a\x11.GetRouteResponse\"\x00\x12=\n" + - "\fGetTimestamp\x12\x14.GetTimestampRequest\x1a\x15.GetTimestampResponse\"\x00\x127\n" + + "\fGetTimestamp\x12\x14.GetTimestampRequest\x1a\x15.GetTimestampResponse\"\x00\x12L\n" + + "\x11ValidateTimestamp\x12\x19.ValidateTimestampRequest\x1a\x1a.ValidateTimestampResponse\"\x00\x127\n" + "\n" + "ListRoutes\x12\x12.ListRoutesRequest\x1a\x13.ListRoutesResponse\"\x00\x127\n" + "\n" + @@ -1348,23 +1499,25 @@ func file_distribution_proto_rawDescGZIP() []byte { } var file_distribution_proto_enumTypes = make([]protoimpl.EnumInfo, 4) -var file_distribution_proto_msgTypes = make([]protoimpl.MessageInfo, 11) +var file_distribution_proto_msgTypes = make([]protoimpl.MessageInfo, 13) var file_distribution_proto_goTypes = []any{ - (RouteState)(0), // 0: RouteState - (SplitJobPhase)(0), // 1: SplitJobPhase - (SplitJobBarrierState)(0), // 2: SplitJobBarrierState - (SplitJobExportPhase)(0), // 3: SplitJobExportPhase - (*GetRouteRequest)(nil), // 4: GetRouteRequest - (*GetRouteResponse)(nil), // 5: GetRouteResponse - (*GetTimestampRequest)(nil), // 6: GetTimestampRequest - (*GetTimestampResponse)(nil), // 7: GetTimestampResponse - (*RouteDescriptor)(nil), // 8: RouteDescriptor - (*SplitJobBracketProgress)(nil), // 9: SplitJobBracketProgress - (*SplitJob)(nil), // 10: SplitJob - (*ListRoutesRequest)(nil), // 11: ListRoutesRequest - (*ListRoutesResponse)(nil), // 12: ListRoutesResponse - (*SplitRangeRequest)(nil), // 13: SplitRangeRequest - (*SplitRangeResponse)(nil), // 14: SplitRangeResponse + (RouteState)(0), // 0: RouteState + (SplitJobPhase)(0), // 1: SplitJobPhase + (SplitJobBarrierState)(0), // 2: SplitJobBarrierState + (SplitJobExportPhase)(0), // 3: SplitJobExportPhase + (*GetRouteRequest)(nil), // 4: GetRouteRequest + (*GetRouteResponse)(nil), // 5: GetRouteResponse + (*GetTimestampRequest)(nil), // 6: GetTimestampRequest + (*GetTimestampResponse)(nil), // 7: GetTimestampResponse + (*ValidateTimestampRequest)(nil), // 8: ValidateTimestampRequest + (*ValidateTimestampResponse)(nil), // 9: ValidateTimestampResponse + (*RouteDescriptor)(nil), // 10: RouteDescriptor + (*SplitJobBracketProgress)(nil), // 11: SplitJobBracketProgress + (*SplitJob)(nil), // 12: SplitJob + (*ListRoutesRequest)(nil), // 13: ListRoutesRequest + (*ListRoutesResponse)(nil), // 14: ListRoutesResponse + (*SplitRangeRequest)(nil), // 15: SplitRangeRequest + (*SplitRangeResponse)(nil), // 16: SplitRangeResponse } var file_distribution_proto_depIdxs = []int32{ 0, // 0: RouteDescriptor.state:type_name -> RouteState @@ -1374,20 +1527,22 @@ var file_distribution_proto_depIdxs = []int32{ 1, // 4: SplitJob.abandon_from_phase:type_name -> SplitJobPhase 2, // 5: SplitJob.cutover_read_fence_state:type_name -> SplitJobBarrierState 2, // 6: SplitJob.target_staged_readiness_state:type_name -> SplitJobBarrierState - 9, // 7: SplitJob.bracket_progress:type_name -> SplitJobBracketProgress - 8, // 8: ListRoutesResponse.routes:type_name -> RouteDescriptor - 8, // 9: SplitRangeResponse.left:type_name -> RouteDescriptor - 8, // 10: SplitRangeResponse.right:type_name -> RouteDescriptor + 11, // 7: SplitJob.bracket_progress:type_name -> SplitJobBracketProgress + 10, // 8: ListRoutesResponse.routes:type_name -> RouteDescriptor + 10, // 9: SplitRangeResponse.left:type_name -> RouteDescriptor + 10, // 10: SplitRangeResponse.right:type_name -> RouteDescriptor 4, // 11: Distribution.GetRoute:input_type -> GetRouteRequest 6, // 12: Distribution.GetTimestamp:input_type -> GetTimestampRequest - 11, // 13: Distribution.ListRoutes:input_type -> ListRoutesRequest - 13, // 14: Distribution.SplitRange:input_type -> SplitRangeRequest - 5, // 15: Distribution.GetRoute:output_type -> GetRouteResponse - 7, // 16: Distribution.GetTimestamp:output_type -> GetTimestampResponse - 12, // 17: Distribution.ListRoutes:output_type -> ListRoutesResponse - 14, // 18: Distribution.SplitRange:output_type -> SplitRangeResponse - 15, // [15:19] is the sub-list for method output_type - 11, // [11:15] is the sub-list for method input_type + 8, // 13: Distribution.ValidateTimestamp:input_type -> ValidateTimestampRequest + 13, // 14: Distribution.ListRoutes:input_type -> ListRoutesRequest + 15, // 15: Distribution.SplitRange:input_type -> SplitRangeRequest + 5, // 16: Distribution.GetRoute:output_type -> GetRouteResponse + 7, // 17: Distribution.GetTimestamp:output_type -> GetTimestampResponse + 9, // 18: Distribution.ValidateTimestamp:output_type -> ValidateTimestampResponse + 14, // 19: Distribution.ListRoutes:output_type -> ListRoutesResponse + 16, // 20: Distribution.SplitRange:output_type -> SplitRangeResponse + 16, // [16:21] is the sub-list for method output_type + 11, // [11:16] is the sub-list for method input_type 11, // [11:11] is the sub-list for extension type_name 11, // [11:11] is the sub-list for extension extendee 0, // [0:11] is the sub-list for field type_name @@ -1404,7 +1559,7 @@ func file_distribution_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_distribution_proto_rawDesc), len(file_distribution_proto_rawDesc)), NumEnums: 4, - NumMessages: 11, + NumMessages: 13, NumExtensions: 0, NumServices: 1, }, diff --git a/proto/distribution.proto b/proto/distribution.proto index 493bfa274..ba6d52fe9 100644 --- a/proto/distribution.proto +++ b/proto/distribution.proto @@ -5,6 +5,7 @@ option go_package = "github.com/bootjp/elastickv/proto"; service Distribution { rpc GetRoute (GetRouteRequest) returns (GetRouteResponse) {} rpc GetTimestamp (GetTimestampRequest) returns (GetTimestampResponse) {} + rpc ValidateTimestamp (ValidateTimestampRequest) returns (ValidateTimestampResponse) {} rpc ListRoutes (ListRoutesRequest) returns (ListRoutesResponse) {} rpc SplitRange (SplitRangeRequest) returns (SplitRangeResponse) {} } @@ -32,6 +33,9 @@ message GetTimestampRequest { // issuance. Once committed, shadow clients return TSO timestamps instead of // legacy HLC candidates, so a rolling cutover cannot reintroduce them. bool activate_cutover = 3; + // ActivatePhaseD closes the legacy compatibility window after every member + // understands the M7 FSM entry. It requires ActivateCutover and is one-way. + bool activate_phase_d = 4; } message GetTimestampResponse { @@ -48,6 +52,21 @@ message GetTimestampResponse { // CutoverActive reports the durable group-0 cutover marker after applying // this request. Shadow nodes switch to the returned TSO timestamp when true. bool cutover_active = 5; + // PhaseDActive reports that legacy per-shard issuance has been retired. + bool phase_d_active = 6; + // PhaseDFloor is the allocation floor captured by the Phase-D marker. + uint64 phase_d_floor = 7; +} + +message ValidateTimestampRequest { + uint64 timestamp = 1; +} + +message ValidateTimestampResponse { + bool valid = 1; + bool phase_d_active = 2; + uint64 phase_d_floor = 3; + uint64 allocation_floor = 4; } enum RouteState { diff --git a/proto/distribution_grpc.pb.go b/proto/distribution_grpc.pb.go index f8d9b82e4..006e7d482 100644 --- a/proto/distribution_grpc.pb.go +++ b/proto/distribution_grpc.pb.go @@ -19,10 +19,11 @@ import ( const _ = grpc.SupportPackageIsVersion9 const ( - Distribution_GetRoute_FullMethodName = "/Distribution/GetRoute" - Distribution_GetTimestamp_FullMethodName = "/Distribution/GetTimestamp" - Distribution_ListRoutes_FullMethodName = "/Distribution/ListRoutes" - Distribution_SplitRange_FullMethodName = "/Distribution/SplitRange" + Distribution_GetRoute_FullMethodName = "/Distribution/GetRoute" + Distribution_GetTimestamp_FullMethodName = "/Distribution/GetTimestamp" + Distribution_ValidateTimestamp_FullMethodName = "/Distribution/ValidateTimestamp" + Distribution_ListRoutes_FullMethodName = "/Distribution/ListRoutes" + Distribution_SplitRange_FullMethodName = "/Distribution/SplitRange" ) // DistributionClient is the client API for Distribution service. @@ -31,6 +32,7 @@ const ( type DistributionClient interface { GetRoute(ctx context.Context, in *GetRouteRequest, opts ...grpc.CallOption) (*GetRouteResponse, error) GetTimestamp(ctx context.Context, in *GetTimestampRequest, opts ...grpc.CallOption) (*GetTimestampResponse, error) + ValidateTimestamp(ctx context.Context, in *ValidateTimestampRequest, opts ...grpc.CallOption) (*ValidateTimestampResponse, error) ListRoutes(ctx context.Context, in *ListRoutesRequest, opts ...grpc.CallOption) (*ListRoutesResponse, error) SplitRange(ctx context.Context, in *SplitRangeRequest, opts ...grpc.CallOption) (*SplitRangeResponse, error) } @@ -63,6 +65,16 @@ func (c *distributionClient) GetTimestamp(ctx context.Context, in *GetTimestampR return out, nil } +func (c *distributionClient) ValidateTimestamp(ctx context.Context, in *ValidateTimestampRequest, opts ...grpc.CallOption) (*ValidateTimestampResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(ValidateTimestampResponse) + err := c.cc.Invoke(ctx, Distribution_ValidateTimestamp_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + func (c *distributionClient) ListRoutes(ctx context.Context, in *ListRoutesRequest, opts ...grpc.CallOption) (*ListRoutesResponse, error) { cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) out := new(ListRoutesResponse) @@ -89,6 +101,7 @@ func (c *distributionClient) SplitRange(ctx context.Context, in *SplitRangeReque type DistributionServer interface { GetRoute(context.Context, *GetRouteRequest) (*GetRouteResponse, error) GetTimestamp(context.Context, *GetTimestampRequest) (*GetTimestampResponse, error) + ValidateTimestamp(context.Context, *ValidateTimestampRequest) (*ValidateTimestampResponse, error) ListRoutes(context.Context, *ListRoutesRequest) (*ListRoutesResponse, error) SplitRange(context.Context, *SplitRangeRequest) (*SplitRangeResponse, error) mustEmbedUnimplementedDistributionServer() @@ -107,6 +120,9 @@ func (UnimplementedDistributionServer) GetRoute(context.Context, *GetRouteReques func (UnimplementedDistributionServer) GetTimestamp(context.Context, *GetTimestampRequest) (*GetTimestampResponse, error) { return nil, status.Error(codes.Unimplemented, "method GetTimestamp not implemented") } +func (UnimplementedDistributionServer) ValidateTimestamp(context.Context, *ValidateTimestampRequest) (*ValidateTimestampResponse, error) { + return nil, status.Error(codes.Unimplemented, "method ValidateTimestamp not implemented") +} func (UnimplementedDistributionServer) ListRoutes(context.Context, *ListRoutesRequest) (*ListRoutesResponse, error) { return nil, status.Error(codes.Unimplemented, "method ListRoutes not implemented") } @@ -170,6 +186,24 @@ func _Distribution_GetTimestamp_Handler(srv interface{}, ctx context.Context, de return interceptor(ctx, in, info, handler) } +func _Distribution_ValidateTimestamp_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(ValidateTimestampRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(DistributionServer).ValidateTimestamp(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: Distribution_ValidateTimestamp_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(DistributionServer).ValidateTimestamp(ctx, req.(*ValidateTimestampRequest)) + } + return interceptor(ctx, in, info, handler) +} + func _Distribution_ListRoutes_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { in := new(ListRoutesRequest) if err := dec(in); err != nil { @@ -221,6 +255,10 @@ var Distribution_ServiceDesc = grpc.ServiceDesc{ MethodName: "GetTimestamp", Handler: _Distribution_GetTimestamp_Handler, }, + { + MethodName: "ValidateTimestamp", + Handler: _Distribution_ValidateTimestamp_Handler, + }, { MethodName: "ListRoutes", Handler: _Distribution_ListRoutes_Handler, From 416c056a6dfa919100dd889c7f4a4850abc3a4b9 Mon Sep 17 00:00:00 2001 From: bootjp Date: Sun, 19 Jul 2026 21:37:30 +0900 Subject: [PATCH 04/14] tso: complete centralized runtime semantics --- kv/tso.go | 45 +++++++++++++++++++++++++++++++++++++-------- 1 file changed, 37 insertions(+), 8 deletions(-) diff --git a/kv/tso.go b/kv/tso.go index c92840db1..2031122d1 100644 --- a/kv/tso.go +++ b/kv/tso.go @@ -199,24 +199,53 @@ func NextTimestampAfterThrough(ctx context.Context, coord Coordinator, startTS u return nextTimestampAfterObserved(ctx, coord, clock, startTS, label) } -// TimestampAllocatorThrough returns the allocator configured behind a -// coordinator or coordinator decorator. +// TimestampAllocatorThrough returns the currently active allocator behind a +// coordinator or coordinator decorator. A runtime allocator in legacy mode is +// intentionally hidden so normal write paths use the coordinator HLC fallback. func TimestampAllocatorThrough(coord Coordinator) (TimestampAllocator, bool) { if provider, ok := coord.(TimestampAllocatorProvider); ok { - alloc := provider.TimestampAllocator() - return alloc, alloc != nil + return resolveTimestampAllocator(provider.TimestampAllocator()) } switch c := coord.(type) { case *Coordinate: - if c != nil && c.tsAllocator != nil { - return c.tsAllocator, true + if c != nil { + return resolveTimestampAllocator(c.tsAllocator) } return nil, false case *ShardedCoordinator: - if c != nil && c.tsAllocator != nil { - return c.tsAllocator, true + if c != nil { + return resolveTimestampAllocator(c.tsAllocator) } return nil, false + default: + alloc, ok := coord.(TimestampAllocator) + if !ok { + return nil, false + } + return resolveTimestampAllocator(alloc) + } +} + +// ConfiguredTimestampAllocatorThrough returns the configured allocator without +// resolving runtime-mode decorators. It is used by long-lived internal servers +// that must keep a DynamicTimestampAllocator reference across later mode-file +// reloads while still falling back to HLC when that allocator reports legacy. +func ConfiguredTimestampAllocatorThrough(coord Coordinator) (TimestampAllocator, bool) { + if provider, ok := coord.(TimestampAllocatorProvider); ok { + alloc := provider.TimestampAllocator() + return alloc, alloc != nil + } + switch c := coord.(type) { + case *Coordinate: + if c == nil || c.tsAllocator == nil { + return nil, false + } + return c.tsAllocator, true + case *ShardedCoordinator: + if c == nil || c.tsAllocator == nil { + return nil, false + } + return c.tsAllocator, true default: alloc, ok := coord.(TimestampAllocator) return alloc, ok From 9b9b7cf1e73cdaf2eb9175cbde8c0fb65ad39dde Mon Sep 17 00:00:00 2001 From: bootjp Date: Sun, 19 Jul 2026 23:24:30 +0900 Subject: [PATCH 05/14] tso: add runtime operations and observability --- Makefile | 47 +- adapter/internal.go | 28 +- adapter/internal_test.go | 52 ++ docs/centralized_tso_operations.md | 115 +++++ ...2026_04_16_implemented_centralized_tso.md} | 100 ++-- .../2026_05_28_implemented_tla_safety_spec.md | 15 +- .../2026_06_23_proposed_scaling_roadmap.md | 44 +- kv/coordinator.go | 6 +- kv/keyviz_label.go | 2 +- kv/sharded_coordinator.go | 14 +- kv/tso.go | 15 + kv/tso_fanout_benchmark_test.go | 140 ++++++ kv/tso_raft.go | 129 ++++- kv/tso_runtime.go | 453 ++++++++++++++++++ kv/tso_runtime_test.go | 407 ++++++++++++++++ main.go | 195 +++++--- main_tso_routing_test.go | 58 ++- monitoring/prometheus/rules/tso-alerts.yml | 103 ++++ monitoring/registry.go | 13 + monitoring/tso.go | 146 ++++++ monitoring/tso_test.go | 48 ++ 21 files changed, 1940 insertions(+), 190 deletions(-) create mode 100644 docs/centralized_tso_operations.md rename docs/design/{2026_04_16_partial_centralized_tso.md => 2026_04_16_implemented_centralized_tso.md} (90%) create mode 100644 kv/tso_fanout_benchmark_test.go create mode 100644 kv/tso_runtime.go create mode 100644 kv/tso_runtime_test.go create mode 100644 monitoring/prometheus/rules/tso-alerts.yml create mode 100644 monitoring/tso.go create mode 100644 monitoring/tso_test.go diff --git a/Makefile b/Makefile index 9dc641ba6..c7e3a8f70 100644 --- a/Makefile +++ b/Makefile @@ -29,35 +29,40 @@ TLA_LIB := ../lib .PHONY: tla-check tla-tools -tla-tools: $(TLA_JAR) - -$(TLA_JAR): +tla-tools: @mkdir -p $(dir $(TLA_JAR)) - @echo "Downloading tla2tools.jar $(TLA_VERSION)..." - @curl -fsSL -o $(TLA_JAR).tmp $(TLA_URL) - @# Prefer sha256sum (GNU coreutils, universal on Linux); fall back to - @# shasum -a 256 (default on macOS). Either yields the same hex - @# digest in the first whitespace-delimited field. - @if command -v sha256sum >/dev/null 2>&1; then \ - actual=$$(sha256sum $(TLA_JAR).tmp | awk '{print $$1}'); \ - elif command -v shasum >/dev/null 2>&1; then \ - actual=$$(shasum -a 256 $(TLA_JAR).tmp | awk '{print $$1}'); \ - else \ - echo "ERROR: neither sha256sum nor shasum is available."; \ - rm -f $(TLA_JAR).tmp; \ - exit 1; \ + @set -eu; \ + sha256_file() { \ + if command -v sha256sum >/dev/null 2>&1; then \ + sha256sum "$$1" | awk '{print $$1}'; \ + elif command -v shasum >/dev/null 2>&1; then \ + shasum -a 256 "$$1" | awk '{print $$1}'; \ + else \ + echo "ERROR: neither sha256sum nor shasum is available." >&2; \ + return 1; \ + fi; \ + }; \ + actual=""; \ + if [ -f "$(TLA_JAR)" ]; then actual=$$(sha256_file "$(TLA_JAR)"); fi; \ + if [ "$$actual" = "$(TLA_SHA256)" ]; then \ + echo "tla2tools.jar ready at $(TLA_JAR) (SHA-256 $(TLA_SHA256))"; \ + exit 0; \ fi; \ + echo "Downloading tla2tools.jar $(TLA_VERSION)..."; \ + trap 'rm -f "$(TLA_JAR).tmp"' EXIT HUP INT TERM; \ + curl -fsSL -o "$(TLA_JAR).tmp" "$(TLA_URL)"; \ + actual=$$(sha256_file "$(TLA_JAR).tmp"); \ if [ "$$actual" != "$(TLA_SHA256)" ]; then \ echo "ERROR: tla2tools.jar SHA-256 mismatch."; \ echo " expected: $(TLA_SHA256)"; \ echo " actual: $$actual"; \ - rm -f $(TLA_JAR).tmp; \ exit 1; \ - fi - @mv $(TLA_JAR).tmp $(TLA_JAR) - @echo "tla2tools.jar ready at $(TLA_JAR) (SHA-256 $(TLA_SHA256))" + fi; \ + mv "$(TLA_JAR).tmp" "$(TLA_JAR)"; \ + trap - EXIT HUP INT TERM; \ + echo "tla2tools.jar ready at $(TLA_JAR) (SHA-256 $(TLA_SHA256))" -tla-check: $(TLA_JAR) +tla-check: tla-tools @# Per-module orchestration lives in scripts/tla-check.sh so adding @# M3..M5 only needs an entry in that script's TLA_MODULES array @# and a `case` line for the gap-invariant string. The script does diff --git a/adapter/internal.go b/adapter/internal.go index bf4279e31..c78bde7ea 100644 --- a/adapter/internal.go +++ b/adapter/internal.go @@ -101,11 +101,17 @@ func (i *Internal) stampTimestamps(ctx context.Context, req *pb.ForwardRequest) func (i *Internal) nextTimestamp(ctx context.Context, label string) (uint64, error) { if i.tsAllocator != nil { ts, err := i.tsAllocator.Next(ctx) - if err != nil { + if err == nil { + return ts, nil + } + if !errors.Is(err, kv.ErrTSOAllocatorRequired) { return 0, errors.Wrap(err, label) } - return ts, nil } + return i.nextLegacyTimestamp(label) +} + +func (i *Internal) nextLegacyTimestamp(label string) (uint64, error) { if i.clock == nil { return 1, nil } @@ -121,15 +127,23 @@ func (i *Internal) nextTimestampAfter(ctx context.Context, min uint64, label str return 0, errors.Wrap(kv.ErrTxnCommitTSRequired, label) } if i.tsAllocator != nil { - return i.nextTimestampAfterFromAllocator(ctx, min, label) + ts, err := i.nextTimestampAfterFromAllocator(ctx, min, label) + if err == nil { + return ts, nil + } + if !errors.Is(err, kv.ErrTSOAllocatorRequired) { + return 0, err + } } + return i.nextLegacyTimestampAfter(min, label) +} + +func (i *Internal) nextLegacyTimestampAfter(min uint64, label string) (uint64, error) { if i.clock == nil { return min + 1, nil } - if i.clock != nil { - i.clock.Observe(min) - } - ts, err := i.nextTimestamp(ctx, label) + i.clock.Observe(min) + ts, err := i.nextLegacyTimestamp(label) if err != nil { return 0, err } diff --git a/adapter/internal_test.go b/adapter/internal_test.go index fcabc5356..a10fbb775 100644 --- a/adapter/internal_test.go +++ b/adapter/internal_test.go @@ -4,9 +4,11 @@ import ( "context" "encoding/binary" "testing" + "time" "github.com/bootjp/elastickv/kv" pb "github.com/bootjp/elastickv/proto" + "github.com/cockroachdb/errors" "github.com/stretchr/testify/require" ) @@ -83,6 +85,35 @@ func TestFillForwardedTxnCommitTS_AssignsCommitTS(t *testing.T) { require.Equal(t, meta.CommitTS, commitTS) } +func TestFillForwardedTxnCommitTS_FallsBackWhenRuntimeAllocatorLegacy(t *testing.T) { + t.Parallel() + + clock := kv.NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(time.Minute).UnixMilli()) + i := &Internal{ + clock: clock, + tsAllocator: internalLegacyRuntimeAllocator{}, + } + startTS := uint64(10) + reqs := []*pb.Request{ + { + IsTxn: true, + Phase: pb.Phase_COMMIT, + Mutations: []*pb.Mutation{ + { + Op: pb.Op_PUT, + Key: []byte(kv.TxnMetaPrefix), + Value: kv.EncodeTxnMeta(kv.TxnMeta{PrimaryKey: []byte("k"), CommitTS: 0}), + }, + }, + }, + } + + commitTS, err := i.fillForwardedTxnCommitTS(context.Background(), reqs, startTS) + require.NoError(t, err) + require.Greater(t, commitTS, startTS) +} + func TestFillForwardedTxnCommitTS_PreservesExistingCommitTS(t *testing.T) { t.Parallel() @@ -175,6 +206,21 @@ func TestFillForwardedTxnCommitTS_StampsCommitTSValueOffset(t *testing.T) { require.Zero(t, reqs[0].Mutations[1].CommitTsValueOffset) } +func TestStampRawTimestamps_FallsBackWhenRuntimeAllocatorLegacy(t *testing.T) { + t.Parallel() + + clock := kv.NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(time.Minute).UnixMilli()) + i := &Internal{ + clock: clock, + tsAllocator: internalLegacyRuntimeAllocator{}, + } + reqs := []*pb.Request{{Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte("k"), Value: []byte("v")}}}} + + require.NoError(t, i.stampRawTimestamps(context.Background(), reqs)) + require.NotZero(t, reqs[0].Ts) +} + func TestFillForwardedTxnCommitTS_PrepareAllowsAlreadyStampedOffsets(t *testing.T) { t.Parallel() @@ -241,3 +287,9 @@ func TestStampTxnTimestamps_UsesSingleTxnStartTS(t *testing.T) { require.Greater(t, meta.CommitTS, uint64(9)) require.Equal(t, meta.CommitTS, commitTS) } + +type internalLegacyRuntimeAllocator struct{} + +func (internalLegacyRuntimeAllocator) Next(context.Context) (uint64, error) { + return 0, errors.WithStack(kv.ErrTSOAllocatorRequired) +} diff --git a/docs/centralized_tso_operations.md b/docs/centralized_tso_operations.md new file mode 100644 index 000000000..c30bb8f8c --- /dev/null +++ b/docs/centralized_tso_operations.md @@ -0,0 +1,115 @@ +# Centralized TSO Operations + +This runbook covers runtime mode changes, production signals, and the rollout +gates for the dedicated group-0 timestamp oracle. The implementation design is +in `docs/design/2026_04_16_implemented_centralized_tso.md`. + +## Runtime configuration + +Start every node with the same atomically replaceable mode file: + +```text +--tsoModeFile=/etc/elastickv/tso-mode +--tsoModeReloadInterval=5s +--tsoBatchSize=256 +``` + +The file contains exactly one mode: `legacy`, `shadow`, `cutover`, or +`phase-d`. `--tsoModeFile` cannot be combined with `--tsoEnabled`, +`--tsoShadowEnabled`, or `--tsoPhaseDEnabled`; those startup-only flags remain +for backward compatibility. Runtime mode reload requires the dedicated group 0. + +Replace the file atomically so a poll cannot observe a truncated value: + +```sh +printf '%s\n' shadow > /etc/elastickv/tso-mode.tmp +mv /etc/elastickv/tso-mode.tmp /etc/elastickv/tso-mode +``` + +An unreadable file, unknown value, backward transition, or skipped phase is +rejected without changing the live allocator. Runtime transitions must be +adjacent and one-way: + +```text +legacy -> shadow -> cutover -> phase-d +``` + +The durable group-0 markers override stale local configuration. Once cutover +or Phase D is committed, a node cannot return to legacy or shadow issuance even +if its mode file moves backward. An allocation already in flight during a +reload may finish under the preceding process mode, but cutover-side requests +still use group 0 and the durable markers never move backward. + +## Rollout gates + +1. Deploy the binary and group 0 while every mode file remains `legacy`. +2. Change every node to `shadow`. Confirm all nodes expose + `elastickv_tso_mode{mode="shadow"} == 1`. +3. Hold shadow mode until `legacy_overlap` remains zero for a full 15-minute + window, allocation errors remain below 1 percent, and allocation p99 remains + below 50 ms. +4. Change nodes to `cutover`. The first production reservation commits the + one-way cutover marker. A node still in shadow observes that marker and + returns the dedicated TSO value. +5. Confirm every binary supports Phase D and every node is on the dedicated + path. Then change nodes to `phase-d`. The first Phase-D reservation commits + its floor marker before returning a new window. + +After the cutover marker commits, rollback means restoring service around the +dedicated TSO path; it never means returning to legacy issuance. Phase D has the +same one-way rule and additionally keeps data-group HLC renewal retired. + +## Metrics and alerts + +| Signal | Meaning | Gate or threshold | +|---|---|---| +| `elastickv_tso_request_duration_seconds` | Local/remote reserve and validation attempt latency, split by outcome | Warning at reserve p99 > 50 ms for 10m; critical at > 200 ms for 5m | +| `elastickv_tso_shadow_comparisons_total` | Accepted candidates, overlap discards, cutover bypasses, and errors | `legacy_overlap` must be zero for 15m before cutover | +| `elastickv_tso_shadow_divergence_timestamps` | Absolute candidate-to-floor distance for accepted and discarded shadow comparisons | Investigate sustained growth before cutover | +| `elastickv_tso_mode` | One-hot process-local mode | More than one mode for 15m warns about a stalled rollout | +| `elastickv_tso_mode_reload_total` | Applied and rejected reloads plus file read/parse failures | Any failure in 10m warns | +| `elastickv_tso_durable_state` | Applied `cutover` and `phase_d` markers | Phase D without cutover is critical | + +The checked rules are in `monitoring/prometheus/rules/tso-alerts.yml`. Allocation +errors warn above 1 percent for 10 minutes and become critical above 10 percent +for 5 minutes. The expressions require at least 0.1 reserve attempts/second so +an idle cluster does not alert on an empty denominator. + +## Write-fanout benchmark + +Run the modeled remote-TSO benchmark with: + +```sh +go test ./kv -run '^$' -bench '^BenchmarkTSOWriteFanout$' -benchmem -benchtime=1s -count=3 +``` + +The harness uses 16 concurrent write coordinators sharing a `BatchAllocator`. +Each operation stamps 1, 3, or 8 fan-out legs, and each refill pays a modeled +1 ms remote-leader plus Raft-commit delay. The following medians were measured +on Go 1.26.5, darwin/arm64, Apple M1 Max: + +| Fan-out | Batch | Wall time/op | p99 | Refills/op | Timestamp writes/s | +|---:|---:|---:|---:|---:|---:| +| 1 | 1 | 1.33 ms | 1,249 ms | 1.000 | 753 | +| 1 | 64 | 24.81 us | 1.323 ms | 0.01563 | 40,309 | +| 1 | 256 | 8.92 us | 0.000292 ms | 0.00391 | 112,081 | +| 3 | 1 | 5.70 ms | 1,233 ms | 3.000 | 526 | +| 3 | 64 | 105.35 us | 3.212 ms | 0.04690 | 28,478 | +| 3 | 256 | 17.30 us | 1.272 ms | 0.01172 | 173,453 | +| 8 | 1 | 10.43 ms | 1,104 ms | 8.000 | 767 | +| 8 | 64 | 157.49 us | 1.316 ms | 0.1251 | 50,796 | +| 8 | 256 | 39.43 us | 1.302 ms | 0.03127 | 202,870 | + +Batch size 1 is intentionally retained as the unamortized baseline and shows +severe waiter tails under contention. The production default of 256 keeps the +modeled fan-out p99 below 4 ms in this harness, including the refill-contention +tail at fan-out 8, and below the 50 ms warning threshold. These are controlled +local measurements, not a substitute for observing the production histograms +during rollout. + +## Intentional non-goals + +This closure does not automate cluster-wide phase orchestration, choose a +dedicated subset of TSO members, or change the HLC ceiling formula. Operators +still advance the shared mode file deployment-by-deployment, and the existing +group membership and clock-floor decisions remain independent future work. diff --git a/docs/design/2026_04_16_partial_centralized_tso.md b/docs/design/2026_04_16_implemented_centralized_tso.md similarity index 90% rename from docs/design/2026_04_16_partial_centralized_tso.md rename to docs/design/2026_04_16_implemented_centralized_tso.md index 979f4dec5..6e3c0c60e 100644 --- a/docs/design/2026_04_16_partial_centralized_tso.md +++ b/docs/design/2026_04_16_implemented_centralized_tso.md @@ -1,13 +1,16 @@ # Centralized Timestamp Oracle (TSO) Design -- Status: Partial — M1-M7 are implemented, including the dedicated group-0 - FSM, leader-routed durable windows, strict term bootstrap, serialized shadow - migration, one-way rolling cutover, durable Phase-D retirement, and - cross-shard SSI timestamp validation. Runtime config reload and production - latency/alerting work remain open. -- Author: bootjp -- Date: 2026-04-16 -- Updated: 2026-07-19 +Status: Implemented +Author: bootjp +Date: 2026-04-16 +Updated: 2026-07-19 + +M1-M8 implement the central subsystem: the dedicated group-0 FSM, +leader-routed durable windows, strict term bootstrap, serialized shadow +migration, one-way cutover and Phase-D retirement, validated cross-shard +timestamps, one-way runtime mode reload, production metrics and alerts, and +write-fanout benchmark evidence. The topology and clock-policy extensions in +Section 9 are intentional non-goals and do not leave central scope incomplete. --- @@ -137,11 +140,20 @@ Implemented: 19. `BatchAllocator` validates cached candidates once Phase D is required. A candidate at/below the Phase-D floor invalidates the entire local window and forces a refill above the marker before any timestamp is returned. - -Remaining: - -1. Runtime config reload for the mode switch; current flags are startup-only. -2. Production benchmark, divergence metrics, and alert thresholds. +20. `--tsoModeFile` polls an atomically replaced `legacy`, `shadow`, `cutover`, + or `phase-d` mode. Live transitions must be adjacent and one-way. Invalid + files, skipped phases, and process-local rollbacks leave the active allocator + unchanged. Applied group-0 markers override a stale local mode immediately + on allocator resolution, before the next poll can run. +21. The production registry exports reserve/validation latency by local or + remote path and outcome, shadow comparison/divergence, process mode, every + reload outcome, and durable marker state. Checked Prometheus rules define + warning and critical latency/error thresholds, persistent reload failure, + mode divergence, and the 15-minute zero-overlap cutover gate. +22. `BenchmarkTSOWriteFanout` models 16 concurrent coordinators, 1/3/8-way + writes, a 1 ms remote TSO commit, and batch sizes 1/64/256. Three-run + evidence supports the existing default batch size 256 and is recorded in + `docs/centralized_tso_operations.md`. ### 1.1 Original Limitation @@ -748,8 +760,10 @@ approach enables a live cutover. ### 7.3 Phase C — Durable TSO Cutover -- The startup flag `--tsoEnabled` switches coordinator issuance to the - leader-routed allocator. Runtime config reload remains future work. +- The backward-compatible startup flag `--tsoEnabled`, or runtime mode file + value `cutover`, switches coordinator issuance to the leader-routed allocator. + A live transition is accepted only after `shadow`; restart recovery may start + directly in cutover when the durable marker already exists. - The first production refill commits the group-0 cutover marker before committing and returning its timestamp window. - The marker is encoded in a versioned TSO-specific envelope whose leading @@ -765,12 +779,13 @@ approach enables a live cutover. - Roll every member to a binary that understands the Phase-D entry and `ValidateTimestamp` RPC. Run all members on the dedicated TSO path before - enabling `--tsoPhaseDEnabled`; an older member would correctly halt on the - unknown control entry, so mixed-version activation is prohibited. -- The first allocation with `--tsoPhaseDEnabled` commits cutover (if needed), - then the Phase-D marker and its pre-Phase-D floor, then a new allocation - window strictly above that floor. The switch is one-way and survives restart - through the TSO V4 snapshot. + enabling `--tsoPhaseDEnabled` or reloading `phase-d`; an older member would + correctly halt on the unknown control entry, so mixed-version activation is + prohibited. +- The first allocation with Phase D requested commits cutover (if needed), then + the Phase-D marker and its pre-Phase-D floor, then a new allocation window + strictly above that floor. The switch is one-way and survives restart through + the TSO V4 snapshot. - `ShardedCoordinator` dynamically stops data-group ceiling proposals and fails closed instead of issuing from its legacy HLC. It continues group-0 renewal only. Shadow mode observes durable cutover and directly uses the @@ -812,6 +827,29 @@ candidate is returned. Once the cutover marker applies, shadow callers stop returning legacy candidates. Independently, every new TSO leader term fences its first window above the strict maximum committed data timestamp. +### 7.6 Runtime Reload and Operational Gates + +`DynamicTimestampAllocator` publishes the process-selected allocator through +an atomic pointer. Before resolving that pointer, it checks the consensus-owned +cutover and Phase-D markers. An applied marker selects the batch/routed TSO path +and enables matching remote confirmation immediately, so a stale `legacy` or +`shadow` file cannot mint a legacy timestamp while waiting for the next poll. +A request that resolved its preceding mode before local marker apply may finish, +but every subsequent resolution observes the one-way durable override. + +`TSORuntimeController` accepts only adjacent one-way live transitions: + +```text +legacy -> shadow -> cutover -> phase-d +``` + +The initial process mode may be later in the sequence for restart recovery. +No-op polls refresh the durable-state gauges, and every failed poll increments +its bounded reload-result counter even when duplicate log messages are +suppressed. Exact rollout gates, alert expressions, atomic file replacement, +and benchmark evidence are maintained in +`docs/centralized_tso_operations.md`. + --- ## 8. Milestones @@ -825,20 +863,24 @@ its first window above the strict maximum committed data timestamp. | M5 — shipped | Preserve the default-group `LocalTSOAllocator` compatibility bridge when group 0 is absent; route coordinator-owned timestamp call sites through the allocator abstraction. | Medium | | M6 — shipped | Run the dedicated group-0 FSM, fence each new TSO leader term above all authoritative data-group commit floors, redirect follower requests to the TSO leader over gRPC, synchronously serialize fail-closed shadow issuance, and commit the one-way rolling cutover marker before production windows. | Low | | M7 — shipped | Commit the durable Phase-D floor marker, preserve V3 snapshots until activation, retire data-shard HLC renewal and legacy/shadow issuance after cutover, activate before read validation while preserving exact applied snapshots through timestamp-bound per-dispatch vouchers, invalidate pre-Phase-D batch windows, and validate unvouched caller-supplied cross-shard SSI timestamps at the group-0 leader from activation onward. | Low | +| M8 — shipped | Atomically reload one-way runtime modes, make durable markers override stale local mode before allocation, expose production latency/divergence/state metrics, install checked alert thresholds, and validate the default batch size with concurrent write-fanout evidence. | Low | --- -## 9. Open Questions +## 9. Resolved Questions and Future Extensions -1. **TSO RTT when TSO leader ≠ write leader:** What batch size minimises tail - latency? Needs benchmarking against realistic write fan-out. +1. **TSO RTT when TSO leader differs from the write leader (resolved):** The + concurrent benchmark covers 1/3/8 fan-out legs with a modeled 1 ms remote + TSO commit. Batch size 1 exposes serialized refill tails, while default 256 + keeps the recorded p99 below the 50 ms warning threshold. -2. **TSO group membership:** Should all cluster nodes join the TSO group, or - should a dedicated subset (e.g. 3 out of N) be used to reduce Raft traffic? +2. **TSO group membership (intentional future extension):** Choosing all nodes + or a dedicated subset is a topology policy outside the central timestamp + subsystem. The implemented routing and fail-closed term fence support either. -3. **Clock floor semantics:** `max(now, ceiling)` vs. `ceiling + 1` — the - stricter form (`ceiling + 1`) guarantees no overlap even if wall clocks - drift, at the cost of one extra millisecond per renewal window. +3. **Clock floor semantics (intentional future extension):** The implementation + retains the existing floor formula. Changing to `ceiling + 1` is an + independent HLC policy decision, not unfinished centralized-TSO scope. 4. **Non-leader TSO requests (resolved):** Followers redirect to the current group-0 leader through `Distribution.GetTimestamp`. Timestamp allocation is 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..265e50f3d 100644 --- a/docs/design/2026_05_28_implemented_tla_safety_spec.md +++ b/docs/design/2026_05_28_implemented_tla_safety_spec.md @@ -661,12 +661,11 @@ does not keep this document in `partial`. mitigation, the doc lifecycle (Section 8.2) ties promotion to specs landing, so a stale `partial` status is a visible signal. -5. **Choice of TSO model.** The HLC spec models the current - per-shard-leader ceiling. The centralized TSO proposal - (`2026_04_16_partial_centralized_tso.md`) would change that. The two - docs are independent; if/when centralized TSO lands, `HLC.tla` (or a - sibling `TSO.tla`) is updated as part of that PR. We do **not** block - this proposal on the TSO decision. +5. **Choice of TSO model.** The HLC spec continues to model the compatibility + per-shard-leader ceiling, while the implemented centralized TSO + (`2026_04_16_implemented_centralized_tso.md`) owns the production ordering + path. A sibling `TSO.tla` remains an independent future extension; runtime + transition safety is pinned by the TSO FSM, routing, reload, and race tests. 6. **TLC tool licensing and binary distribution.** `tla2tools.jar` is MIT-licensed but not in any package manager we already depend on. @@ -716,8 +715,8 @@ does not keep this document in `partial`. - `distribution/engine.go`, `distribution/catalog.go`, `distribution/watcher.go` — route catalog. - `docs/architecture_overview.md` — system-level diagrams. -- `docs/design/2026_04_16_partial_centralized_tso.md` — the TSO proposal - that this spec is independent of (Section 9, risk 5). +- `docs/design/2026_04_16_implemented_centralized_tso.md` — the implemented TSO + design that this spec is independent of (Section 9, risk 5). - Diego Ongaro's Raft TLA+ specification — reference for the abstract `Raft.tla` interface. - CockroachDB and TiDB MVCC / HLC TLA+ models — public prior art for diff --git a/docs/design/2026_06_23_proposed_scaling_roadmap.md b/docs/design/2026_06_23_proposed_scaling_roadmap.md index fe74eda05..4d92db147 100644 --- a/docs/design/2026_06_23_proposed_scaling_roadmap.md +++ b/docs/design/2026_06_23_proposed_scaling_roadmap.md @@ -202,38 +202,18 @@ memory each group's private cache/memtable pins. ### (d) Cross-group transactions at scale -- **Per-group HLC vs centralized TSO.** Today the ceiling is renewed only on - the default group (§1.5); the TSO doc - (`docs/design/2026_04_16_partial_centralized_tso.md`) has shipped the - near-term fix (renew on all led groups, in parallel — its M1) and still - leaves the dedicated TSO Raft group with a batch allocator open. - **Assessment of whether TSO is still needed once groups multiply:** the - near-term per-group fix (TSO doc §6) was *necessary regardless* — it is the - minimum correctness fix the moment a node leads a non-default group it is - member of, and it has landed before multi-node multi-group bootstrap (b) - makes that topology reachable. The *full* dedicated-TSO component (TSO doc - M6–M7) is a larger component, but the need for a shared ordering source is - not a throughput/amortization question once cross-node cross-group - transactions are possible: with the per-group fix in place, every node's - timestamps are self-monotonic, but **global** monotonicity across coordinator - nodes is still not guaranteed without a shared oracle (TSO doc §6 - "Guarantee"). Cross-group transactions (`kv/transaction.go`, - `kv/txn_codec.go`) whose timestamps can be allocated by different ingress / - coordinator nodes need a single ordering source for OCC commit-ts - comparability, regardless of where the participating groups' current leaders - live. `LeaderProxy.Commit` / `Internal.Forward` preserve non-zero timestamps, - so leader co-location does not prove single-clock allocation. (The shared-snapshot - invariant — every operation in one txn reading at the *same* `startTS` — is - already upheld: `nextStartTS` allocates one `startTS` for the whole txn and - propagates it via `reqs.StartTS` to every participating group; the gap is the - cross-*coordinator* comparability of the per-txn `commitTS`, not the per-txn - `startTS`. See OQ-1.) So: per-group renewal fix is in-scope-soon and - load-bearing for one node leading multiple groups; before enabling any - cross-group transaction mode in which more than one coordinator node can issue - `startTS` / `commitTS`, the roadmap must either pull the dedicated TSO group - forward or land a narrower single-oracle bridge that pins cross-group - timestamp allocation to one designated leader. Until such txns are enabled, - the per-group fix plus the shared-`*HLC` property remains adequate. +- **Per-group HLC vs centralized TSO.** The centralized TSO design + (`docs/design/2026_04_16_implemented_centralized_tso.md`) has shipped M1-M8: + all-led-group compatibility renewal, the dedicated group-0 FSM, leader-routed + durable windows, one-way cutover and Phase D, runtime mode reload, and + operational gates. The shared ordering source is available for cross-node + cross-group transactions. Deployments that remain in `legacy` still have + only per-node monotonicity; before enabling multiple coordinator nodes for + cross-group issuance, they must complete the documented `shadow -> cutover` + sequence. `LeaderProxy.Commit` / `Internal.Forward` preserve non-zero + timestamps, and one `startTS` remains shared by every transaction participant. + Phase D additionally validates a caller-supplied cross-shard `StartTS` at the + group-0 leader before commit allocation. ### (e) Cluster size / membership diff --git a/kv/coordinator.go b/kv/coordinator.go index af91a09fd..eda72056f 100644 --- a/kv/coordinator.go +++ b/kv/coordinator.go @@ -1012,11 +1012,11 @@ func (c *Coordinate) allocateTimestampAfter(ctx context.Context, label string, m if min == ^uint64(0) { return 0, errors.Wrap(ErrTxnCommitTSRequired, label) } - if c.tsAllocator != nil { + if allocator, ok := resolveTimestampAllocator(c.tsAllocator); ok { if min > 0 { - return nextTimestampAfterFromAllocator(ctx, c.tsAllocator, min, label) + return nextTimestampAfterFromAllocator(ctx, allocator, min, label) } - return nextTimestampFromAllocator(ctx, c.tsAllocator, label) + return nextTimestampFromAllocator(ctx, allocator, label) } if c.clock == nil { return 0, errors.Wrap(ErrTSOClockNil, label) diff --git a/kv/keyviz_label.go b/kv/keyviz_label.go index ba40e30d5..09a3629c3 100644 --- a/kv/keyviz_label.go +++ b/kv/keyviz_label.go @@ -91,7 +91,7 @@ func (c keyVizLabeledCoordinator) RaftLeaderForKey(key []byte) string { func (c keyVizLabeledCoordinator) Clock() *HLC { return c.inner.Clock() } func (c keyVizLabeledCoordinator) TimestampAllocator() TimestampAllocator { - alloc, _ := TimestampAllocatorThrough(c.inner) + alloc, _ := ConfiguredTimestampAllocatorThrough(c.inner) return alloc } diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index 3705fcf4c..4a54f3b3e 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -1629,11 +1629,11 @@ func (c *ShardedCoordinator) allocateTimestampAfter(ctx context.Context, label s if min == ^uint64(0) { return 0, errors.Wrap(ErrTxnCommitTSRequired, label) } - if c.tsAllocator != nil { + if allocator, ok := resolveTimestampAllocator(c.tsAllocator); ok { if min > 0 { - return nextTimestampAfterFromAllocator(ctx, c.tsAllocator, min, label) + return nextTimestampAfterFromAllocator(ctx, allocator, min, label) } - return nextTimestampFromAllocator(ctx, c.tsAllocator, label) + return nextTimestampFromAllocator(ctx, allocator, label) } if c.tsoCutoverState != nil && c.tsoCutoverState.PhaseDActive() { return 0, errors.Wrap(ErrTSOAllocatorRequired, @@ -1663,7 +1663,8 @@ func (c *ShardedCoordinator) NextAfter(ctx context.Context, min uint64) (uint64, } func (c *ShardedCoordinator) nextTxnTSAfter(ctx context.Context, startTS uint64) (uint64, error) { - if c.clock == nil && c.tsAllocator == nil { + _, allocatorConfigured := resolveTimestampAllocator(c.tsAllocator) + if c.clock == nil && !allocatorConfigured { nextTS := startTS + 1 if nextTS == 0 { return 0, nil @@ -1725,7 +1726,8 @@ func (c *ShardedCoordinator) nextStartTS(ctx context.Context, elems []*Elem[OP]) if c.clock != nil && maxTS > 0 { c.clock.Observe(maxTS) } - if c.clock == nil && c.tsAllocator == nil { + _, allocatorConfigured := resolveTimestampAllocator(c.tsAllocator) + if c.clock == nil && !allocatorConfigured { return maxTS + 1, nil } ts, err := c.allocateTimestampAfter(ctx, "allocate sharded startTS", maxTS) @@ -2195,7 +2197,7 @@ func (c *ShardedCoordinator) rawLogs(ctx context.Context, reqs *OperationGroup[O } func (c *ShardedCoordinator) rawLogTimestamp(ctx context.Context) (uint64, error) { - if c.tsAllocator != nil { + if _, ok := resolveTimestampAllocator(c.tsAllocator); ok { return 0, nil } return c.allocateTimestamp(ctx, "allocate sharded raw log ts") diff --git a/kv/tso.go b/kv/tso.go index 2031122d1..0f0be8181 100644 --- a/kv/tso.go +++ b/kv/tso.go @@ -252,6 +252,21 @@ func ConfiguredTimestampAllocatorThrough(coord Coordinator) (TimestampAllocator, } } +type currentTimestampAllocatorProvider interface { + currentTimestampAllocator() TimestampAllocator +} + +func resolveTimestampAllocator(alloc TimestampAllocator) (TimestampAllocator, bool) { + for range 4 { + provider, ok := alloc.(currentTimestampAllocatorProvider) + if !ok { + return alloc, alloc != nil + } + alloc = provider.currentTimestampAllocator() + } + return alloc, alloc != nil +} + func coordinatorTimestampAllocator(coord Coordinator) (TimestampAllocator, bool) { return TimestampAllocatorThrough(coord) } diff --git a/kv/tso_fanout_benchmark_test.go b/kv/tso_fanout_benchmark_test.go new file mode 100644 index 000000000..00e759c7a --- /dev/null +++ b/kv/tso_fanout_benchmark_test.go @@ -0,0 +1,140 @@ +package kv + +import ( + "context" + "fmt" + "sort" + "sync" + "sync/atomic" + "testing" + "time" +) + +// BenchmarkTSOWriteFanout models 16 concurrent write coordinators sharing one +// process-local BatchAllocator. Each logical operation stamps every fan-out +// leg, while a refill pays a 1 ms remote leader plus Raft commit delay. +func BenchmarkTSOWriteFanout(b *testing.B) { + const ( + writers = 16 + tsoRTT = time.Millisecond + ) + for _, fanout := range []int{1, 3, 8} { + for _, batchSize := range []int{1, 64, 256} { + name := fmt.Sprintf("writers=%d/fanout=%d/batch=%d/rtt=1ms", writers, fanout, batchSize) + b.Run(name, func(b *testing.B) { + runTSOWriteFanoutBenchmark(b, writers, fanout, batchSize, tsoRTT) + }) + } + } +} + +func runTSOWriteFanoutBenchmark(b *testing.B, writers, fanout, batchSize int, rtt time.Duration) { + backend := &benchmarkTSOAllocator{rtt: rtt} + allocator, err := NewBatchAllocator(backend, batchSize) + if err != nil { + b.Fatal(err) + } + latencies := make([]int64, b.N) + var next atomic.Uint64 + var firstErr atomic.Pointer[benchmarkTSOError] + ctx := context.Background() + b.ResetTimer() + started := time.Now() + var wg sync.WaitGroup + for range writers { + wg.Add(1) + go runTSOWriteFanoutWorker(ctx, allocator, fanout, latencies, &next, &firstErr, &wg) + } + wg.Wait() + elapsed := time.Since(started) + b.StopTimer() + reportTSOWriteFanoutMetrics(b, fanout, latencies, backend, elapsed, firstErr.Load()) +} + +func runTSOWriteFanoutWorker( + ctx context.Context, + allocator TimestampAllocator, + fanout int, + latencies []int64, + next *atomic.Uint64, + firstErr *atomic.Pointer[benchmarkTSOError], + wg *sync.WaitGroup, +) { + defer wg.Done() + for { + index := next.Add(1) - 1 + if index >= uint64(len(latencies)) { //nolint:gosec // benchmark sizes are non-negative. + return + } + opStarted := time.Now() + for range fanout { + if _, err := allocator.Next(ctx); err != nil { + firstErr.CompareAndSwap(nil, &benchmarkTSOError{err: err}) + return + } + } + latencies[index] = time.Since(opStarted).Nanoseconds() + } +} + +func reportTSOWriteFanoutMetrics( + b *testing.B, + fanout int, + latencies []int64, + backend *benchmarkTSOAllocator, + elapsed time.Duration, + recorded *benchmarkTSOError, +) { + if recorded != nil { + b.Fatal(recorded.err) + } + sort.Slice(latencies, func(i, j int) bool { return latencies[i] < latencies[j] }) + b.ReportMetric(percentileDuration(latencies, 50).Seconds()*1000, "p50-ms") + b.ReportMetric(percentileDuration(latencies, 99).Seconds()*1000, "p99-ms") + b.ReportMetric(float64(backend.refills.Load())/float64(b.N), "refills/op") + b.ReportMetric(float64(b.N*fanout)/elapsed.Seconds(), "writes/s") +} + +func percentileDuration(sorted []int64, percentile int) time.Duration { + if len(sorted) == 0 { + return 0 + } + index := (len(sorted)*percentile + 99) / 100 + if index > 0 { + index-- + } + return time.Duration(sorted[index]) +} + +type benchmarkTSOError struct { + err error +} + +type benchmarkTSOAllocator struct { + rtt time.Duration + next atomic.Uint64 + refills atomic.Uint64 +} + +func (a *benchmarkTSOAllocator) Next(ctx context.Context) (uint64, error) { + return a.NextBatch(ctx, 1) +} + +func (a *benchmarkTSOAllocator) NextBatch(ctx context.Context, n int) (uint64, error) { + timer := time.NewTimer(a.rtt) + defer timer.Stop() + select { + case <-ctx.Done(): + return 0, ctx.Err() + case <-timer.C: + } + a.refills.Add(1) + end := a.next.Add(uint64(n)) //nolint:gosec // benchmark batch sizes are positive. + return end - uint64(n) + 1, nil //nolint:gosec // benchmark batch sizes are positive. +} + +func (a *benchmarkTSOAllocator) IsLeader() bool { return false } + +func (a *benchmarkTSOAllocator) RunLeaseRenewal(ctx context.Context) { + <-ctx.Done() +} diff --git a/kv/tso_raft.go b/kv/tso_raft.go index 3975f5733..63f06d637 100644 --- a/kv/tso_raft.go +++ b/kv/tso_raft.go @@ -5,6 +5,7 @@ import ( stderrors "errors" "log/slog" "sync" + "sync/atomic" "time" "github.com/bootjp/elastickv/internal/raftengine" @@ -389,29 +390,36 @@ type LeaderRoutedTSOAllocator struct { remoteRequest tsoRemoteRequest retryBudget time.Duration retryInterval time.Duration - activate bool - activatePhaseD bool + activation atomic.Uint32 clock *HLC remoteValidate tsoRemoteValidation + observer TSOObserver } +const ( + tsoActivationInactive uint32 = iota + tsoActivationCutover + tsoActivationPhaseD +) + type LeaderRoutedTSOAllocatorOption func(*LeaderRoutedTSOAllocator) func WithTSOCutoverActivation() LeaderRoutedTSOAllocatorOption { - return func(a *LeaderRoutedTSOAllocator) { a.activate = true } + return func(a *LeaderRoutedTSOAllocator) { a.setActivation(true, false) } } func WithTSOPhaseDActivation() LeaderRoutedTSOAllocatorOption { - return func(a *LeaderRoutedTSOAllocator) { - a.activate = true - a.activatePhaseD = true - } + return func(a *LeaderRoutedTSOAllocator) { a.setActivation(true, true) } } func WithTSORoutedClock(clock *HLC) LeaderRoutedTSOAllocatorOption { return func(a *LeaderRoutedTSOAllocator) { a.clock = clock } } +func WithTSOObserver(observer TSOObserver) LeaderRoutedTSOAllocatorOption { + return func(a *LeaderRoutedTSOAllocator) { a.observer = observer } +} + func NewLeaderRoutedTSOAllocator( local TSOAllocator, leader raftengine.LeaderView, @@ -459,7 +467,8 @@ func (a *LeaderRoutedTSOAllocator) NextBatchAfter(ctx context.Context, n int, mi } func (a *LeaderRoutedTSOAllocator) nextBatchAfter(ctx context.Context, n int, min uint64) (uint64, error) { - reservation, err := a.nextReservation(ctx, n, min, a.activate, a.activatePhaseD, a.activate) + activate, activatePhaseD := a.activationState() + reservation, err := a.nextReservation(ctx, n, min, activate, activatePhaseD, activate) if err != nil { return 0, err } @@ -524,19 +533,28 @@ func (a *LeaderRoutedTSOAllocator) tryBatch( ) (TSOReservation, error) { var empty TSOReservation if a.local.IsLeader() { - return a.tryLocalBatch(ctx, n, min, activate, activatePhaseD, requireReservation) + started := time.Now() + reservation, err := a.tryLocalBatch(ctx, n, min, activate, activatePhaseD, requireReservation) + a.observeRequest("reserve", "local", err, time.Since(started)) + return reservation, err } addr := leaderAddrFromEngine(a.leader) if addr == "" { - return empty, errors.WithStack(ErrLeaderNotFound) + err := errors.WithStack(ErrLeaderNotFound) + a.observeRequest("reserve", "remote", err, 0) + return empty, err } + started := time.Now() reservation, err := a.remoteRequest(ctx, addr, n, min, activate, activatePhaseD) if err != nil { + a.observeRequest("reserve", "remote", err, time.Since(started)) return empty, err } if err := validateTSOReservation(reservation, n, min, requireReservation); err != nil { + a.observeRequest("reserve", "remote", err, time.Since(started)) return empty, err } + a.observeRequest("reserve", "remote", nil, time.Since(started)) return reservation, nil } @@ -629,13 +647,18 @@ func (a *LeaderRoutedTSOAllocator) ValidateDurableTimestamp(ctx context.Context, if !ok { return errors.WithStack(ErrTSOProtocolUnsupported) } + started := time.Now() lastErr = errors.WithStack(validator.ValidateDurableTimestamp(ctx, timestamp)) + a.observeRequest("validate", "local", lastErr, time.Since(started)) } else { addr := leaderAddrFromEngine(a.leader) if addr == "" { lastErr = errors.WithStack(ErrLeaderNotFound) + a.observeRequest("validate", "remote", lastErr, 0) } else { + started := time.Now() lastErr = a.remoteValidate(ctx, addr, timestamp) + a.observeRequest("validate", "remote", lastErr, time.Since(started)) } } if lastErr == nil { @@ -679,7 +702,8 @@ func (a *LeaderRoutedTSOAllocator) PhaseDActive() bool { } func (a *LeaderRoutedTSOAllocator) PhaseDRequired() bool { - return a != nil && (a.activatePhaseD || a.PhaseDActive()) + _, activatePhaseD := a.activationState() + return a != nil && (activatePhaseD || a.PhaseDActive()) } func (a *LeaderRoutedTSOAllocator) observeReservation(reservation TSOReservation) { @@ -688,6 +712,65 @@ func (a *LeaderRoutedTSOAllocator) observeReservation(reservation TSOReservation } end := reservation.Base + uint64(reservation.Count) - 1 //nolint:gosec // validated reservation. a.clock.Observe(end) + if a.observer != nil { + a.observer.ObserveTSODurableState(reservation.CutoverActive, reservation.PhaseDActive) + } +} + +// setActivation atomically publishes the durable markers future reservations +// must request. In-flight reservations may complete under the prior mode; all +// such modes still use group 0, and committed markers remain one-way. +func (a *LeaderRoutedTSOAllocator) setActivation(cutover, phaseD bool) { + if a == nil { + return + } + state := tsoActivationInactive + if cutover || phaseD { + state = tsoActivationCutover + } + if phaseD { + state = tsoActivationPhaseD + } + a.activation.Store(state) +} + +// promoteActivation advances the requested durable marker without allowing a +// concurrent observer to lower an already-requested Phase D activation. +func (a *LeaderRoutedTSOAllocator) promoteActivation(cutover, phaseD bool) { + if a == nil { + return + } + target := tsoActivationInactive + if cutover || phaseD { + target = tsoActivationCutover + } + if phaseD { + target = tsoActivationPhaseD + } + for current := a.activation.Load(); current < target; current = a.activation.Load() { + if a.activation.CompareAndSwap(current, target) { + return + } + } +} + +func (a *LeaderRoutedTSOAllocator) activationState() (bool, bool) { + if a == nil { + return false, false + } + state := a.activation.Load() + return state >= tsoActivationCutover, state >= tsoActivationPhaseD +} + +func (a *LeaderRoutedTSOAllocator) observeRequest(operation, path string, err error, duration time.Duration) { + if a == nil || a.observer == nil { + return + } + outcome := "success" + if err != nil { + outcome = "error" + } + a.observer.ObserveTSORequest(operation, path, outcome, duration) } func validateRoutedTSOWindow(base uint64, n int, min uint64) error { @@ -794,10 +877,11 @@ func nonNilTSOContext(ctx context.Context) context.Context { // and retried; once the durable cutover marker is active, the allocator returns // the reserved TSO timestamp directly so rolling restarts cannot mix sources. type ShadowTimestampAllocator struct { - legacy *HLC - shadow TSOShadowReservationAllocator - cutover tsoCutoverState - log *slog.Logger + legacy *HLC + shadow TSOShadowReservationAllocator + cutover tsoCutoverState + log *slog.Logger + observer TSOObserver } type ShadowTimestampAllocatorOption func(*ShadowTimestampAllocator) @@ -806,6 +890,10 @@ func WithTSOShadowCutoverState(state tsoCutoverState) ShadowTimestampAllocatorOp return func(a *ShadowTimestampAllocator) { a.cutover = state } } +func WithTSOShadowObserver(observer TSOObserver) ShadowTimestampAllocatorOption { + return func(a *ShadowTimestampAllocator) { a.observer = observer } +} + func NewShadowTimestampAllocator( legacy *HLC, shadow TSOShadowReservationAllocator, @@ -848,6 +936,7 @@ func (a *ShadowTimestampAllocator) NextAfter(ctx context.Context, min uint64) (u func (a *ShadowTimestampAllocator) nextAfter(ctx context.Context, min uint64) (uint64, error) { ctx = nonNilTSOContext(ctx) if a.cutover != nil && a.cutover.CutoverActive() { + a.observeShadowComparison("cutover_active", 0) return a.nextDedicatedAfter(ctx, min) } a.legacy.Observe(min) @@ -861,6 +950,7 @@ func (a *ShadowTimestampAllocator) nextAfter(ctx context.Context, min uint64) (u } reservation, err := a.shadow.ValidateShadowTimestamp(ctx, legacyTS) if err != nil { + a.observeShadowComparison("error", 0) a.log.ErrorContext(ctx, "tso shadow allocation failed", slog.Uint64("legacy_ts", legacyTS), slog.Any("err", err), @@ -869,11 +959,14 @@ func (a *ShadowTimestampAllocator) nextAfter(ctx context.Context, min uint64) (u } a.legacy.Observe(reservation.Base) if reservation.CutoverActive { + a.observeShadowComparison("cutover_active", 0) return reservation.Base, nil } if legacyTS > reservation.PreviousAllocationFloor { + a.observeShadowComparison("legacy_ahead", legacyTS-reservation.PreviousAllocationFloor) return legacyTS, nil } + a.observeShadowComparison("legacy_overlap", reservation.PreviousAllocationFloor-legacyTS) a.log.WarnContext(ctx, "tso shadow discarded overlapping legacy timestamp", slog.Uint64("legacy_ts", legacyTS), slog.Uint64("previous_tso_floor", reservation.PreviousAllocationFloor), @@ -911,3 +1004,9 @@ func (a *ShadowTimestampAllocator) PhaseDRequired() bool { func (a *ShadowTimestampAllocator) Close() error { return nil } + +func (a *ShadowTimestampAllocator) observeShadowComparison(result string, divergence uint64) { + if a != nil && a.observer != nil { + a.observer.ObserveTSOShadowComparison(result, divergence) + } +} diff --git a/kv/tso_runtime.go b/kv/tso_runtime.go new file mode 100644 index 000000000..fb341b724 --- /dev/null +++ b/kv/tso_runtime.go @@ -0,0 +1,453 @@ +package kv + +import ( + "context" + "log/slog" + "os" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/cockroachdb/errors" +) + +var ( + ErrInvalidTSOMode = errors.New("tso: invalid runtime mode") + ErrTSOModeRollback = errors.New("tso: runtime mode rollback is prohibited") + ErrUnsafeTSOModeTransition = errors.New("tso: runtime mode transition skips a required phase") +) + +// TSOMode is the process-local timestamp issuance mode. Its ordering is part +// of the migration contract: runtime reload may only move to the next value. +type TSOMode uint32 + +const ( + TSOModeLegacy TSOMode = iota + TSOModeShadow + TSOModeCutover + TSOModePhaseD +) + +func ParseTSOMode(raw string) (TSOMode, error) { + switch strings.ToLower(strings.TrimSpace(raw)) { + case "legacy": + return TSOModeLegacy, nil + case "shadow": + return TSOModeShadow, nil + case "cutover": + return TSOModeCutover, nil + case "phase-d", "phase_d", "phased": + return TSOModePhaseD, nil + default: + return TSOModeLegacy, errors.Wrapf(ErrInvalidTSOMode, "%q", strings.TrimSpace(raw)) + } +} + +func (m TSOMode) String() string { + switch m { + case TSOModeLegacy: + return "legacy" + case TSOModeShadow: + return "shadow" + case TSOModeCutover: + return "cutover" + case TSOModePhaseD: + return "phase-d" + default: + return "invalid" + } +} + +func validTSOMode(mode TSOMode) bool { + return mode <= TSOModePhaseD +} + +// TSOObserver is the bounded-cardinality operational surface for production +// allocation latency, shadow divergence, durable state, and mode reloads. +type TSOObserver interface { + ObserveTSORequest(operation, path, outcome string, duration time.Duration) + ObserveTSOShadowComparison(result string, divergence uint64) + ObserveTSOMode(mode string) + ObserveTSOModeReload(result string) + ObserveTSODurableState(cutoverActive, phaseDActive bool) +} + +type timestampAllocatorSlot struct { + allocator TimestampAllocator +} + +// DynamicTimestampAllocator atomically publishes an optional allocator. A nil +// current allocator deliberately means "use the coordinator's legacy HLC +// path"; callers must resolve it through TimestampAllocatorThrough rather than +// calling Next directly. +type DynamicTimestampAllocator struct { + current atomic.Pointer[timestampAllocatorSlot] + durableState TSORuntimeState + durableAllocator TimestampAllocator + routed *LeaderRoutedTSOAllocator +} + +func NewDynamicTimestampAllocator(initial TimestampAllocator) *DynamicTimestampAllocator { + d := &DynamicTimestampAllocator{} + d.store(initial) + return d +} + +func (d *DynamicTimestampAllocator) store(allocator TimestampAllocator) { + if d == nil { + return + } + if allocator == nil { + d.current.Store(nil) + return + } + d.current.Store(×tampAllocatorSlot{allocator: allocator}) +} + +func (d *DynamicTimestampAllocator) currentTimestampAllocator() TimestampAllocator { + if d == nil { + return nil + } + if allocator := d.durableTimestampAllocator(); allocator != nil { + return allocator + } + slot := d.current.Load() + if slot == nil { + return nil + } + return slot.allocator +} + +// configureDurableOverride installs immutable references before the dynamic +// allocator is published. An applied one-way marker must override a stale mode +// file immediately rather than waiting for the next reload poll. +func (d *DynamicTimestampAllocator) configureDurableOverride( + state TSORuntimeState, + allocator TimestampAllocator, + routed *LeaderRoutedTSOAllocator, +) { + d.durableState = state + d.durableAllocator = allocator + d.routed = routed +} + +func (d *DynamicTimestampAllocator) durableTimestampAllocator() TimestampAllocator { + if d.durableState == nil || d.durableAllocator == nil { + return nil + } + phaseD := d.durableState.PhaseDActive() + if !phaseD && !d.durableState.CutoverActive() { + return nil + } + d.routed.promoteActivation(true, phaseD) + return d.durableAllocator +} + +func (d *DynamicTimestampAllocator) Next(ctx context.Context) (uint64, error) { + allocator := d.currentTimestampAllocator() + if allocator == nil { + return 0, errors.WithStack(ErrTSOAllocatorRequired) + } + timestamp, err := allocator.Next(ctx) + return timestamp, errors.Wrap(err, "dynamic tso next") +} + +func (d *DynamicTimestampAllocator) NextAfter(ctx context.Context, min uint64) (uint64, error) { + allocator := d.currentTimestampAllocator() + if allocator == nil { + return 0, errors.WithStack(ErrTSOAllocatorRequired) + } + if after, ok := allocator.(TimestampAfterAllocator); ok { + timestamp, err := after.NextAfter(ctx, min) + return timestamp, errors.Wrap(err, "dynamic tso next after") + } + return nextTimestampAfterFromAllocator(ctx, allocator, min, "dynamic tso next after") +} + +func (d *DynamicTimestampAllocator) ValidateDurableTimestamp(ctx context.Context, timestamp uint64) error { + allocator := d.currentTimestampAllocator() + if allocator == nil { + return errors.WithStack(ErrTSOAllocatorRequired) + } + validator, ok := allocator.(DurableTimestampValidator) + if !ok { + return errors.WithStack(ErrTSOProtocolUnsupported) + } + return errors.WithStack(validator.ValidateDurableTimestamp(ctx, timestamp)) +} + +func (d *DynamicTimestampAllocator) PhaseDActive() bool { + state, ok := d.currentTimestampAllocator().(TSOPhaseDState) + return ok && state.PhaseDActive() +} + +func (d *DynamicTimestampAllocator) PhaseDRequired() bool { + state, ok := d.currentTimestampAllocator().(TSOPhaseDState) + return ok && state.PhaseDRequired() +} + +func (d *DynamicTimestampAllocator) Invalidate() { + invalidateTimestampWindow(d.currentTimestampAllocator()) +} + +// TSORuntimeState is the consensus-owned one-way migration state. +type TSORuntimeState interface { + CutoverActive() bool + PhaseDActive() bool +} + +type TSORuntimeControllerConfig struct { + Clock *HLC + Routed *LeaderRoutedTSOAllocator + State TSORuntimeState + BatchSize int + InitialMode TSOMode + Logger *slog.Logger + Observer TSOObserver +} + +// TSORuntimeController owns the process-local mode while treating group-0's +// durable markers as authoritative. InitialMode may start at any phase for a +// restart, but ApplyMode requires adjacent, one-way runtime transitions. +type TSORuntimeController struct { + dynamic *DynamicTimestampAllocator + shadow *ShadowTimestampAllocator + batch *BatchAllocator + routed *LeaderRoutedTSOAllocator + state TSORuntimeState + observer TSOObserver + + mu sync.Mutex + mode atomic.Uint32 +} + +func NewTSORuntimeController(cfg TSORuntimeControllerConfig) (*TSORuntimeController, error) { + if cfg.Clock == nil { + return nil, errors.WithStack(ErrTSOClockNil) + } + if cfg.Routed == nil { + return nil, errors.WithStack(ErrTSOAllocatorRequired) + } + if cfg.State == nil { + return nil, errors.WithStack(ErrTSOStateRequired) + } + if !validTSOMode(cfg.InitialMode) { + return nil, errors.Wrapf(ErrInvalidTSOMode, "%d", cfg.InitialMode) + } + if cfg.Logger == nil { + cfg.Logger = slog.Default() + } + shadow, err := NewShadowTimestampAllocator( + cfg.Clock, + cfg.Routed, + cfg.Logger, + WithTSOShadowCutoverState(cfg.State), + WithTSOShadowObserver(cfg.Observer), + ) + if err != nil { + return nil, errors.Wrap(err, "configure runtime tso shadow allocator") + } + batch, err := NewBatchAllocator(cfg.Routed, cfg.BatchSize) + if err != nil { + return nil, errors.Wrap(err, "configure runtime tso batch allocator") + } + dynamic := NewDynamicTimestampAllocator(nil) + dynamic.configureDurableOverride(cfg.State, batch, cfg.Routed) + controller := &TSORuntimeController{ + dynamic: dynamic, + shadow: shadow, + batch: batch, + routed: cfg.Routed, + state: cfg.State, + observer: cfg.Observer, + } + controller.installMode(controller.effectiveMode(cfg.InitialMode)) + return controller, nil +} + +func (c *TSORuntimeController) Allocator() *DynamicTimestampAllocator { + if c == nil { + return nil + } + return c.dynamic +} + +func (c *TSORuntimeController) CurrentMode() TSOMode { + if c == nil { + return TSOModeLegacy + } + return TSOMode(c.mode.Load()) +} + +func (c *TSORuntimeController) EffectiveMode(requested TSOMode) TSOMode { + if c == nil { + return requested + } + return c.effectiveMode(requested) +} + +func (c *TSORuntimeController) effectiveMode(requested TSOMode) TSOMode { + if c.state != nil && c.state.PhaseDActive() { + return TSOModePhaseD + } + if c.state != nil && c.state.CutoverActive() && requested < TSOModeCutover { + return TSOModeCutover + } + return requested +} + +func (c *TSORuntimeController) ApplyMode(requested TSOMode) error { + if c == nil { + return errors.WithStack(ErrTSOAllocatorRequired) + } + if !validTSOMode(requested) { + c.observeReload("rejected") + return errors.Wrapf(ErrInvalidTSOMode, "%d", requested) + } + c.mu.Lock() + defer c.mu.Unlock() + target := c.effectiveMode(requested) + durableFloor := c.effectiveMode(TSOModeLegacy) + current := c.CurrentMode() + if target < current { + c.observeReload("rejected") + return errors.Wrapf(ErrTSOModeRollback, "%s -> %s", current, target) + } + if target == current { + c.observeDurableState() + return nil + } + if target > current+1 && target > durableFloor+1 { + c.observeReload("rejected") + return errors.Wrapf(ErrUnsafeTSOModeTransition, "%s -> %s", current, target) + } + c.installMode(target) + c.observeReload("applied") + return nil +} + +func (c *TSORuntimeController) installMode(mode TSOMode) { + switch mode { + case TSOModeLegacy: + c.routed.setActivation(false, false) + c.dynamic.store(nil) + case TSOModeShadow: + c.routed.setActivation(false, false) + c.dynamic.store(c.shadow) + case TSOModeCutover: + c.routed.setActivation(true, false) + c.batch.Invalidate() + c.dynamic.store(c.batch) + case TSOModePhaseD: + c.routed.setActivation(true, true) + c.batch.Invalidate() + c.dynamic.store(c.batch) + } + c.mode.Store(uint32(mode)) + if c.observer != nil { + c.observer.ObserveTSOMode(mode.String()) + } + c.observeDurableState() +} + +func (c *TSORuntimeController) observeReload(result string) { + if c != nil && c.observer != nil { + c.observer.ObserveTSOModeReload(result) + } +} + +func (c *TSORuntimeController) observeDurableState() { + if c != nil && c.observer != nil && c.state != nil { + c.observer.ObserveTSODurableState(c.state.CutoverActive(), c.state.PhaseDActive()) + } +} + +func ReadTSOModeFile(path string) (TSOMode, error) { + raw, err := os.ReadFile(path) + if err != nil { + return TSOModeLegacy, errors.Wrap(err, "read tso mode file") + } + mode, err := ParseTSOMode(string(raw)) + return mode, errors.Wrap(err, "parse tso mode file") +} + +// RunTSOModeFileReload polls an atomically replaceable plain-text mode file. +// Invalid reads or transitions leave the current allocator untouched. +func RunTSOModeFileReload( + ctx context.Context, + path string, + interval time.Duration, + controller *TSORuntimeController, + logger *slog.Logger, +) error { + if controller == nil { + return errors.WithStack(ErrTSOAllocatorRequired) + } + if strings.TrimSpace(path) == "" { + return nil + } + if interval <= 0 { + return errors.New("tso mode reload interval must be positive") + } + if logger == nil { + logger = slog.Default() + } + ctx = nonNilTSOContext(ctx) + ticker := time.NewTicker(interval) + defer ticker.Stop() + lastError := "" + for { + select { + case <-ctx.Done(): + return nil + case <-ticker.C: + err := reloadTSOModeFile(path, controller) + lastError = reportTSOModeReloadError(ctx, path, lastError, err, controller, logger) + } + } +} + +func reloadTSOModeFile(path string, controller *TSORuntimeController) error { + mode, err := ReadTSOModeFile(path) + if err != nil { + return err + } + return errors.WithStack(controller.ApplyMode(mode)) +} + +func reportTSOModeReloadError( + ctx context.Context, + path string, + lastError string, + err error, + controller *TSORuntimeController, + logger *slog.Logger, +) string { + if err == nil { + return "" + } + if !errors.Is(err, ErrTSOModeRollback) && !errors.Is(err, ErrUnsafeTSOModeTransition) { + controller.observeReload(classifyTSOModeReloadError(err)) + } + if lastError == err.Error() { + return lastError + } + logger.ErrorContext(ctx, "tso mode reload rejected", + slog.String("path", path), + slog.String("current_mode", controller.CurrentMode().String()), + slog.Any("err", err), + ) + return err.Error() +} + +func classifyTSOModeReloadError(err error) string { + switch { + case errors.Is(err, ErrInvalidTSOMode): + return "parse_error" + case errors.Is(err, ErrTSOModeRollback), errors.Is(err, ErrUnsafeTSOModeTransition): + return "rejected" + default: + return "read_error" + } +} diff --git a/kv/tso_runtime_test.go b/kv/tso_runtime_test.go new file mode 100644 index 000000000..3ed8b39ab --- /dev/null +++ b/kv/tso_runtime_test.go @@ -0,0 +1,407 @@ +package kv + +import ( + "context" + "io" + "log/slog" + "os" + "path/filepath" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/bootjp/elastickv/distribution" + "github.com/cockroachdb/errors" + "github.com/stretchr/testify/require" +) + +func TestParseTSOMode(t *testing.T) { + t.Parallel() + for _, tc := range []struct { + raw string + want TSOMode + }{ + {raw: "legacy\n", want: TSOModeLegacy}, + {raw: "SHADOW", want: TSOModeShadow}, + {raw: "cutover", want: TSOModeCutover}, + {raw: "phase-d", want: TSOModePhaseD}, + {raw: "phase_d", want: TSOModePhaseD}, + } { + mode, err := ParseTSOMode(tc.raw) + require.NoError(t, err) + require.Equal(t, tc.want, mode) + } + _, err := ParseTSOMode("disabled") + require.ErrorIs(t, err, ErrInvalidTSOMode) +} + +func TestDynamicTimestampAllocatorPreservesLegacyFallbackUntilActivated(t *testing.T) { + t.Parallel() + clock := NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(time.Minute).UnixMilli()) + dynamic := NewDynamicTimestampAllocator(nil) + coord := NewShardedCoordinator(distribution.NewEngine(), nil, 1, clock, nil). + WithTSOAllocator(dynamic) + + allocator, ok := TimestampAllocatorThrough(coord) + require.False(t, ok) + require.Nil(t, allocator) + legacyTS, err := NextTimestampThrough(context.Background(), coord, "legacy dynamic fallback") + require.NoError(t, err) + + dedicated := &phaseDTestAllocator{next: legacyTS + 100} + dynamic.store(dedicated) + allocator, ok = TimestampAllocatorThrough(coord) + require.True(t, ok) + require.Same(t, dedicated, allocator) + dedicatedTS, err := NextTimestampThrough(context.Background(), coord, "active dynamic allocator") + require.NoError(t, err) + require.Equal(t, legacyTS+100, dedicatedTS) +} + +func TestDynamicTimestampAllocatorRejectsValidationWithoutAllocator(t *testing.T) { + t.Parallel() + dynamic := NewDynamicTimestampAllocator(nil) + require.ErrorIs(t, + dynamic.ValidateDurableTimestamp(context.Background(), 1), + ErrTSOAllocatorRequired, + ) +} + +func TestTSORuntimeControllerEnforcesAdjacentOneWayTransitions(t *testing.T) { + t.Parallel() + controller, state := newTestTSORuntimeController(t) + require.Equal(t, TSOModeLegacy, controller.CurrentMode()) + require.Nil(t, controller.Allocator().currentTimestampAllocator()) + + err := controller.ApplyMode(TSOModeCutover) + require.ErrorIs(t, err, ErrUnsafeTSOModeTransition) + require.Equal(t, TSOModeLegacy, controller.CurrentMode()) + + require.NoError(t, controller.ApplyMode(TSOModeShadow)) + require.Equal(t, TSOModeShadow, controller.CurrentMode()) + require.ErrorIs(t, controller.ApplyMode(TSOModeLegacy), ErrTSOModeRollback) + + require.NoError(t, controller.ApplyMode(TSOModeCutover)) + _, err = controller.Allocator().Next(context.Background()) + require.NoError(t, err) + require.True(t, state.CutoverActive()) + + require.NoError(t, controller.ApplyMode(TSOModePhaseD)) + controller.Allocator().Invalidate() + _, err = controller.Allocator().Next(context.Background()) + require.NoError(t, err) + require.True(t, state.PhaseDActive()) + require.NoError(t, controller.ApplyMode(TSOModeCutover)) + require.Equal(t, TSOModePhaseD, controller.CurrentMode()) +} + +func TestTSORuntimeControllerDurableStateOverridesStaleMode(t *testing.T) { + t.Parallel() + state := &runtimeTSOState{} + state.cutover.Store(true) + controller, _ := newTestTSORuntimeControllerWithState(t, state) + require.Equal(t, TSOModeCutover, controller.CurrentMode()) + require.NotNil(t, controller.Allocator().currentTimestampAllocator()) + require.NoError(t, controller.ApplyMode(TSOModeShadow)) + require.Equal(t, TSOModeCutover, controller.CurrentMode()) + + state.phaseD.Store(true) + require.NoError(t, controller.ApplyMode(TSOModeLegacy)) + require.Equal(t, TSOModePhaseD, controller.CurrentMode()) +} + +func TestTSORuntimeControllerDurableStateMaySkipLocalPhases(t *testing.T) { + t.Parallel() + state := &runtimeTSOState{} + controller, _ := newTestTSORuntimeControllerWithState(t, state) + + state.cutover.Store(true) + require.NoError(t, controller.ApplyMode(TSOModeLegacy)) + require.Equal(t, TSOModeCutover, controller.CurrentMode()) + + state.phaseD.Store(true) + require.NoError(t, controller.ApplyMode(TSOModeLegacy)) + require.Equal(t, TSOModePhaseD, controller.CurrentMode()) +} + +func TestTSORuntimeControllerAllowsPhaseDAfterDurableCutoverOverride(t *testing.T) { + t.Parallel() + state := &runtimeTSOState{} + controller, _ := newTestTSORuntimeControllerWithState(t, state) + require.Equal(t, TSOModeLegacy, controller.CurrentMode()) + + state.cutover.Store(true) + require.NoError(t, controller.ApplyMode(TSOModePhaseD)) + require.Equal(t, TSOModePhaseD, controller.CurrentMode()) +} + +func TestTSORuntimeControllerDurableCutoverOverridesLegacyBeforeReload(t *testing.T) { + t.Parallel() + controller, state := newTestTSORuntimeController(t) + require.Equal(t, TSOModeLegacy, controller.CurrentMode()) + + state.cutover.Store(true) + coord := NewShardedCoordinator(distribution.NewEngine(), nil, 1, controller.shadow.legacy, nil). + WithTSOAllocator(controller.Allocator()) + allocator, ok := TimestampAllocatorThrough(coord) + require.True(t, ok) + require.Same(t, controller.batch, allocator) + activateCutover, activatePhaseD := controller.routed.activationState() + require.True(t, activateCutover) + require.False(t, activatePhaseD) + + _, err := NextTimestampThrough(context.Background(), coord, "durable cutover before reload") + require.NoError(t, err) + require.Equal(t, TSOModeLegacy, controller.CurrentMode(), "mode poll has not run") +} + +func TestDurableCutoverObservationDoesNotLowerPhaseDActivation(t *testing.T) { + t.Parallel() + controller, state := newTestTSORuntimeController(t) + controller.routed.setActivation(true, true) + state.cutover.Store(true) + + require.Same(t, controller.batch, controller.Allocator().currentTimestampAllocator()) + activateCutover, activatePhaseD := controller.routed.activationState() + require.True(t, activateCutover) + require.True(t, activatePhaseD) +} + +func TestReloadTSOModeFileRefreshesDurableStateWithoutModeChange(t *testing.T) { + t.Parallel() + state := &runtimeTSOState{} + observer := &recordingRuntimeTSOObserver{} + controller, _ := newTestTSORuntimeControllerWithObserver(t, TSOModeCutover, state, observer) + path := filepath.Join(t.TempDir(), "tso-mode") + writeTSOModeFile(t, path, "cutover\n") + observer.durableCalls = 0 + + state.cutover.Store(true) + require.NoError(t, reloadTSOModeFile(path, controller)) + require.Equal(t, 1, observer.durableCalls) + require.True(t, observer.cutoverActive) + require.False(t, observer.phaseDActive) +} + +func TestReportTSOModeReloadErrorCountsEveryFailedPoll(t *testing.T) { + t.Parallel() + observer := &recordingRuntimeTSOObserver{} + controller, _ := newTestTSORuntimeControllerWithObserver( + t, + TSOModeLegacy, + &runtimeTSOState{}, + observer, + ) + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + err := errors.Wrap(ErrInvalidTSOMode, "parse tso mode file") + + lastError := reportTSOModeReloadError(context.Background(), "mode", "", err, controller, logger) + reportTSOModeReloadError(context.Background(), "mode", lastError, err, controller, logger) + require.Equal(t, 2, observer.reloads["parse_error"]) +} + +func TestRunTSOModeFileReloadAppliesSafeSequence(t *testing.T) { + t.Parallel() + controller, _ := newTestTSORuntimeController(t) + path := filepath.Join(t.TempDir(), "tso-mode") + writeTSOModeFile(t, path, "legacy\n") + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { + done <- RunTSOModeFileReload(ctx, path, time.Millisecond, controller, slog.Default()) + }() + t.Cleanup(func() { + cancel() + require.NoError(t, <-done) + }) + + writeTSOModeFile(t, path, "shadow\n") + require.Eventually(t, func() bool { + return controller.CurrentMode() == TSOModeShadow + }, time.Second, time.Millisecond) + + writeTSOModeFile(t, path, "phase-d\n") + require.Never(t, func() bool { + return controller.CurrentMode() == TSOModePhaseD + }, 20*time.Millisecond, time.Millisecond) + require.Equal(t, TSOModeShadow, controller.CurrentMode()) + + writeTSOModeFile(t, path, "cutover\n") + require.Eventually(t, func() bool { + return controller.CurrentMode() == TSOModeCutover + }, time.Second, time.Millisecond) + writeTSOModeFile(t, path, "phase-d\n") + require.Eventually(t, func() bool { + return controller.CurrentMode() == TSOModePhaseD + }, time.Second, time.Millisecond) + + writeTSOModeFile(t, path, "legacy\n") + require.Never(t, func() bool { + return controller.CurrentMode() != TSOModePhaseD + }, 20*time.Millisecond, time.Millisecond) +} + +func writeTSOModeFile(t *testing.T, path, contents string) { + t.Helper() + tmp := path + ".tmp" + require.NoError(t, os.WriteFile(tmp, []byte(contents), 0o600)) + require.NoError(t, os.Rename(tmp, path)) +} + +func TestTSORuntimeControllerConcurrentAllocationsDuringReload(t *testing.T) { + controller, _ := newTestTSORuntimeController(t) + clock := NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(time.Minute).UnixMilli()) + coord := NewShardedCoordinator(distribution.NewEngine(), nil, 1, clock, nil). + WithTSOAllocator(controller.Allocator()) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + var wg sync.WaitGroup + errs := make(chan error, 32) + for range 16 { + wg.Add(1) + go func() { + defer wg.Done() + for range 200 { + if _, err := NextTimestampThrough(ctx, coord, "runtime reload race"); err != nil { + errs <- err + return + } + } + }() + } + require.NoError(t, controller.ApplyMode(TSOModeShadow)) + require.NoError(t, controller.ApplyMode(TSOModeCutover)) + require.NoError(t, controller.ApplyMode(TSOModePhaseD)) + wg.Wait() + close(errs) + for err := range errs { + require.NoError(t, err) + } +} + +func newTestTSORuntimeController(t *testing.T) (*TSORuntimeController, *runtimeTSOState) { + t.Helper() + return newTestTSORuntimeControllerWithState(t, &runtimeTSOState{}) +} + +func newTestTSORuntimeControllerWithState( + t *testing.T, + state *runtimeTSOState, +) (*TSORuntimeController, *runtimeTSOState) { + return newTestTSORuntimeControllerWithObserver(t, TSOModeLegacy, state, nil) +} + +func newTestTSORuntimeControllerWithObserver( + t *testing.T, + initial TSOMode, + state *runtimeTSOState, + observer TSOObserver, +) (*TSORuntimeController, *runtimeTSOState) { + t.Helper() + clock := NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(time.Minute).UnixMilli()) + local := &runtimeTSOAllocator{state: state} + routed, err := NewLeaderRoutedTSOAllocator(local, &recordingTSOEngine{}, WithTSORoutedClock(clock)) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, routed.Close()) }) + controller, err := NewTSORuntimeController(TSORuntimeControllerConfig{ + Clock: clock, + Routed: routed, + State: state, + BatchSize: 8, + InitialMode: initial, + Observer: observer, + }) + require.NoError(t, err) + return controller, state +} + +type recordingRuntimeTSOObserver struct { + reloads map[string]int + durableCalls int + cutoverActive bool + phaseDActive bool +} + +func (o *recordingRuntimeTSOObserver) ObserveTSORequest(string, string, string, time.Duration) {} +func (o *recordingRuntimeTSOObserver) ObserveTSOShadowComparison(string, uint64) {} +func (o *recordingRuntimeTSOObserver) ObserveTSOMode(string) {} + +func (o *recordingRuntimeTSOObserver) ObserveTSOModeReload(result string) { + if o.reloads == nil { + o.reloads = make(map[string]int) + } + o.reloads[result]++ +} + +func (o *recordingRuntimeTSOObserver) ObserveTSODurableState(cutoverActive, phaseDActive bool) { + o.durableCalls++ + o.cutoverActive = cutoverActive + o.phaseDActive = phaseDActive +} + +type runtimeTSOState struct { + cutover atomic.Bool + phaseD atomic.Bool +} + +func (s *runtimeTSOState) CutoverActive() bool { return s.cutover.Load() } +func (s *runtimeTSOState) PhaseDActive() bool { return s.phaseD.Load() } + +type runtimeTSOAllocator struct { + state *runtimeTSOState + next atomic.Uint64 + mu sync.Mutex +} + +func (a *runtimeTSOAllocator) Next(ctx context.Context) (uint64, error) { + return a.NextBatch(ctx, 1) +} + +func (a *runtimeTSOAllocator) NextBatch(ctx context.Context, n int) (uint64, error) { + reservation, err := a.ReserveBatchAfter(ctx, n, 0, false, false) + return reservation.Base, err +} + +func (a *runtimeTSOAllocator) NextBatchAfter(ctx context.Context, n int, min uint64) (uint64, error) { + reservation, err := a.ReserveBatchAfter(ctx, n, min, false, false) + return reservation.Base, err +} + +func (a *runtimeTSOAllocator) ReserveBatchAfter( + _ context.Context, + n int, + min uint64, + activateCutover bool, + activatePhaseD bool, +) (TSOReservation, error) { + a.mu.Lock() + defer a.mu.Unlock() + if activateCutover || activatePhaseD { + a.state.cutover.Store(true) + } + if activatePhaseD { + a.state.phaseD.Store(true) + } + previous := a.next.Load() + base := max(previous+1, min+1) + end := base + uint64(n) - 1 //nolint:gosec // test uses positive bounded n. + a.next.Store(end) + return TSOReservation{ + Base: base, + Count: n, + PreviousAllocationFloor: previous, + CutoverActive: a.state.CutoverActive(), + PhaseDActive: a.state.PhaseDActive(), + }, nil +} + +func (a *runtimeTSOAllocator) ValidateDurableTimestamp(context.Context, uint64) error { return nil } +func (a *runtimeTSOAllocator) PhaseDActive() bool { return a.state.PhaseDActive() } +func (a *runtimeTSOAllocator) PhaseDRequired() bool { return a.state.PhaseDActive() } +func (a *runtimeTSOAllocator) IsLeader() bool { return true } +func (a *runtimeTSOAllocator) RunLeaseRenewal(ctx context.Context) { <-ctx.Done() } diff --git a/main.go b/main.go index 5f515a91e..af302f1ff 100644 --- a/main.go +++ b/main.go @@ -55,6 +55,7 @@ const ( etcdMaxSizePerMsg = 1 << 20 etcdMaxInflightMsg = 1024 defaultTSOBatchSize = 256 + defaultTSOReload = 5 * time.Second defaultFilesystemRootMode = 0o755 defaultFilesystemPlacementScanInterval = 30 * time.Second @@ -128,6 +129,8 @@ var ( tsoShadowEnabled = flag.Bool("tsoShadowEnabled", false, "Serialize legacy HLC issuance through the dedicated TSO; fail closed on TSO errors and switch to TSO values after cutover") tsoPhaseDEnabled = flag.Bool("tsoPhaseDEnabled", false, "Close the centralized-TSO compatibility window after every group-0 member understands the M7 marker") tsoBatchSize = flag.Int("tsoBatchSize", defaultTSOBatchSize, "Timestamp batch size used by TSO cutover and shadow validation") + tsoModeFile = flag.String("tsoModeFile", "", "Path to a runtime-reloadable TSO mode file containing legacy, shadow, cutover, or phase-d") + tsoModeReloadInterval = flag.Duration("tsoModeReloadInterval", defaultTSOReload, "Polling interval for tsoModeFile") leaderBalance = flag.Bool("leaderBalance", false, "Enable automatic count-based Raft-group leader balancing on the default-group leader") leaderBalanceInterval = flag.Duration("leaderBalanceInterval", defaultLeaderBalanceInterval, "Interval between leader-balance scheduler evaluations") leaderBalanceGroupCooldown = flag.Duration("leaderBalanceGroupCooldown", defaultLeaderBalanceGroupCooldown, "Minimum time before the scheduler can move the same raft group again") @@ -515,7 +518,12 @@ func run() error { WithKeyVizLabelsEnabled(*keyvizLabelsEnabled). WithAllShardGroups(dataGroupIDs(cfg.groups)...). WithPartitionResolver(buildSQSPartitionResolver(cfg.sqsFifoPartitionMap)) - tsoWiring, err := configureCoordinatorTSO(coordinate, shardGroups, shardStore) + tsoWiring, err := configureCoordinatorTSOWithObserver( + coordinate, + shardGroups, + shardStore, + metricsRegistry.TSOObserver(), + ) if err != nil { return err } @@ -535,6 +543,7 @@ func run() error { sqsAdvertisesHTFIFO(), slog.Default()) cleanup.Add(leadershipRefusalDeregister) eg, runCtx := errgroup.WithContext(ctx) + startTSOModeReloader(runCtx, eg, tsoWiring) startRaftEngineLifecycleWatchers(runCtx, eg, runtimes) // setupDistributionCatalog + the Stage 7a process-start registration // gate are bundled so run() has a single startup-fault path: a @@ -622,6 +631,21 @@ func run() error { return nil } +func startTSOModeReloader(ctx context.Context, eg *errgroup.Group, wiring coordinatorTSOWiring) { + if eg == nil || wiring.runtimeController == nil || strings.TrimSpace(*tsoModeFile) == "" { + return + } + eg.Go(func() error { + return kv.RunTSOModeFileReload( + ctx, + *tsoModeFile, + *tsoModeReloadInterval, + wiring.runtimeController, + slog.Default(), + ) + }) +} + func startRaftEngineLifecycleWatchers(ctx context.Context, eg *errgroup.Group, runtimes []*raftGroupRuntime) { for _, rt := range runtimes { if rt == nil { @@ -2060,17 +2084,12 @@ func configurationContainsMember(configuration raftengine.Configuration, localID } type coordinatorTSOWiring struct { - serverAllocator kv.TSOAllocator - routedAllocator *kv.LeaderRoutedTSOAllocator - shadowAllocator *kv.ShadowTimestampAllocator + serverAllocator kv.TSOAllocator + routedAllocator *kv.LeaderRoutedTSOAllocator + runtimeController *kv.TSORuntimeController } func (w coordinatorTSOWiring) Close() error { - if w.shadowAllocator != nil { - if err := w.shadowAllocator.Close(); err != nil { - return errors.Wrap(err, "close TSO shadow validator") - } - } if w.routedAllocator == nil { return nil } @@ -2087,38 +2106,98 @@ func configureCoordinatorTSO( coordinate *kv.ShardedCoordinator, shardGroups map[uint64]*kv.ShardGroup, floorProviders ...kv.TSOCutoverFloorProvider, +) (coordinatorTSOWiring, error) { + var floorProvider kv.TSOCutoverFloorProvider + if len(floorProviders) > 0 { + floorProvider = floorProviders[0] + } + return configureCoordinatorTSOWithObserver(coordinate, shardGroups, floorProvider, nil) +} + +func configureCoordinatorTSOWithObserver( + coordinate *kv.ShardedCoordinator, + shardGroups map[uint64]*kv.ShardGroup, + floorProvider kv.TSOCutoverFloorProvider, + observer kv.TSOObserver, ) (coordinatorTSOWiring, error) { var wiring coordinatorTSOWiring if coordinate == nil { return wiring, errors.Wrap(kv.ErrTSOCoordinatorNil, "configure tso allocator") } - if *tsoEnabled && *tsoShadowEnabled { - return wiring, errors.New("--tsoEnabled and --tsoShadowEnabled are mutually exclusive") - } - if *tsoPhaseDEnabled && !*tsoEnabled { - return wiring, errors.New("--tsoPhaseDEnabled requires --tsoEnabled") + mode, err := configuredTSOMode() + if err != nil { + return wiring, err } tsoGroup, dedicated := shardGroups[dedicatedTSORaftGroupID] if !dedicated { - return configureLegacyCoordinatorTSO(coordinate) + return configureLegacyCoordinatorTSO(coordinate, mode) } - var floorProvider kv.TSOCutoverFloorProvider - if len(floorProviders) > 0 { - floorProvider = floorProviders[0] + return configureDedicatedCoordinatorTSO(coordinate, tsoGroup, floorProvider, mode, observer) +} + +func configuredTSOMode() (kv.TSOMode, error) { + if err := validateTSOModeFlags(); err != nil { + return kv.TSOModeLegacy, err + } + if strings.TrimSpace(*tsoModeFile) != "" { + mode, err := kv.ReadTSOModeFile(*tsoModeFile) + return mode, errors.Wrap(err, "load initial tso runtime mode") + } + return tsoModeFromLegacyFlags(), nil +} + +func validateTSOModeFlags() error { + if err := validateLegacyTSOModeFlags(); err != nil { + return err + } + if strings.TrimSpace(*tsoModeFile) == "" { + return nil + } + if *tsoEnabled || *tsoShadowEnabled || *tsoPhaseDEnabled { + return errors.New("--tsoModeFile cannot be combined with tsoEnabled, tsoShadowEnabled, or tsoPhaseDEnabled") + } + if *tsoModeReloadInterval <= 0 { + return errors.New("--tsoModeReloadInterval must be positive") } - return configureDedicatedCoordinatorTSO(coordinate, tsoGroup, floorProvider) + return nil } -func configureLegacyCoordinatorTSO(coordinate *kv.ShardedCoordinator) (coordinatorTSOWiring, error) { +func validateLegacyTSOModeFlags() error { + if *tsoEnabled && *tsoShadowEnabled { + return errors.New("--tsoEnabled and --tsoShadowEnabled are mutually exclusive") + } + if *tsoPhaseDEnabled && !*tsoEnabled { + return errors.New("--tsoPhaseDEnabled requires --tsoEnabled") + } + return nil +} + +func tsoModeFromLegacyFlags() kv.TSOMode { + switch { + case *tsoPhaseDEnabled: + return kv.TSOModePhaseD + case *tsoEnabled: + return kv.TSOModeCutover + case *tsoShadowEnabled: + return kv.TSOModeShadow + default: + return kv.TSOModeLegacy + } +} + +func configureLegacyCoordinatorTSO(coordinate *kv.ShardedCoordinator, mode kv.TSOMode) (coordinatorTSOWiring, error) { var wiring coordinatorTSOWiring - if *tsoPhaseDEnabled { + if strings.TrimSpace(*tsoModeFile) != "" { + return wiring, errors.Wrap(kv.ErrTSOGroupRequired, "configure tso mode reload") + } + if mode == kv.TSOModePhaseD { return wiring, errors.Wrap(kv.ErrTSOGroupRequired, "configure tso phase D") } - if *tsoShadowEnabled { + if mode == kv.TSOModeShadow { return wiring, errors.Wrap(kv.ErrTSOGroupRequired, "configure tso shadow allocator") } - if !*tsoEnabled { + if mode == kv.TSOModeLegacy { return wiring, nil } // Preserve the pre-group-0 bridge for deployments that enable TSO @@ -2139,6 +2218,8 @@ func configureDedicatedCoordinatorTSO( coordinate *kv.ShardedCoordinator, tsoGroup *kv.ShardGroup, floorProvider kv.TSOCutoverFloorProvider, + mode kv.TSOMode, + observer kv.TSOObserver, ) (coordinatorTSOWiring, error) { var wiring coordinatorTSOWiring local, err := kv.NewRaftTSOAllocator( @@ -2152,64 +2233,46 @@ func configureDedicatedCoordinatorTSO( wiring.serverAllocator = local coordinate.WithTimestampGroup(dedicatedTSORaftGroupID) coordinate.WithTSOCutoverState(tsoGroup.TSOState) - if !dedicatedTSOAllocatorRequired(tsoGroup.TSOState) { + if !dedicatedTSOAllocatorRequired(tsoGroup.TSOState, mode) { return wiring, nil } - routedOpts := dedicatedTSORoutingOptions(coordinate.Clock()) + routedOpts := dedicatedTSORoutingOptions(coordinate.Clock(), observer) routed, err := kv.NewLeaderRoutedTSOAllocator(local, tsoGroup.Engine, routedOpts...) if err != nil { return wiring, errors.Wrap(err, "configure tso leader routing") } wiring.routedAllocator = routed - return installDedicatedCoordinatorAllocator(coordinate, tsoGroup.TSOState, wiring, routed) + controller, err := kv.NewTSORuntimeController(kv.TSORuntimeControllerConfig{ + Clock: coordinate.Clock(), + Routed: routed, + State: tsoGroup.TSOState, + BatchSize: *tsoBatchSize, + InitialMode: mode, + Logger: slog.Default(), + Observer: observer, + }) + if err != nil { + _ = wiring.Close() + return coordinatorTSOWiring{}, errors.Wrap(err, "configure tso runtime controller") + } + wiring.runtimeController = controller + coordinate.WithTSOAllocator(controller.Allocator()) + return wiring, nil } -func dedicatedTSOAllocatorRequired(state *kv.TSOStateMachine) bool { - return *tsoEnabled || *tsoShadowEnabled || (state != nil && state.CutoverActive()) +func dedicatedTSOAllocatorRequired(state *kv.TSOStateMachine, mode kv.TSOMode) bool { + return strings.TrimSpace(*tsoModeFile) != "" || mode != kv.TSOModeLegacy || (state != nil && state.CutoverActive()) } -func dedicatedTSORoutingOptions(clock *kv.HLC) []kv.LeaderRoutedTSOAllocatorOption { +func dedicatedTSORoutingOptions(clock *kv.HLC, observer kv.TSOObserver) []kv.LeaderRoutedTSOAllocatorOption { opts := []kv.LeaderRoutedTSOAllocatorOption{kv.WithTSORoutedClock(clock)} - if *tsoEnabled { - opts = append(opts, kv.WithTSOCutoverActivation()) - } - if *tsoPhaseDEnabled { - opts = append(opts, kv.WithTSOPhaseDActivation()) + if observer != nil { + opts = append(opts, kv.WithTSOObserver(observer)) } return opts } -func installDedicatedCoordinatorAllocator( - coordinate *kv.ShardedCoordinator, - state *kv.TSOStateMachine, - wiring coordinatorTSOWiring, - routed *kv.LeaderRoutedTSOAllocator, -) (coordinatorTSOWiring, error) { - if *tsoShadowEnabled && (state == nil || !state.CutoverActive()) { - shadow, err := kv.NewShadowTimestampAllocator( - coordinate.Clock(), - routed, - slog.Default(), - kv.WithTSOShadowCutoverState(state), - ) - if err != nil { - _ = wiring.Close() - return coordinatorTSOWiring{}, errors.Wrap(err, "configure tso shadow allocator") - } - wiring.shadowAllocator = shadow - coordinate.WithTSOAllocator(shadow) - return wiring, nil - } - batch, err := kv.NewBatchAllocator(routed, *tsoBatchSize) - if err != nil { - _ = wiring.Close() - return coordinatorTSOWiring{}, errors.Wrap(err, "configure tso batch allocator") - } - coordinate.WithTSOAllocator(batch) - return wiring, nil -} - type hlcLeaseRenewalBlocker interface { SetHLCLeaseRenewalBlocker(func() bool) } @@ -2396,7 +2459,7 @@ func (c startupGatedCoordinator) Clock() *kv.HLC { } func (c startupGatedCoordinator) TimestampAllocator() kv.TimestampAllocator { - alloc, _ := kv.TimestampAllocatorThrough(c.inner) + alloc, _ := kv.ConfiguredTimestampAllocatorThrough(c.inner) return alloc } @@ -2812,7 +2875,7 @@ func startRaftServers( } func internalTimestampOptions(coordinate kv.Coordinator) []adapter.InternalOption { - if alloc, ok := kv.TimestampAllocatorThrough(coordinate); ok { + if alloc, ok := kv.ConfiguredTimestampAllocatorThrough(coordinate); ok { return []adapter.InternalOption{adapter.WithInternalTimestampAllocator(alloc)} } return nil diff --git a/main_tso_routing_test.go b/main_tso_routing_test.go index a2527fa5f..db4e9dcc3 100644 --- a/main_tso_routing_test.go +++ b/main_tso_routing_test.go @@ -3,6 +3,8 @@ package main import ( "context" "fmt" + "os" + "path/filepath" "sync/atomic" "testing" "time" @@ -139,7 +141,6 @@ func TestConfigureCoordinatorTSORestoresDurableCutoverWithoutFlags(t *testing.T) t.Cleanup(func() { require.NoError(t, wiring.Close()) }) require.NotNil(t, wiring.serverAllocator) require.NotNil(t, wiring.routedAllocator) - require.Nil(t, wiring.shadowAllocator) require.True(t, coord.IsTimestampLeader()) require.True(t, fsm.CutoverActive()) @@ -158,7 +159,8 @@ func TestConfigureCoordinatorTSODurableCutoverOverridesShadowFlag(t *testing.T) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, wiring.Close()) }) require.NotNil(t, wiring.routedAllocator) - require.Nil(t, wiring.shadowAllocator, "durable cutover cannot return to legacy shadow issuance") + require.Equal(t, kv.TSOModeCutover, wiring.runtimeController.CurrentMode(), + "durable cutover cannot return to legacy shadow issuance") require.True(t, coord.IsTimestampLeader()) require.True(t, fsm.CutoverActive()) @@ -179,10 +181,56 @@ func TestInternalTimestampOptionsPreservesTSOThroughStartupGate(t *testing.T) { require.Same(t, allocator, got) require.Len(t, internalTimestampOptions(gated), 1) + runtimeAllocator := kv.NewDynamicTimestampAllocator(nil) + runtimeCoord := newMainTSOCoordinator(kv.NewHLC(), nil).WithTSOAllocator(runtimeAllocator) + runtimeGated := startupGatedCoordinator{inner: runtimeCoord, gate: &startupPublicKVGate{}} + got, ok = kv.TimestampAllocatorThrough(runtimeGated) + require.False(t, ok) + require.Nil(t, got) + configured, ok := kv.ConfiguredTimestampAllocatorThrough(runtimeGated) + require.True(t, ok) + require.Same(t, runtimeAllocator, configured) + require.Len(t, internalTimestampOptions(runtimeGated), 1) + legacy := startupGatedCoordinator{inner: newMainTSOCoordinator(kv.NewHLC(), nil)} require.Empty(t, internalTimestampOptions(legacy)) } +func TestConfigureCoordinatorTSOModeFileStartsRuntimeController(t *testing.T) { + setTSOModeFlags(t, false, false) + path := filepath.Join(t.TempDir(), "tso-mode") + require.NoError(t, os.WriteFile(path, []byte("shadow\n"), 0o600)) + *tsoModeFile = path + *tsoModeReloadInterval = time.Millisecond + clock := kv.NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(time.Minute).UnixMilli()) + fsm := kv.NewTSOStateMachine(clock) + engine := &mainTSOEngine{state: raftengine.StateLeader, tsoState: fsm} + groups := map[uint64]*kv.ShardGroup{ + dedicatedTSORaftGroupID: {Engine: engine, TSOState: fsm}, + } + coord := newMainTSOCoordinator(clock, groups) + + wiring, err := configureCoordinatorTSO(coord, groups, mainTSOFloorProvider{}) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, wiring.Close()) }) + require.NotNil(t, wiring.runtimeController) + require.Equal(t, kv.TSOModeShadow, wiring.runtimeController.CurrentMode()) + _, configured := kv.TimestampAllocatorThrough(coord) + require.True(t, configured) +} + +func TestConfigureCoordinatorTSOModeFileRequiresDedicatedGroup(t *testing.T) { + setTSOModeFlags(t, false, false) + path := filepath.Join(t.TempDir(), "tso-mode") + require.NoError(t, os.WriteFile(path, []byte("legacy\n"), 0o600)) + *tsoModeFile = path + coord := newMainTSOCoordinator(kv.NewHLC(), nil) + + _, err := configureCoordinatorTSO(coord, nil) + require.ErrorIs(t, err, kv.ErrTSOGroupRequired) +} + func newActiveMainTSOCutover(t *testing.T) (*kv.HLC, *kv.TSOStateMachine, *mainTSOEngine, map[uint64]*kv.ShardGroup) { t.Helper() clock := kv.NewHLC() @@ -205,15 +253,21 @@ func setTSOModeFlags(t *testing.T, enabled, shadow bool) { oldShadow := *tsoShadowEnabled oldPhaseD := *tsoPhaseDEnabled oldBatchSize := *tsoBatchSize + oldModeFile := *tsoModeFile + oldReloadInterval := *tsoModeReloadInterval *tsoEnabled = enabled *tsoShadowEnabled = shadow *tsoPhaseDEnabled = false *tsoBatchSize = 8 + *tsoModeFile = "" + *tsoModeReloadInterval = defaultTSOReload t.Cleanup(func() { *tsoEnabled = oldEnabled *tsoShadowEnabled = oldShadow *tsoPhaseDEnabled = oldPhaseD *tsoBatchSize = oldBatchSize + *tsoModeFile = oldModeFile + *tsoModeReloadInterval = oldReloadInterval }) } diff --git a/monitoring/prometheus/rules/tso-alerts.yml b/monitoring/prometheus/rules/tso-alerts.yml new file mode 100644 index 000000000..bbee49800 --- /dev/null +++ b/monitoring/prometheus/rules/tso-alerts.yml @@ -0,0 +1,103 @@ +groups: + - name: elastickv-tso + rules: + - alert: ElasticKVTSOAllocationErrorRateHigh + expr: | + ( + sum(rate(elastickv_tso_request_duration_seconds_count{operation="reserve",outcome="error"}[5m])) + / + clamp_min(sum(rate(elastickv_tso_request_duration_seconds_count{operation="reserve"}[5m])), 0.001) + ) > 0.01 + and + sum(rate(elastickv_tso_request_duration_seconds_count{operation="reserve"}[5m])) > 0.1 + for: 10m + labels: + severity: warning + annotations: + summary: TSO allocation error rate exceeds 1 percent + description: Dedicated TSO reserve attempts have exceeded a 1 percent error rate for 10 minutes. + + - alert: ElasticKVTSOAllocationErrorRateCritical + expr: | + ( + sum(rate(elastickv_tso_request_duration_seconds_count{operation="reserve",outcome="error"}[5m])) + / + clamp_min(sum(rate(elastickv_tso_request_duration_seconds_count{operation="reserve"}[5m])), 0.001) + ) > 0.10 + and + sum(rate(elastickv_tso_request_duration_seconds_count{operation="reserve"}[5m])) > 0.1 + for: 5m + labels: + severity: critical + annotations: + summary: TSO allocation error rate exceeds 10 percent + description: Dedicated TSO reserve attempts have exceeded a 10 percent error rate for 5 minutes; writes may be failing closed. + + - alert: ElasticKVTSOAllocationLatencyHigh + expr: | + histogram_quantile( + 0.99, + sum by (le) ( + rate(elastickv_tso_request_duration_seconds_bucket{operation="reserve",outcome="success"}[5m]) + ) + ) > 0.05 + for: 10m + labels: + severity: warning + annotations: + summary: TSO allocation p99 exceeds 50 milliseconds + description: Successful dedicated TSO reserve attempts have exceeded 50 ms p99 for 10 minutes. + + - alert: ElasticKVTSOAllocationLatencyCritical + expr: | + histogram_quantile( + 0.99, + sum by (le) ( + rate(elastickv_tso_request_duration_seconds_bucket{operation="reserve",outcome="success"}[5m]) + ) + ) > 0.20 + for: 5m + labels: + severity: critical + annotations: + summary: TSO allocation p99 exceeds 200 milliseconds + description: Successful dedicated TSO reserve attempts have exceeded 200 ms p99 for 5 minutes. + + - alert: ElasticKVTSOShadowOverlapDetected + expr: sum(increase(elastickv_tso_shadow_comparisons_total{result="legacy_overlap"}[15m])) > 0 + for: 5m + labels: + severity: warning + annotations: + summary: Legacy timestamps overlap the durable TSO floor + description: At least one shadow candidate was discarded in the last 15 minutes. Keep every node in shadow mode until this remains zero for a full 15-minute window. + + - alert: ElasticKVTSOModeReloadRejected + expr: sum(increase(elastickv_tso_mode_reload_total{result=~"rejected|parse_error|read_error"}[10m])) > 0 + for: 1m + labels: + severity: warning + annotations: + summary: A runtime TSO mode reload was rejected + description: A node could not read or parse its mode file, or the requested transition violated the one-way phase order. + + - alert: ElasticKVTSOModeDivergence + expr: count(count by (mode) (elastickv_tso_mode == 1)) > 1 + for: 15m + labels: + severity: warning + annotations: + summary: Cluster nodes remain in different TSO modes + description: More than one process-local TSO mode has remained active for 15 minutes. Finish the rolling transition or investigate a rejected reload. + + - alert: ElasticKVTSOPhaseDWithoutCutover + expr: | + max(elastickv_tso_durable_state{marker="phase_d"}) + > + max(elastickv_tso_durable_state{marker="cutover"}) + for: 1m + labels: + severity: critical + annotations: + summary: Durable TSO marker ordering is invalid + description: Phase D is visible without the prerequisite durable cutover marker. Stop rollout and inspect group-0 state. diff --git a/monitoring/registry.go b/monitoring/registry.go index 27acdb10c..568bfeeb6 100644 --- a/monitoring/registry.go +++ b/monitoring/registry.go @@ -29,6 +29,8 @@ type Registry struct { hlcObserver *HLCObserver coldStart *ColdStartMetrics coldStartObs *ColdStartObserver + tso *TSOMetrics + tsoObserver *TSOObserver } // NewRegistry builds a registry with constant labels that identify the local node. @@ -59,6 +61,8 @@ func NewRegistry(nodeID string, nodeAddress string) *Registry { r.hlcObserver = newHLCObserver(r.hlc) r.coldStart = newColdStartMetrics(registerer) r.coldStartObs = newColdStartObserver(r.coldStart) + r.tso = newTSOMetrics(registerer) + r.tsoObserver = newTSOObserver(r.tso) return r } @@ -279,3 +283,12 @@ func (r *Registry) HLCObserver() *HLCObserver { } return r.hlcObserver } + +// TSOObserver returns the shared runtime TSO observer. The kv package consumes +// it through its narrow observer interface, keeping Prometheus ownership here. +func (r *Registry) TSOObserver() *TSOObserver { + if r == nil { + return nil + } + return r.tsoObserver +} diff --git a/monitoring/tso.go b/monitoring/tso.go new file mode 100644 index 000000000..535db2a40 --- /dev/null +++ b/monitoring/tso.go @@ -0,0 +1,146 @@ +package monitoring + +import ( + "time" + + "github.com/prometheus/client_golang/prometheus" +) + +var tsoRequestDurationBuckets = []float64{ + 0.0001, 0.00025, 0.0005, 0.001, 0.0025, 0.005, + 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 5, +} + +const ( + tsoDivergenceBucketFactor = 16 + tsoDivergenceBucketCount = 12 +) + +type TSOMetrics struct { + requestDuration *prometheus.HistogramVec + shadowComparison *prometheus.CounterVec + shadowDivergence *prometheus.HistogramVec + mode *prometheus.GaugeVec + reloads *prometheus.CounterVec + durableState *prometheus.GaugeVec +} + +func newTSOMetrics(registerer prometheus.Registerer) *TSOMetrics { + m := &TSOMetrics{ + requestDuration: prometheus.NewHistogramVec( + prometheus.HistogramOpts{ + Name: "elastickv_tso_request_duration_seconds", + Help: "End-to-end duration of dedicated TSO reserve and validation attempts by local or remote leader path and outcome.", + Buckets: tsoRequestDurationBuckets, + }, + []string{"operation", "path", "outcome"}, + ), + shadowComparison: prometheus.NewCounterVec( + prometheus.CounterOpts{ + Name: "elastickv_tso_shadow_comparisons_total", + Help: "Shadow migration comparisons by bounded result: legacy_ahead, legacy_overlap, cutover_active, or error.", + }, + []string{"result"}, + ), + shadowDivergence: prometheus.NewHistogramVec( + prometheus.HistogramOpts{ + Name: "elastickv_tso_shadow_divergence_timestamps", + Help: "Absolute timestamp distance between a shadow legacy candidate and the preceding durable TSO allocation floor.", + Buckets: prometheus.ExponentialBuckets(1, tsoDivergenceBucketFactor, tsoDivergenceBucketCount), + }, + []string{"result"}, + ), + mode: prometheus.NewGaugeVec( + prometheus.GaugeOpts{ + Name: "elastickv_tso_mode", + Help: "Current process-local TSO mode as a one-hot gauge.", + }, + []string{"mode"}, + ), + reloads: prometheus.NewCounterVec( + prometheus.CounterOpts{ + Name: "elastickv_tso_mode_reload_total", + Help: "Runtime TSO mode reload outcomes.", + }, + []string{"result"}, + ), + durableState: prometheus.NewGaugeVec( + prometheus.GaugeOpts{ + Name: "elastickv_tso_durable_state", + Help: "Consensus-owned one-way TSO marker state for cutover and phase_d.", + }, + []string{"marker"}, + ), + } + if registerer != nil { + registerer.MustRegister( + m.requestDuration, + m.shadowComparison, + m.shadowDivergence, + m.mode, + m.reloads, + m.durableState, + ) + } + return m +} + +type TSOObserver struct { + metrics *TSOMetrics +} + +func newTSOObserver(metrics *TSOMetrics) *TSOObserver { + return &TSOObserver{metrics: metrics} +} + +func (o *TSOObserver) ObserveTSORequest(operation, path, outcome string, duration time.Duration) { + if o == nil || o.metrics == nil { + return + } + o.metrics.requestDuration.WithLabelValues(operation, path, outcome).Observe(duration.Seconds()) +} + +func (o *TSOObserver) ObserveTSOShadowComparison(result string, divergence uint64) { + if o == nil || o.metrics == nil { + return + } + o.metrics.shadowComparison.WithLabelValues(result).Inc() + if result == "legacy_ahead" || result == "legacy_overlap" { + o.metrics.shadowDivergence.WithLabelValues(result).Observe(float64(divergence)) + } +} + +func (o *TSOObserver) ObserveTSOMode(mode string) { + if o == nil || o.metrics == nil { + return + } + for _, candidate := range []string{"legacy", "shadow", "cutover", "phase-d"} { + value := 0.0 + if candidate == mode { + value = 1 + } + o.metrics.mode.WithLabelValues(candidate).Set(value) + } +} + +func (o *TSOObserver) ObserveTSOModeReload(result string) { + if o == nil || o.metrics == nil { + return + } + o.metrics.reloads.WithLabelValues(result).Inc() +} + +func (o *TSOObserver) ObserveTSODurableState(cutoverActive, phaseDActive bool) { + if o == nil || o.metrics == nil { + return + } + o.metrics.durableState.WithLabelValues("cutover").Set(boolMetricValue(cutoverActive)) + o.metrics.durableState.WithLabelValues("phase_d").Set(boolMetricValue(phaseDActive)) +} + +func boolMetricValue(value bool) float64 { + if value { + return 1 + } + return 0 +} diff --git a/monitoring/tso_test.go b/monitoring/tso_test.go new file mode 100644 index 000000000..7c220c156 --- /dev/null +++ b/monitoring/tso_test.go @@ -0,0 +1,48 @@ +package monitoring + +import ( + "testing" + "time" + + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/testutil" + dto "github.com/prometheus/client_model/go" + "github.com/stretchr/testify/require" +) + +func TestTSOObserverExportsOperationalMetrics(t *testing.T) { + t.Parallel() + reg := prometheus.NewRegistry() + metrics := newTSOMetrics(reg) + observer := newTSOObserver(metrics) + + observer.ObserveTSORequest("reserve", "remote", "success", 3*time.Millisecond) + observer.ObserveTSORequest("reserve", "remote", "error", 20*time.Millisecond) + observer.ObserveTSOShadowComparison("legacy_ahead", 7) + observer.ObserveTSOShadowComparison("legacy_overlap", 3) + observer.ObserveTSOMode("shadow") + observer.ObserveTSOModeReload("applied") + observer.ObserveTSODurableState(true, false) + + require.Equal(t, 1.0, testutil.ToFloat64(metrics.shadowComparison.WithLabelValues("legacy_ahead"))) + require.Equal(t, 1.0, testutil.ToFloat64(metrics.shadowComparison.WithLabelValues("legacy_overlap"))) + require.Equal(t, 1.0, testutil.ToFloat64(metrics.mode.WithLabelValues("shadow"))) + require.Equal(t, 0.0, testutil.ToFloat64(metrics.mode.WithLabelValues("legacy"))) + require.Equal(t, 1.0, testutil.ToFloat64(metrics.reloads.WithLabelValues("applied"))) + require.Equal(t, 1.0, testutil.ToFloat64(metrics.durableState.WithLabelValues("cutover"))) + require.Equal(t, 0.0, testutil.ToFloat64(metrics.durableState.WithLabelValues("phase_d"))) + + families, err := reg.Gather() + require.NoError(t, err) + require.True(t, hasMetricFamily(families, "elastickv_tso_request_duration_seconds")) + require.True(t, hasMetricFamily(families, "elastickv_tso_shadow_divergence_timestamps")) +} + +func hasMetricFamily(families []*dto.MetricFamily, name string) bool { + for _, family := range families { + if family.GetName() == name { + return true + } + } + return false +} From 229c2069f72c138b55d344da004ae300b89c5b1d Mon Sep 17 00:00:00 2001 From: bootjp Date: Thu, 23 Jul 2026 17:19:07 +0900 Subject: [PATCH 06/14] proto: reserve removed raft status fields --- proto/service.pb.go | 4 ++-- proto/service.proto | 2 ++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/proto/service.pb.go b/proto/service.pb.go index a70a3434b..bf6cd8ff9 100644 --- a/proto/service.pb.go +++ b/proto/service.pb.go @@ -2377,7 +2377,7 @@ const file_service_proto_rawDesc = "" + "\bstart_ts\x18\x01 \x01(\x04R\astartTs\",\n" + "\x10RollbackResponse\x12\x18\n" + "\asuccess\x18\x01 \x01(\bR\asuccess\"\x18\n" + - "\x16RaftAdminStatusRequest\"\xa2\x03\n" + + "\x16RaftAdminStatusRequest\"\xae\x03\n" + "\x17RaftAdminStatusResponse\x12%\n" + "\x05state\x18\x01 \x01(\x0e2\x0f.RaftAdminStateR\x05state\x12\x1b\n" + "\tleader_id\x18\x02 \x01(\tR\bleaderId\x12%\n" + @@ -2391,7 +2391,7 @@ const file_service_proto_rawDesc = "" + "fsmPending\x12\x1b\n" + "\tnum_peers\x18\n" + " \x01(\x04R\bnumPeers\x12,\n" + - "\x12last_contact_nanos\x18\v \x01(\x03R\x10lastContactNanos\"\x1f\n" + + "\x12last_contact_nanos\x18\v \x01(\x03R\x10lastContactNanosJ\x04\b\f\x10\rJ\x04\b\r\x10\x0e\"\x1f\n" + "\x1dRaftAdminConfigurationRequest\"W\n" + "\x0fRaftAdminMember\x12\x0e\n" + "\x02id\x18\x01 \x01(\tR\x02id\x12\x18\n" + diff --git a/proto/service.proto b/proto/service.proto index 6b01e6012..815e67648 100644 --- a/proto/service.proto +++ b/proto/service.proto @@ -187,6 +187,8 @@ enum RaftAdminState { message RaftAdminStatusRequest {} message RaftAdminStatusResponse { + reserved 12, 13; + RaftAdminState state = 1; string leader_id = 2; string leader_address = 3; From 917f421b4c1f6f36b90d612f59315d6818867941 Mon Sep 17 00:00:00 2001 From: bootjp Date: Thu, 23 Jul 2026 18:20:40 +0900 Subject: [PATCH 07/14] tso: preserve runtime allocator handles --- kv/keyviz_label.go | 5 ++++ kv/tso.go | 12 +++++++++ kv/tso_runtime_test.go | 34 +++++++++++++++++++++++++ main.go | 6 +++++ main_tso_routing_test.go | 55 ++++++++++++++++++++++++++++++++++++++++ 5 files changed, 112 insertions(+) diff --git a/kv/keyviz_label.go b/kv/keyviz_label.go index 09a3629c3..7a1fe5daf 100644 --- a/kv/keyviz_label.go +++ b/kv/keyviz_label.go @@ -91,6 +91,11 @@ func (c keyVizLabeledCoordinator) RaftLeaderForKey(key []byte) string { func (c keyVizLabeledCoordinator) Clock() *HLC { return c.inner.Clock() } func (c keyVizLabeledCoordinator) TimestampAllocator() TimestampAllocator { + alloc, _ := TimestampAllocatorThrough(c.inner) + return alloc +} + +func (c keyVizLabeledCoordinator) ConfiguredTimestampAllocator() TimestampAllocator { alloc, _ := ConfiguredTimestampAllocatorThrough(c.inner) return alloc } diff --git a/kv/tso.go b/kv/tso.go index 0f0be8181..99e67b37a 100644 --- a/kv/tso.go +++ b/kv/tso.go @@ -48,6 +48,14 @@ type TimestampAllocatorProvider interface { TimestampAllocator() TimestampAllocator } +// ConfiguredTimestampAllocatorProvider lets coordinator decorators expose the +// configured allocator without resolving runtime-mode wrappers. Long-lived +// services use this to keep a DynamicTimestampAllocator reference while the +// active mode is still legacy. +type ConfiguredTimestampAllocatorProvider interface { + ConfiguredTimestampAllocator() TimestampAllocator +} + type TimestampAfterAllocator interface { NextAfter(ctx context.Context, min uint64) (uint64, error) } @@ -231,6 +239,10 @@ func TimestampAllocatorThrough(coord Coordinator) (TimestampAllocator, bool) { // that must keep a DynamicTimestampAllocator reference across later mode-file // reloads while still falling back to HLC when that allocator reports legacy. func ConfiguredTimestampAllocatorThrough(coord Coordinator) (TimestampAllocator, bool) { + if provider, ok := coord.(ConfiguredTimestampAllocatorProvider); ok { + alloc := provider.ConfiguredTimestampAllocator() + return alloc, alloc != nil + } if provider, ok := coord.(TimestampAllocatorProvider); ok { alloc := provider.TimestampAllocator() return alloc, alloc != nil diff --git a/kv/tso_runtime_test.go b/kv/tso_runtime_test.go index 3ed8b39ab..c2b794170 100644 --- a/kv/tso_runtime_test.go +++ b/kv/tso_runtime_test.go @@ -60,6 +60,26 @@ func TestDynamicTimestampAllocatorPreservesLegacyFallbackUntilActivated(t *testi require.Equal(t, legacyTS+100, dedicatedTS) } +func TestConfiguredTimestampAllocatorThroughPrefersConfiguredProvider(t *testing.T) { + t.Parallel() + + active := &phaseDTestAllocator{next: 10} + configured := NewDynamicTimestampAllocator(nil) + provider := configuredAllocatorTestProvider{ + Coordinator: NewShardedCoordinator(distribution.NewEngine(), nil, 1, NewHLC(), nil), + active: active, + configured: configured, + } + + got, ok := TimestampAllocatorThrough(provider) + require.True(t, ok) + require.Same(t, active, got) + + got, ok = ConfiguredTimestampAllocatorThrough(provider) + require.True(t, ok) + require.Same(t, configured, got) +} + func TestDynamicTimestampAllocatorRejectsValidationWithoutAllocator(t *testing.T) { t.Parallel() dynamic := NewDynamicTimestampAllocator(nil) @@ -349,6 +369,20 @@ type runtimeTSOState struct { phaseD atomic.Bool } +type configuredAllocatorTestProvider struct { + Coordinator + active TimestampAllocator + configured TimestampAllocator +} + +func (p configuredAllocatorTestProvider) TimestampAllocator() TimestampAllocator { + return p.active +} + +func (p configuredAllocatorTestProvider) ConfiguredTimestampAllocator() TimestampAllocator { + return p.configured +} + func (s *runtimeTSOState) CutoverActive() bool { return s.cutover.Load() } func (s *runtimeTSOState) PhaseDActive() bool { return s.phaseD.Load() } diff --git a/main.go b/main.go index af302f1ff..dad62e9f7 100644 --- a/main.go +++ b/main.go @@ -2417,6 +2417,7 @@ var _ kv.LeaseReadableCoordinator = (*startupGatedCoordinator)(nil) var _ kv.AllGroupsLeaseReadableCoordinator = (*startupGatedCoordinator)(nil) var _ kv.GroupRoutableCoordinator = (*startupGatedCoordinator)(nil) var _ kv.TimestampAllocatorProvider = (*startupGatedCoordinator)(nil) +var _ kv.ConfiguredTimestampAllocatorProvider = (*startupGatedCoordinator)(nil) var _ kv.AppliedReadTimestampVoucher = (*startupGatedCoordinator)(nil) func (c startupGatedCoordinator) Dispatch(ctx context.Context, reqs *kv.OperationGroup[kv.OP]) (*kv.CoordinateResponse, error) { @@ -2459,6 +2460,11 @@ func (c startupGatedCoordinator) Clock() *kv.HLC { } func (c startupGatedCoordinator) TimestampAllocator() kv.TimestampAllocator { + alloc, _ := kv.TimestampAllocatorThrough(c.inner) + return alloc +} + +func (c startupGatedCoordinator) ConfiguredTimestampAllocator() kv.TimestampAllocator { alloc, _ := kv.ConfiguredTimestampAllocatorThrough(c.inner) return alloc } diff --git a/main_tso_routing_test.go b/main_tso_routing_test.go index db4e9dcc3..d56044534 100644 --- a/main_tso_routing_test.go +++ b/main_tso_routing_test.go @@ -9,10 +9,12 @@ import ( "testing" "time" + "github.com/bootjp/elastickv/adapter" "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/internal/raftengine" "github.com/bootjp/elastickv/keyviz" "github.com/bootjp/elastickv/kv" + pb "github.com/bootjp/elastickv/proto" "github.com/stretchr/testify/require" ) @@ -196,6 +198,46 @@ func TestInternalTimestampOptionsPreservesTSOThroughStartupGate(t *testing.T) { require.Empty(t, internalTimestampOptions(legacy)) } +func TestInternalForwardUsesRuntimeAllocatorAfterModeReload(t *testing.T) { + t.Parallel() + + setTSOModeFlags(t, false, false) + path := filepath.Join(t.TempDir(), "tso-mode") + require.NoError(t, os.WriteFile(path, []byte("legacy\n"), 0o600)) + *tsoModeFile = path + + clock := kv.NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(time.Minute).UnixMilli()) + fsm := kv.NewTSOStateMachine(clock) + engine := &mainTSOEngine{state: raftengine.StateLeader, tsoState: fsm} + groups := map[uint64]*kv.ShardGroup{ + dedicatedTSORaftGroupID: {Engine: engine, TSOState: fsm}, + } + coord := newMainTSOCoordinator(clock, groups) + wiring, err := configureCoordinatorTSO(coord, groups, mainTSOFloorProvider{}) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, wiring.Close()) }) + require.Equal(t, kv.TSOModeLegacy, wiring.runtimeController.CurrentMode()) + + gated := startupGatedCoordinator{inner: coord, gate: &startupPublicKVGate{}} + opts := internalTimestampOptions(gated) + require.Len(t, opts, 1) + require.NoError(t, wiring.runtimeController.ApplyMode(kv.TSOModeShadow)) + require.NoError(t, wiring.runtimeController.ApplyMode(kv.TSOModeCutover)) + + txn := &mainForwardTxn{} + internal := adapter.NewInternalWithEngine(txn, engine, nil, nil, opts...) + reqs := []*pb.Request{{ + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte("k"), Value: []byte("v")}}, + }} + resp, err := internal.Forward(context.Background(), &pb.ForwardRequest{Requests: reqs}) + require.NoError(t, err) + require.True(t, resp.GetSuccess()) + require.Len(t, txn.requests, 1) + require.Greater(t, txn.requests[0].GetTs(), uint64(1), + "timestamp 1 would mean the long-lived Internal server lost the runtime allocator and used its nil-clock fallback") +} + func TestConfigureCoordinatorTSOModeFileStartsRuntimeController(t *testing.T) { setTSOModeFlags(t, false, false) path := filepath.Join(t.TempDir(), "tso-mode") @@ -291,6 +333,19 @@ func (a *mainTimestampAllocator) Next(context.Context) (uint64, error) { return a.next, nil } +type mainForwardTxn struct { + requests []*pb.Request +} + +func (t *mainForwardTxn) Commit(_ context.Context, reqs []*pb.Request) (*kv.TransactionResponse, error) { + t.requests = reqs + return &kv.TransactionResponse{CommitIndex: 1}, nil +} + +func (t *mainForwardTxn) Abort(context.Context, []*pb.Request) (*kv.TransactionResponse, error) { + return &kv.TransactionResponse{}, nil +} + func (mainTSOFloorProvider) GlobalCommittedTimestampFloor(context.Context) (uint64, error) { return 0, nil } From 2402e770c465c654be0893bb466d1bf81af7e375 Mon Sep 17 00:00:00 2001 From: bootjp Date: Thu, 23 Jul 2026 18:58:29 +0900 Subject: [PATCH 08/14] tso: cover runtime allocator reload --- main_tso_routing_test.go | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/main_tso_routing_test.go b/main_tso_routing_test.go index d56044534..7d0ea5ec6 100644 --- a/main_tso_routing_test.go +++ b/main_tso_routing_test.go @@ -222,11 +222,13 @@ func TestInternalForwardUsesRuntimeAllocatorAfterModeReload(t *testing.T) { gated := startupGatedCoordinator{inner: coord, gate: &startupPublicKVGate{}} opts := internalTimestampOptions(gated) require.Len(t, opts, 1) - require.NoError(t, wiring.runtimeController.ApplyMode(kv.TSOModeShadow)) - require.NoError(t, wiring.runtimeController.ApplyMode(kv.TSOModeCutover)) txn := &mainForwardTxn{} internal := adapter.NewInternalWithEngine(txn, engine, nil, nil, opts...) + + require.NoError(t, wiring.runtimeController.ApplyMode(kv.TSOModeShadow)) + require.NoError(t, wiring.runtimeController.ApplyMode(kv.TSOModeCutover)) + reqs := []*pb.Request{{ Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte("k"), Value: []byte("v")}}, }} From 360418a58df180ef7b1290e13fc63f75e4e59908 Mon Sep 17 00:00:00 2001 From: bootjp Date: Thu, 23 Jul 2026 20:54:22 +0900 Subject: [PATCH 09/14] tso: bind applied read vouchers --- adapter/distribution_server_test.go | 10 ++- adapter/dynamodb_item_write.go | 22 ++++-- adapter/dynamodb_onephase_dedup_test.go | 17 +++++ adapter/redis_list_dedup_test.go | 60 +++++++++++++++ adapter/redis_retry_test.go | 23 ++++++ adapter/redis_stream_cmds.go | 97 ++++++++++++++----------- adapter/s3.go | 1 + adapter/s3_hlc_fence_test.go | 9 ++- kv/coordinator.go | 2 +- kv/keyviz_label.go | 4 +- kv/sharded_coordinator.go | 35 ++++++--- kv/sharded_coordinator_txn_test.go | 14 +++- kv/tso.go | 57 +++++++++++---- main.go | 4 +- 14 files changed, 263 insertions(+), 92 deletions(-) diff --git a/adapter/distribution_server_test.go b/adapter/distribution_server_test.go index 3aec7ea42..e1a88cf93 100644 --- a/adapter/distribution_server_test.go +++ b/adapter/distribution_server_test.go @@ -3,6 +3,7 @@ package adapter import ( "context" "encoding/binary" + stderrors "errors" "net" "testing" "time" @@ -1113,6 +1114,10 @@ func (s *distributionCoordinatorStub) TimestampAllocator() kv.TimestampAllocator return s.allocator } +func (s *distributionCoordinatorStub) VouchAppliedReadTimestamp(uint64, kv.AppliedReadTimestampVoucherRef) error { + return nil +} + func newDistributionCoordinatorStub(st store.MVCCStore, leader bool) *distributionCoordinatorStub { return &distributionCoordinatorStub{ store: st, @@ -1376,9 +1381,12 @@ func (a *distributionTSOAllocator) ValidateDurableTimestamp(_ context.Context, t if a.err != nil { return a.err } - if !a.phaseD || timestamp <= a.phaseDFloor || timestamp > a.base { + if !a.phaseD || timestamp == 0 || timestamp > a.base { return kv.ErrTSOTimestampInvalid } + if timestamp <= a.phaseDFloor { + return stderrors.Join(kv.ErrTSOTimestampInvalid, kv.ErrTSOTimestampPrePhaseD) + } return nil } diff --git a/adapter/dynamodb_item_write.go b/adapter/dynamodb_item_write.go index bee55d2fb..c08861f9c 100644 --- a/adapter/dynamodb_item_write.go +++ b/adapter/dynamodb_item_write.go @@ -168,7 +168,7 @@ func (d *DynamoDBServer) retryItemWriteWithGenerationLegacy( return plan, nil } plan.req.StartTS = readTS - if err = d.commitItemWrite(ctx, plan.req); err != nil { + if err = d.commitItemWrite(ctx, readTimestamp, plan.req); err != nil { if !isRetryableTransactWriteError(err) { return nil, errors.WithStack(err) } @@ -208,6 +208,10 @@ type reusableItemWrite struct { // — the write set was built once from attempt 1's read — so plan is also the // correct value to return when the FSM dedup no-ops the apply (R1). plan *itemWritePlan + // readTimestamp carries the Phase-D dispatch capability that validates the + // reused StartTS. Retaining only the numeric StartTS would leave retries + // unable to re-vouch the same applied read watermark. + readTimestamp kv.ReadTimestamp // commitTS is the most recent dispatched commit_ts for this write set; the // next retry passes it as PrevCommitTS so the FSM probes exactly the attempt // that might have landed. @@ -291,13 +295,14 @@ func (d *DynamoDBServer) itemWriteFirstAttempt( } plan.req.StartTS = readTS plan.req.CommitTS = commitTS - if dispErr := d.commitItemWrite(ctx, plan.req); dispErr != nil { + if dispErr := d.commitItemWrite(ctx, readTimestamp, plan.req); dispErr != nil { // dispErr is already wrapped by commitItemWrite; return it raw. if isRetryableTransactWriteError(dispErr) { return nil, &reusableItemWrite{ - plan: plan, - commitTS: commitTS, - probeKey: kv.PrimaryKeyForElems(plan.req.Elems), + plan: plan, + readTimestamp: readTimestamp, + commitTS: commitTS, + probeKey: kv.PrimaryKeyForElems(plan.req.Elems), }, dispErr } return nil, nil, dispErr @@ -319,7 +324,7 @@ func (d *DynamoDBServer) itemWriteReuseAttempt( } pending.plan.req.CommitTS = commitTS pending.plan.req.PrevCommitTS = pending.commitTS - dispErr := d.commitItemWrite(ctx, pending.plan.req) + dispErr := d.commitItemWrite(ctx, pending.readTimestamp, pending.plan.req) if dispErr == nil { return d.finishItemWriteAttempt(ctx, tableName, pending.plan) } @@ -441,8 +446,9 @@ func (d *DynamoDBServer) preparePutItemWrite(ctx context.Context, in putItemInpu }, nil } -func (d *DynamoDBServer) commitItemWrite(ctx context.Context, req *kv.OperationGroup[kv.OP]) error { - _, err := d.coordinator.Dispatch(ctx, req) +func (d *DynamoDBServer) commitItemWrite(ctx context.Context, readTimestamp kv.ReadTimestamp, req *kv.OperationGroup[kv.OP]) error { + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + _, err := kv.DispatchWithReadTimestamp(dispatchCtx, d.coordinator, req) if err != nil { return errors.WithStack(err) } diff --git a/adapter/dynamodb_onephase_dedup_test.go b/adapter/dynamodb_onephase_dedup_test.go index ce7b9c7d4..1137c0d1a 100644 --- a/adapter/dynamodb_onephase_dedup_test.go +++ b/adapter/dynamodb_onephase_dedup_test.go @@ -133,6 +133,23 @@ func TestItemWriteDedup_PriorAttemptDidNotLand_Applies(t *testing.T) { require.Equal(t, 0, coord.probeNoOps, "nothing landed, so the probe must miss and the reuse applies") } +func TestItemWriteDedup_PhaseDVouchesReuse(t *testing.T) { + t.Parallel() + ctx := context.Background() + st := store.NewMVCCStore() + coord := newPhaseDDedupTestCoordinator(st, 1, false) + 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, 2, coord.dispatches) + require.Equal(t, uint64(2), coord.vouches.Load(), "first attempt and reused write set must each reserve a Phase-D dispatch voucher") +} + // TestItemWriteDedup_SelfInflictedReuseConflict_ReturnsSuccess: attempt 1 // pre-rejects, the reuse then LANDS but surfaces WriteConflict (self-inflicted // conflict under churn). The adapter-side self-conflict guard probes the reuse's diff --git a/adapter/redis_list_dedup_test.go b/adapter/redis_list_dedup_test.go index 428bbe110..6a3d8b4a7 100644 --- a/adapter/redis_list_dedup_test.go +++ b/adapter/redis_list_dedup_test.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "errors" + "sync/atomic" "testing" "github.com/bootjp/elastickv/kv" @@ -54,6 +55,18 @@ type dedupTestCoordinator struct { beforeDispatch func(n int) } +type phaseDDedupTestCoordinator struct { + *dedupTestCoordinator + alloc *phaseDDedupAllocator + vouches atomic.Uint64 +} + +type phaseDDedupAllocator struct { + next atomic.Uint64 + validateCalls atomic.Uint64 + phaseDFloor uint64 +} + func newDedupTestCoordinator(st store.MVCCStore, ambiguousDispatch int, lands bool) *dedupTestCoordinator { return &dedupTestCoordinator{ occAdapterCoordinator: newOCCAdapterCoordinator(st), @@ -62,6 +75,53 @@ func newDedupTestCoordinator(st store.MVCCStore, ambiguousDispatch int, lands bo } } +func newPhaseDDedupTestCoordinator(st store.MVCCStore, ambiguousDispatch int, lands bool) *phaseDDedupTestCoordinator { + return &phaseDDedupTestCoordinator{ + dedupTestCoordinator: newDedupTestCoordinator(st, ambiguousDispatch, lands), + alloc: &phaseDDedupAllocator{phaseDFloor: 10}, + } +} + +func (c *phaseDDedupTestCoordinator) TimestampAllocator() kv.TimestampAllocator { + return c.alloc +} + +func (c *phaseDDedupTestCoordinator) VouchAppliedReadTimestamp(uint64, kv.AppliedReadTimestampVoucherRef) error { + c.vouches.Add(1) + return nil +} + +func (a *phaseDDedupAllocator) Next(ctx context.Context) (uint64, error) { + return a.NextAfter(ctx, 0) +} + +func (a *phaseDDedupAllocator) NextAfter(_ context.Context, min uint64) (uint64, error) { + for { + cur := a.next.Load() + next := cur + 1 + if next <= min { + next = min + 1 + } + if a.next.CompareAndSwap(cur, next) { + return next, nil + } + } +} + +func (a *phaseDDedupAllocator) ValidateDurableTimestamp(_ context.Context, timestamp uint64) error { + a.validateCalls.Add(1) + if timestamp == 0 { + return kv.ErrTSOTimestampInvalid + } + if timestamp <= a.phaseDFloor { + return errors.Join(kv.ErrTSOTimestampInvalid, kv.ErrTSOTimestampPrePhaseD) + } + return nil +} + +func (a *phaseDDedupAllocator) PhaseDActive() bool { return true } +func (a *phaseDDedupAllocator) PhaseDRequired() bool { return true } + func (c *dedupTestCoordinator) Dispatch(ctx context.Context, req *kv.OperationGroup[kv.OP]) (*kv.CoordinateResponse, error) { c.dispatches++ n := c.dispatches diff --git a/adapter/redis_retry_test.go b/adapter/redis_retry_test.go index dd5eab7f5..4817d8ee4 100644 --- a/adapter/redis_retry_test.go +++ b/adapter/redis_retry_test.go @@ -352,6 +352,29 @@ func TestRedisXAddDedupsLandedWireWriteConflict(t *testing.T) { require.Equal(t, int64(1), meta.Length, "the generated XADD entry must not be appended twice") } +func TestRedisXAddDedupPhaseDVouchesReuse(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + coord := newPhaseDDedupTestCoordinator(st, 1, false) + srv := &RedisServer{ + store: st, + coordinator: coord, + scriptCache: map[string]string{}, + onePhaseTxnDedup: true, + } + conn := &recordingConn{} + + srv.xadd(conn, redcon.Command{Args: [][]byte{ + []byte(cmdXAdd), []byte("retry:stream"), []byte("*"), []byte("field"), []byte("value"), + }}) + + require.Empty(t, conn.err) + require.NotEmpty(t, conn.bulk) + require.Equal(t, 2, coord.dispatches) + require.Equal(t, uint64(2), coord.vouches.Load(), "first attempt and reused write set must each reserve a Phase-D dispatch voucher") +} + func TestRedisXAddDedupDisabledDoesNotReplayLandedWireConflict(t *testing.T) { t.Parallel() diff --git a/adapter/redis_stream_cmds.go b/adapter/redis_stream_cmds.go index 0cfea772a..0023c6705 100644 --- a/adapter/redis_stream_cmds.go +++ b/adapter/redis_stream_cmds.go @@ -231,11 +231,11 @@ func (r *RedisServer) xadd(conn redcon.Conn, cmd redcon.Command) { } func (r *RedisServer) xaddTxn(ctx context.Context, key []byte, req xaddRequest) (string, error) { - id, readTS, elems, err := r.prepareXAdd(ctx, key, req) + id, readTimestamp, elems, err := r.prepareXAdd(ctx, key, req) if err != nil { return "", err } - return id, r.dispatchAndSignalStream(ctx, true, readTS, elems, key) + return id, r.dispatchAndSignalStreamWithReadTimestamp(ctx, true, readTimestamp, elems, key) } // prepareXAdd fixes one XADD attempt's stream ID and exact write set at a @@ -246,38 +246,39 @@ func (r *RedisServer) prepareXAdd( ctx context.Context, key []byte, req xaddRequest, -) (string, uint64, []*kv.Elem[kv.OP], error) { - readTS, err := r.xaddReadTimestamp(ctx, key) +) (string, kv.ReadTimestamp, []*kv.Elem[kv.OP], error) { + readTimestamp, err := r.xaddReadTimestamp(ctx, key) if err != nil { - return "", 0, nil, err + return "", kv.ReadTimestamp{}, nil, err } + readTS := readTimestamp.Timestamp() typ, err := r.streamTypeForXAdd(ctx, key, readTS) if err != nil { - return "", 0, nil, err + return "", kv.ReadTimestamp{}, nil, err } legacyCleanup, meta, metaFound, err := r.streamWriteBase(ctx, key, readTS) if err != nil { - return "", 0, nil, err + return "", kv.ReadTimestamp{}, nil, err } legacyCleanup, meta, metaFound, err = r.streamCleanupForExpiredRecreate( ctx, key, readTS, typ, legacyCleanup, meta, metaFound) if err != nil { - return "", 0, nil, err + return "", kv.ReadTimestamp{}, nil, err } id, parsedID, err := resolveXAddID(meta, metaFound, req.id) if err != nil { - return "", 0, nil, err + return "", kv.ReadTimestamp{}, nil, err } if err := xaddEnforceMaxWideColumn(key, meta.Length, req.maxLen); err != nil { - return "", 0, nil, err + return "", kv.ReadTimestamp{}, nil, err } entryValue, err := marshalStreamEntry(newRedisStreamEntry(id, req.fields)) if err != nil { - return "", 0, nil, err + return "", kv.ReadTimestamp{}, nil, err } // Capacity hint covers: stream cleanup Dels + one entry Put + one meta @@ -296,7 +297,7 @@ func (r *RedisServer) prepareXAdd( nextLen, trim, err := r.xaddTrimIfNeeded(ctx, key, readTS, req.maxLen, meta.Length+1) if err != nil { - return "", 0, nil, err + return "", kv.ReadTimestamp{}, nil, err } elems = append(elems, trim...) elems = appendMaxLenZeroSelfDel(elems, req.maxLen, key, parsedID) @@ -308,11 +309,11 @@ func (r *RedisServer) prepareXAdd( ExpireAt: meta.ExpireAt, }) if err != nil { - return "", 0, nil, cockerrors.WithStack(err) + return "", kv.ReadTimestamp{}, nil, cockerrors.WithStack(err) } elems = append(elems, &kv.Elem[kv.OP]{Op: kv.Put, Key: store.StreamMetaKey(key), Value: metaBytes}) - return id, readTS, elems, nil + return id, readTimestamp, elems, nil } // reusableXAdd captures the exact ID and write set from an XADD attempt. A @@ -320,11 +321,12 @@ func (r *RedisServer) prepareXAdd( // already landed instead of resolving "*" against newer stream metadata and // appending a duplicate entry. type reusableXAdd struct { - elems []*kv.Elem[kv.OP] - startTS uint64 - commitTS uint64 - probeKey []byte - id string + elems []*kv.Elem[kv.OP] + readTimestamp kv.ReadTimestamp + startTS uint64 + commitTS uint64 + probeKey []byte + id string } // xaddTxnWithDedup retries XADD only through exact write-set reuse. This is the @@ -365,16 +367,18 @@ func (r *RedisServer) firstXAddAttempt( key []byte, req xaddRequest, ) (string, *reusableXAdd, error) { - id, readTS, elems, err := r.prepareXAdd(ctx, key, req) + id, readTimestamp, elems, err := r.prepareXAdd(ctx, key, req) if err != nil { return "", nil, err } + readTS := readTimestamp.Timestamp() startTS := normalizeStartTS(readTS) commitTS, err := r.nextCommitTSAfter(ctx, startTS, "redis xadd first-attempt: allocate commitTS") if err != nil { return "", nil, cockerrors.WithStack(err) } - _, dispErr := r.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + attemptCtx := readTimestamp.WithDispatchVoucher(ctx) + _, dispErr := kv.DispatchWithReadTimestamp(attemptCtx, r.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, CommitTS: commitTS, @@ -392,11 +396,12 @@ func (r *RedisServer) firstXAddAttempt( return "", nil, cockerrors.WithStack(dispErr) } return "", &reusableXAdd{ - elems: elems, - startTS: startTS, - commitTS: commitTS, - probeKey: kv.PrimaryKeyForElems(elems), - id: id, + elems: elems, + readTimestamp: readTimestamp, + startTS: startTS, + commitTS: commitTS, + probeKey: kv.PrimaryKeyForElems(elems), + id: id, }, cockerrors.WithStack(dispErr) } @@ -413,7 +418,8 @@ func (r *RedisServer) dispatchXAddReuse( if allocErr != nil { return "", false, cockerrors.WithStack(allocErr) } - _, dispErr := r.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + attemptCtx := pending.readTimestamp.WithDispatchVoucher(ctx) + _, dispErr := kv.DispatchWithReadTimestamp(attemptCtx, r.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: pending.startTS, CommitTS: commitTS, @@ -475,37 +481,40 @@ func (r *RedisServer) streamCleanupForExpiredRecreate( return cleanup, store.StreamMeta{}, false, nil } -func (r *RedisServer) xaddReadTimestamp(ctx context.Context, key []byte) (uint64, error) { - readTS, err := r.beginTxnStartTS(ctx, "redis xadd: begin read timestamp") +func (r *RedisServer) xaddReadTimestamp(ctx context.Context, key []byte) (kv.ReadTimestamp, error) { + readTimestamp, err := r.beginTxnReadTimestamp(ctx, "redis xadd: begin read timestamp") if err != nil { - return 0, cockerrors.WithStack(err) + return kv.ReadTimestamp{}, cockerrors.WithStack(err) } + readTS := readTimestamp.Timestamp() typ, err := r.keyTypeAtExpect(ctx, key, readTS, redisTypeStream) if err != nil { - return 0, err + return kv.ReadTimestamp{}, err } if typ != redisTypeNone && typ != redisTypeStream { - return 0, wrongTypeError() + return kv.ReadTimestamp{}, wrongTypeError() } - return readTS, nil + return readTimestamp, nil } -// dispatchAndSignalStream dispatches the elems through the coordinator -// and, on success, wakes any XREAD BLOCK waiter on the same node. -// dispatchElems blocks until the FSM applies locally, so by the time -// Signal fires the new entries are visible at the readTS the woken -// waiter will pick on its next iteration. Pulled out of xaddTxn so the -// parent function stays under the cyclop budget — the signal step -// would otherwise add an extra branch on the dispatch error path. -func (r *RedisServer) dispatchAndSignalStream( +func (r *RedisServer) dispatchAndSignalStreamWithReadTimestamp( ctx context.Context, isTxn bool, - startTS uint64, + readTimestamp kv.ReadTimestamp, elems []*kv.Elem[kv.OP], streamKey []byte, ) error { - if err := r.dispatchElems(ctx, isTxn, startTS, elems); err != nil { - return err + if len(elems) == 0 { + return nil + } + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + _, err := kv.DispatchWithReadTimestamp(dispatchCtx, r.coordinator, &kv.OperationGroup[kv.OP]{ + IsTxn: isTxn, + StartTS: normalizeStartTS(readTimestamp.Timestamp()), + Elems: elems, + }) + if err != nil { + return cockerrors.WithStack(err) } r.streamWaiters.Signal(streamKey) return nil diff --git a/adapter/s3.go b/adapter/s3.go index bcb5b0355..b70d31da0 100644 --- a/adapter/s3.go +++ b/adapter/s3.go @@ -2447,6 +2447,7 @@ func (s *S3Server) beginTxnReadTimestamp(ctx context.Context, readTS uint64, lab if readTS == ^uint64(0) { if alloc, ok := kv.TimestampAllocatorThrough(s.coordinator); ok { if phaseD, phaseDOK := alloc.(kv.TSOPhaseDState); phaseDOK && (phaseD.PhaseDRequired() || phaseD.PhaseDActive()) { + readTS = 1 readTimestamp, err := kv.BeginReadTimestampThrough(ctx, s.coordinator, readTS, label) return readTimestamp, errors.WithStack(err) } diff --git a/adapter/s3_hlc_fence_test.go b/adapter/s3_hlc_fence_test.go index 348a48bb7..b5b473d2c 100644 --- a/adapter/s3_hlc_fence_test.go +++ b/adapter/s3_hlc_fence_test.go @@ -70,7 +70,7 @@ func TestS3BeginTxnReadTimestampPhaseDPreservesAppliedWatermark(t *testing.T) { require.Zero(t, allocator.count, "an applied Phase-D read watermark must not allocate ahead of Raft apply") } -func TestS3BeginTxnReadTimestampPhaseDRejectsLatestSentinel(t *testing.T) { +func TestS3BeginTxnReadTimestampPhaseDNormalizesEmptySnapshotSentinel(t *testing.T) { t.Parallel() allocator := &distributionTSOAllocator{base: 100, phaseD: true, phaseDFloor: 10} @@ -78,9 +78,10 @@ func TestS3BeginTxnReadTimestampPhaseDRejectsLatestSentinel(t *testing.T) { coord.allocator = allocator srv := &S3Server{coordinator: coord} - _, err := srv.beginTxnReadTimestamp(context.Background(), ^uint64(0), "test") - require.ErrorIs(t, err, kv.ErrTSOTimestampInvalid) - require.Zero(t, allocator.count, "the latest sentinel must fail closed instead of allocating an unapplied read timestamp") + readTimestamp, err := srv.beginTxnReadTimestamp(context.Background(), ^uint64(0), "test") + require.NoError(t, err) + require.Equal(t, uint64(1), readTimestamp.Timestamp()) + require.Zero(t, allocator.count, "an empty S3 snapshot must not allocate ahead of Raft apply") } // TestS3NextTxnCommitTSFailsClosedOnExpiredCeiling verifies that diff --git a/kv/coordinator.go b/kv/coordinator.go index af91a09fd..beec5a599 100644 --- a/kv/coordinator.go +++ b/kv/coordinator.go @@ -243,7 +243,7 @@ var _ AppliedReadTimestampVoucher = (*Coordinate)(nil) // VouchAppliedReadTimestamp is a no-op for the single-group coordinator. The // sharded coordinator consumes vouchers before cross-group StartTS validation. -func (c *Coordinate) VouchAppliedReadTimestamp(uint64) error { +func (c *Coordinate) VouchAppliedReadTimestamp(uint64, AppliedReadTimestampVoucherRef) error { return nil } diff --git a/kv/keyviz_label.go b/kv/keyviz_label.go index ba40e30d5..5d112f545 100644 --- a/kv/keyviz_label.go +++ b/kv/keyviz_label.go @@ -95,12 +95,12 @@ func (c keyVizLabeledCoordinator) TimestampAllocator() TimestampAllocator { return alloc } -func (c keyVizLabeledCoordinator) VouchAppliedReadTimestamp(timestamp uint64) error { +func (c keyVizLabeledCoordinator) VouchAppliedReadTimestamp(timestamp uint64, ref AppliedReadTimestampVoucherRef) error { voucher, ok := c.inner.(AppliedReadTimestampVoucher) if !ok { return errors.WithStack(ErrTSOProtocolUnsupported) } - return errors.WithStack(voucher.VouchAppliedReadTimestamp(timestamp)) + return errors.WithStack(voucher.VouchAppliedReadTimestamp(timestamp, ref)) } func (c keyVizLabeledCoordinator) LeaseRead(ctx context.Context) (uint64, error) { diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index 3705fcf4c..af1143e0a 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -392,7 +392,7 @@ type ShardedCoordinator struct { PhaseDActive() bool } appliedReadVoucherMu sync.Mutex - appliedReadVouchers map[uint64]uint64 + appliedReadVouchers map[appliedReadVoucherKey]uint64 // allShardGroupIDs, when configured, is the explicit set of data groups // that whole-keyspace operations must visit. It lets callers keep // non-data groups (for example a reserved timestamp group) in c.groups @@ -433,6 +433,11 @@ type ShardedCoordinator struct { registrationGate *RegistrationGate } +type appliedReadVoucherKey struct { + timestamp uint64 + ref AppliedReadTimestampVoucherRef +} + // RegistrationGate carries the Stage 7a §4.1 registration-before- // first-write barrier into the coordinator. main.go owns the // encryption StateCache and the registration goroutine and supplies @@ -487,40 +492,46 @@ func (c *ShardedCoordinator) TimestampAllocator() TimestampAllocator { // VouchAppliedReadTimestamp records one use of an audited adapter watermark. // The bounded map prevents abandoned requests from growing process memory. -func (c *ShardedCoordinator) VouchAppliedReadTimestamp(timestamp uint64) error { - if c == nil || timestamp == 0 || timestamp == ^uint64(0) { +func (c *ShardedCoordinator) VouchAppliedReadTimestamp(timestamp uint64, ref AppliedReadTimestampVoucherRef) error { + if c == nil || timestamp == 0 || timestamp == ^uint64(0) || ref.id == 0 { return errors.WithStack(ErrTSOTimestampInvalid) } + key := appliedReadVoucherKey{timestamp: timestamp, ref: ref} c.appliedReadVoucherMu.Lock() defer c.appliedReadVoucherMu.Unlock() if c.appliedReadVouchers == nil { - c.appliedReadVouchers = make(map[uint64]uint64) + c.appliedReadVouchers = make(map[appliedReadVoucherKey]uint64) } - if uses, ok := c.appliedReadVouchers[timestamp]; ok { - c.appliedReadVouchers[timestamp] = uses + 1 + if uses, ok := c.appliedReadVouchers[key]; ok { + c.appliedReadVouchers[key] = uses + 1 return nil } if len(c.appliedReadVouchers) >= maxAppliedReadTimestampVouches { return errors.WithStack(ErrTSOReadVoucherLimit) } - c.appliedReadVouchers[timestamp] = 1 + c.appliedReadVouchers[key] = 1 return nil } -func (c *ShardedCoordinator) consumeAppliedReadTimestampVoucher(timestamp uint64) bool { +func (c *ShardedCoordinator) consumeAppliedReadTimestampVoucher(ctx context.Context, timestamp uint64) bool { if c == nil { return false } + ref, ok := appliedReadTimestampVoucherRefFromContext(ctx, timestamp) + if !ok { + return false + } + key := appliedReadVoucherKey{timestamp: timestamp, ref: ref} c.appliedReadVoucherMu.Lock() defer c.appliedReadVoucherMu.Unlock() - uses, ok := c.appliedReadVouchers[timestamp] + uses, ok := c.appliedReadVouchers[key] if !ok { return false } if uses <= 1 { - delete(c.appliedReadVouchers, timestamp) + delete(c.appliedReadVouchers, key) } else { - c.appliedReadVouchers[timestamp] = uses - 1 + c.appliedReadVouchers[key] = uses - 1 } return true } @@ -1246,7 +1257,7 @@ func (c *ShardedCoordinator) validateCallerSuppliedTxnStart(ctx context.Context, if !callerSupplied { return nil } - if c.consumeAppliedReadTimestampVoucher(startTS) || singleShard { + if c.consumeAppliedReadTimestampVoucher(ctx, startTS) || singleShard { return nil } return c.validateCrossShardReadTimestamp(ctx, startTS) diff --git a/kv/sharded_coordinator_txn_test.go b/kv/sharded_coordinator_txn_test.go index 82826272a..d2bd878d8 100644 --- a/kv/sharded_coordinator_txn_test.go +++ b/kv/sharded_coordinator_txn_test.go @@ -241,7 +241,7 @@ func TestShardedCoordinatorDispatchTxn_PhaseDAcceptsValidatedCallerStartTS(t *te require.Equal(t, uint64(100), g2Txn.requests[0].Ts) } -func TestShardedCoordinatorDispatchTxn_PhaseDAcceptsVouchedAppliedWatermarkOnce(t *testing.T) { +func TestShardedCoordinatorDispatchTxn_PhaseDAcceptsBoundAppliedWatermark(t *testing.T) { t.Parallel() prePhaseDErr := errors.Join(ErrTSOTimestampInvalid, ErrTSOTimestampPrePhaseD) coord, g1Txn, g2Txn, alloc := newPhaseDCrossShardCoordinator(t, prePhaseDErr) @@ -261,15 +261,23 @@ func TestShardedCoordinatorDispatchTxn_PhaseDAcceptsVouchedAppliedWatermarkOnce( }, } } + _, err = coord.Dispatch(context.Background(), request()) + require.ErrorIs(t, err, ErrTSOTimestampPrePhaseD, "a numeric timestamp without the bound capability must not steal a voucher") + require.Equal(t, uint64(2), alloc.validateCalls.Load()) + require.Empty(t, g1Txn.requests) + require.Empty(t, g2Txn.requests) + + ctx := readTS.WithDispatchVoucher(context.Background()) + _, err = DispatchWithReadTimestamp(ctx, coord, request()) require.NoError(t, err) - require.Equal(t, uint64(1), alloc.validateCalls.Load(), "voucher must bypass numeric Phase-D validation") + require.Equal(t, uint64(2), alloc.validateCalls.Load(), "bound voucher must bypass numeric Phase-D validation") require.Len(t, g1Txn.requests, 2) require.Len(t, g2Txn.requests, 2) _, err = coord.Dispatch(context.Background(), request()) require.ErrorIs(t, err, ErrTSOTimestampPrePhaseD) - require.Equal(t, uint64(2), alloc.validateCalls.Load(), "voucher must be single-use") + require.Equal(t, uint64(3), alloc.validateCalls.Load(), "voucher must be bound and single-use") } func TestDispatchWithReadTimestampVouchesEveryBoundDispatch(t *testing.T) { diff --git a/kv/tso.go b/kv/tso.go index c92840db1..4538c2e59 100644 --- a/kv/tso.go +++ b/kv/tso.go @@ -69,7 +69,13 @@ type TSOPhaseDState interface { // It is a process-local capability used only to distinguish audited adapter // snapshots from arbitrary caller-supplied StartTS values during Phase D. type AppliedReadTimestampVoucher interface { - VouchAppliedReadTimestamp(uint64) error + VouchAppliedReadTimestamp(uint64, AppliedReadTimestampVoucherRef) error +} + +// AppliedReadTimestampVoucherRef is an opaque process-local dispatch +// capability. Only this package can mint a non-zero ref. +type AppliedReadTimestampVoucherRef struct { + id uint64 } // ReadTimestamp is the adapter-side result of beginning a transaction snapshot. @@ -86,9 +92,19 @@ func (t ReadTimestamp) Timestamp() uint64 { return t.timestamp } +func (t ReadTimestamp) voucherRef() (AppliedReadTimestampVoucherRef, bool) { + if t.voucher == nil || t.voucher.ref.id == 0 { + return AppliedReadTimestampVoucherRef{}, false + } + return t.voucher.ref, true +} + +var nextAppliedReadDispatchVoucherID atomic.Uint64 + type appliedReadDispatchVoucher struct { mu sync.Mutex prepared uint64 + ref AppliedReadTimestampVoucherRef } type appliedReadDispatchVoucherContextKey struct{} @@ -100,10 +116,21 @@ func (t ReadTimestamp) WithDispatchVoucher(ctx context.Context) context.Context return context.WithValue(nonNilTSOContext(ctx), appliedReadDispatchVoucherContextKey{}, t) } +func appliedReadTimestampVoucherRefFromContext(ctx context.Context, timestamp uint64) (AppliedReadTimestampVoucherRef, bool) { + if ctx == nil { + return AppliedReadTimestampVoucherRef{}, false + } + readTimestamp, ok := ctx.Value(appliedReadDispatchVoucherContextKey{}).(ReadTimestamp) + if !ok || readTimestamp.timestamp != timestamp { + return AppliedReadTimestampVoucherRef{}, false + } + return readTimestamp.voucherRef() +} + // DispatchWithReadTimestamp dispatches an OCC operation under the applied-read -// capability bound by ReadTimestamp.WithDispatchVoucher. The first dispatch -// consumes the voucher reserved by BeginReadTimestampThrough; every subsequent -// dispatch reserves one additional use immediately before dispatching. +// capability bound by ReadTimestamp.WithDispatchVoucher. Each dispatch reserves +// one token tied to that bound capability immediately before dispatching, so a +// same-valued StartTS from another request cannot consume it. func DispatchWithReadTimestamp( ctx context.Context, coord Coordinator, @@ -128,18 +155,22 @@ func DispatchWithReadTimestamp( return resp, errors.WithStack(err) } +func newAppliedReadDispatchVoucher() *appliedReadDispatchVoucher { + id := nextAppliedReadDispatchVoucherID.Add(1) + if id == 0 { + id = nextAppliedReadDispatchVoucherID.Add(1) + } + return &appliedReadDispatchVoucher{ref: AppliedReadTimestampVoucherRef{id: id}} +} + func (v *appliedReadDispatchVoucher) prepare(coord Coordinator, timestamp uint64) error { v.mu.Lock() defer v.mu.Unlock() - if v.prepared == 0 { - v.prepared = 1 - return nil - } voucher, ok := coord.(AppliedReadTimestampVoucher) if !ok { return errors.WithStack(ErrTSOProtocolUnsupported) } - if err := voucher.VouchAppliedReadTimestamp(timestamp); err != nil { + if err := voucher.VouchAppliedReadTimestamp(timestamp, v.ref); err != nil { return errors.WithStack(err) } v.prepared++ @@ -261,7 +292,7 @@ func BeginReadTimestampThrough( } readTimestamp := ReadTimestamp{timestamp: legacyTimestamp} if vouched { - readTimestamp.voucher = &appliedReadDispatchVoucher{} + readTimestamp.voucher = newAppliedReadDispatchVoucher() } return readTimestamp, nil } @@ -296,13 +327,9 @@ func validateAppliedReadTimestamp( if !errors.Is(err, ErrTSOTimestampPrePhaseD) { return false, errors.Wrap(err, label) } - voucher, ok := coord.(AppliedReadTimestampVoucher) - if !ok { + if _, ok := coord.(AppliedReadTimestampVoucher); !ok { return false, errors.Wrap(ErrTSOProtocolUnsupported, label+": applied read voucher unavailable") } - if err := voucher.VouchAppliedReadTimestamp(timestamp); err != nil { - return false, errors.Wrap(err, label) - } return true, nil } diff --git a/main.go b/main.go index 5f515a91e..91fc75256 100644 --- a/main.go +++ b/main.go @@ -2400,12 +2400,12 @@ func (c startupGatedCoordinator) TimestampAllocator() kv.TimestampAllocator { return alloc } -func (c startupGatedCoordinator) VouchAppliedReadTimestamp(timestamp uint64) error { +func (c startupGatedCoordinator) VouchAppliedReadTimestamp(timestamp uint64, ref kv.AppliedReadTimestampVoucherRef) error { voucher, ok := c.inner.(kv.AppliedReadTimestampVoucher) if !ok { return errors.WithStack(kv.ErrTSOProtocolUnsupported) } - return errors.WithStack(voucher.VouchAppliedReadTimestamp(timestamp)) + return errors.WithStack(voucher.VouchAppliedReadTimestamp(timestamp, ref)) } func (c startupGatedCoordinator) LeaseRead(ctx context.Context) (uint64, error) { From f6414537c0ae6f8cb5bfebe58764a7b18c5071e7 Mon Sep 17 00:00:00 2001 From: bootjp Date: Thu, 23 Jul 2026 21:01:48 +0900 Subject: [PATCH 10/14] proto: reserve retired raft status fields --- proto/service.pb.go | 4 ++-- proto/service.proto | 2 ++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/proto/service.pb.go b/proto/service.pb.go index a70a3434b..bf6cd8ff9 100644 --- a/proto/service.pb.go +++ b/proto/service.pb.go @@ -2377,7 +2377,7 @@ const file_service_proto_rawDesc = "" + "\bstart_ts\x18\x01 \x01(\x04R\astartTs\",\n" + "\x10RollbackResponse\x12\x18\n" + "\asuccess\x18\x01 \x01(\bR\asuccess\"\x18\n" + - "\x16RaftAdminStatusRequest\"\xa2\x03\n" + + "\x16RaftAdminStatusRequest\"\xae\x03\n" + "\x17RaftAdminStatusResponse\x12%\n" + "\x05state\x18\x01 \x01(\x0e2\x0f.RaftAdminStateR\x05state\x12\x1b\n" + "\tleader_id\x18\x02 \x01(\tR\bleaderId\x12%\n" + @@ -2391,7 +2391,7 @@ const file_service_proto_rawDesc = "" + "fsmPending\x12\x1b\n" + "\tnum_peers\x18\n" + " \x01(\x04R\bnumPeers\x12,\n" + - "\x12last_contact_nanos\x18\v \x01(\x03R\x10lastContactNanos\"\x1f\n" + + "\x12last_contact_nanos\x18\v \x01(\x03R\x10lastContactNanosJ\x04\b\f\x10\rJ\x04\b\r\x10\x0e\"\x1f\n" + "\x1dRaftAdminConfigurationRequest\"W\n" + "\x0fRaftAdminMember\x12\x0e\n" + "\x02id\x18\x01 \x01(\tR\x02id\x12\x18\n" + diff --git a/proto/service.proto b/proto/service.proto index 6b01e6012..815e67648 100644 --- a/proto/service.proto +++ b/proto/service.proto @@ -187,6 +187,8 @@ enum RaftAdminState { message RaftAdminStatusRequest {} message RaftAdminStatusResponse { + reserved 12, 13; + RaftAdminState state = 1; string leader_id = 2; string leader_address = 3; From c8e3adaccc0beccb3406f9a197561d6cb16c0719 Mon Sep 17 00:00:00 2001 From: bootjp Date: Thu, 23 Jul 2026 21:24:55 +0900 Subject: [PATCH 11/14] adapter: align timestamp validation status test --- adapter/distribution_server_test.go | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/adapter/distribution_server_test.go b/adapter/distribution_server_test.go index e1a88cf93..64d30d41d 100644 --- a/adapter/distribution_server_test.go +++ b/adapter/distribution_server_test.go @@ -174,7 +174,7 @@ func TestDistributionServerValidateTimestamp(t *testing.T) { require.Equal(t, uint64(501), resp.GetAllocationFloor()) _, err = s.ValidateTimestamp(context.Background(), &pb.ValidateTimestampRequest{Timestamp: 499}) - require.Equal(t, codes.InvalidArgument, status.Code(err)) + require.Equal(t, codes.OutOfRange, status.Code(err)) } func TestDistributionServerValidateTimestamp_LeaderRoutedRPC(t *testing.T) { @@ -197,7 +197,9 @@ func TestDistributionServerValidateTimestamp_LeaderRoutedRPC(t *testing.T) { t.Cleanup(func() { require.NoError(t, routed.Close()) }) require.NoError(t, routed.ValidateDurableTimestamp(context.Background(), 700)) - require.ErrorIs(t, routed.ValidateDurableTimestamp(context.Background(), 699), kv.ErrTSOTimestampInvalid) + err = routed.ValidateDurableTimestamp(context.Background(), 699) + require.ErrorIs(t, err, kv.ErrTSOTimestampInvalid) + require.ErrorIs(t, err, kv.ErrTSOTimestampPrePhaseD) } func TestDistributionServerGetTimestamp_RejectsFollower(t *testing.T) { From 580c5b58f0ea7a1bd7e1fa089898f7f3b79c2bd4 Mon Sep 17 00:00:00 2001 From: bootjp Date: Thu, 23 Jul 2026 21:33:09 +0900 Subject: [PATCH 12/14] tso: preserve applied-read vouchers through gated dispatch --- adapter/redis_compat_helpers.go | 17 +++++++ adapter/redis_strings.go | 5 +- adapter/s3_put_object.go | 18 +++---- kv/coordinator.go | 4 ++ kv/keyviz_label.go | 6 +++ kv/sharded_coordinator.go | 21 +++++++++ kv/sharded_coordinator_txn_test.go | 75 ++++++++++++++++++++++++++++++ kv/tso.go | 26 +++++++++-- kv/tso_raft.go | 11 +++-- kv/tso_raft_test.go | 38 +++++++++++++++ main.go | 6 +++ 11 files changed, 209 insertions(+), 18 deletions(-) diff --git a/adapter/redis_compat_helpers.go b/adapter/redis_compat_helpers.go index 292787f2c..19e6d9891 100644 --- a/adapter/redis_compat_helpers.go +++ b/adapter/redis_compat_helpers.go @@ -778,6 +778,23 @@ func (r *RedisServer) dispatchElems(ctx context.Context, isTxn bool, startTS uin return errors.WithStack(err) } +func (r *RedisServer) dispatchReadTimestampElems(ctx context.Context, isTxn bool, readTimestamp kv.ReadTimestamp, elems []*kv.Elem[kv.OP]) error { + if len(elems) == 0 { + return nil + } + startTS := readTimestamp.Timestamp() + if startTS == ^uint64(0) { + startTS = 0 + } + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + _, err := kv.DispatchWithReadTimestamp(dispatchCtx, r.coordinator, &kv.OperationGroup[kv.OP]{ + IsTxn: isTxn, + StartTS: startTS, + Elems: elems, + }) + return errors.WithStack(err) +} + // readRedisStringAt reads a Redis string value, trying the prefixed key first // and falling back to the bare key for legacy data written before the // !redis|str| prefix migration. Returns the decoded user value and the diff --git a/adapter/redis_strings.go b/adapter/redis_strings.go index a61ce4066..7a5492685 100644 --- a/adapter/redis_strings.go +++ b/adapter/redis_strings.go @@ -522,10 +522,11 @@ func (r *RedisServer) delLocal(keys [][]byte) (int, error) { err := r.retryRedisWrite(ctx, func() error { elems := []*kv.Elem[kv.OP]{} nextRemoved := 0 - readTS, err := r.beginTxnStartTS(ctx, "redis del: begin read timestamp") + readTimestamp, err := r.beginTxnReadTimestamp(ctx, "redis del: begin read timestamp") if err != nil { return errors.WithStack(err) } + readTS := readTimestamp.Timestamp() for _, key := range keys { keyElems, existed, err := r.deleteLogicalKeyElems(ctx, key, readTS) if err != nil { @@ -536,7 +537,7 @@ func (r *RedisServer) delLocal(keys [][]byte) (int, error) { } elems = append(elems, keyElems...) } - if err := r.dispatchElems(ctx, true, readTS, elems); err != nil { + if err := r.dispatchReadTimestampElems(ctx, true, readTimestamp, elems); err != nil { return err } removed = nextRemoved diff --git a/adapter/s3_put_object.go b/adapter/s3_put_object.go index 8b883d99f..8019aac38 100644 --- a/adapter/s3_put_object.go +++ b/adapter/s3_put_object.go @@ -10,12 +10,13 @@ import ( ) type s3PutObjectState struct { - startTS uint64 - meta *s3BucketMeta - headKey []byte - previous *s3ObjectManifest - uploadID string - readPin *kv.ActiveTimestampToken + startTS uint64 + readTimestamp kv.ReadTimestamp + meta *s3BucketMeta + headKey []byte + previous *s3ObjectManifest + uploadID string + readPin *kv.ActiveTimestampToken } func (s *S3Server) prepareS3PutObject(ctx context.Context, request *http.Request, bucket, objectKey string) (*s3PutObjectState, error) { @@ -26,7 +27,7 @@ func (s *S3Server) prepareS3PutObject(ctx context.Context, request *http.Request } readTS = readTimestamp.Timestamp() startTS := readTS - state := &s3PutObjectState{startTS: startTS, readPin: s.pinReadTS(readTS)} + state := &s3PutObjectState{startTS: startTS, readTimestamp: readTimestamp, readPin: s.pinReadTS(readTS)} prepared := false defer func() { if !prepared { @@ -111,7 +112,8 @@ func (s *S3Server) commitS3PutObject(ctx context.Context, request *http.Request, if err != nil { return false, errors.WithStack(err) } - _, err = s.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := state.readTimestamp.WithDispatchVoucher(ctx) + _, err = kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: state.startTS, CommitTS: commitTS, diff --git a/kv/coordinator.go b/kv/coordinator.go index beec5a599..0a478ce7d 100644 --- a/kv/coordinator.go +++ b/kv/coordinator.go @@ -240,6 +240,7 @@ type Coordinate struct { var _ Coordinator = (*Coordinate)(nil) var _ AppliedReadTimestampVoucher = (*Coordinate)(nil) +var _ AppliedReadTimestampVoucherRevoker = (*Coordinate)(nil) // VouchAppliedReadTimestamp is a no-op for the single-group coordinator. The // sharded coordinator consumes vouchers before cross-group StartTS validation. @@ -247,6 +248,9 @@ func (c *Coordinate) VouchAppliedReadTimestamp(uint64, AppliedReadTimestampVouch return nil } +// RevokeAppliedReadTimestamp is a no-op for the single-group coordinator. +func (c *Coordinate) RevokeAppliedReadTimestamp(uint64, AppliedReadTimestampVoucherRef) {} + type Coordinator interface { Dispatch(ctx context.Context, reqs *OperationGroup[OP]) (*CoordinateResponse, error) IsLeader() bool diff --git a/kv/keyviz_label.go b/kv/keyviz_label.go index 5d112f545..96264d69c 100644 --- a/kv/keyviz_label.go +++ b/kv/keyviz_label.go @@ -103,6 +103,12 @@ func (c keyVizLabeledCoordinator) VouchAppliedReadTimestamp(timestamp uint64, re return errors.WithStack(voucher.VouchAppliedReadTimestamp(timestamp, ref)) } +func (c keyVizLabeledCoordinator) RevokeAppliedReadTimestamp(timestamp uint64, ref AppliedReadTimestampVoucherRef) { + if revoker, ok := c.inner.(AppliedReadTimestampVoucherRevoker); ok { + revoker.RevokeAppliedReadTimestamp(timestamp, ref) + } +} + func (c keyVizLabeledCoordinator) LeaseRead(ctx context.Context) (uint64, error) { if lr, ok := c.inner.(LeaseReadableCoordinator); ok { idx, err := lr.LeaseRead(ctx) diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index af1143e0a..8b0aad0f5 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -513,6 +513,27 @@ func (c *ShardedCoordinator) VouchAppliedReadTimestamp(timestamp uint64, ref App return nil } +// RevokeAppliedReadTimestamp removes one prepared voucher that did not reach +// ShardedCoordinator dispatch validation, for example because an outer +// coordinator decorator rejected the dispatch first. +func (c *ShardedCoordinator) RevokeAppliedReadTimestamp(timestamp uint64, ref AppliedReadTimestampVoucherRef) { + if c == nil || timestamp == 0 || timestamp == ^uint64(0) || ref.id == 0 { + return + } + key := appliedReadVoucherKey{timestamp: timestamp, ref: ref} + c.appliedReadVoucherMu.Lock() + defer c.appliedReadVoucherMu.Unlock() + uses, ok := c.appliedReadVouchers[key] + if !ok { + return + } + if uses <= 1 { + delete(c.appliedReadVouchers, key) + return + } + c.appliedReadVouchers[key] = uses - 1 +} + func (c *ShardedCoordinator) consumeAppliedReadTimestampVoucher(ctx context.Context, timestamp uint64) bool { if c == nil { return false diff --git a/kv/sharded_coordinator_txn_test.go b/kv/sharded_coordinator_txn_test.go index d2bd878d8..6ef124de2 100644 --- a/kv/sharded_coordinator_txn_test.go +++ b/kv/sharded_coordinator_txn_test.go @@ -315,6 +315,34 @@ func TestDispatchWithReadTimestampVouchesEveryBoundDispatch(t *testing.T) { require.Equal(t, uint64(2), alloc.validateCalls.Load()) } +func TestDispatchWithReadTimestampRevokesVoucherWhenOuterGateRejects(t *testing.T) { + t.Parallel() + prePhaseDErr := errors.Join(ErrTSOTimestampInvalid, ErrTSOTimestampPrePhaseD) + coord, _, _, _ := newPhaseDCrossShardCoordinator(t, prePhaseDErr) + gateErr := errors.New("startup gate rejected dispatch") + gated := phaseDGateCoordinator{inner: coord, err: gateErr} + + readTimestamp, err := BeginReadTimestampThrough(context.Background(), gated, 10, "vouch gated applied watermark") + require.NoError(t, err) + _, err = DispatchWithReadTimestamp( + readTimestamp.WithDispatchVoucher(context.Background()), + gated, + &OperationGroup[OP]{ + IsTxn: true, + StartTS: readTimestamp.Timestamp(), + Elems: []*Elem[OP]{ + {Op: Put, Key: []byte("b"), Value: []byte("v1")}, + {Op: Put, Key: []byte("x"), Value: []byte("v2")}, + }, + }, + ) + require.ErrorIs(t, err, gateErr) + + coord.appliedReadVoucherMu.Lock() + defer coord.appliedReadVoucherMu.Unlock() + require.Empty(t, coord.appliedReadVouchers) +} + func TestReadTimestampVoucherBindingShadowsParentCapability(t *testing.T) { prePhaseDErr := errors.Join(ErrTSOTimestampInvalid, ErrTSOTimestampPrePhaseD) coord, _, _, alloc := newPhaseDCrossShardCoordinator(t, prePhaseDErr) @@ -424,6 +452,53 @@ func newPhaseDCrossShardCoordinator( return coord, g1Txn, g2Txn, alloc } +type phaseDGateCoordinator struct { + inner *ShardedCoordinator + err error +} + +func (c phaseDGateCoordinator) Dispatch(context.Context, *OperationGroup[OP]) (*CoordinateResponse, error) { + return nil, c.err +} + +func (c phaseDGateCoordinator) IsLeader() bool { return c.inner.IsLeader() } + +func (c phaseDGateCoordinator) VerifyLeader(ctx context.Context) error { + return c.inner.VerifyLeader(ctx) +} + +func (c phaseDGateCoordinator) LinearizableRead(ctx context.Context) (uint64, error) { + return c.inner.LinearizableRead(ctx) +} + +func (c phaseDGateCoordinator) RaftLeader() string { return c.inner.RaftLeader() } + +func (c phaseDGateCoordinator) IsLeaderForKey(key []byte) bool { + return c.inner.IsLeaderForKey(key) +} + +func (c phaseDGateCoordinator) VerifyLeaderForKey(ctx context.Context, key []byte) error { + return c.inner.VerifyLeaderForKey(ctx, key) +} + +func (c phaseDGateCoordinator) RaftLeaderForKey(key []byte) string { + return c.inner.RaftLeaderForKey(key) +} + +func (c phaseDGateCoordinator) Clock() *HLC { return c.inner.Clock() } + +func (c phaseDGateCoordinator) TimestampAllocator() TimestampAllocator { + return c.inner.TimestampAllocator() +} + +func (c phaseDGateCoordinator) VouchAppliedReadTimestamp(timestamp uint64, ref AppliedReadTimestampVoucherRef) error { + return c.inner.VouchAppliedReadTimestamp(timestamp, ref) +} + +func (c phaseDGateCoordinator) RevokeAppliedReadTimestamp(timestamp uint64, ref AppliedReadTimestampVoucherRef) { + c.inner.RevokeAppliedReadTimestamp(timestamp, ref) +} + func TestShardedCoordinatorDispatchTxn_SingleShardUsesOnePhase(t *testing.T) { t.Parallel() diff --git a/kv/tso.go b/kv/tso.go index 4538c2e59..16aa8b53a 100644 --- a/kv/tso.go +++ b/kv/tso.go @@ -72,6 +72,13 @@ type AppliedReadTimestampVoucher interface { VouchAppliedReadTimestamp(uint64, AppliedReadTimestampVoucherRef) error } +// AppliedReadTimestampVoucherRevoker removes an unused voucher registration. +// Decorators that forward VouchAppliedReadTimestamp should forward revocation +// too so pre-dispatch gates cannot leak inner-coordinator voucher entries. +type AppliedReadTimestampVoucherRevoker interface { + RevokeAppliedReadTimestamp(uint64, AppliedReadTimestampVoucherRef) +} + // AppliedReadTimestampVoucherRef is an opaque process-local dispatch // capability. Only this package can mint a non-zero ref. type AppliedReadTimestampVoucherRef struct { @@ -148,9 +155,11 @@ func DispatchWithReadTimestamp( if reqs == nil || reqs.StartTS != readTimestamp.timestamp { return nil, errors.WithStack(ErrTSOTimestampInvalid) } - if err := readTimestamp.voucher.prepare(coord, readTimestamp.timestamp); err != nil { + revoke, err := readTimestamp.voucher.prepare(coord, readTimestamp.timestamp) + if err != nil { return nil, err } + defer revoke() resp, err := coord.Dispatch(ctx, reqs) return resp, errors.WithStack(err) } @@ -163,18 +172,25 @@ func newAppliedReadDispatchVoucher() *appliedReadDispatchVoucher { return &appliedReadDispatchVoucher{ref: AppliedReadTimestampVoucherRef{id: id}} } -func (v *appliedReadDispatchVoucher) prepare(coord Coordinator, timestamp uint64) error { +func (v *appliedReadDispatchVoucher) prepare(coord Coordinator, timestamp uint64) (func(), error) { v.mu.Lock() defer v.mu.Unlock() voucher, ok := coord.(AppliedReadTimestampVoucher) if !ok { - return errors.WithStack(ErrTSOProtocolUnsupported) + return nil, errors.WithStack(ErrTSOProtocolUnsupported) } if err := voucher.VouchAppliedReadTimestamp(timestamp, v.ref); err != nil { - return errors.WithStack(err) + return nil, errors.WithStack(err) } v.prepared++ - return nil + revoke := func() {} + if revoker, ok := coord.(AppliedReadTimestampVoucherRevoker); ok { + ref := v.ref + revoke = func() { + revoker.RevokeAppliedReadTimestamp(timestamp, ref) + } + } + return revoke, nil } type tsoBatchAfterAllocator interface { diff --git a/kv/tso_raft.go b/kv/tso_raft.go index 3975f5733..dcfe236af 100644 --- a/kv/tso_raft.go +++ b/kv/tso_raft.go @@ -197,7 +197,12 @@ func (a *RaftTSOAllocator) prepareLeaderTermReservation( activateCutover bool, activatePhaseD bool, ) (uint64, error) { - termFloor, err := a.termCommitFloor(ctx, term) + // The ordinary per-term cache protects all reservations served by this TSO + // leader term. Phase-D activation is a one-way cutover boundary, so it must + // fence legacy data-group commits that may have landed after this term's + // first non-Phase-D reservation. + forceFreshFloor := activatePhaseD && !a.state.PhaseDActive() + termFloor, err := a.termCommitFloor(ctx, term, forceFreshFloor) if err != nil { return 0, err } @@ -232,8 +237,8 @@ func (a *RaftTSOAllocator) activateDurableMarkers( return a.commitPhaseD(ctx, engine, previousFloor) } -func (a *RaftTSOAllocator) termCommitFloor(ctx context.Context, term uint64) (uint64, error) { - if a.initializedTerm == term { +func (a *RaftTSOAllocator) termCommitFloor(ctx context.Context, term uint64, forceFresh bool) (uint64, error) { + if a.initializedTerm == term && !forceFresh { return a.state.AllocationFloor(), nil } floor, err := a.floorProvider.GlobalCommittedTimestampFloor(ctx) diff --git a/kv/tso_raft_test.go b/kv/tso_raft_test.go index 56215a38f..c33b9ee80 100644 --- a/kv/tso_raft_test.go +++ b/kv/tso_raft_test.go @@ -259,6 +259,38 @@ func TestRaftTSOAllocatorCommitsPhaseDBeforeWindowAndValidatesRange(t *testing.T require.ErrorIs(t, alloc.ValidateDurableTimestamp(context.Background(), end+1), ErrTSOTimestampInvalid) } +func TestRaftTSOAllocatorResamplesCommitFloorWhenActivatingPhaseDInInitializedTerm(t *testing.T) { + clock := NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) + fsm := NewTSOStateMachine(clock) + engine := &recordingTSOEngine{ + state: raftengine.StateLeader, + leader: raftengine.LeaderInfo{Address: "self"}, + term: 1, + apply: applyTSOTestFSM(fsm), + } + provider := &recordingTSOFloorProvider{} + alloc, err := NewRaftTSOAllocator( + &ShardGroup{Engine: engine, TSOState: fsm}, + clock, + WithTSOCutoverFloorProvider(provider), + ) + require.NoError(t, err) + + _, err = alloc.Next(context.Background()) + require.NoError(t, err) + require.Equal(t, 1, provider.callCount()) + + laterDataFloor := fsm.AllocationFloor() + 1_000 + provider.setFloor(laterDataFloor) + reservation, err := alloc.ReserveBatchAfter(context.Background(), testTSOBatchSize, 0, true, true) + require.NoError(t, err) + require.Equal(t, 2, provider.callCount(), "Phase-D activation must resample the data commit floor even within an initialized term") + require.Equal(t, laterDataFloor, reservation.PreviousAllocationFloor) + require.Equal(t, laterDataFloor, reservation.PhaseDFloor) + require.Greater(t, reservation.Base, laterDataFloor) +} + func TestRaftTSOAllocatorFencesAboveRestoredPhaseDFloor(t *testing.T) { clock := NewHLC() clock.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) @@ -754,3 +786,9 @@ func (p *recordingTSOFloorProvider) callCount() int { defer p.mu.Unlock() return p.calls } + +func (p *recordingTSOFloorProvider) setFloor(floor uint64) { + p.mu.Lock() + defer p.mu.Unlock() + p.floor = floor +} diff --git a/main.go b/main.go index 91fc75256..e2f968a61 100644 --- a/main.go +++ b/main.go @@ -2408,6 +2408,12 @@ func (c startupGatedCoordinator) VouchAppliedReadTimestamp(timestamp uint64, ref return errors.WithStack(voucher.VouchAppliedReadTimestamp(timestamp, ref)) } +func (c startupGatedCoordinator) RevokeAppliedReadTimestamp(timestamp uint64, ref kv.AppliedReadTimestampVoucherRef) { + if revoker, ok := c.inner.(kv.AppliedReadTimestampVoucherRevoker); ok { + revoker.RevokeAppliedReadTimestamp(timestamp, ref) + } +} + func (c startupGatedCoordinator) LeaseRead(ctx context.Context) (uint64, error) { return kv.LeaseReadThrough(c.inner, ctx) //nolint:wrapcheck // Pass through coordinator errors unchanged. } From 874a14f389eaba571502ad6e1023545e70590c86 Mon Sep 17 00:00:00 2001 From: bootjp Date: Thu, 23 Jul 2026 22:19:57 +0900 Subject: [PATCH 13/14] tso: preserve applied-read vouchers across adapters --- adapter/distribution_server.go | 9 ++-- adapter/distribution_server_test.go | 5 ++- adapter/dynamodb_transact.go | 19 +++++--- adapter/redis_lua_context.go | 41 +++++++++-------- adapter/s3.go | 3 +- adapter/s3_multipart_complete.go | 3 +- adapter/sqs_fifo.go | 6 ++- adapter/sqs_messages.go | 68 ++++++++++++++++------------- adapter/sqs_messages_batch.go | 2 +- 9 files changed, 91 insertions(+), 65 deletions(-) diff --git a/adapter/distribution_server.go b/adapter/distribution_server.go index 2e8ed35e6..de5356bcb 100644 --- a/adapter/distribution_server.go +++ b/adapter/distribution_server.go @@ -370,7 +370,7 @@ func (s *DistributionServer) SplitRange(ctx context.Context, req *pb.SplitRangeR } left, right := splitCatalogRoutes(parent, splitKey, leftID, rightID, 0) - saved, err := s.saveSplitResultViaCoordinator(ctx, snapshot.ReadTS, req.GetExpectedCatalogVersion(), parent.RouteID, left, right) + saved, err := s.saveSplitResultViaCoordinator(ctx, readTimestamp, req.GetExpectedCatalogVersion(), parent.RouteID, left, right) if err != nil { return nil, err } @@ -448,7 +448,7 @@ func (s *DistributionServer) verifyCatalogLeader(ctx context.Context) error { func (s *DistributionServer) saveSplitResultViaCoordinator( ctx context.Context, - readTS uint64, + readTimestamp kv.ReadTimestamp, expectedVersion uint64, parentID uint64, left distribution.RouteDescriptor, @@ -467,10 +467,11 @@ func (s *DistributionServer) saveSplitResultViaCoordinator( if err != nil { return distribution.CatalogSnapshot{}, grpcStatusErrorf(codes.Internal, "build split mutations: %v", err) } - resp, err := s.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + resp, err := kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, &kv.OperationGroup[kv.OP]{ Elems: ops, IsTxn: true, - StartTS: readTS, + StartTS: readTimestamp.Timestamp(), }) if err != nil { return distribution.CatalogSnapshot{}, grpcStatusErrorf(codes.Internal, "commit split mutations: %v", err) diff --git a/adapter/distribution_server_test.go b/adapter/distribution_server_test.go index 64d30d41d..5cf6df728 100644 --- a/adapter/distribution_server_test.go +++ b/adapter/distribution_server_test.go @@ -785,7 +785,7 @@ func TestDistributionServerSplitRange_PhaseDReadsAtValidatedAppliedWatermark(t * base: legacyFloor + 100, leader: true, phaseD: true, - phaseDFloor: legacyFloor - 1, + phaseDFloor: legacyFloor, } coordinator := newDistributionCoordinatorStub(baseStore, true) coordinator.allocator = allocator @@ -802,6 +802,7 @@ func TestDistributionServerSplitRange_PhaseDReadsAtValidatedAppliedWatermark(t * }) require.NoError(t, err) require.Equal(t, legacyFloor, coordinator.lastStartTS) + require.Equal(t, 1, coordinator.vouchCalls) } func TestDistributionServerSplitRange_UsesPersistentNextRouteID(t *testing.T) { @@ -1110,6 +1111,7 @@ type distributionCoordinatorStub struct { asyncApplyDone chan error asyncApplyDelay time.Duration dispatchCalls int + vouchCalls int } func (s *distributionCoordinatorStub) TimestampAllocator() kv.TimestampAllocator { @@ -1117,6 +1119,7 @@ func (s *distributionCoordinatorStub) TimestampAllocator() kv.TimestampAllocator } func (s *distributionCoordinatorStub) VouchAppliedReadTimestamp(uint64, kv.AppliedReadTimestampVoucherRef) error { + s.vouchCalls++ return nil } diff --git a/adapter/dynamodb_transact.go b/adapter/dynamodb_transact.go index 3d5cd82af..1d482be5c 100644 --- a/adapter/dynamodb_transact.go +++ b/adapter/dynamodb_transact.go @@ -695,14 +695,15 @@ func (d *DynamoDBServer) transactWriteItemsWithRetry(ctx context.Context, in tra func (d *DynamoDBServer) runTransactWriteAttempt( ctx context.Context, - reqs *kv.OperationGroup[kv.OP], + reqs *preparedTransactWriteItemsRequest, generations map[string]uint64, cleanupKeys [][]byte, ) (bool, error, error) { - if len(reqs.Elems) == 0 { + if len(reqs.group.Elems) == 0 { return true, nil, nil } - if _, err := d.coordinator.Dispatch(ctx, reqs); err != nil { + dispatchCtx := reqs.readTimestamp.WithDispatchVoucher(ctx) + if _, err := kv.DispatchWithReadTimestamp(dispatchCtx, d.coordinator, reqs.group); err != nil { wrapped := errors.WithStack(err) if !isRetryableTransactWriteError(err) { return false, nil, wrapped @@ -723,7 +724,12 @@ func (d *DynamoDBServer) runTransactWriteAttempt( return false, nil, nil } -func (d *DynamoDBServer) buildTransactWriteItemsRequest(ctx context.Context, in transactWriteItemsInput) (*kv.OperationGroup[kv.OP], map[string]uint64, [][]byte, error) { +type preparedTransactWriteItemsRequest struct { + readTimestamp kv.ReadTimestamp + group *kv.OperationGroup[kv.OP] +} + +func (d *DynamoDBServer) buildTransactWriteItemsRequest(ctx context.Context, in transactWriteItemsInput) (*preparedTransactWriteItemsRequest, map[string]uint64, [][]byte, error) { tableNames, err := collectTransactWriteTableNames(in) if err != nil { return nil, nil, nil, err @@ -756,7 +762,10 @@ func (d *DynamoDBServer) buildTransactWriteItemsRequest(ctx context.Context, in return nil, nil, nil, err } } - return reqs, tableGenerations, cleanup, nil + return &preparedTransactWriteItemsRequest{ + readTimestamp: readTimestamp, + group: reqs, + }, tableGenerations, cleanup, nil } // processTransactWriteItem validates and plans a single item within a diff --git a/adapter/redis_lua_context.go b/adapter/redis_lua_context.go index e906e9b10..b9beeb479 100644 --- a/adapter/redis_lua_context.go +++ b/adapter/redis_lua_context.go @@ -16,9 +16,10 @@ import ( ) type luaScriptContext struct { - server *RedisServer - startTS uint64 - readPin *kv.ActiveTimestampToken + server *RedisServer + startTS uint64 + readTimestamp kv.ReadTimestamp + readPin *kv.ActiveTimestampToken // ctx is the request-scoped context captured at newLuaScriptContext // time. New cmd* handlers should propagate it into store operations @@ -265,21 +266,22 @@ func newLuaScriptContext(ctx context.Context, server *RedisServer) (*luaScriptCo } startTS := readTimestamp.Timestamp() return &luaScriptContext{ - server: server, - startTS: startTS, - readPin: server.pinReadTS(startTS), - ctx: ctx, - touched: map[string]struct{}{}, - deleted: map[string]bool{}, - everDeleted: map[string]bool{}, - negativeType: map[string]bool{}, - strings: map[string]*luaStringState{}, - lists: map[string]*luaListState{}, - hashes: map[string]*luaHashState{}, - sets: map[string]*luaSetState{}, - zsets: map[string]*luaZSetState{}, - streams: map[string]*luaStreamState{}, - ttls: map[string]*luaTTLState{}, + server: server, + startTS: startTS, + readTimestamp: readTimestamp, + readPin: server.pinReadTS(startTS), + ctx: ctx, + touched: map[string]struct{}{}, + deleted: map[string]bool{}, + everDeleted: map[string]bool{}, + negativeType: map[string]bool{}, + strings: map[string]*luaStringState{}, + lists: map[string]*luaListState{}, + hashes: map[string]*luaHashState{}, + sets: map[string]*luaSetState{}, + zsets: map[string]*luaZSetState{}, + streams: map[string]*luaStreamState{}, + ttls: map[string]*luaTTLState{}, }, nil } @@ -3338,7 +3340,8 @@ func (c *luaScriptContext) commit() error { return nil } - if _, err := c.server.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := c.readTimestamp.WithDispatchVoucher(ctx) + if _, err := kv.DispatchWithReadTimestamp(dispatchCtx, c.server.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: luaCommitFloor(c.startTS), CommitTS: commitTS, diff --git a/adapter/s3.go b/adapter/s3.go index b70d31da0..cc1250261 100644 --- a/adapter/s3.go +++ b/adapter/s3.go @@ -1204,7 +1204,8 @@ func (s *S3Server) createMultipartUpload(w http.ResponseWriter, r *http.Request, writeS3InternalError(w, err) return } - if _, err := s.coordinator.Dispatch(r.Context(), &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(r.Context()) + if _, err := kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, CommitTS: commitTS, diff --git a/adapter/s3_multipart_complete.go b/adapter/s3_multipart_complete.go index acc6fe67a..3f9e08502 100644 --- a/adapter/s3_multipart_complete.go +++ b/adapter/s3_multipart_complete.go @@ -207,7 +207,8 @@ func (s *S3Server) commitS3MultipartCompletionAttempt(ctx context.Context, compl if err != nil { return nil, errors.WithStack(err) } - _, err = s.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + _, err = kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, CommitTS: commitTS, diff --git a/adapter/sqs_fifo.go b/adapter/sqs_fifo.go index b4445adeb..d41b40ef0 100644 --- a/adapter/sqs_fifo.go +++ b/adapter/sqs_fifo.go @@ -174,8 +174,9 @@ func (s *SQSServer) sendFifoMessage( in sqsSendMessageInput, dedupID string, delay int64, - readTS uint64, + readTimestamp kv.ReadTimestamp, ) (map[string]string, bool, error) { + readTS := readTimestamp.Timestamp() // HT-FIFO: hash the MessageGroupId once at the entry point so // every key built in this transaction (data, vis, byage, dedup, // group-lock, sequence) lands in the same partition. partitionFor @@ -248,7 +249,8 @@ func (s *SQSServer) sendFifoMessage( {Op: kv.Put, Key: seqKey, Value: []byte(strconv.FormatUint(nextSeq, 10))}, }, } - if _, err := s.coordinator.Dispatch(ctx, req); err != nil { + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + if _, err := kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, req); err != nil { if isRetryableTransactWriteError(err) { return nil, true, nil } diff --git a/adapter/sqs_messages.go b/adapter/sqs_messages.go index 89ad6dc40..ca18825ec 100644 --- a/adapter/sqs_messages.go +++ b/adapter/sqs_messages.go @@ -1622,18 +1622,20 @@ func (s *SQSServer) deleteMessageWithRetry(ctx context.Context, queueName string backoff := transactRetryInitialBackoff deadline := time.Now().Add(transactRetryMaxDuration) for range transactRetryMaxAttempts { - meta, rec, dataKey, readTS, outcome, err := s.loadMessageForDelete(ctx, queueName, handle) + meta, rec, dataKey, readTimestamp, outcome, err := s.loadMessageForDelete(ctx, queueName, handle) if err != nil { return err } if outcome == sqsDeleteNoOp { return nil } + readTS := readTimestamp.Timestamp() req, err := s.buildDeleteOps(ctx, queueName, meta, handle, rec, dataKey, readTS) if err != nil { return err } - if _, err := s.coordinator.Dispatch(ctx, req); err == nil { + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + if _, err := kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, req); err == nil { // Hot-partition observability (§11 PR 7): record the // successful delete on the partitioned commit branch // only. Legacy queues stay off the metric. @@ -1697,8 +1699,9 @@ const ( // outcome for AWS-compatible DeleteMessage semantics: structural errors // propagate; missing records and token mismatches on an otherwise-valid // queue return sqsDeleteNoOp; matching tokens return sqsDeleteProceed -// with the loaded record. The readTS it took the snapshot at is -// returned so the caller can pass it as StartTS on the OCC dispatch, +// with the loaded record. The ReadTimestamp it took the snapshot at is +// returned so the caller can pass it as StartTS on the OCC dispatch and +// bind any Phase-D applied-read voucher to that dispatch, // pinning the read-write conflict detection window. // // The caller-supplied QueueUrl is cross-checked against the handle's @@ -1707,41 +1710,41 @@ const ( // to a different (or recreated) queue and we reject it as a structural // error — silently succeeding would let misrouted deletes ack messages // that cannot possibly be deleted on this queue. -func (s *SQSServer) loadMessageForDelete(ctx context.Context, queueName string, handle *decodedReceiptHandle) (*sqsQueueMeta, *sqsMessageRecord, []byte, uint64, sqsDeleteOutcome, error) { +func (s *SQSServer) loadMessageForDelete(ctx context.Context, queueName string, handle *decodedReceiptHandle) (*sqsQueueMeta, *sqsMessageRecord, []byte, kv.ReadTimestamp, sqsDeleteOutcome, error) { readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs delete message: begin read timestamp") if err != nil { - return nil, nil, nil, 0, sqsDeleteProceed, errors.WithStack(err) + return nil, nil, nil, kv.ReadTimestamp{}, sqsDeleteProceed, errors.WithStack(err) } readTS := readTimestamp.Timestamp() meta, exists, err := s.loadQueueMetaAt(ctx, queueName, readTS) if err != nil { - return nil, nil, nil, readTS, sqsDeleteProceed, errors.WithStack(err) + return nil, nil, nil, readTimestamp, sqsDeleteProceed, errors.WithStack(err) } if !exists { - return nil, nil, nil, readTS, sqsDeleteProceed, newSQSAPIError(http.StatusBadRequest, sqsErrQueueDoesNotExist, "queue does not exist") + return nil, nil, nil, readTimestamp, sqsDeleteProceed, newSQSAPIError(http.StatusBadRequest, sqsErrQueueDoesNotExist, "queue does not exist") } if meta.Generation != handle.QueueGeneration { - return nil, nil, nil, readTS, sqsDeleteProceed, newSQSAPIError(http.StatusBadRequest, sqsErrReceiptHandleInvalid, "receipt handle does not belong to this queue") + return nil, nil, nil, readTimestamp, sqsDeleteProceed, newSQSAPIError(http.StatusBadRequest, sqsErrReceiptHandleInvalid, "receipt handle does not belong to this queue") } if err := validateReceiptHandleVersion(meta, handle); err != nil { - return nil, nil, nil, readTS, sqsDeleteProceed, newSQSAPIError(http.StatusBadRequest, sqsErrReceiptHandleInvalid, "receipt handle is not valid for this queue") + return nil, nil, nil, readTimestamp, sqsDeleteProceed, newSQSAPIError(http.StatusBadRequest, sqsErrReceiptHandleInvalid, "receipt handle is not valid for this queue") } dataKey := sqsMsgDataKeyDispatch(meta, queueName, handle.Partition, handle.QueueGeneration, handle.MessageIDHex) raw, err := s.store.GetAt(ctx, dataKey, readTS) if err != nil { if errors.Is(err, store.ErrKeyNotFound) { - return meta, nil, nil, readTS, sqsDeleteNoOp, nil + return meta, nil, nil, readTimestamp, sqsDeleteNoOp, nil } - return nil, nil, nil, readTS, sqsDeleteProceed, errors.WithStack(err) + return nil, nil, nil, readTimestamp, sqsDeleteProceed, errors.WithStack(err) } rec, err := decodeSQSMessageRecord(raw) if err != nil { - return nil, nil, nil, readTS, sqsDeleteProceed, errors.WithStack(err) + return nil, nil, nil, readTimestamp, sqsDeleteProceed, errors.WithStack(err) } if !bytes.Equal(rec.CurrentReceiptToken, handle.ReceiptToken) { - return meta, nil, nil, readTS, sqsDeleteNoOp, nil + return meta, nil, nil, readTimestamp, sqsDeleteNoOp, nil } - return meta, rec, dataKey, readTS, sqsDeleteProceed, nil + return meta, rec, dataKey, readTimestamp, sqsDeleteProceed, nil } func (s *SQSServer) changeMessageVisibility(w http.ResponseWriter, r *http.Request) { @@ -1795,10 +1798,11 @@ func (s *SQSServer) changeVisibilityWithRetry(ctx context.Context, queueName str backoff := transactRetryInitialBackoff deadline := time.Now().Add(transactRetryMaxDuration) for range transactRetryMaxAttempts { - meta, rec, dataKey, readTS, apiErr := s.loadAndVerifyMessage(ctx, queueName, handle) + meta, rec, dataKey, readTimestamp, apiErr := s.loadAndVerifyMessage(ctx, queueName, handle) if apiErr != nil { return apiErr } + readTS := readTimestamp.Timestamp() now := time.Now().UnixMilli() if rec.VisibleAtMillis <= now { return newSQSAPIError(http.StatusBadRequest, sqsErrMessageNotInflight, "message is not currently in flight") @@ -1826,7 +1830,8 @@ func (s *SQSServer) changeVisibilityWithRetry(ctx context.Context, queueName str {Op: kv.Put, Key: dataKey, Value: recordBytes}, }, } - if _, err := s.coordinator.Dispatch(ctx, req); err == nil { + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + if _, err := kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, req); err == nil { return nil } else if !isRetryableTransactWriteError(err) { return errors.WithStack(err) @@ -1855,9 +1860,10 @@ func (s *SQSServer) parseQueueAndReceipt(queueUrl, receiptHandle string) (string // loadAndVerifyMessage reads the data record for the given handle and // verifies that the receipt token matches the current one on record. -// Returns the record, its key, the snapshot timestamp the read ran at, -// or a typed SQS error. Callers use the snapshot as StartTS on the -// OCC dispatch so concurrent commits cannot slip past ReadKeys. +// Returns the record, its key, the ReadTimestamp the read ran at, or a +// typed SQS error. Callers use the snapshot as StartTS on the OCC +// dispatch and bind its Phase-D applied-read voucher so concurrent commits +// cannot slip past ReadKeys. // // The caller-supplied QueueUrl is cross-checked against the handle's // embedded queue_generation, mirroring loadMessageForDelete: an @@ -1865,41 +1871,41 @@ func (s *SQSServer) parseQueueAndReceipt(queueUrl, receiptHandle string) (string // cleans them up, so a handle from a deleted / recreated queue must // be rejected with ReceiptHandleIsInvalid instead of silently // mutating the orphan record. -func (s *SQSServer) loadAndVerifyMessage(ctx context.Context, queueName string, handle *decodedReceiptHandle) (*sqsQueueMeta, *sqsMessageRecord, []byte, uint64, error) { +func (s *SQSServer) loadAndVerifyMessage(ctx context.Context, queueName string, handle *decodedReceiptHandle) (*sqsQueueMeta, *sqsMessageRecord, []byte, kv.ReadTimestamp, error) { readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs change message visibility: begin read timestamp") if err != nil { - return nil, nil, nil, 0, errors.WithStack(err) + return nil, nil, nil, kv.ReadTimestamp{}, errors.WithStack(err) } readTS := readTimestamp.Timestamp() meta, exists, err := s.loadQueueMetaAt(ctx, queueName, readTS) if err != nil { - return nil, nil, nil, readTS, errors.WithStack(err) + return nil, nil, nil, readTimestamp, errors.WithStack(err) } if !exists { - return nil, nil, nil, readTS, newSQSAPIError(http.StatusBadRequest, sqsErrQueueDoesNotExist, "queue does not exist") + return nil, nil, nil, readTimestamp, newSQSAPIError(http.StatusBadRequest, sqsErrQueueDoesNotExist, "queue does not exist") } if meta.Generation != handle.QueueGeneration { - return nil, nil, nil, readTS, newSQSAPIError(http.StatusBadRequest, sqsErrReceiptHandleInvalid, "receipt handle does not belong to this queue") + return nil, nil, nil, readTimestamp, newSQSAPIError(http.StatusBadRequest, sqsErrReceiptHandleInvalid, "receipt handle does not belong to this queue") } if err := validateReceiptHandleVersion(meta, handle); err != nil { - return nil, nil, nil, readTS, newSQSAPIError(http.StatusBadRequest, sqsErrReceiptHandleInvalid, "receipt handle is not valid for this queue") + return nil, nil, nil, readTimestamp, newSQSAPIError(http.StatusBadRequest, sqsErrReceiptHandleInvalid, "receipt handle is not valid for this queue") } dataKey := sqsMsgDataKeyDispatch(meta, queueName, handle.Partition, handle.QueueGeneration, handle.MessageIDHex) raw, err := s.store.GetAt(ctx, dataKey, readTS) if err != nil { if errors.Is(err, store.ErrKeyNotFound) { - return nil, nil, nil, readTS, newSQSAPIError(http.StatusBadRequest, sqsErrReceiptHandleInvalid, "message not found") + return nil, nil, nil, readTimestamp, newSQSAPIError(http.StatusBadRequest, sqsErrReceiptHandleInvalid, "message not found") } - return nil, nil, nil, readTS, errors.WithStack(err) + return nil, nil, nil, readTimestamp, errors.WithStack(err) } rec, err := decodeSQSMessageRecord(raw) if err != nil { - return nil, nil, nil, readTS, errors.WithStack(err) + return nil, nil, nil, readTimestamp, errors.WithStack(err) } if !bytes.Equal(rec.CurrentReceiptToken, handle.ReceiptToken) { - return nil, nil, nil, readTS, newSQSAPIError(http.StatusBadRequest, sqsErrInvalidReceiptHandle, "receipt handle token does not match") + return nil, nil, nil, readTimestamp, newSQSAPIError(http.StatusBadRequest, sqsErrInvalidReceiptHandle, "receipt handle token does not match") } - return meta, rec, dataKey, readTS, nil + return meta, rec, dataKey, readTimestamp, nil } // ------------------------ small helpers ------------------------ diff --git a/adapter/sqs_messages_batch.go b/adapter/sqs_messages_batch.go index 1afbe51e6..4ca825682 100644 --- a/adapter/sqs_messages_batch.go +++ b/adapter/sqs_messages_batch.go @@ -380,7 +380,7 @@ func (s *SQSServer) runFifoSendWithRetry( if err != nil { return nil, err } - resp, retry, err := s.sendFifoMessage(ctx, queueName, meta, in, dedupID, delay, readTS) + resp, retry, err := s.sendFifoMessage(ctx, queueName, meta, in, dedupID, delay, readTimestamp) if err != nil { return nil, err } From 5587fd53c162ff9f168ba88e274155a66bdc6cca Mon Sep 17 00:00:00 2001 From: bootjp Date: Thu, 23 Jul 2026 23:25:38 +0900 Subject: [PATCH 14/14] adapter: preserve Phase D read vouchers --- adapter/redis_lua_context.go | 2 +- adapter/redis_lua_phase_d_test.go | 26 ++++++++++++ adapter/s3.go | 3 +- adapter/s3_admin.go | 3 +- adapter/s3_admin_objects.go | 3 +- adapter/s3_hlc_fence_test.go | 66 +++++++++++++++++++++++++++++++ 6 files changed, 99 insertions(+), 4 deletions(-) create mode 100644 adapter/redis_lua_phase_d_test.go diff --git a/adapter/redis_lua_context.go b/adapter/redis_lua_context.go index b9beeb479..8274cb199 100644 --- a/adapter/redis_lua_context.go +++ b/adapter/redis_lua_context.go @@ -260,7 +260,7 @@ func newLuaScriptContext(ctx context.Context, server *RedisServer) (*luaScriptCo if _, err := kv.LeaseReadThrough(server.coordinator, ctx); err != nil { return nil, errors.WithStack(err) } - readTimestamp, err := kv.BeginReadTimestampThrough(ctx, server.coordinator, server.readTS(), "redis lua: begin read timestamp") + readTimestamp, err := server.beginTxnReadTimestamp(ctx, "redis lua: begin read timestamp") if err != nil { return nil, errors.WithStack(err) } diff --git a/adapter/redis_lua_phase_d_test.go b/adapter/redis_lua_phase_d_test.go new file mode 100644 index 000000000..73d02bf81 --- /dev/null +++ b/adapter/redis_lua_phase_d_test.go @@ -0,0 +1,26 @@ +package adapter + +import ( + "context" + "testing" + + "github.com/bootjp/elastickv/store" + "github.com/stretchr/testify/require" +) + +func TestRedisLuaBeginReadTimestampPhaseDNormalizesEmptySnapshotSentinel(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + allocator := &distributionTSOAllocator{base: 100, phaseD: true, phaseDFloor: 10} + coord := newDistributionCoordinatorStub(st, true) + coord.allocator = allocator + server := NewRedisServer(nil, "", st, coord, nil, nil) + + scriptCtx, err := newLuaScriptContext(context.Background(), server) + require.NoError(t, err) + defer scriptCtx.Close() + + require.Equal(t, uint64(1), scriptCtx.startTS) + require.Zero(t, allocator.count, "an empty Redis Lua snapshot must not allocate ahead of Raft apply") +} diff --git a/adapter/s3.go b/adapter/s3.go index cc1250261..ef95cb577 100644 --- a/adapter/s3.go +++ b/adapter/s3.go @@ -646,7 +646,8 @@ func (s *S3Server) createBucket(w http.ResponseWriter, r *http.Request, bucket s {Op: kv.Put, Key: s3keys.BucketGenerationKey(bucket), Value: encodeS3Generation(nextGeneration)}, }, } - _, err = s.coordinator.Dispatch(r.Context(), req) + dispatchCtx := readTimestamp.WithDispatchVoucher(r.Context()) + _, err = kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, req) return errors.WithStack(err) }) if err != nil { diff --git a/adapter/s3_admin.go b/adapter/s3_admin.go index fe72e8b3a..227e74ee0 100644 --- a/adapter/s3_admin.go +++ b/adapter/s3_admin.go @@ -295,7 +295,8 @@ func (s *S3Server) adminCreateBucketTxn(ctx context.Context, principal AdminPrin if err != nil { return nil, errors.WithStack(err) } - _, err = s.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + _, err = kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, CommitTS: commitTS, diff --git a/adapter/s3_admin_objects.go b/adapter/s3_admin_objects.go index 118e0092a..4dbc8f648 100644 --- a/adapter/s3_admin_objects.go +++ b/adapter/s3_admin_objects.go @@ -347,7 +347,8 @@ func (s *S3Server) adminPutObjectStream(ctx context.Context, bucket, key string, s.cleanupManifestBlobs(ctx, bucket, meta.Generation, key, uploadedManifest()) return nil, 0, errors.WithStack(err) } - if _, err := s.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + if _, err := kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, CommitTS: commitTS, diff --git a/adapter/s3_hlc_fence_test.go b/adapter/s3_hlc_fence_test.go index b5b473d2c..ff7735a4a 100644 --- a/adapter/s3_hlc_fence_test.go +++ b/adapter/s3_hlc_fence_test.go @@ -2,6 +2,9 @@ package adapter import ( "context" + "net/http" + "net/http/httptest" + "strings" "testing" "time" @@ -84,6 +87,69 @@ func TestS3BeginTxnReadTimestampPhaseDNormalizesEmptySnapshotSentinel(t *testing require.Zero(t, allocator.count, "an empty S3 snapshot must not allocate ahead of Raft apply") } +func TestS3CreateBucketPhaseDBindsReadVoucher(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + allocator := &distributionTSOAllocator{base: 100, phaseD: true, phaseDFloor: 10} + coord := newDistributionCoordinatorStub(st, true) + coord.allocator = allocator + server := NewS3Server(nil, "", st, coord, nil) + + req := httptest.NewRequest(http.MethodPut, "/voucher-bucket", nil) + rec := httptest.NewRecorder() + server.createBucket(rec, req, "voucher-bucket") + + require.Equalf(t, http.StatusOK, rec.Code, "body=%s", rec.Body.String()) + require.Equal(t, 1, coord.vouchCalls, "create bucket must carry the applied-read voucher into dispatch") + require.Equal(t, uint64(1), coord.lastStartTS) +} + +func TestS3AdminCreateBucketPhaseDBindsReadVoucher(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + allocator := &distributionTSOAllocator{base: 100, phaseD: true, phaseDFloor: 10} + coord := newDistributionCoordinatorStub(st, true) + coord.allocator = allocator + server := NewS3Server(nil, "", st, coord, nil) + + _, err := server.AdminCreateBucket(context.Background(), fullAdminBucketsPrincipal(), "admin-voucher-bucket", s3AclPrivate) + + require.NoError(t, err) + require.Equal(t, 1, coord.vouchCalls, "admin create bucket must carry the applied-read voucher into dispatch") + require.Equal(t, uint64(1), coord.lastStartTS) +} + +func TestS3AdminPutObjectPhaseDBindsReadVoucher(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + const generation = uint64(1) + bucketMeta, err := encodeS3BucketMeta(&s3BucketMeta{ + BucketName: "admin-voucher-bucket", + Generation: generation, + CreatedAtHLC: 1, + Region: s3DefaultRegion, + Owner: "AKIA_FULL", + Acl: s3AclPrivate, + }) + require.NoError(t, err) + require.NoError(t, st.PutAt(ctx, s3keys.BucketMetaKey("admin-voucher-bucket"), bucketMeta, 1, 0)) + + allocator := &distributionTSOAllocator{base: 100, phaseD: true, phaseDFloor: 10} + coord := newDistributionCoordinatorStub(st, true) + coord.allocator = allocator + server := NewS3Server(nil, "", st, coord, nil) + + err = server.AdminPutObject(ctx, fullAdminBucketsPrincipal(), "admin-voucher-bucket", "key.txt", strings.NewReader(""), "text/plain") + + require.NoError(t, err) + require.Equal(t, 1, coord.vouchCalls, "admin put object must carry the applied-read voucher into final dispatch") + require.Equal(t, uint64(1), coord.lastStartTS) +} + // TestS3NextTxnCommitTSFailsClosedOnExpiredCeiling verifies that // nextTxnCommitTS surfaces ErrCeilingExpired through the // NextFenced() it calls after Observe(startTS). This is the