From 9182f378f75401617f38cfae9d80b92305c10c1b Mon Sep 17 00:00:00 2001 From: bootjp Date: Sun, 19 Jul 2026 21:20:45 +0900 Subject: [PATCH] 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 {