From 97a42ba72c18205fcb84a43ad4ca0cedc48efa31 Mon Sep 17 00:00:00 2001 From: bootjp Date: Sun, 19 Jul 2026 21:15:25 +0900 Subject: [PATCH 1/2] 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 2/2] 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 {