diff --git a/adapter/distribution_server.go b/adapter/distribution_server.go index 072294d78..a13fd28fd 100644 --- a/adapter/distribution_server.go +++ b/adapter/distribution_server.go @@ -21,15 +21,16 @@ 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 - watchInterval time.Duration - watchLeader func() bool - fsObserver DistributionFilesystemObserver - reloadRetry struct { + mu sync.Mutex + engine *distribution.Engine + catalog *distribution.CatalogStore + coordinator kv.Coordinator + timestampAllocator kv.TSOAllocator + readTracker *kv.ActiveTimestampTracker + watchInterval time.Duration + watchLeader func() bool + fsObserver DistributionFilesystemObserver + reloadRetry struct { attempts int interval time.Duration } @@ -53,6 +54,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 @@ -154,10 +164,181 @@ 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, activatePhaseD, err := timestampActivationValues(req) + if err != nil { + return nil, err + } + if s.timestampAllocator == nil { + return s.legacyTimestampResponse(count, minTimestamp, activateCutover, activatePhaseD) + } + + reservation, err := s.allocateTimestampReservation(ctx, count, minTimestamp, activateCutover, activatePhaseD) + 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, + PhaseDActive: reservation.PhaseDActive, + PhaseDFloor: reservation.PhaseDFloor, + }, nil +} + +func timestampActivationValues(req *pb.GetTimestampRequest) (bool, bool, error) { + activateCutover := req != nil && req.GetActivateCutover() + activatePhaseD := req != nil && req.GetActivatePhaseD() + if activatePhaseD && !activateCutover { + return false, false, errors.WithStack(status.Error(codes.InvalidArgument, + "phase D activation requires cutover activation")) + } + return activateCutover, activatePhaseD, nil +} + +func (s *DistributionServer) legacyTimestampResponse( + count int, + minTimestamp uint64, + activateCutover bool, + activatePhaseD bool, +) (*pb.GetTimestampResponse, error) { + if count != 1 || minTimestamp != 0 || activateCutover || activatePhaseD { + return nil, errors.WithStack(status.Error(codes.FailedPrecondition, "dedicated TSO allocator is not configured")) + } + if s.engine == nil { + return nil, errors.WithStack(status.Error(codes.Unavailable, "distribution engine is not configured")) + } + return &pb.GetTimestampResponse{Timestamp: s.engine.NextTimestamp()}, nil +} + +// ValidateTimestamp verifies a read/start timestamp against the durable M7 +// allocation range. The local allocator performs the group-0 leader fence; +// followers reject so clients re-resolve rather than validating stale state. +func (s *DistributionServer) ValidateTimestamp( + ctx context.Context, + req *pb.ValidateTimestampRequest, +) (*pb.ValidateTimestampResponse, error) { + if req == nil || req.GetTimestamp() == 0 { + return nil, errors.WithStack(status.Error(codes.InvalidArgument, "timestamp is required")) + } + validator, ok := s.timestampAllocator.(kv.DurableTimestampValidator) + if !ok { + return nil, errors.WithStack(status.Error(codes.FailedPrecondition, + "dedicated TSO timestamp validation is not configured")) + } + if err := validator.ValidateDurableTimestamp(ctx, req.GetTimestamp()); err != nil { + return nil, timestampRPCError(err) + } + resp := &pb.ValidateTimestampResponse{Valid: true, PhaseDActive: true} + if state, ok := s.timestampAllocator.(interface { + PhaseDFloor() uint64 + AllocationFloor() uint64 + }); ok { + resp.PhaseDFloor = state.PhaseDFloor() + resp.AllocationFloor = state.AllocationFloor() + } + return resp, nil +} + +func (s *DistributionServer) allocateTimestampReservation( + ctx context.Context, + count int, + minTimestamp uint64, + activateCutover bool, + activatePhaseD bool, +) (kv.TSOReservation, error) { + if allocator, ok := s.timestampAllocator.(kv.TSOReservationAllocator); ok { + reservation, err := allocator.ReserveBatchAfter(ctx, count, minTimestamp, activateCutover, activatePhaseD) + return reservation, errors.Wrap(err, "reserve dedicated TSO window") + } + reservation := kv.TSOReservation{Count: count} + switch { + case activateCutover || activatePhaseD: + 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())) + case errors.Is(err, kv.ErrTSOTimestampPrePhaseD): + return errors.WithStack(status.Error(codes.OutOfRange, err.Error())) + case errors.Is(err, kv.ErrTSOTimestampInvalid): + return errors.WithStack(status.Error(codes.InvalidArgument, err.Error())) + case errors.Is(err, kv.ErrTSOPhaseDInactive): + return errors.WithStack(status.Error(codes.FailedPrecondition, err.Error())) + default: + return errors.WithStack(status.Error(codes.Unavailable, err.Error())) + } } // ListRoutes returns all durable routes from catalog storage. @@ -342,7 +523,16 @@ func (s *DistributionServer) SplitRange(ctx context.Context, req *pb.SplitRangeR return nil, err } - snapshot, err := s.loadCatalogSnapshot(ctx) + readTimestamp, err := kv.BeginReadTimestampThrough( + ctx, + s.coordinator, + s.catalog.LatestCommitTS(), + "distribution split range: begin read timestamp", + ) + if err != nil { + return nil, grpcStatusErrorf(codes.Internal, "begin catalog split snapshot: %v", err) + } + snapshot, err := s.loadCatalogSnapshotAt(ctx, readTimestamp.Timestamp()) if err != nil { return nil, err } @@ -352,15 +542,8 @@ func (s *DistributionServer) SplitRange(ctx context.Context, req *pb.SplitRangeR return nil, err } - parent, found := findRouteByID(snapshot.Routes, req.GetRouteId()) - if !found { - return nil, grpcStatusError(codes.NotFound, errDistributionUnknownRoute.Error()) - } - - rawSplitKey := req.GetSplitKey() - splitKey := distribution.CloneBytes(fskeys.NormalizeSplitBoundary(kv.RouteKey(rawSplitKey))) - if err := validateSplitKey(parent, splitKey); err != nil { - s.observeFilePinnedHotspotIfNeeded(rawSplitKey, splitKey, err) + parent, splitKey, err := s.prepareSplitRange(snapshot.Routes, req) + if err != nil { return nil, err } @@ -370,7 +553,7 @@ func (s *DistributionServer) SplitRange(ctx context.Context, req *pb.SplitRangeR } left, right := splitCatalogRoutes(parent, splitKey, leftID, rightID, 0) - saved, err := s.saveSplitResultViaCoordinator(ctx, snapshot.ReadTS, req.GetExpectedCatalogVersion(), parent.RouteID, left, right) + saved, err := s.saveSplitResultViaCoordinator(ctx, readTimestamp, req.GetExpectedCatalogVersion(), parent.RouteID, left, right) if err != nil { return nil, err } @@ -389,6 +572,20 @@ func (s *DistributionServer) SplitRange(ctx context.Context, req *pb.SplitRangeR }, nil } +func (s *DistributionServer) prepareSplitRange(routes []distribution.RouteDescriptor, req *pb.SplitRangeRequest) (distribution.RouteDescriptor, []byte, error) { + parent, found := findRouteByID(routes, req.GetRouteId()) + if !found { + return distribution.RouteDescriptor{}, nil, grpcStatusError(codes.NotFound, errDistributionUnknownRoute.Error()) + } + rawSplitKey := req.GetSplitKey() + splitKey := distribution.CloneBytes(fskeys.NormalizeSplitBoundary(kv.RouteKey(rawSplitKey))) + if err := validateSplitKey(parent, splitKey); err != nil { + s.observeFilePinnedHotspotIfNeeded(rawSplitKey, splitKey, err) + return distribution.RouteDescriptor{}, nil, err + } + return parent, splitKey, nil +} + func (s *DistributionServer) pinReadTS(ts uint64) *kv.ActiveTimestampToken { if s == nil || s.readTracker == nil { return nil @@ -434,7 +631,7 @@ func (s *DistributionServer) verifyCatalogLeader(ctx context.Context) error { func (s *DistributionServer) saveSplitResultViaCoordinator( ctx context.Context, - readTS uint64, + readTimestamp kv.ReadTimestamp, expectedVersion uint64, parentID uint64, left distribution.RouteDescriptor, @@ -448,6 +645,7 @@ func (s *DistributionServer) saveSplitResultViaCoordinator( return distribution.CatalogSnapshot{}, grpcStatusError(codes.Internal, errDistributionRouteIDOverflow.Error()) } nextRouteID := right.RouteID + 1 + readTS := readTimestamp.Timestamp() commitTS, err := kv.NextTimestampAfterThrough(ctx, s.coordinator, readTS, "split range: allocate commitTS") if err != nil { return distribution.CatalogSnapshot{}, grpcStatusErrorf(codes.Internal, "allocate split commit timestamp: %v", err) @@ -477,7 +675,8 @@ func (s *DistributionServer) saveSplitResultViaCoordinator( return distribution.CatalogSnapshot{}, grpcStatusErrorf(codes.Internal, "convert catalog delta mutations: %v", err) } ops = append(ops, deltaOps...) - resp, err := s.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + resp, err := kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, &kv.OperationGroup[kv.OP]{ Elems: ops, IsTxn: true, StartTS: readTS, @@ -590,6 +789,17 @@ func (s *DistributionServer) loadCatalogSnapshot(ctx context.Context) (distribut return snapshot, nil } +func (s *DistributionServer) loadCatalogSnapshotAt(ctx context.Context, readTS uint64) (distribution.CatalogSnapshot, error) { + if s.catalog == nil { + return distribution.CatalogSnapshot{}, grpcStatusError(codes.FailedPrecondition, errDistributionCatalogNotConfigured.Error()) + } + snapshot, err := s.catalog.SnapshotAt(ctx, readTS) + if err != nil { + return distribution.CatalogSnapshot{}, grpcStatusErrorf(codes.Internal, "load route catalog: %v", err) + } + return snapshot, nil +} + func (s *DistributionServer) loadCatalogSnapshotAtVersion( ctx context.Context, readTS uint64, diff --git a/adapter/distribution_server_test.go b/adapter/distribution_server_test.go index ce268cddc..f39399928 100644 --- a/adapter/distribution_server_test.go +++ b/adapter/distribution_server_test.go @@ -3,15 +3,19 @@ package adapter import ( "context" "encoding/binary" + stderrors "errors" + "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 +77,244 @@ 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_CommitsPhaseDAndReturnsFloor(t *testing.T) { + t.Parallel() + + alloc := &distributionTSOAllocator{base: 501, previousFloor: 499, leader: true} + s := NewDistributionServer( + distribution.NewEngine(), + nil, + WithDistributionTimestampAllocator(alloc), + ) + resp, err := s.GetTimestamp(context.Background(), &pb.GetTimestampRequest{ + Count: 1, + MinTimestamp: 500, + ActivateCutover: true, + ActivatePhaseD: true, + }) + require.NoError(t, err) + require.True(t, alloc.activate) + require.True(t, alloc.activatePhaseD) + require.True(t, resp.GetCutoverActive()) + require.True(t, resp.GetPhaseDActive()) + require.Equal(t, uint64(499), resp.GetPhaseDFloor()) +} + +func TestDistributionServerGetTimestamp_RejectsPhaseDWithoutCutover(t *testing.T) { + t.Parallel() + + s := NewDistributionServer( + distribution.NewEngine(), + nil, + WithDistributionTimestampAllocator(&distributionTSOAllocator{leader: true}), + ) + _, err := s.GetTimestamp(context.Background(), &pb.GetTimestampRequest{ActivatePhaseD: true}) + require.Equal(t, codes.InvalidArgument, status.Code(err)) +} + +func TestDistributionServerValidateTimestamp(t *testing.T) { + t.Parallel() + + alloc := &distributionTSOAllocator{ + base: 501, + leader: true, + phaseD: true, + phaseDFloor: 499, + } + s := NewDistributionServer( + distribution.NewEngine(), + nil, + WithDistributionTimestampAllocator(alloc), + ) + resp, err := s.ValidateTimestamp(context.Background(), &pb.ValidateTimestampRequest{Timestamp: 500}) + require.NoError(t, err) + require.True(t, resp.GetValid()) + require.True(t, resp.GetPhaseDActive()) + require.Equal(t, uint64(499), resp.GetPhaseDFloor()) + require.Equal(t, uint64(501), resp.GetAllocationFloor()) + + _, err = s.ValidateTimestamp(context.Background(), &pb.ValidateTimestampRequest{Timestamp: 499}) + require.Equal(t, codes.OutOfRange, status.Code(err)) +} + +func TestDistributionServerValidateTimestamp_LeaderRoutedRPC(t *testing.T) { + serverAlloc := &distributionTSOAllocator{ + base: 701, + leader: true, + phaseD: true, + phaseDFloor: 699, + } + addr := serveDistributionTestServer(t, NewDistributionServer( + distribution.NewEngine(), + nil, + WithDistributionTimestampAllocator(serverAlloc), + )) + routed, err := kv.NewLeaderRoutedTSOAllocator( + &distributionTSOAllocator{leader: false}, + distributionLeaderView{addr: addr}, + ) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, routed.Close()) }) + + require.NoError(t, routed.ValidateDurableTimestamp(context.Background(), 700)) + err = routed.ValidateDurableTimestamp(context.Background(), 699) + require.ErrorIs(t, err, kv.ErrTSOTimestampInvalid) + require.ErrorIs(t, err, kv.ErrTSOTimestampPrePhaseD) +} + +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() @@ -524,6 +766,54 @@ func TestDistributionServerSplitRange_UsesCoordinatorForCatalogWrites(t *testing require.Equal(t, coordinator.lastCommitTS, changes.Deltas[0].Mutations[2].Route.SplitAtHLC) } +func TestDistributionServerSplitRange_PhaseDReadsAtValidatedAppliedWatermark(t *testing.T) { + t.Parallel() + + ctx := context.Background() + baseStore := store.NewMVCCStore() + catalog := distribution.NewCatalogStore(baseStore) + saved, err := catalog.Save(ctx, 0, []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte(""), + End: []byte("m"), + GroupID: 1, + State: distribution.RouteStateActive, + }, + { + RouteID: 2, + Start: []byte("m"), + End: nil, + GroupID: 2, + State: distribution.RouteStateActive, + }, + }) + require.NoError(t, err) + legacyFloor := catalog.LatestCommitTS() + allocator := &distributionTSOAllocator{ + base: legacyFloor + 100, + leader: true, + phaseD: true, + phaseDFloor: legacyFloor, + } + coordinator := newDistributionCoordinatorStub(baseStore, true) + coordinator.allocator = allocator + s := NewDistributionServer( + distribution.NewEngine(), + catalog, + WithDistributionCoordinator(coordinator), + ) + + _, err = s.SplitRange(ctx, &pb.SplitRangeRequest{ + ExpectedCatalogVersion: saved.Version, + RouteId: 1, + SplitKey: []byte("g"), + }) + require.NoError(t, err) + require.Equal(t, legacyFloor, coordinator.lastStartTS) + require.Equal(t, 1, coordinator.vouchCalls) +} + func TestDistributionServerSplitRange_UsesPersistentNextRouteID(t *testing.T) { t.Parallel() @@ -819,6 +1109,7 @@ func seededDistributionServerWithoutCoordinator(t *testing.T) (*DistributionServ type distributionCoordinatorStub struct { store store.MVCCStore + allocator kv.TimestampAllocator leader bool clock *kv.HLC nextTS uint64 @@ -829,6 +1120,16 @@ type distributionCoordinatorStub struct { asyncApplyDone chan error asyncApplyDelay time.Duration dispatchCalls int + vouchCalls int +} + +func (s *distributionCoordinatorStub) TimestampAllocator() kv.TimestampAllocator { + return s.allocator +} + +func (s *distributionCoordinatorStub) VouchAppliedReadTimestamp(uint64, kv.AppliedReadTimestampVoucherRef) error { + s.vouchCalls++ + return nil } func newDistributionCoordinatorStub(st store.MVCCStore, leader bool) *distributionCoordinatorStub { @@ -1004,6 +1305,137 @@ 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 + activatePhaseD bool + cutover bool + phaseD bool + phaseDFloor uint64 +} + +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, + activatePhaseD bool, +) (kv.TSOReservation, error) { + a.count = n + a.min = min + a.activate = activate + a.activatePhaseD = activatePhaseD + if a.err != nil { + return kv.TSOReservation{}, a.err + } + if activate { + a.cutover = true + } + if activatePhaseD { + a.phaseD = true + a.phaseDFloor = a.previousFloor + } + return kv.TSOReservation{ + Base: a.base, + Count: n, + PreviousAllocationFloor: a.previousFloor, + CutoverActive: a.cutover, + PhaseDActive: a.phaseD, + PhaseDFloor: a.phaseDFloor, + }, nil +} + +func (a *distributionTSOAllocator) ValidateDurableTimestamp(_ context.Context, timestamp uint64) error { + if a.err != nil { + return a.err + } + if !a.phaseD || timestamp == 0 || timestamp > a.base { + return kv.ErrTSOTimestampInvalid + } + if timestamp <= a.phaseDFloor { + return stderrors.Join(kv.ErrTSOTimestampInvalid, kv.ErrTSOTimestampPrePhaseD) + } + return nil +} + +func (a *distributionTSOAllocator) PhaseDFloor() uint64 { return a.phaseDFloor } + +func (a *distributionTSOAllocator) AllocationFloor() uint64 { return a.base } + +func (a *distributionTSOAllocator) PhaseDActive() bool { return a.phaseD } + +func (a *distributionTSOAllocator) PhaseDRequired() bool { return a.phaseD } + +func (a *distributionTSOAllocator) IsLeader() bool { return a.leader } + +func (a *distributionTSOAllocator) RunLeaseRenewal(ctx context.Context) { + <-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/dynamodb_item_write.go b/adapter/dynamodb_item_write.go index a5aecda19..c08861f9c 100644 --- a/adapter/dynamodb_item_write.go +++ b/adapter/dynamodb_item_write.go @@ -155,7 +155,11 @@ func (d *DynamoDBServer) retryItemWriteWithGenerationLegacy( backoff := transactRetryInitialBackoff deadline := time.Now().Add(transactRetryMaxDuration) for range transactRetryMaxAttempts { - readTS := d.nextTxnReadTS() + readTimestamp, err := d.beginTxnReadTimestamp(ctx, "dynamodb item-write legacy: begin read timestamp") + if err != nil { + return nil, errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() plan, err := prepare(readTS) if err != nil { return nil, err @@ -164,7 +168,7 @@ func (d *DynamoDBServer) retryItemWriteWithGenerationLegacy( return plan, nil } plan.req.StartTS = readTS - if err = d.commitItemWrite(ctx, plan.req); err != nil { + if err = d.commitItemWrite(ctx, readTimestamp, plan.req); err != nil { if !isRetryableTransactWriteError(err) { return nil, errors.WithStack(err) } @@ -204,6 +208,10 @@ type reusableItemWrite struct { // — the write set was built once from attempt 1's read — so plan is also the // correct value to return when the FSM dedup no-ops the apply (R1). plan *itemWritePlan + // readTimestamp carries the Phase-D dispatch capability that validates the + // reused StartTS. Retaining only the numeric StartTS would leave retries + // unable to re-vouch the same applied read watermark. + readTimestamp kv.ReadTimestamp // commitTS is the most recent dispatched commit_ts for this write set; the // next retry passes it as PrevCommitTS so the FSM probes exactly the attempt // that might have landed. @@ -266,7 +274,11 @@ func (d *DynamoDBServer) itemWriteFirstAttempt( tableName string, prepare func(readTS uint64) (*itemWritePlan, error), ) (*itemWritePlan, *reusableItemWrite, error) { - readTS := d.nextTxnReadTS() + readTimestamp, err := d.beginTxnReadTimestamp(ctx, "dynamodb item-write: begin read timestamp") + if err != nil { + return nil, nil, errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() plan, err := prepare(readTS) if err != nil { return nil, nil, err @@ -283,13 +295,14 @@ func (d *DynamoDBServer) itemWriteFirstAttempt( } plan.req.StartTS = readTS plan.req.CommitTS = commitTS - if dispErr := d.commitItemWrite(ctx, plan.req); dispErr != nil { + if dispErr := d.commitItemWrite(ctx, readTimestamp, plan.req); dispErr != nil { // dispErr is already wrapped by commitItemWrite; return it raw. if isRetryableTransactWriteError(dispErr) { return nil, &reusableItemWrite{ - plan: plan, - commitTS: commitTS, - probeKey: kv.PrimaryKeyForElems(plan.req.Elems), + plan: plan, + readTimestamp: readTimestamp, + commitTS: commitTS, + probeKey: kv.PrimaryKeyForElems(plan.req.Elems), }, dispErr } return nil, nil, dispErr @@ -311,7 +324,7 @@ func (d *DynamoDBServer) itemWriteReuseAttempt( } pending.plan.req.CommitTS = commitTS pending.plan.req.PrevCommitTS = pending.commitTS - dispErr := d.commitItemWrite(ctx, pending.plan.req) + dispErr := d.commitItemWrite(ctx, pending.readTimestamp, pending.plan.req) if dispErr == nil { return d.finishItemWriteAttempt(ctx, tableName, pending.plan) } @@ -433,8 +446,9 @@ func (d *DynamoDBServer) preparePutItemWrite(ctx context.Context, in putItemInpu }, nil } -func (d *DynamoDBServer) commitItemWrite(ctx context.Context, req *kv.OperationGroup[kv.OP]) error { - _, err := d.coordinator.Dispatch(ctx, req) +func (d *DynamoDBServer) commitItemWrite(ctx context.Context, readTimestamp kv.ReadTimestamp, req *kv.OperationGroup[kv.OP]) error { + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + _, err := kv.DispatchWithReadTimestamp(dispatchCtx, d.coordinator, req) if err != nil { return errors.WithStack(err) } diff --git a/adapter/dynamodb_locks.go b/adapter/dynamodb_locks.go index 8bbb3698f..b68db8000 100644 --- a/adapter/dynamodb_locks.go +++ b/adapter/dynamodb_locks.go @@ -115,6 +115,11 @@ func (d *DynamoDBServer) nextTxnReadTS() uint64 { return maxTS } +func (d *DynamoDBServer) beginTxnReadTimestamp(ctx context.Context, label string) (kv.ReadTimestamp, error) { + readTimestamp, err := kv.BeginReadTimestampThrough(ctx, d.coordinator, d.nextTxnReadTS(), label) + return readTimestamp, errors.WithStack(err) +} + func (d *DynamoDBServer) pinReadTS(ts uint64) *kv.ActiveTimestampToken { if d == nil || d.readTracker == nil { return &kv.ActiveTimestampToken{} diff --git a/adapter/dynamodb_migration.go b/adapter/dynamodb_migration.go index 6fcaff1d9..86b715de8 100644 --- a/adapter/dynamodb_migration.go +++ b/adapter/dynamodb_migration.go @@ -21,12 +21,11 @@ func (d *DynamoDBServer) ensureLegacyTableMigrationLocked(ctx context.Context, t backoff := transactRetryInitialBackoff deadline := time.Now().Add(transactRetryMaxDuration) for range transactRetryMaxAttempts { - readTS := d.nextTxnReadTS() - schema, exists, err := d.loadTableSchemaAt(ctx, tableName, readTS) + readTimestamp, schema, migrationRequired, err := d.legacyMigrationSnapshot(ctx, tableName) if err != nil { return errors.WithStack(err) } - if !exists || !schema.needsLegacyKeyMigration() { + if !migrationRequired { return nil } // Admin read-only callers (AdminScanTable) must not trigger @@ -39,7 +38,7 @@ func (d *DynamoDBServer) ensureLegacyTableMigrationLocked(ctx context.Context, t "table requires a one-time legacy-key migration before admin read endpoints are available; migrate via the SigV4 surface first") } if !schema.usesOrderedKeyEncoding() { - err = d.startLegacyTableKeyMigration(ctx, schema, readTS) + err = d.startLegacyTableKeyMigration(ctx, schema, readTimestamp) } else { err = d.migrateLegacyTableGeneration(ctx, schema) } @@ -57,14 +56,30 @@ func (d *DynamoDBServer) ensureLegacyTableMigrationLocked(ctx context.Context, t return newDynamoAPIError(http.StatusInternalServerError, dynamoErrInternal, "legacy table migration retry attempts exhausted") } +func (d *DynamoDBServer) legacyMigrationSnapshot( + ctx context.Context, + tableName string, +) (kv.ReadTimestamp, *dynamoTableSchema, bool, error) { + readTimestamp, err := d.beginTxnReadTimestamp(ctx, "dynamodb legacy-table migration: begin read timestamp") + if err != nil { + return kv.ReadTimestamp{}, nil, false, errors.WithStack(err) + } + schema, exists, err := d.loadTableSchemaAt(ctx, tableName, readTimestamp.Timestamp()) + if err != nil { + return kv.ReadTimestamp{}, nil, false, errors.WithStack(err) + } + return readTimestamp, schema, exists && schema.needsLegacyKeyMigration(), nil +} + func (d *DynamoDBServer) startLegacyTableKeyMigration( ctx context.Context, schema *dynamoTableSchema, - readTS uint64, + readTimestamp kv.ReadTimestamp, ) error { if schema == nil || schema.usesOrderedKeyEncoding() { return nil } + readTS := readTimestamp.Timestamp() nextGeneration, err := d.nextTableGenerationAt(ctx, schema.TableName, readTS) if err != nil { return err @@ -151,7 +166,11 @@ func (d *DynamoDBServer) migrateLegacyItem( backoff := transactRetryInitialBackoff deadline := time.Now().Add(transactRetryMaxDuration) for range transactRetryMaxAttempts { - readTS := d.nextTxnReadTS() + readTimestamp, err := d.beginTxnReadTimestamp(ctx, "dynamodb legacy-item migration: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() req, done, err := d.buildLegacyMigrationRequest(ctx, targetSchema, sourceSchema, targetKey, sourceKey, readTS) if err != nil { return err @@ -292,7 +311,11 @@ func (d *DynamoDBServer) finalizeLegacyTableMigration(ctx context.Context, schem if err != nil { return errors.WithStack(err) } - readTS := d.nextTxnReadTS() + readTimestamp, err := d.beginTxnReadTimestamp(ctx, "dynamodb finalize migration: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() req := &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: readTS, diff --git a/adapter/dynamodb_onephase_dedup_test.go b/adapter/dynamodb_onephase_dedup_test.go index ce7b9c7d4..1137c0d1a 100644 --- a/adapter/dynamodb_onephase_dedup_test.go +++ b/adapter/dynamodb_onephase_dedup_test.go @@ -133,6 +133,23 @@ func TestItemWriteDedup_PriorAttemptDidNotLand_Applies(t *testing.T) { require.Equal(t, 0, coord.probeNoOps, "nothing landed, so the probe must miss and the reuse applies") } +func TestItemWriteDedup_PhaseDVouchesReuse(t *testing.T) { + t.Parallel() + ctx := context.Background() + st := store.NewMVCCStore() + coord := newPhaseDDedupTestCoordinator(st, 1, false) + schema, server := newDedupItemWriteServer(st, coord, true) + seedDedupItem(t, st, schema, "1", "2") + + plan, err := server.updateItemWithRetry(ctx, appendListInput()) + require.NoError(t, err) + require.NotNil(t, plan) + + require.Equal(t, []string{"1", "2", "3"}, readListValues(t, server, schema)) + require.Equal(t, 2, coord.dispatches) + require.Equal(t, uint64(2), coord.vouches.Load(), "first attempt and reused write set must each reserve a Phase-D dispatch voucher") +} + // TestItemWriteDedup_SelfInflictedReuseConflict_ReturnsSuccess: attempt 1 // pre-rejects, the reuse then LANDS but surfaces WriteConflict (self-inflicted // conflict under churn). The adapter-side self-conflict guard probes the reuse's diff --git a/adapter/dynamodb_schema.go b/adapter/dynamodb_schema.go index 0ffa0f917..2100a4768 100644 --- a/adapter/dynamodb_schema.go +++ b/adapter/dynamodb_schema.go @@ -183,7 +183,11 @@ func (d *DynamoDBServer) createTableWithRetry(ctx context.Context, tableName str backoff := transactRetryInitialBackoff deadline := time.Now().Add(transactRetryMaxDuration) for range transactRetryMaxAttempts { - readTS := d.nextTxnReadTS() + readTimestamp, err := d.beginTxnReadTimestamp(ctx, "dynamodb create table: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() exists, err := d.tableExistsAt(ctx, tableName, readTS) if err != nil { return err @@ -199,6 +203,7 @@ func (d *DynamoDBServer) createTableWithRetry(ctx context.Context, tableName str if err != nil { return err } + req.StartTS = readTS if _, err := d.coordinator.Dispatch(ctx, req); err == nil { return nil } @@ -293,7 +298,11 @@ func (d *DynamoDBServer) deleteTableWithRetry(ctx context.Context, tableName str backoff := transactRetryInitialBackoff deadline := time.Now().Add(transactRetryMaxDuration) for range transactRetryMaxAttempts { - readTS := d.nextTxnReadTS() + readTimestamp, err := d.beginTxnReadTimestamp(ctx, "dynamodb delete table: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() schema, exists, err := d.loadTableSchemaAt(ctx, tableName, readTS) if err != nil { return errors.WithStack(err) @@ -304,7 +313,7 @@ func (d *DynamoDBServer) deleteTableWithRetry(ctx context.Context, tableName str req := &kv.OperationGroup[kv.OP]{ IsTxn: true, - StartTS: 0, + StartTS: readTS, Elems: []*kv.Elem[kv.OP]{ {Op: kv.Del, Key: dynamoTableMetaKey(tableName)}, }, diff --git a/adapter/dynamodb_transact.go b/adapter/dynamodb_transact.go index 3f58a688a..1d482be5c 100644 --- a/adapter/dynamodb_transact.go +++ b/adapter/dynamodb_transact.go @@ -695,14 +695,15 @@ func (d *DynamoDBServer) transactWriteItemsWithRetry(ctx context.Context, in tra func (d *DynamoDBServer) runTransactWriteAttempt( ctx context.Context, - reqs *kv.OperationGroup[kv.OP], + reqs *preparedTransactWriteItemsRequest, generations map[string]uint64, cleanupKeys [][]byte, ) (bool, error, error) { - if len(reqs.Elems) == 0 { + if len(reqs.group.Elems) == 0 { return true, nil, nil } - if _, err := d.coordinator.Dispatch(ctx, reqs); err != nil { + dispatchCtx := reqs.readTimestamp.WithDispatchVoucher(ctx) + if _, err := kv.DispatchWithReadTimestamp(dispatchCtx, d.coordinator, reqs.group); err != nil { wrapped := errors.WithStack(err) if !isRetryableTransactWriteError(err) { return false, nil, wrapped @@ -723,7 +724,12 @@ func (d *DynamoDBServer) runTransactWriteAttempt( return false, nil, nil } -func (d *DynamoDBServer) buildTransactWriteItemsRequest(ctx context.Context, in transactWriteItemsInput) (*kv.OperationGroup[kv.OP], map[string]uint64, [][]byte, error) { +type preparedTransactWriteItemsRequest struct { + readTimestamp kv.ReadTimestamp + group *kv.OperationGroup[kv.OP] +} + +func (d *DynamoDBServer) buildTransactWriteItemsRequest(ctx context.Context, in transactWriteItemsInput) (*preparedTransactWriteItemsRequest, map[string]uint64, [][]byte, error) { tableNames, err := collectTransactWriteTableNames(in) if err != nil { return nil, nil, nil, err @@ -733,7 +739,11 @@ func (d *DynamoDBServer) buildTransactWriteItemsRequest(ctx context.Context, in return nil, nil, nil, err } } - readTS := d.nextTxnReadTS() + readTimestamp, err := d.beginTxnReadTimestamp(ctx, "dynamodb transact-write: begin read timestamp") + if err != nil { + return nil, nil, nil, errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() reqs := &kv.OperationGroup[kv.OP]{ IsTxn: true, // Keep transaction start aligned with the snapshot used to evaluate @@ -752,7 +762,10 @@ func (d *DynamoDBServer) buildTransactWriteItemsRequest(ctx context.Context, in return nil, nil, nil, err } } - return reqs, tableGenerations, cleanup, nil + return &preparedTransactWriteItemsRequest{ + readTimestamp: readTimestamp, + group: reqs, + }, tableGenerations, cleanup, nil } // processTransactWriteItem validates and plans a single item within a diff --git a/adapter/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/adapter/redis_collection_ttl.go b/adapter/redis_collection_ttl.go index 7aadfa945..b492a3385 100644 --- a/adapter/redis_collection_ttl.go +++ b/adapter/redis_collection_ttl.go @@ -148,6 +148,23 @@ func (r *RedisServer) dispatchCollectionExpire( return true, r.dispatchElems(ctx, true, readTS, elems) } +func (r *RedisServer) dispatchCollectionExpireReadTimestamp( + ctx context.Context, + key []byte, + readTimestamp kv.ReadTimestamp, + typ redisValueType, + expireAt time.Time, +) (bool, error) { + readTS := readTimestamp.Timestamp() + ttlMs := redisExpireAtMillis(expireAt) + elems, ok, err := r.collectionExpireElems(ctx, key, readTS, typ, ttlMs) + if err != nil || !ok { + return ok, err + } + elems = append(elems, &kv.Elem[kv.OP]{Op: kv.Put, Key: redisTTLKey(key), Value: encodeRedisTTL(expireAt)}) + return true, r.dispatchReadTimestampElems(ctx, readTimestamp, elems) +} + func (r *RedisServer) collectionExpireElems( ctx context.Context, key []byte, diff --git a/adapter/redis_compat_helpers.go b/adapter/redis_compat_helpers.go index 016b423c6..94b4d86d5 100644 --- a/adapter/redis_compat_helpers.go +++ b/adapter/redis_compat_helpers.go @@ -778,6 +778,23 @@ func (r *RedisServer) dispatchElems(ctx context.Context, isTxn bool, startTS uin return errors.WithStack(err) } +func (r *RedisServer) dispatchReadTimestampElems(ctx context.Context, readTimestamp kv.ReadTimestamp, elems []*kv.Elem[kv.OP]) error { + if len(elems) == 0 { + return nil + } + startTS := readTimestamp.Timestamp() + if startTS == ^uint64(0) { + startTS = 0 + } + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + _, err := kv.DispatchWithReadTimestamp(dispatchCtx, r.coordinator, &kv.OperationGroup[kv.OP]{ + IsTxn: true, + StartTS: startTS, + Elems: elems, + }) + return errors.WithStack(err) +} + // readRedisStringAt reads a Redis string value, trying the prefixed key first // and falling back to the bare key for legacy data written before the // !redis|str| prefix migration. Returns the decoded user value and the @@ -1278,14 +1295,15 @@ func (r *RedisServer) listValuesAt(ctx context.Context, key []byte, readTS uint6 return r.fetchListRange(ctx, key, meta, 0, meta.Len-1, readTS) } -func (r *RedisServer) rewriteListTxn(ctx context.Context, key []byte, readTS uint64, values []string) error { +func (r *RedisServer) rewriteListTxn(ctx context.Context, key []byte, readTimestamp kv.ReadTimestamp, values []string) error { + readTS := readTimestamp.Timestamp() elems, _, err := r.deleteLogicalKeyElems(ctx, key, readTS) if err != nil { return err } if len(values) == 0 { - return r.dispatchElems(ctx, true, readTS, elems) + return r.dispatchReadTimestampElems(ctx, readTimestamp, elems) } rawValues := make([][]byte, 0, len(values)) @@ -1302,7 +1320,8 @@ func (r *RedisServer) rewriteListTxn(ctx context.Context, key []byte, readTS uin return err } elems = append(elems, ops...) - _, err = r.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + _, err = kv.DispatchWithReadTimestamp(dispatchCtx, r.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, CommitTS: commitTS, diff --git a/adapter/redis_delta_compactor.go b/adapter/redis_delta_compactor.go index b90ed2fa0..44282ff94 100644 --- a/adapter/redis_delta_compactor.go +++ b/adapter/redis_delta_compactor.go @@ -249,7 +249,15 @@ func (c *DeltaCompactor) compactUrgentKey(ctx context.Context, req urgentCompact func (c *DeltaCompactor) compactUrgentKeyBatch(ctx context.Context, req urgentCompactionRequest, h *collectionDeltaHandler, prefix, end []byte) (int, bool) { // Use a fresh readTS each iteration so we observe the committed state from // the previous compaction pass and do not re-scan already-deleted delta keys. - readTS := snapshotTS(c.coord.Clock(), c.st) + readTimestamp, err := kv.BeginReadTimestampThrough(ctx, c.coord, snapshotTS(c.coord.Clock(), c.st), + "redis delta compactor urgent: begin read timestamp") + if err != nil { + c.logger.WarnContext(ctx, "delta compactor urgent: timestamp allocation failed", + "type", req.typeName, "key", string(req.userKey), "error", err) + return 0, true + } + ctx = readTimestamp.WithDispatchVoucher(ctx) + readTS := readTimestamp.Timestamp() // Scan one extra beyond MaxDeltaScanLimit to detect whether more remain. kvs, err := c.st.ScanAt(ctx, prefix, end, store.MaxDeltaScanLimit+1, readTS) @@ -305,9 +313,15 @@ func (c *DeltaCompactor) SyncOnce(ctx context.Context) error { "backoff_until", until) return nil } - readTS := snapshotTS(c.coord.Clock(), c.st) tickCtx, cancel := context.WithTimeout(ctx, c.timeout) defer cancel() + readTimestamp, err := kv.BeginReadTimestampThrough(tickCtx, c.coord, snapshotTS(c.coord.Clock(), c.st), + "redis delta compactor: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } + tickCtx = readTimestamp.WithDispatchVoucher(tickCtx) + readTS := readTimestamp.Timestamp() combined := c.compactBackgroundHandlers(tickCtx, readTS) if tickCtx.Err() == nil { @@ -604,7 +618,7 @@ func (c *DeltaCompactor) dispatchCompaction(ctx context.Context, readTS uint64, if err != nil { return errors.WithStack(err) } - _, err = c.coord.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + _, err = kv.DispatchWithReadTimestamp(ctx, c.coord, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: normalizeStartTS(readTS), CommitTS: commitTS, diff --git a/adapter/redis_expire_cmds.go b/adapter/redis_expire_cmds.go index a40c8adc0..4e094ec66 100644 --- a/adapter/redis_expire_cmds.go +++ b/adapter/redis_expire_cmds.go @@ -47,29 +47,24 @@ func (r *RedisServer) getdel(conn redcon.Conn, cmd redcon.Command) { defer cancel() var v []byte err := r.retryRedisWrite(ctx, func() error { - readTS := r.readTS() - typ, err := r.keyTypeAt(ctx, key, readTS) + readTimestamp, err := r.beginTxnReadTimestamp(ctx, "redis getdel: begin read timestamp") + if err != nil { + return cockerrors.WithStack(err) + } + readTS := readTimestamp.Timestamp() + raw, exists, err := r.getdelValueAt(ctx, key, readTS) if err != nil { return err } - if typ == redisTypeNone { + if !exists { v = nil return nil } - if typ != redisTypeString { - return wrongTypeError() - } - raw, _, err := r.readRedisStringAt(key, readTS) - if err != nil { - // Key may have expired or been deleted between type check and read. - v = nil - return nil //nolint:nilerr // treat not-found/expired as nil value - } elems, _, err := r.deleteLogicalKeyElems(ctx, key, readTS) if err != nil { return err } - if err := r.dispatchElems(ctx, true, readTS, elems); err != nil { + if err := r.dispatchReadTimestampElems(ctx, readTimestamp, elems); err != nil { return err } v = raw @@ -86,6 +81,25 @@ func (r *RedisServer) getdel(conn redcon.Conn, cmd redcon.Command) { conn.WriteBulk(v) } +func (r *RedisServer) getdelValueAt(ctx context.Context, key []byte, readTS uint64) ([]byte, bool, error) { + typ, err := r.keyTypeAt(ctx, key, readTS) + if err != nil { + return nil, false, err + } + if typ == redisTypeNone { + return nil, false, nil + } + if typ != redisTypeString { + return nil, false, wrongTypeError() + } + raw, _, err := r.readRedisStringAt(key, readTS) + if err != nil { + // Key may have expired or been deleted between type check and read. + return nil, false, nil //nolint:nilerr // treat not-found/expired as nil value + } + return raw, true, nil +} + // SETNX key value — set if not exists, returns 1 on success, 0 on failure func (r *RedisServer) setnx(conn redcon.Conn, cmd redcon.Command) { if r.proxyToLeader(conn, cmd, cmd.Args[1]) { @@ -185,25 +199,29 @@ func parseExpireTTL(raw []byte) (int64, error) { return ttl, nil } -func (r *RedisServer) prepareExpire(key []byte, nxOnly bool) (uint64, bool, error) { - readTS := r.readTS() - exists, err := r.logicalExistsAt(context.Background(), key, readTS) +func (r *RedisServer) prepareExpireReadTimestamp(ctx context.Context, key []byte, nxOnly bool) (kv.ReadTimestamp, bool, error) { + readTimestamp, err := r.beginTxnReadTimestamp(ctx, "redis expire: begin read timestamp") + if err != nil { + return kv.ReadTimestamp{}, false, cockerrors.WithStack(err) + } + readTS := readTimestamp.Timestamp() + exists, err := r.logicalExistsAt(ctx, key, readTS) if err != nil { - return 0, false, err + return kv.ReadTimestamp{}, false, err } if !exists { - return readTS, false, nil + return readTimestamp, false, nil } if !nxOnly { - return readTS, true, nil + return readTimestamp, true, nil } - currentTTL, err := r.ttlAt(context.Background(), key, readTS) + currentTTL, err := r.ttlAt(ctx, key, readTS) if err != nil { - return 0, false, err + return kv.ReadTimestamp{}, false, err } - return readTS, !hasActiveTTL(currentTTL, time.Now()), nil + return readTimestamp, !hasActiveTTL(currentTTL, time.Now()), nil } func (r *RedisServer) setExpire(conn redcon.Conn, cmd redcon.Command, unit time.Duration) { @@ -253,21 +271,22 @@ func (r *RedisServer) setExpire(conn redcon.Conn, cmd redcon.Command, unit time. // then re-invokes doSetExpire with a fresh readTS, providing OCC safety without // an explicit mutex. Leadership is verified by coordinator.Dispatch itself. func (r *RedisServer) doSetExpire(ctx context.Context, key []byte, ttl int64, expireAt time.Time, nxOnly bool) (int, error) { - readTS, eligible, err := r.prepareExpire(key, nxOnly) + readTimestamp, eligible, err := r.prepareExpireReadTimestamp(ctx, key, nxOnly) if err != nil { return 0, err } if !eligible { return 0, nil } + readTS := readTimestamp.Timestamp() if ttl <= 0 { - return r.expireDeleteKey(ctx, key, readTS) + return r.expireDeleteKeyReadTimestamp(ctx, key, readTimestamp) } typ, err := r.rawKeyTypeAt(ctx, key, readTS) if err != nil { return 0, err } - applied, err := r.dispatchExpireForType(ctx, key, readTS, typ, expireAt) + applied, err := r.dispatchExpireForReadTimestamp(ctx, key, readTimestamp, typ, expireAt) if err != nil || !applied { return 0, err } @@ -301,6 +320,31 @@ func (r *RedisServer) dispatchExpireForType( return true, r.dispatchElems(ctx, true, readTS, elems) } +func (r *RedisServer) dispatchExpireForReadTimestamp( + ctx context.Context, + key []byte, + readTimestamp kv.ReadTimestamp, + typ redisValueType, + expireAt time.Time, +) (bool, error) { + readTS := readTimestamp.Timestamp() + if typ == redisTypeString { + plain, err := r.isPlainRedisString(ctx, key, readTS) + if err != nil { + return false, err + } + if plain { + return r.dispatchStringExpireReadTimestamp(ctx, key, readTimestamp, expireAt) + } + return r.dispatchHLLExpireReadTimestamp(ctx, key, readTimestamp, expireAt) + } + if isNonStringCollectionType(typ) { + return r.dispatchCollectionExpireReadTimestamp(ctx, key, readTimestamp, typ, expireAt) + } + elems := []*kv.Elem[kv.OP]{{Op: kv.Put, Key: redisTTLKey(key), Value: encodeRedisTTL(expireAt)}} + return true, r.dispatchReadTimestampElems(ctx, readTimestamp, elems) +} + // isPlainRedisString distinguishes a plain Redis string (stored under // !redis|str| or, for legacy data, the bare key) from a HyperLogLog // (stored under !redis|hll|), both of which rawKeyTypeAt reports as @@ -321,12 +365,13 @@ func (r *RedisServer) isPlainRedisString(ctx context.Context, key []byte, readTS return legacy, nil } -func (r *RedisServer) expireDeleteKey(ctx context.Context, key []byte, readTS uint64) (int, error) { +func (r *RedisServer) expireDeleteKeyReadTimestamp(ctx context.Context, key []byte, readTimestamp kv.ReadTimestamp) (int, error) { + readTS := readTimestamp.Timestamp() elems, existed, err := r.deleteLogicalKeyElems(ctx, key, readTS) if err != nil { return 0, err } - if err := r.dispatchElems(ctx, true, readTS, elems); err != nil { + if err := r.dispatchReadTimestampElems(ctx, readTimestamp, elems); err != nil { return 0, err } if existed { @@ -359,6 +404,23 @@ func (r *RedisServer) dispatchStringExpire(ctx context.Context, key []byte, read return true, r.dispatchElems(ctx, true, readTS, elems) } +func (r *RedisServer) dispatchStringExpireReadTimestamp(ctx context.Context, key []byte, readTimestamp kv.ReadTimestamp, expireAt time.Time) (bool, error) { + readTS := readTimestamp.Timestamp() + userValue, _, readErr := r.readRedisStringAt(key, readTS) + if readErr != nil { + if cockerrors.Is(readErr, store.ErrKeyNotFound) { + return false, nil + } + return false, cockerrors.WithStack(readErr) + } + encoded := encodeRedisStr(userValue, &expireAt) + elems := []*kv.Elem[kv.OP]{ + {Op: kv.Put, Key: redisStrKey(key), Value: encoded}, + {Op: kv.Put, Key: redisTTLKey(key), Value: encodeRedisTTL(expireAt)}, + } + return true, r.dispatchReadTimestampElems(ctx, readTimestamp, elems) +} + func (r *RedisServer) dispatchHLLExpire(ctx context.Context, key []byte, readTS uint64, expireAt time.Time) (bool, error) { raw, err := r.store.GetAt(ctx, redisHLLKey(key), readTS) if err != nil { @@ -381,3 +443,27 @@ func (r *RedisServer) dispatchHLLExpire(ctx context.Context, key []byte, readTS } return true, r.dispatchElems(ctx, true, readTS, elems) } + +func (r *RedisServer) dispatchHLLExpireReadTimestamp(ctx context.Context, key []byte, readTimestamp kv.ReadTimestamp, expireAt time.Time) (bool, error) { + readTS := readTimestamp.Timestamp() + raw, err := r.store.GetAt(ctx, redisHLLKey(key), readTS) + if err != nil { + if cockerrors.Is(err, store.ErrKeyNotFound) { + return false, nil + } + return false, cockerrors.WithStack(err) + } + value, _, _, err := decodeRedisHLL(raw) + if err != nil { + return false, err + } + encoded, err := encodeRedisHLL(value, &expireAt) + if err != nil { + return false, err + } + elems := []*kv.Elem[kv.OP]{ + {Op: kv.Put, Key: redisHLLKey(key), Value: encoded}, + {Op: kv.Put, Key: redisTTLKey(key), Value: encodeRedisTTL(expireAt)}, + } + return true, r.dispatchReadTimestampElems(ctx, readTimestamp, elems) +} diff --git a/adapter/redis_hash_cmds.go b/adapter/redis_hash_cmds.go index 9943d6ef5..df7d8eb9a 100644 --- a/adapter/redis_hash_cmds.go +++ b/adapter/redis_hash_cmds.go @@ -162,7 +162,11 @@ func (r *RedisServer) applyHashFieldPairs(key []byte, args [][]byte) (int, error defer cancel() var added int err := r.retryRedisWrite(ctx, func() error { - readTS := r.readTS() + readTimestamp, err := r.beginTxnReadTimestamp(ctx, "redis hash field write: begin read timestamp") + if err != nil { + return cockerrors.WithStack(err) + } + readTS := readTimestamp.Timestamp() typ, err := r.keyTypeOrEmptyAt(ctx, key, readTS, redisTypeHash) if err != nil { return err @@ -209,7 +213,8 @@ func (r *RedisServer) applyHashFieldPairs(key []byte, args [][]byte) (int, error return nil } - _, dispatchErr := r.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + _, dispatchErr := kv.DispatchWithReadTimestamp(dispatchCtx, r.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, CommitTS: commitTS, @@ -408,7 +413,8 @@ func (r *RedisServer) hdel(conn redcon.Conn, cmd redcon.Command) { } // hdelWideColumn deletes the given fields from the wide-column hash and emits a negative delta. -func (r *RedisServer) hdelWideColumn(ctx context.Context, key []byte, fields [][]byte, readTS uint64) (int, error) { +func (r *RedisServer) hdelWideColumn(ctx context.Context, key []byte, fields [][]byte, readTimestamp kv.ReadTimestamp) (int, error) { + readTS := readTimestamp.Timestamp() delElems, removed, err := r.resolveHashFieldDelElems(ctx, key, fields, readTS) if err != nil { return 0, err @@ -428,7 +434,8 @@ func (r *RedisServer) hdelWideColumn(ctx context.Context, key []byte, fields [][ Key: store.HashMetaDeltaKey(key, commitTS, 0), Value: deltaVal, }) - _, dispatchErr := r.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + _, dispatchErr := kv.DispatchWithReadTimestamp(dispatchCtx, r.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, CommitTS: commitTS, @@ -472,7 +479,11 @@ func (r *RedisServer) resolveHashFieldDelElems(ctx context.Context, key []byte, } func (r *RedisServer) hdelTxn(ctx context.Context, key []byte, fields [][]byte) (int, error) { - readTS := r.readTS() + readTimestamp, err := r.beginTxnReadTimestamp(ctx, "redis hash field delete: begin read timestamp") + if err != nil { + return 0, cockerrors.WithStack(err) + } + readTS := readTimestamp.Timestamp() typ, err := r.keyTypeAtExpect(ctx, key, readTS, redisTypeHash) if err != nil { return 0, err @@ -492,7 +503,7 @@ func (r *RedisServer) hdelTxn(ctx context.Context, key []byte, fields [][]byte) return 0, cockerrors.WithStack(err) } if len(wideKVs) > 0 { - return r.hdelWideColumn(ctx, key, fields, readTS) + return r.hdelWideColumn(ctx, key, fields, readTimestamp) } // Legacy blob path. @@ -504,7 +515,7 @@ func (r *RedisServer) hdelTxn(ctx context.Context, key []byte, fields [][]byte) if removed == 0 { return 0, nil } - return removed, r.persistHashTxn(ctx, key, readTS, value) + return removed, r.persistHashReadTimestampTxn(ctx, key, readTimestamp, value) } func removeHashFields(value redisHashValue, fields [][]byte) int { @@ -553,6 +564,39 @@ func (r *RedisServer) persistHashTxn(ctx context.Context, key []byte, readTS uin return r.dispatchElems(ctx, true, readTS, elems) } +func (r *RedisServer) persistHashReadTimestampTxn(ctx context.Context, key []byte, readTimestamp kv.ReadTimestamp, value redisHashValue) error { + readTS := readTimestamp.Timestamp() + if len(value) == 0 { + elems, _, err := r.deleteLogicalKeyElems(ctx, key, readTS) + if err != nil { + return err + } + return r.dispatchReadTimestampElems(ctx, readTimestamp, elems) + } + ttlMs, expired, err := legacyTTLMillisForMigrationAt(ctx, r.store, key, readTS) + if err != nil { + return err + } + if expired { + return r.dispatchReadTimestampElems(ctx, readTimestamp, legacyExpiredCollectionCleanupElems(key, redisHashKey(key))) + } + elems := make([]*kv.Elem[kv.OP], 0, len(value)+1) + for field, val := range value { + elems = append(elems, &kv.Elem[kv.OP]{ + Op: kv.Put, + Key: store.HashFieldKey(key, []byte(field)), + Value: []byte(val), + }) + } + elems = append(elems, &kv.Elem[kv.OP]{ + Op: kv.Put, + Key: store.HashMetaKey(key), + Value: store.MarshalHashMeta(store.HashMeta{Len: int64(len(value)), ExpireAt: ttlMs}), + }) + elems = append(elems, &kv.Elem[kv.OP]{Op: kv.Del, Key: redisHashKey(key)}) + return r.dispatchReadTimestampElems(ctx, readTimestamp, elems) +} + func (r *RedisServer) hexists(conn redcon.Conn, cmd redcon.Command) { if r.proxyToLeader(conn, cmd, cmd.Args[1]) { return @@ -729,7 +773,8 @@ func (r *RedisServer) readHashFieldInt(ctx context.Context, key, field []byte, r // hincrbyWithMigration handles the HINCRBY case where a legacy JSON blob must be migrated // atomically with the increment operation. -func (r *RedisServer) hincrbyWithMigration(ctx context.Context, key, fieldKey []byte, readTS, commitTS uint64, current int64, isNewField bool, typ redisValueType, increment int64, cleanupElems []*kv.Elem[kv.OP]) (int64, error) { +func (r *RedisServer) hincrbyWithMigration(ctx context.Context, key, fieldKey []byte, readTimestamp kv.ReadTimestamp, commitTS uint64, current int64, isNewField bool, typ redisValueType, increment int64, cleanupElems []*kv.Elem[kv.OP]) (int64, error) { + readTS := readTimestamp.Timestamp() migrationElems, migErr := r.buildHashLegacyMigrationElems(ctx, key, readTS) if migErr != nil { return 0, migErr @@ -751,7 +796,8 @@ func (r *RedisServer) hincrbyWithMigration(ctx context.Context, key, fieldKey [] }, ) } - _, dispatchErr := r.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + _, dispatchErr := kv.DispatchWithReadTimestamp(dispatchCtx, r.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: normalizeStartTS(readTS), CommitTS: commitTS, @@ -762,7 +808,11 @@ func (r *RedisServer) hincrbyWithMigration(ctx context.Context, key, fieldKey [] } func (r *RedisServer) hincrbyTxn(ctx context.Context, key, field []byte, increment int64) (int64, error) { - readTS := r.readTS() + readTimestamp, err := r.beginTxnReadTimestamp(ctx, "redis hash increment: begin read timestamp") + if err != nil { + return 0, cockerrors.WithStack(err) + } + readTS := readTimestamp.Timestamp() typ, err := r.keyTypeOrEmptyAt(ctx, key, readTS, redisTypeHash) if err != nil { return 0, err @@ -787,7 +837,7 @@ func (r *RedisServer) hincrbyTxn(ctx context.Context, key, field []byte, increme // If a legacy blob exists, migrate it atomically with the increment. if len(legacyValue) > 0 { - return r.hincrbyWithMigration(ctx, key, fieldKey, readTS, commitTS, current, isNewField, typ, increment, cleanupElems) + return r.hincrbyWithMigration(ctx, key, fieldKey, readTimestamp, commitTS, current, isNewField, typ, increment, cleanupElems) } current += increment @@ -806,7 +856,8 @@ func (r *RedisServer) hincrbyTxn(ctx context.Context, key, field []byte, increme }, ) } - _, dispatchErr := r.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + _, dispatchErr := kv.DispatchWithReadTimestamp(dispatchCtx, r.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, CommitTS: commitTS, @@ -831,7 +882,11 @@ func (r *RedisServer) incrLegacy(conn redcon.Conn, cmd redcon.Command) { defer cancel() var current int64 if err := r.retryRedisWrite(ctx, func() error { - readTS := r.readTS() + readTimestamp, err := r.beginTxnReadTimestamp(ctx, "redis incr: begin read timestamp") + if err != nil { + return cockerrors.WithStack(err) + } + readTS := readTimestamp.Timestamp() typ, err := r.keyTypeAt(ctx, cmd.Args[1], readTS) if err != nil { return err @@ -867,7 +922,7 @@ func (r *RedisServer) incrLegacy(conn redcon.Conn, cmd redcon.Command) { // cannot later expire a now-persistent key. elems = append(elems, &kv.Elem[kv.OP]{Op: kv.Del, Key: redisTTLKey(cmd.Args[1])}) } - return r.dispatchElems(ctx, true, readTS, elems) + return r.dispatchReadTimestampElems(ctx, readTimestamp, elems) }); err != nil { writeRedisError(conn, err) return diff --git a/adapter/redis_list_dedup_test.go b/adapter/redis_list_dedup_test.go index 428bbe110..6a3d8b4a7 100644 --- a/adapter/redis_list_dedup_test.go +++ b/adapter/redis_list_dedup_test.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "errors" + "sync/atomic" "testing" "github.com/bootjp/elastickv/kv" @@ -54,6 +55,18 @@ type dedupTestCoordinator struct { beforeDispatch func(n int) } +type phaseDDedupTestCoordinator struct { + *dedupTestCoordinator + alloc *phaseDDedupAllocator + vouches atomic.Uint64 +} + +type phaseDDedupAllocator struct { + next atomic.Uint64 + validateCalls atomic.Uint64 + phaseDFloor uint64 +} + func newDedupTestCoordinator(st store.MVCCStore, ambiguousDispatch int, lands bool) *dedupTestCoordinator { return &dedupTestCoordinator{ occAdapterCoordinator: newOCCAdapterCoordinator(st), @@ -62,6 +75,53 @@ func newDedupTestCoordinator(st store.MVCCStore, ambiguousDispatch int, lands bo } } +func newPhaseDDedupTestCoordinator(st store.MVCCStore, ambiguousDispatch int, lands bool) *phaseDDedupTestCoordinator { + return &phaseDDedupTestCoordinator{ + dedupTestCoordinator: newDedupTestCoordinator(st, ambiguousDispatch, lands), + alloc: &phaseDDedupAllocator{phaseDFloor: 10}, + } +} + +func (c *phaseDDedupTestCoordinator) TimestampAllocator() kv.TimestampAllocator { + return c.alloc +} + +func (c *phaseDDedupTestCoordinator) VouchAppliedReadTimestamp(uint64, kv.AppliedReadTimestampVoucherRef) error { + c.vouches.Add(1) + return nil +} + +func (a *phaseDDedupAllocator) Next(ctx context.Context) (uint64, error) { + return a.NextAfter(ctx, 0) +} + +func (a *phaseDDedupAllocator) NextAfter(_ context.Context, min uint64) (uint64, error) { + for { + cur := a.next.Load() + next := cur + 1 + if next <= min { + next = min + 1 + } + if a.next.CompareAndSwap(cur, next) { + return next, nil + } + } +} + +func (a *phaseDDedupAllocator) ValidateDurableTimestamp(_ context.Context, timestamp uint64) error { + a.validateCalls.Add(1) + if timestamp == 0 { + return kv.ErrTSOTimestampInvalid + } + if timestamp <= a.phaseDFloor { + return errors.Join(kv.ErrTSOTimestampInvalid, kv.ErrTSOTimestampPrePhaseD) + } + return nil +} + +func (a *phaseDDedupAllocator) PhaseDActive() bool { return true } +func (a *phaseDDedupAllocator) PhaseDRequired() bool { return true } + func (c *dedupTestCoordinator) Dispatch(ctx context.Context, req *kv.OperationGroup[kv.OP]) (*kv.CoordinateResponse, error) { c.dispatches++ n := c.dispatches diff --git a/adapter/redis_lists.go b/adapter/redis_lists.go index 1f7a6fcf2..d692bbf2e 100644 --- a/adapter/redis_lists.go +++ b/adapter/redis_lists.go @@ -99,7 +99,11 @@ func (r *RedisServer) listPushCore(ctx context.Context, key []byte, values [][]b var newLen int64 err := r.retryRedisWrite(ctx, func() error { - readTS := r.readTS() + readTimestamp, err := r.beginTxnReadTimestamp(ctx, "redis list push: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() meta, metaExists, typ, cleanupElems, err := r.listPushSnapshot(ctx, key, readTS) if err != nil { return err @@ -122,7 +126,8 @@ func (r *RedisServer) listPushCore(ctx context.Context, key []byte, values [][]b } // Dispatch with the pre-allocated commitTS. - _, dispErr := r.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + _, dispErr := kv.DispatchWithReadTimestamp(dispatchCtx, r.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, CommitTS: commitTS, @@ -189,7 +194,8 @@ type reusableListPush struct { // boundary-touching commit fires WriteConflict and the adapter drops // pending → recomputes. Empty when attempt 1 read an empty list (no // boundary to fence; the OCC on the write key suffices for that case). - readKeys [][]byte + readKeys [][]byte + readTimestamp kv.ReadTimestamp } // dispatchListPushReuse runs one iteration of the option-2 reuse path: @@ -212,7 +218,8 @@ func (r *RedisServer) dispatchListPushReuse(ctx context.Context, key []byte, pen if allocErr != nil { return 0, false, errors.WithStack(allocErr) } - _, dispErr := r.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := pending.readTimestamp.WithDispatchVoucher(ctx) + _, dispErr := kv.DispatchWithReadTimestamp(dispatchCtx, r.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: pending.startTS, CommitTS: commitTS, @@ -380,19 +387,19 @@ func (r *RedisServer) listPushCoreWithDedup(ctx context.Context, key []byte, val var newLen int64 var pending *reusableListPush err := r.retryRedisWrite(ctx, func() error { - if pending != nil { - length, drop, dispErr := r.dispatchListPushReuse(ctx, key, pending) - if drop { - pending = nil - } - if dispErr != nil { - return dispErr + length, handled, err := r.tryPendingListPush(ctx, key, &pending) + if handled { + if err != nil { + return err } newLen = length return nil } - - readTS := r.readTS() + readTimestamp, err := r.beginTxnReadTimestamp(ctx, "redis list push: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() meta, metaExists, typ, cleanupElems, err := r.listPushSnapshot(ctx, key, readTS) if err != nil { return err @@ -416,7 +423,8 @@ func (r *RedisServer) listPushCoreWithDedup(ctx context.Context, key []byte, val } boundaryReads := listPushReadKeys(key, meta, typ, metaExists) - _, dispErr := r.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + _, dispErr := kv.DispatchWithReadTimestamp(dispatchCtx, r.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, CommitTS: commitTS, @@ -442,11 +450,12 @@ func (r *RedisServer) listPushCoreWithDedup(ctx context.Context, key []byte, val // idempotency cache) and out of scope for this design. if isRetryableRedisTxnErr(dispErr) { pending = &reusableListPush{ - ops: ops, - startTS: startTS, - commitTS: commitTS, - length: updatedMeta.Len, - readKeys: boundaryReads, + ops: ops, + startTS: startTS, + commitTS: commitTS, + length: updatedMeta.Len, + readKeys: boundaryReads, + readTimestamp: readTimestamp, } } return errors.WithStack(dispErr) @@ -454,6 +463,21 @@ func (r *RedisServer) listPushCoreWithDedup(ctx context.Context, key []byte, val return newLen, err } +func (r *RedisServer) tryPendingListPush( + ctx context.Context, + key []byte, + pending **reusableListPush, +) (int64, bool, error) { + if *pending == nil { + return 0, false, nil + } + length, drop, err := r.dispatchListPushReuse(ctx, key, *pending) + if drop { + *pending = nil + } + return length, true, err +} + func (r *RedisServer) listRPush(ctx context.Context, key []byte, values [][]byte) (int64, error) { return r.listPushCore(ctx, key, values, r.buildRPushOps) } @@ -593,7 +617,8 @@ func (r *RedisServer) checkListKeyType(ctx context.Context, key []byte, readTS u // listPopClaimOnce executes one attempt of a pop-with-claim transaction. // Returns (nil, nil) for a missing key or an empty list, and the popped // values otherwise. -func (r *RedisServer) listPopClaimOnce(ctx context.Context, key []byte, count int, left bool, readTS uint64) ([]string, error) { +func (r *RedisServer) listPopClaimOnce(ctx context.Context, key []byte, count int, left bool, readTimestamp kv.ReadTimestamp) ([]string, error) { + readTS := readTimestamp.Timestamp() found, typeErr := r.checkListKeyType(ctx, key, readTS) if typeErr != nil || !found { return nil, typeErr @@ -618,7 +643,7 @@ func (r *RedisServer) listPopClaimOnce(ctx context.Context, key []byte, count in return nil, buildErr } - if err := r.commitListPop(ctx, key, elems, n, left, readTS); err != nil { + if err := r.commitListPop(ctx, key, elems, n, left, readTimestamp); err != nil { return nil, err } return values, nil @@ -628,7 +653,8 @@ func (r *RedisServer) listPopClaimOnce(ctx context.Context, key []byte, count in // and dispatches the pop transaction. Extracted from listPopClaimOnce // so that function stays under the cyclop ceiling after the HLC-4 // (iii) NextFenced fence added a new error branch (PR #867 Phase 2b). -func (r *RedisServer) commitListPop(ctx context.Context, key []byte, elems []*kv.Elem[kv.OP], n int64, left bool, readTS uint64) error { +func (r *RedisServer) commitListPop(ctx context.Context, key []byte, elems []*kv.Elem[kv.OP], n int64, left bool, readTimestamp kv.ReadTimestamp) error { + readTS := readTimestamp.Timestamp() // n is the number of sequence positions claimed (including any holes). // HeadDelta and LenDelta must use n, not len(values), so that Head // advances past holes and the metadata stays consistent with Tail. @@ -651,7 +677,8 @@ func (r *RedisServer) commitListPop(ctx context.Context, key []byte, elems []*kv }, ) - if _, dispErr := r.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + if _, dispErr := kv.DispatchWithReadTimestamp(dispatchCtx, r.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, CommitTS: commitTS, @@ -675,7 +702,11 @@ func (r *RedisServer) listPopClaim(ctx context.Context, key []byte, count int, l var popped []string err := r.retryRedisWrite(ctx, func() error { - result, popErr := r.listPopClaimOnce(ctx, key, count, left, r.readTS()) + readTimestamp, readErr := r.beginTxnReadTimestamp(ctx, "redis list pop: begin read timestamp") + if readErr != nil { + return errors.WithStack(readErr) + } + result, popErr := r.listPopClaimOnce(ctx, key, count, left, readTimestamp) if popErr != nil { return popErr } @@ -824,9 +855,14 @@ func (r *RedisServer) ltrim(conn redcon.Conn, cmd redcon.Command) { } ctx, cancel := context.WithTimeout(context.Background(), redisDispatchTimeout) defer cancel() + key := cmd.Args[1] if err := r.retryRedisWrite(ctx, func() error { - readTS := r.readTS() - typ, err := r.keyTypeAt(ctx, cmd.Args[1], readTS) + readTimestamp, err := r.beginTxnReadTimestamp(ctx, "redis ltrim: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() + typ, err := r.keyTypeAt(ctx, key, readTS) if err != nil { return err } @@ -836,16 +872,11 @@ func (r *RedisServer) ltrim(conn redcon.Conn, cmd redcon.Command) { if typ != redisTypeList { return wrongTypeError() } - current, err := r.listValuesAt(ctx, cmd.Args[1], readTS) + current, err := r.listValuesAt(ctx, key, readTS) if err != nil { return err } - s, e := normalizeRankRange(start, stop, len(current)) - trimmed := []string{} - if e >= s { - trimmed = append(trimmed, current[s:e+1]...) - } - return r.rewriteListTxn(ctx, cmd.Args[1], readTS, trimmed) + return r.rewriteListTxn(ctx, key, readTimestamp, trimListValues(current, start, stop)) }); err != nil { writeRedisError(conn, err) return @@ -853,6 +884,14 @@ func (r *RedisServer) ltrim(conn redcon.Conn, cmd redcon.Command) { conn.WriteString("OK") } +func trimListValues(current []string, start, stop int) []string { + s, e := normalizeRankRange(start, stop, len(current)) + if e < s { + return []string{} + } + return append([]string{}, current[s:e+1]...) +} + func (r *RedisServer) lindex(conn redcon.Conn, cmd redcon.Command) { if r.proxyToLeader(conn, cmd, cmd.Args[1]) { return diff --git a/adapter/redis_lua_context.go b/adapter/redis_lua_context.go index 439378d84..b39afed97 100644 --- a/adapter/redis_lua_context.go +++ b/adapter/redis_lua_context.go @@ -16,9 +16,10 @@ import ( ) type luaScriptContext struct { - server *RedisServer - startTS uint64 - readPin *kv.ActiveTimestampToken + server *RedisServer + startTS uint64 + readTimestamp kv.ReadTimestamp + readPin *kv.ActiveTimestampToken // ctx is the request-scoped context captured at newLuaScriptContext // time. New cmd* handlers should propagate it into store operations @@ -273,24 +274,29 @@ func newLuaScriptContext(ctx context.Context, server *RedisServer) (*luaScriptCo if _, err := kv.LeaseReadThrough(server.coordinator, ctx); err != nil { return nil, errors.WithStack(err) } - startTS := server.readTS() + readTimestamp, err := server.beginTxnReadTimestamp(ctx, "redis lua: begin read timestamp") + if err != nil { + return nil, errors.WithStack(err) + } + startTS := readTimestamp.Timestamp() return &luaScriptContext{ - server: server, - startTS: startTS, - readPin: server.pinReadTS(startTS), - ctx: ctx, - touched: map[string]struct{}{}, - readKeys: map[string][]byte{}, - deleted: map[string]bool{}, - everDeleted: map[string]bool{}, - negativeType: map[string]bool{}, - strings: map[string]*luaStringState{}, - lists: map[string]*luaListState{}, - hashes: map[string]*luaHashState{}, - sets: map[string]*luaSetState{}, - zsets: map[string]*luaZSetState{}, - streams: map[string]*luaStreamState{}, - ttls: map[string]*luaTTLState{}, + server: server, + startTS: startTS, + readTimestamp: readTimestamp, + readPin: server.pinReadTS(startTS), + ctx: ctx, + touched: map[string]struct{}{}, + readKeys: map[string][]byte{}, + deleted: map[string]bool{}, + everDeleted: map[string]bool{}, + negativeType: map[string]bool{}, + strings: map[string]*luaStringState{}, + lists: map[string]*luaListState{}, + hashes: map[string]*luaHashState{}, + sets: map[string]*luaSetState{}, + zsets: map[string]*luaZSetState{}, + streams: map[string]*luaStreamState{}, + ttls: map[string]*luaTTLState{}, }, nil } @@ -3568,7 +3574,8 @@ func (c *luaScriptContext) commit() error { return nil } - if _, err := c.server.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := c.readTimestamp.WithDispatchVoucher(ctx) + if _, err := kv.DispatchWithReadTimestamp(dispatchCtx, c.server.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: luaCommitFloor(c.startTS), CommitTS: commitTS, diff --git a/adapter/redis_lua_phase_d_test.go b/adapter/redis_lua_phase_d_test.go new file mode 100644 index 000000000..73d02bf81 --- /dev/null +++ b/adapter/redis_lua_phase_d_test.go @@ -0,0 +1,26 @@ +package adapter + +import ( + "context" + "testing" + + "github.com/bootjp/elastickv/store" + "github.com/stretchr/testify/require" +) + +func TestRedisLuaBeginReadTimestampPhaseDNormalizesEmptySnapshotSentinel(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + allocator := &distributionTSOAllocator{base: 100, phaseD: true, phaseDFloor: 10} + coord := newDistributionCoordinatorStub(st, true) + coord.allocator = allocator + server := NewRedisServer(nil, "", st, coord, nil, nil) + + scriptCtx, err := newLuaScriptContext(context.Background(), server) + require.NoError(t, err) + defer scriptCtx.Close() + + require.Equal(t, uint64(1), scriptCtx.startTS) + require.Zero(t, allocator.count, "an empty Redis Lua snapshot must not allocate ahead of Raft apply") +} diff --git a/adapter/redis_retry_test.go b/adapter/redis_retry_test.go index 7dff88b7e..539df8797 100644 --- a/adapter/redis_retry_test.go +++ b/adapter/redis_retry_test.go @@ -352,6 +352,29 @@ func TestRedisXAddDedupsLandedWireWriteConflict(t *testing.T) { require.Equal(t, int64(1), meta.Length, "the generated XADD entry must not be appended twice") } +func TestRedisXAddDedupPhaseDVouchesReuse(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + coord := newPhaseDDedupTestCoordinator(st, 1, false) + srv := &RedisServer{ + store: st, + coordinator: coord, + scriptCache: map[string]string{}, + onePhaseTxnDedup: true, + } + conn := &recordingConn{} + + srv.xadd(conn, redcon.Command{Args: [][]byte{ + []byte(cmdXAdd), []byte("retry:stream"), []byte("*"), []byte("field"), []byte("value"), + }}) + + require.Empty(t, conn.err) + require.NotEmpty(t, conn.bulk) + require.Equal(t, 2, coord.dispatches) + require.Equal(t, uint64(2), coord.vouches.Load(), "first attempt and reused write set must each reserve a Phase-D dispatch voucher") +} + func TestRedisXAddDedupDisabledDoesNotReplayLandedWireConflict(t *testing.T) { t.Parallel() diff --git a/adapter/redis_set_cmds.go b/adapter/redis_set_cmds.go index d29aba6e1..c5beb2487 100644 --- a/adapter/redis_set_cmds.go +++ b/adapter/redis_set_cmds.go @@ -199,9 +199,10 @@ func sortedExactSetMembers(existing map[string]struct{}) []string { return out } -func (r *RedisServer) persistExactSetMembersTxn(ctx context.Context, kind string, key []byte, readTS uint64, members map[string]struct{}) error { +func (r *RedisServer) persistExactSetMembersTxn(ctx context.Context, kind string, key []byte, readTimestamp kv.ReadTimestamp, members map[string]struct{}) error { + readTS := readTimestamp.Timestamp() if kind != setKind { - return r.persistHLLMembersTxn(ctx, key, readTS, members) + return r.persistHLLMembersTxn(ctx, key, readTimestamp, members) } // Wide-column set: full rewrite (used when the whole state is available). if len(members) == 0 { @@ -209,7 +210,7 @@ func (r *RedisServer) persistExactSetMembersTxn(ctx context.Context, kind string if err != nil { return err } - return r.dispatchElems(ctx, true, readTS, elems) + return r.dispatchReadTimestampElems(ctx, readTimestamp, elems) } elems := make([]*kv.Elem[kv.OP], 0, len(members)+setWideColOverhead) for member := range members { @@ -226,10 +227,11 @@ func (r *RedisServer) persistExactSetMembersTxn(ctx context.Context, kind string }) // Remove legacy blob if present. elems = append(elems, &kv.Elem[kv.OP]{Op: kv.Del, Key: redisSetKey(key)}) - return r.dispatchElems(ctx, true, readTS, elems) + return r.dispatchReadTimestampElems(ctx, readTimestamp, elems) } -func (r *RedisServer) persistHLLMembersTxn(ctx context.Context, key []byte, readTS uint64, members map[string]struct{}) error { +func (r *RedisServer) persistHLLMembersTxn(ctx context.Context, key []byte, readTimestamp kv.ReadTimestamp, members map[string]struct{}) error { + readTS := readTimestamp.Timestamp() // HLL keeps a single anchor payload, now wrapped in an inline-TTL envelope // while preserving legacy payload reads during migration. if len(members) == 0 { @@ -237,7 +239,7 @@ func (r *RedisServer) persistHLLMembersTxn(ctx context.Context, key []byte, read if err != nil { return err } - return r.dispatchElems(ctx, true, readTS, elems) + return r.dispatchReadTimestampElems(ctx, readTimestamp, elems) } ttl, err := r.ttlAt(ctx, key, readTS) if err != nil { @@ -251,7 +253,7 @@ func (r *RedisServer) persistHLLMembersTxn(ctx context.Context, key []byte, read return err } elems := []*kv.Elem[kv.OP]{{Op: kv.Put, Key: redisHLLKey(key), Value: payload}} - return r.dispatchElems(ctx, true, readTS, appendHLLScanIndexElem(elems, key, ttl)) + return r.dispatchReadTimestampElems(ctx, readTimestamp, appendHLLScanIndexElem(elems, key, ttl)) } func appendHLLScanIndexElem(elems []*kv.Elem[kv.OP], key []byte, ttl *time.Time) []*kv.Elem[kv.OP] { @@ -278,7 +280,11 @@ func applySetMemberMutation(elems []*kv.Elem[kv.OP], memberKey []byte, exists, a func (r *RedisServer) mutateExactSetLegacy(conn redcon.Conn, ctx context.Context, kind string, key []byte, members [][]byte, add bool) { var changed int if err := r.retryRedisWrite(ctx, func() error { - readTS := r.readTS() + readTimestamp, err := r.beginTxnReadTimestamp(ctx, "redis set mutation: begin read timestamp") + if err != nil { + return cockerrors.WithStack(err) + } + readTS := readTimestamp.Timestamp() if err := r.validateExactSetKind(kind, key, readTS); err != nil { return err } @@ -291,7 +297,7 @@ func (r *RedisServer) mutateExactSetLegacy(conn redcon.Conn, ctx context.Context if changed == 0 { return nil } - return r.persistExactSetMembersTxn(ctx, kind, key, readTS, existing) + return r.persistExactSetMembersTxn(ctx, kind, key, readTimestamp, existing) }); err != nil { writeRedisError(conn, err) return @@ -303,49 +309,26 @@ func (r *RedisServer) mutateExactSetLegacy(conn redcon.Conn, ctx context.Context func (r *RedisServer) mutateExactSetWide(conn redcon.Conn, ctx context.Context, key []byte, members [][]byte, add bool) { var changed int if err := r.retryRedisWrite(ctx, func() error { - readTS := r.readTS() - typ, err := r.keyTypeOrEmptyAt(ctx, key, readTS, redisTypeSet) + readTimestamp, err := r.beginTxnReadTimestamp(ctx, "redis set mutation: begin read timestamp") if err != nil { - return err + return cockerrors.WithStack(err) } - cleanupElems, migrationElems, legacyMemberBase, expiredRecreate, err := r.setWideMutationBase(ctx, key, readTS, typ) + readTS := readTimestamp.Timestamp() + prepared, err := r.prepareExactSetWideMutation(ctx, key, members, add, readTS) if err != nil { return err } - - startTS := normalizeStartTS(readTS) - commitTS, err := r.nextCommitTSAfter(ctx, startTS, "mutateExactSetWide: allocate commitTS") - if err != nil { - return cockerrors.WithStack(err) - } - - elems := make([]*kv.Elem[kv.OP], 0, len(cleanupElems)+len(migrationElems)+len(members)+setWideColOverhead) - elems = append(elems, cleanupElems...) - elems = append(elems, migrationElems...) - - var lenDelta int64 - var mutErr error - elems, changed, lenDelta, mutErr = r.applySetMemberMutations(ctx, key, members, add, readTS, elems, legacyMemberBase, expiredRecreate) - if mutErr != nil { - return mutErr - } - - if changed == 0 && len(migrationElems) == 0 && len(cleanupElems) == 0 { + changed = prepared.changed + if prepared.skip { return nil } - - elems = appendSetDeltaElems(elems, key, lenDelta, commitTS) - - if len(elems) == 0 { - return nil - } - - _, dispatchErr := r.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + _, dispatchErr := kv.DispatchWithReadTimestamp(dispatchCtx, r.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, - StartTS: startTS, - CommitTS: commitTS, - ReadKeys: redisTxnWideCreateReadKeys(key, typ, redisTxnWideSetFenceKey), - Elems: elems, + StartTS: prepared.startTS, + CommitTS: prepared.commitTS, + ReadKeys: redisTxnWideCreateReadKeys(key, prepared.typ, redisTxnWideSetFenceKey), + Elems: prepared.elems, }) return cockerrors.WithStack(dispatchErr) }); err != nil { @@ -355,6 +338,50 @@ func (r *RedisServer) mutateExactSetWide(conn redcon.Conn, ctx context.Context, conn.WriteInt(changed) } +type exactSetWideMutation struct { + typ redisValueType + startTS uint64 + commitTS uint64 + changed int + elems []*kv.Elem[kv.OP] + skip bool +} + +func (r *RedisServer) prepareExactSetWideMutation( + ctx context.Context, + key []byte, + members [][]byte, + add bool, + readTS uint64, +) (exactSetWideMutation, error) { + var prepared exactSetWideMutation + typ, err := r.keyTypeOrEmptyAt(ctx, key, readTS, redisTypeSet) + if err != nil { + return prepared, err + } + cleanupElems, migrationElems, legacyMemberBase, expiredRecreate, err := r.setWideMutationBase(ctx, key, readTS, typ) + if err != nil { + return prepared, err + } + startTS := normalizeStartTS(readTS) + commitTS, err := r.nextCommitTSAfter(ctx, startTS, "mutateExactSetWide: allocate commitTS") + if err != nil { + return prepared, cockerrors.WithStack(err) + } + elems := make([]*kv.Elem[kv.OP], 0, len(cleanupElems)+len(migrationElems)+len(members)+setWideColOverhead) + elems = append(elems, cleanupElems...) + elems = append(elems, migrationElems...) + elems, changed, lenDelta, err := r.applySetMemberMutations(ctx, key, members, add, readTS, elems, legacyMemberBase, expiredRecreate) + if err != nil { + return prepared, err + } + elems = appendSetDeltaElems(elems, key, lenDelta, commitTS) + return exactSetWideMutation{ + typ: typ, startTS: startTS, commitTS: commitTS, changed: changed, elems: elems, + skip: (changed == 0 && len(migrationElems) == 0 && len(cleanupElems) == 0) || len(elems) == 0, + }, nil +} + func (r *RedisServer) setWideMutationBase( ctx context.Context, key []byte, @@ -718,7 +745,11 @@ func (r *RedisServer) pfadd(conn redcon.Conn, cmd redcon.Command) { defer cancel() var changed int if err := r.retryRedisWrite(ctx, func() error { - readTS := r.readTS() + readTimestamp, err := r.beginTxnReadTimestamp(ctx, "redis pfadd: begin read timestamp") + if err != nil { + return cockerrors.WithStack(err) + } + readTS := readTimestamp.Timestamp() if err := r.validateExactSetKind(hllKind, cmd.Args[1], readTS); err != nil { return err } @@ -733,7 +764,7 @@ func (r *RedisServer) pfadd(conn redcon.Conn, cmd redcon.Command) { return nil } - return r.persistExactSetMembersTxn(ctx, hllKind, cmd.Args[1], readTS, existing) + return r.persistExactSetMembersTxn(ctx, hllKind, cmd.Args[1], readTimestamp, existing) }); err != nil { writeRedisError(conn, err) return diff --git a/adapter/redis_stream_cmds.go b/adapter/redis_stream_cmds.go index 52efed53d..628811ead 100644 --- a/adapter/redis_stream_cmds.go +++ b/adapter/redis_stream_cmds.go @@ -232,11 +232,11 @@ func (r *RedisServer) xadd(conn redcon.Conn, cmd redcon.Command) { } func (r *RedisServer) xaddTxn(ctx context.Context, key []byte, req xaddRequest) (string, error) { - id, readTS, elems, err := r.prepareXAdd(ctx, key, req) + id, readTimestamp, elems, err := r.prepareXAdd(ctx, key, req) if err != nil { return "", err } - return id, r.dispatchAndSignalStream(ctx, true, readTS, elems, key) + return id, r.dispatchAndSignalStreamWithReadTimestamp(ctx, true, readTimestamp, elems, key) } // prepareXAdd fixes one XADD attempt's stream ID and exact write set at a @@ -247,35 +247,24 @@ func (r *RedisServer) prepareXAdd( ctx context.Context, key []byte, req xaddRequest, -) (string, uint64, []*kv.Elem[kv.OP], error) { - readTS := r.readTS() - typ, err := r.streamTypeForXAdd(ctx, key, readTS) +) (string, kv.ReadTimestamp, []*kv.Elem[kv.OP], error) { + readTimestamp, readTS, legacyCleanup, meta, metaFound, err := r.prepareXAddBase(ctx, key) if err != nil { - return "", 0, nil, err - } - - legacyCleanup, meta, metaFound, err := r.streamWriteBase(ctx, key, readTS) - if err != nil { - return "", 0, nil, err - } - legacyCleanup, meta, metaFound, err = r.streamCleanupForExpiredRecreate( - ctx, key, readTS, typ, legacyCleanup, meta, metaFound) - if err != nil { - return "", 0, nil, err + return "", kv.ReadTimestamp{}, nil, err } id, parsedID, err := resolveXAddID(meta, metaFound, req.id) if err != nil { - return "", 0, nil, err + return "", kv.ReadTimestamp{}, nil, err } if err := xaddEnforceMaxWideColumn(key, meta.Length, req.maxLen); err != nil { - return "", 0, nil, err + return "", kv.ReadTimestamp{}, nil, err } entryValue, err := marshalStreamEntry(newRedisStreamEntry(id, req.fields)) if err != nil { - return "", 0, nil, err + return "", kv.ReadTimestamp{}, nil, err } // Capacity hint covers: stream cleanup Dels + one entry Put + one meta @@ -294,7 +283,7 @@ func (r *RedisServer) prepareXAdd( nextMeta, trim, err := r.xaddTrimIfNeeded(ctx, key, readTS, req.maxLen, meta) if err != nil { - return "", 0, nil, err + return "", kv.ReadTimestamp{}, nil, err } elems = append(elems, trim...) elems = appendMaxLenZeroSelfDel(elems, req.maxLen, key, parsedID) @@ -306,11 +295,36 @@ func (r *RedisServer) prepareXAdd( nextMeta.LastSeq = parsedID.seq metaBytes, err := store.MarshalStreamMeta(nextMeta) if err != nil { - return "", 0, nil, cockerrors.WithStack(err) + return "", kv.ReadTimestamp{}, nil, cockerrors.WithStack(err) } elems = append(elems, &kv.Elem[kv.OP]{Op: kv.Put, Key: store.StreamMetaKey(key), Value: metaBytes}) - return id, readTS, elems, nil + return id, readTimestamp, elems, nil +} + +func (r *RedisServer) prepareXAddBase( + ctx context.Context, + key []byte, +) (kv.ReadTimestamp, uint64, []*kv.Elem[kv.OP], store.StreamMeta, bool, error) { + readTimestamp, err := r.xaddReadTimestamp(ctx, key) + if err != nil { + return kv.ReadTimestamp{}, 0, nil, store.StreamMeta{}, false, err + } + readTS := readTimestamp.Timestamp() + typ, err := r.streamTypeForXAdd(ctx, key, readTS) + if err != nil { + return kv.ReadTimestamp{}, 0, nil, store.StreamMeta{}, false, err + } + legacyCleanup, meta, metaFound, err := r.streamWriteBase(ctx, key, readTS) + if err != nil { + return kv.ReadTimestamp{}, 0, nil, store.StreamMeta{}, false, err + } + legacyCleanup, meta, metaFound, err = r.streamCleanupForExpiredRecreate( + ctx, key, readTS, typ, legacyCleanup, meta, metaFound) + if err != nil { + return kv.ReadTimestamp{}, 0, nil, store.StreamMeta{}, false, err + } + return readTimestamp, readTS, legacyCleanup, meta, metaFound, nil } // reusableXAdd captures the exact ID and write set from an XADD attempt. A @@ -318,11 +332,12 @@ func (r *RedisServer) prepareXAdd( // already landed instead of resolving "*" against newer stream metadata and // appending a duplicate entry. type reusableXAdd struct { - elems []*kv.Elem[kv.OP] - startTS uint64 - commitTS uint64 - probeKey []byte - id string + elems []*kv.Elem[kv.OP] + readTimestamp kv.ReadTimestamp + startTS uint64 + commitTS uint64 + probeKey []byte + id string } // xaddTxnWithDedup retries XADD only through exact write-set reuse. This is the @@ -363,16 +378,18 @@ func (r *RedisServer) firstXAddAttempt( key []byte, req xaddRequest, ) (string, *reusableXAdd, error) { - id, readTS, elems, err := r.prepareXAdd(ctx, key, req) + id, readTimestamp, elems, err := r.prepareXAdd(ctx, key, req) if err != nil { return "", nil, err } + readTS := readTimestamp.Timestamp() startTS := normalizeStartTS(readTS) commitTS, err := r.nextCommitTSAfter(ctx, startTS, "redis xadd first-attempt: allocate commitTS") if err != nil { return "", nil, cockerrors.WithStack(err) } - _, dispErr := r.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + attemptCtx := readTimestamp.WithDispatchVoucher(ctx) + _, dispErr := kv.DispatchWithReadTimestamp(attemptCtx, r.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, CommitTS: commitTS, @@ -390,11 +407,12 @@ func (r *RedisServer) firstXAddAttempt( return "", nil, cockerrors.WithStack(dispErr) } return "", &reusableXAdd{ - elems: elems, - startTS: startTS, - commitTS: commitTS, - probeKey: kv.PrimaryKeyForElems(elems), - id: id, + elems: elems, + readTimestamp: readTimestamp, + startTS: startTS, + commitTS: commitTS, + probeKey: kv.PrimaryKeyForElems(elems), + id: id, }, cockerrors.WithStack(dispErr) } @@ -411,7 +429,8 @@ func (r *RedisServer) dispatchXAddReuse( if allocErr != nil { return "", false, cockerrors.WithStack(allocErr) } - _, dispErr := r.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + attemptCtx := pending.readTimestamp.WithDispatchVoucher(ctx) + _, dispErr := kv.DispatchWithReadTimestamp(attemptCtx, r.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: pending.startTS, CommitTS: commitTS, @@ -473,22 +492,40 @@ func (r *RedisServer) streamCleanupForExpiredRecreate( return cleanup, store.StreamMeta{}, false, nil } -// dispatchAndSignalStream dispatches the elems through the coordinator -// and, on success, wakes any XREAD BLOCK waiter on the same node. -// dispatchElems blocks until the FSM applies locally, so by the time -// Signal fires the new entries are visible at the readTS the woken -// waiter will pick on its next iteration. Pulled out of xaddTxn so the -// parent function stays under the cyclop budget — the signal step -// would otherwise add an extra branch on the dispatch error path. -func (r *RedisServer) dispatchAndSignalStream( +func (r *RedisServer) xaddReadTimestamp(ctx context.Context, key []byte) (kv.ReadTimestamp, error) { + readTimestamp, err := r.beginTxnReadTimestamp(ctx, "redis xadd: begin read timestamp") + if err != nil { + return kv.ReadTimestamp{}, cockerrors.WithStack(err) + } + readTS := readTimestamp.Timestamp() + typ, err := r.keyTypeAtExpect(ctx, key, readTS, redisTypeStream) + if err != nil { + return kv.ReadTimestamp{}, err + } + if typ != redisTypeNone && typ != redisTypeStream { + return kv.ReadTimestamp{}, wrongTypeError() + } + return readTimestamp, nil +} + +func (r *RedisServer) dispatchAndSignalStreamWithReadTimestamp( ctx context.Context, isTxn bool, - startTS uint64, + readTimestamp kv.ReadTimestamp, elems []*kv.Elem[kv.OP], streamKey []byte, ) error { - if err := r.dispatchElems(ctx, isTxn, startTS, elems); err != nil { - return err + if len(elems) == 0 { + return nil + } + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + _, err := kv.DispatchWithReadTimestamp(dispatchCtx, r.coordinator, &kv.OperationGroup[kv.OP]{ + IsTxn: isTxn, + StartTS: normalizeStartTS(readTimestamp.Timestamp()), + Elems: elems, + }) + if err != nil { + return cockerrors.WithStack(err) } r.streamWaiters.Signal(streamKey) return nil @@ -820,12 +857,8 @@ func (r *RedisServer) streamTypeForXAdd(ctx context.Context, key []byte, readTS } } -// flushLegacyCleanupOnTrimNoOp commits the legacy-blob Del + meta Put -// for an XTRIM whose length is already under maxLen. Without this -// flush a subsequent read would still find the stale legacy blob. -// Returns 0 removed entries; callers use that directly. -func (r *RedisServer) flushLegacyCleanupOnTrimNoOp( - ctx context.Context, readTS uint64, key []byte, +func (r *RedisServer) flushLegacyCleanupOnTrimNoOpReadTimestamp( + ctx context.Context, readTimestamp kv.ReadTimestamp, key []byte, meta store.StreamMeta, legacyCleanup []*kv.Elem[kv.OP], ) (int, error) { if len(legacyCleanup) == 0 { @@ -838,11 +871,15 @@ func (r *RedisServer) flushLegacyCleanupOnTrimNoOp( elems := make([]*kv.Elem[kv.OP], 0, len(legacyCleanup)+1) elems = append(elems, legacyCleanup...) elems = append(elems, &kv.Elem[kv.OP]{Op: kv.Put, Key: store.StreamMetaKey(key), Value: metaBytes}) - return 0, r.dispatchElems(ctx, true, readTS, elems) + return 0, r.dispatchReadTimestampElems(ctx, readTimestamp, elems) } func (r *RedisServer) xtrimTxn(ctx context.Context, key []byte, maxLen int) (int, error) { - readTS := r.readTS() + readTimestamp, err := r.beginTxnReadTimestamp(ctx, "redis xtrim: begin read timestamp") + if err != nil { + return 0, cockerrors.WithStack(err) + } + readTS := readTimestamp.Timestamp() proceed, err := r.streamTypeForWrite(ctx, key, readTS) if err != nil || !proceed { return 0, err @@ -854,7 +891,7 @@ func (r *RedisServer) xtrimTxn(ctx context.Context, key []byte, maxLen int) (int } if meta.Length <= int64(maxLen) { - return r.flushLegacyCleanupOnTrimNoOp(ctx, readTS, key, meta, legacyCleanup) + return r.flushLegacyCleanupOnTrimNoOpReadTimestamp(ctx, readTimestamp, key, meta, legacyCleanup) } // Cap the trim request at maxWideColumnItems so a single XTRIM cannot @@ -891,7 +928,7 @@ func (r *RedisServer) xtrimTxn(ctx context.Context, key []byte, maxLen int) (int return 0, cockerrors.WithStack(err) } elems = append(elems, &kv.Elem[kv.OP]{Op: kv.Put, Key: store.StreamMetaKey(key), Value: metaBytes}) - return actualRemoved, r.dispatchElems(ctx, true, readTS, elems) + return actualRemoved, r.dispatchReadTimestampElems(ctx, readTimestamp, elems) } func (r *RedisServer) xrange(conn redcon.Conn, cmd redcon.Command) { diff --git a/adapter/redis_strings.go b/adapter/redis_strings.go index a01fa41a1..3236fe983 100644 --- a/adapter/redis_strings.go +++ b/adapter/redis_strings.go @@ -130,7 +130,8 @@ func (r *RedisServer) loadRedisSetState(ctx context.Context, key []byte, readTS return state, nil } -func (r *RedisServer) replaceWithStringTxn(ctx context.Context, key, value []byte, ttl *time.Time, typ redisValueType, readTS uint64) error { +func (r *RedisServer) replaceWithStringReadTimestampTxn(ctx context.Context, key, value []byte, ttl *time.Time, typ redisValueType, readTimestamp kv.ReadTimestamp) error { + readTS := readTimestamp.Timestamp() var elems []*kv.Elem[kv.OP] if isNonStringCollectionType(typ) { delElems, _, err := r.deleteLogicalKeyElems(ctx, key, readTS) @@ -139,22 +140,24 @@ func (r *RedisServer) replaceWithStringTxn(ctx context.Context, key, value []byt } elems = append(elems, delElems...) } - // Embed TTL in the string value; write !redis|ttl| as a secondary scan index. encoded := encodeRedisStr(bytes.Clone(value), ttl) elems = append(elems, &kv.Elem[kv.OP]{Op: kv.Put, Key: redisStrKey(key), Value: encoded}) if ttl != nil { elems = append(elems, &kv.Elem[kv.OP]{Op: kv.Put, Key: redisTTLKey(key), Value: encodeRedisTTL(*ttl)}) } else { - // Clear any prior scan index so a persistent string is not later expired. elems = append(elems, &kv.Elem[kv.OP]{Op: kv.Del, Key: redisTTLKey(key)}) } - return r.dispatchElems(ctx, true, readTS, elems) + return r.dispatchReadTimestampElems(ctx, readTimestamp, elems) } func (r *RedisServer) executeSet(ctx context.Context, key, value []byte, opts redisSetOptions) (redisSetExecution, error) { var result redisSetExecution err := r.retryRedisWrite(ctx, func() error { - readTS := r.readTS() + readTimestamp, err := r.beginTxnReadTimestamp(ctx, "redis set string: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() state, err := r.loadRedisSetState(ctx, key, readTS, opts.returnOld) if err != nil { return err @@ -169,7 +172,7 @@ func (r *RedisServer) executeSet(ctx context.Context, key, value []byte, opts re return wrongTypeError() } // Use rawTyp for cleanup so expired-but-lingering internal keys are deleted. - if err := r.replaceWithStringTxn(ctx, key, value, opts.ttl, state.rawTyp, readTS); err != nil { + if err := r.replaceWithStringReadTimestampTxn(ctx, key, value, opts.ttl, state.rawTyp, readTimestamp); err != nil { return err } result = redisSetExecution{state: state, wroteOldBulk: opts.returnOld} @@ -519,7 +522,11 @@ func (r *RedisServer) delLocal(keys [][]byte) (int, error) { err := r.retryRedisWrite(ctx, func() error { elems := []*kv.Elem[kv.OP]{} nextRemoved := 0 - readTS := r.readTS() + readTimestamp, err := r.beginTxnReadTimestamp(ctx, "redis del: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() for _, key := range keys { keyElems, existed, err := r.deleteLogicalKeyElems(ctx, key, readTS) if err != nil { @@ -530,7 +537,7 @@ func (r *RedisServer) delLocal(keys [][]byte) (int, error) { } elems = append(elems, keyElems...) } - if err := r.dispatchElems(ctx, true, readTS, elems); err != nil { + if err := r.dispatchReadTimestampElems(ctx, readTimestamp, elems); err != nil { return err } removed = nextRemoved diff --git a/adapter/redis_txn.go b/adapter/redis_txn.go index a38001517..76b1515cd 100644 --- a/adapter/redis_txn.go +++ b/adapter/redis_txn.go @@ -1577,7 +1577,7 @@ func (t *txnContext) commit() error { CommitTS: prepared.commitTS, ReadKeys: prepared.readKeys, } - if _, err := t.server.coordinator.Dispatch(prepared.ctx, group); err != nil { + if _, err := kv.DispatchWithReadTimestamp(prepared.ctx, t.server.coordinator, group); err != nil { return errors.WithStack(err) } return nil @@ -2379,13 +2379,18 @@ func (r *RedisServer) runTransactionDirect(queue []redcon.Command) ([]redisResul var results []redisResult err := r.retryRedisWrite(dispatchCtx, func() error { - startTS := r.txnStartTS() + readTimestamp, err := r.beginTxnReadTimestamp(dispatchCtx, "redis exec: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } + attemptCtx := readTimestamp.WithDispatchVoucher(dispatchCtx) + startTS := readTimestamp.Timestamp() readPin := r.pinReadTS(startTS) defer readPin.Release() txn := &txnContext{ server: r, - ctx: dispatchCtx, + ctx: attemptCtx, working: map[string]*txnValue{}, replacers: map[string]*stringReplacement{}, listStates: map[string]*listTxnState{}, @@ -2412,7 +2417,7 @@ func (r *RedisServer) runTransactionDirect(queue []redcon.Command) ([]redisResul nextResults = append(nextResults, res) } - if err := txn.validateReadSet(dispatchCtx); err != nil { + if err := txn.validateReadSet(attemptCtx); err != nil { return err } if err := txn.commit(); err != nil { @@ -2446,11 +2451,12 @@ func (r *RedisServer) runTransactionDirect(queue []redcon.Command) ([]redisResul // results are only returned when reuse actually represents the // outcome of attempt 1's intent. type reusableExecTxn struct { - elems []*kv.Elem[kv.OP] - startTS uint64 - commitTS uint64 - readKeys [][]byte - results []redisResult + elems []*kv.Elem[kv.OP] + startTS uint64 + commitTS uint64 + readKeys [][]byte + results []redisResult + readTimestamp kv.ReadTimestamp } // dispatchExecReuse runs one iteration of the option-2 reuse path for @@ -2468,6 +2474,7 @@ type reusableExecTxn struct { // is the current length" question; the client-visible result IS the // cached results array. func (r *RedisServer) dispatchExecReuse(ctx context.Context, pending *reusableExecTxn) (results []redisResult, drop bool, err error) { + ctx = pending.readTimestamp.WithDispatchVoucher(ctx) // gemini PR-A HIGH: persistence-grade commit_ts allocation must honor the // HLC-4 physical-ceiling fence (see kv/hlc.go NextFenced + the TLA proof // at tla/hlc/MCHLC_gap.cfg). Clock().Next() bypasses the ceiling and @@ -2478,7 +2485,7 @@ func (r *RedisServer) dispatchExecReuse(ctx context.Context, pending *reusableEx if allocErr != nil { return nil, false, errors.WithStack(allocErr) } - _, dispErr := r.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + _, dispErr := kv.DispatchWithReadTimestamp(ctx, r.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: pending.startTS, CommitTS: commitTS, @@ -2593,7 +2600,12 @@ func (r *RedisServer) runTransactionWithDedup(queue []redcon.Command) ([]redisRe // from runTransactionWithDedup to keep that loop under the cyclop // budget; the dedup rationale lives there. func (r *RedisServer) firstExecAttempt(dispatchCtx context.Context, queue []redcon.Command) ([]redisResult, *reusableExecTxn, error) { - startTS := r.txnStartTS() + readTimestamp, err := r.beginTxnReadTimestamp(dispatchCtx, "redis exec: begin read timestamp") + if err != nil { + return nil, nil, errors.WithStack(err) + } + dispatchCtx = readTimestamp.WithDispatchVoucher(dispatchCtx) + startTS := readTimestamp.Timestamp() readPin := r.pinReadTS(startTS) defer readPin.Release() @@ -2647,7 +2659,7 @@ func (r *RedisServer) firstExecAttempt(dispatchCtx context.Context, queue []redc CommitTS: prepared.commitTS, ReadKeys: prepared.readKeys, } - if _, dispErr := r.coordinator.Dispatch(prepared.ctx, group); dispErr != nil { + if _, dispErr := kv.DispatchWithReadTimestamp(prepared.ctx, r.coordinator, group); dispErr != nil { // Preserve the exact attempt for a forwarded conflict only after // restoring its typed form. runTransactionWithDedup can then reuse this // write set instead of replaying the EXEC body from a new snapshot. @@ -2660,11 +2672,12 @@ func (r *RedisServer) firstExecAttempt(dispatchCtx context.Context, queue []redc // escape to the client are out of scope for this loop. if isRetryableRedisTxnErr(dispErr) { return nil, &reusableExecTxn{ - elems: prepared.elems, - startTS: txn.startTS, - commitTS: prepared.commitTS, - readKeys: prepared.readKeys, - results: nextResults, + elems: prepared.elems, + startTS: txn.startTS, + commitTS: prepared.commitTS, + readKeys: prepared.readKeys, + results: nextResults, + readTimestamp: readTimestamp, }, errors.WithStack(dispErr) } return nil, nil, errors.WithStack(dispErr) @@ -2710,6 +2723,11 @@ func (r *RedisServer) txnStartTS() uint64 { return maxTS } +func (r *RedisServer) beginTxnReadTimestamp(ctx context.Context, label string) (kv.ReadTimestamp, error) { + readTimestamp, err := kv.BeginReadTimestampThrough(ctx, r.coordinator, r.txnStartTS(), label) + return readTimestamp, errors.WithStack(err) +} + func (r *RedisServer) writeResults(conn redcon.Conn, results []redisResult) { conn.WriteArray(len(results)) for _, res := range results { diff --git a/adapter/redis_zset_cmds.go b/adapter/redis_zset_cmds.go index cfc8be8a4..f9757d53b 100644 --- a/adapter/redis_zset_cmds.go +++ b/adapter/redis_zset_cmds.go @@ -600,7 +600,11 @@ func (r *RedisServer) applyZAddPair(ctx context.Context, key []byte, p zaddPair, } func (r *RedisServer) zaddTxn(ctx context.Context, key []byte, flags zaddFlags, pairs []zaddPair) (int, error) { - readTS := r.readTS() + readTimestamp, err := r.beginTxnReadTimestamp(ctx, "redis zadd: begin read timestamp") + if err != nil { + return 0, cockerrors.WithStack(err) + } + readTS := readTimestamp.Timestamp() base, err := r.prepareZSetWriteBase(ctx, key, readTS, len(pairs)) if err != nil { return 0, err @@ -645,7 +649,7 @@ func (r *RedisServer) zaddTxn(ctx context.Context, key []byte, flags zaddFlags, }) } - return added, r.dispatchAndSignalZSet(ctx, readTS, commitTS, elems, key, + return added, r.dispatchAndSignalZSet(ctx, readTimestamp, commitTS, elems, key, redisTxnWideCreateReadKeys(key, base.typ, redisTxnWideZSetFenceKey)) } @@ -700,12 +704,15 @@ func (r *RedisServer) prepareZSetWriteBase(ctx context.Context, key []byte, read // dispatch error path. func (r *RedisServer) dispatchAndSignalZSet( ctx context.Context, - readTS, commitTS uint64, + readTimestamp kv.ReadTimestamp, + commitTS uint64, elems []*kv.Elem[kv.OP], zsetKey []byte, readKeys [][]byte, ) error { - _, err := r.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + readTS := readTimestamp.Timestamp() + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + _, err := kv.DispatchWithReadTimestamp(dispatchCtx, r.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: normalizeStartTS(readTS), CommitTS: commitTS, @@ -722,7 +729,11 @@ func (r *RedisServer) dispatchAndSignalZSet( // zincrbyTxn performs one attempt of ZINCRBY in wide-column format. // Returns the new score after applying increment. func (r *RedisServer) zincrbyTxn(ctx context.Context, key []byte, member string, increment float64) (float64, error) { - readTS := r.readTS() + readTimestamp, err := r.beginTxnReadTimestamp(ctx, "redis zincrby: begin read timestamp") + if err != nil { + return 0, cockerrors.WithStack(err) + } + readTS := readTimestamp.Timestamp() base, err := r.prepareZSetWriteBase(ctx, key, readTS, 0) if err != nil { return 0, err @@ -763,7 +774,7 @@ func (r *RedisServer) zincrbyTxn(ctx context.Context, key []byte, member string, Value: deltaVal, }) } - if err := r.dispatchAndSignalZSet(ctx, readTS, commitTS, elems, key, + if err := r.dispatchAndSignalZSet(ctx, readTimestamp, commitTS, elems, key, redisTxnWideCreateReadKeys(key, base.typ, redisTxnWideZSetFenceKey)); err != nil { return 0, err } @@ -843,13 +854,14 @@ func removeZSetMembers(members map[string]float64, rawMembers [][]byte) []redisZ return removed } -func (r *RedisServer) persistZSetEntriesTxn(ctx context.Context, key []byte, readTS uint64, entries []redisZSetEntry) error { +func (r *RedisServer) persistZSetEntriesTxn(ctx context.Context, key []byte, readTimestamp kv.ReadTimestamp, entries []redisZSetEntry) error { + readTS := readTimestamp.Timestamp() if len(entries) == 0 { elems, _, err := r.deleteLogicalKeyElems(ctx, key, readTS) if err != nil { return err } - return r.dispatchElems(ctx, true, readTS, elems) + return r.dispatchReadTimestampElems(ctx, readTimestamp, elems) } memberPrefix := store.ZSetMemberScanPrefix(key) @@ -883,7 +895,8 @@ func (r *RedisServer) persistZSetEntriesTxn(ctx context.Context, key []byte, rea Key: store.ZSetMetaDeltaKey(key, commitTS, 0), Value: deltaVal, }) - _, dispatchErr := r.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + _, dispatchErr := kv.DispatchWithReadTimestamp(dispatchCtx, r.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, CommitTS: commitTS, @@ -892,25 +905,26 @@ func (r *RedisServer) persistZSetEntriesTxn(ctx context.Context, key []byte, rea }) return cockerrors.WithStack(dispatchErr) } - return r.dispatchElems(ctx, true, readTS, elems) + return r.dispatchReadTimestampElems(ctx, readTimestamp, elems) } payload, err := marshalZSetValue(redisZSetValue{Entries: entries}) if err != nil { return err } - return r.dispatchElems(ctx, true, readTS, []*kv.Elem[kv.OP]{ + return r.dispatchReadTimestampElems(ctx, readTimestamp, []*kv.Elem[kv.OP]{ {Op: kv.Put, Key: redisZSetKey(key), Value: payload}, }) } -func (r *RedisServer) persistZSetRemovalsTxn(ctx context.Context, key []byte, readTS uint64, removed, remaining []redisZSetEntry) error { +func (r *RedisServer) persistZSetRemovalsTxn(ctx context.Context, key []byte, readTimestamp kv.ReadTimestamp, removed, remaining []redisZSetEntry) error { + readTS := readTimestamp.Timestamp() if len(remaining) == 0 { elems, _, err := r.deleteLogicalKeyElems(ctx, key, readTS) if err != nil { return err } - return r.dispatchElems(ctx, true, readTS, elems) + return r.dispatchReadTimestampElems(ctx, readTimestamp, elems) } memberPrefix := store.ZSetMemberScanPrefix(key) memberEnd := store.PrefixScanEnd(memberPrefix) @@ -919,7 +933,7 @@ func (r *RedisServer) persistZSetRemovalsTxn(ctx context.Context, key []byte, re return cockerrors.WithStack(err) } if len(probeKVs) == 0 { - return r.persistZSetEntriesTxn(ctx, key, readTS, remaining) + return r.persistZSetEntriesTxn(ctx, key, readTimestamp, remaining) } startTS := normalizeStartTS(readTS) commitTS, err := r.nextCommitTSAfter(ctx, startTS, "persistZSetRemovalsTxn: allocate commitTS") @@ -941,7 +955,8 @@ func (r *RedisServer) persistZSetRemovalsTxn(ctx context.Context, key []byte, re Value: deltaVal, }) elems = append(elems, redisTxnWideZSetFenceElem(key)) - _, dispatchErr := r.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + _, dispatchErr := kv.DispatchWithReadTimestamp(dispatchCtx, r.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, CommitTS: commitTS, @@ -1024,7 +1039,11 @@ func (r *RedisServer) zremWithTypeProbe(conn redcon.Conn, cmd redcon.Command, fa defer cancel() var removed int if err := r.retryRedisWrite(ctx, func() error { - readTS := r.readTS() + readTimestamp, err := r.beginTxnReadTimestamp(ctx, "redis zrem: begin read timestamp") + if err != nil { + return cockerrors.WithStack(err) + } + readTS := readTimestamp.Timestamp() typ, err := r.zremTypeAt(ctx, cmd.Args[1], readTS, fastMiss) if err != nil { return err @@ -1046,7 +1065,7 @@ func (r *RedisServer) zremWithTypeProbe(conn redcon.Conn, cmd redcon.Command, fa if removed == 0 { return nil } - return r.persistZSetRemovalsTxn(ctx, cmd.Args[1], readTS, removedEntries, zsetMapToEntries(members)) + return r.persistZSetRemovalsTxn(ctx, cmd.Args[1], readTimestamp, removedEntries, zsetMapToEntries(members)) }); err != nil { writeRedisError(conn, err) return @@ -1095,22 +1114,19 @@ func (r *RedisServer) zremrangebyrank(conn redcon.Conn, cmd redcon.Command) { defer cancel() var removed int if err := r.retryRedisWrite(ctx, func() error { - readTS := r.readTS() - typ, err := r.keyTypeAtExpect(ctx, cmd.Args[1], readTS, redisTypeZSet) + readTimestamp, err := r.beginTxnReadTimestamp(ctx, "redis zremrangebyrank: begin read timestamp") + if err != nil { + return cockerrors.WithStack(err) + } + readTS := readTimestamp.Timestamp() + value, exists, err := r.zsetMutationValueAt(ctx, cmd.Args[1], readTS) if err != nil { return err } - if typ == redisTypeNone { + if !exists { removed = 0 return nil } - if typ != redisTypeZSet { - return wrongTypeError() - } - value, _, err := r.loadZSetAt(ctx, cmd.Args[1], readTS) - if err != nil { - return err - } s, e := normalizeRankRange(start, stop, len(value.Entries)) if e < s { removed = 0 @@ -1120,7 +1136,7 @@ func (r *RedisServer) zremrangebyrank(conn redcon.Conn, cmd redcon.Command) { remaining = append(remaining, value.Entries[e+1:]...) removedEntries := append([]redisZSetEntry(nil), value.Entries[s:e+1]...) removed = len(removedEntries) - return r.persistZSetRemovalsTxn(ctx, cmd.Args[1], readTS, removedEntries, remaining) + return r.persistZSetRemovalsTxn(ctx, cmd.Args[1], readTimestamp, removedEntries, remaining) }); err != nil { writeRedisError(conn, err) return @@ -1128,6 +1144,25 @@ func (r *RedisServer) zremrangebyrank(conn redcon.Conn, cmd redcon.Command) { conn.WriteInt(removed) } +func (r *RedisServer) zsetMutationValueAt( + ctx context.Context, + key []byte, + readTS uint64, +) (redisZSetValue, bool, error) { + typ, err := r.keyTypeAtExpect(ctx, key, readTS, redisTypeZSet) + if err != nil { + return redisZSetValue{}, false, err + } + if typ == redisTypeNone { + return redisZSetValue{}, false, nil + } + if typ != redisTypeZSet { + return redisZSetValue{}, false, wrongTypeError() + } + value, _, err := r.loadZSetAt(ctx, key, readTS) + return value, true, err +} + // tryBZPopMinWithMode runs one BZPOPMIN attempt against key. The // fast flag selects keyTypeAtExpectFast (no slow-path fallback, no // wrongType detection) when true; the caller MUST guarantee that the @@ -1140,16 +1175,20 @@ func (r *RedisServer) tryBZPopMinWithMode(key []byte, fast bool) (*bzpopminResul defer cancel() var result *bzpopminResult err := r.retryRedisWrite(ctx, func() error { - readTS := r.readTS() + readTimestamp, err := r.beginTxnReadTimestamp(ctx, "redis bzpopmin: begin read timestamp") + if err != nil { + return cockerrors.WithStack(err) + } + readTS := readTimestamp.Timestamp() var typ redisValueType - var err error + var typeErr error if fast { - typ, err = r.keyTypeAtExpectFast(ctx, key, readTS, redisTypeZSet) + typ, typeErr = r.keyTypeAtExpectFast(ctx, key, readTS, redisTypeZSet) } else { - typ, err = r.keyTypeAtExpect(ctx, key, readTS, redisTypeZSet) + typ, typeErr = r.keyTypeAtExpect(ctx, key, readTS, redisTypeZSet) } - if err != nil { - return err + if typeErr != nil { + return typeErr } if typ == redisTypeNone { result = nil @@ -1178,7 +1217,7 @@ func (r *RedisServer) tryBZPopMinWithMode(key []byte, fast bool) (*bzpopminResul } isWide := len(probeKVs) > 0 - if err := r.persistBZPopMinResult(ctx, key, readTS, popped, remaining, isWide); err != nil { + if err := r.persistBZPopMinResult(ctx, key, readTimestamp, popped, remaining, isWide); err != nil { return err } result = &bzpopminResult{key: key, entry: popped} @@ -1187,13 +1226,14 @@ func (r *RedisServer) tryBZPopMinWithMode(key []byte, fast bool) (*bzpopminResul return result, err } -func (r *RedisServer) persistBZPopMinResult(ctx context.Context, key []byte, readTS uint64, popped redisZSetEntry, remaining []redisZSetEntry, isWide bool) error { +func (r *RedisServer) persistBZPopMinResult(ctx context.Context, key []byte, readTimestamp kv.ReadTimestamp, popped redisZSetEntry, remaining []redisZSetEntry, isWide bool) error { + readTS := readTimestamp.Timestamp() if len(remaining) == 0 { elems, _, err := r.deleteLogicalKeyElems(ctx, key, readTS) if err != nil { return err } - return r.dispatchElems(ctx, true, readTS, elems) + return r.dispatchReadTimestampElems(ctx, readTimestamp, elems) } if isWide { // Wide-column: delete the popped member key + score index, emit delta -1. @@ -1209,7 +1249,8 @@ func (r *RedisServer) persistBZPopMinResult(ctx context.Context, key []byte, rea {Op: kv.Put, Key: store.ZSetMetaDeltaKey(key, commitTS, 0), Value: deltaVal}, redisTxnWideZSetFenceElem(key), } - _, dispatchErr := r.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + _, dispatchErr := kv.DispatchWithReadTimestamp(dispatchCtx, r.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, CommitTS: commitTS, @@ -1223,7 +1264,7 @@ func (r *RedisServer) persistBZPopMinResult(ctx context.Context, key []byte, rea if err != nil { return err } - return r.dispatchElems(ctx, true, readTS, []*kv.Elem[kv.OP]{ + return r.dispatchReadTimestampElems(ctx, readTimestamp, []*kv.Elem[kv.OP]{ {Op: kv.Put, Key: redisZSetKey(key), Value: payload}, }) } diff --git a/adapter/s3.go b/adapter/s3.go index 1893af37e..3d210da2a 100644 --- a/adapter/s3.go +++ b/adapter/s3.go @@ -626,10 +626,12 @@ func (s *S3Server) createBucket(w http.ResponseWriter, r *http.Request, bucket s } err := s.retryS3Mutation(r.Context(), func() error { readTS := s.readTS() - startTS, err := s.txnStartTS(r.Context(), readTS) + readTimestamp, err := s.beginTxnReadTimestamp(r.Context(), readTS, "s3 create bucket: begin read timestamp") if err != nil { return errors.Wrap(err, "s3: allocate startTS for mutation") } + readTS = readTimestamp.Timestamp() + startTS := readTS readPin := s.pinReadTS(readTS) defer readPin.Release() @@ -674,7 +676,8 @@ func (s *S3Server) createBucket(w http.ResponseWriter, r *http.Request, bucket s {Op: kv.Put, Key: s3keys.BucketGenerationKey(bucket), Value: encodeS3Generation(nextGeneration)}, }, } - _, err = s.coordinator.Dispatch(r.Context(), req) + dispatchCtx := readTimestamp.WithDispatchVoucher(r.Context()) + _, err = kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, req) return errors.WithStack(err) }) if err != nil { @@ -705,10 +708,12 @@ func (s *S3Server) deleteBucket(w http.ResponseWriter, r *http.Request, bucket s var deletedGeneration uint64 err := s.retryS3Mutation(r.Context(), func() error { readTS := s.readTS() - startTS, err := s.txnStartTS(r.Context(), readTS) + readTimestamp, err := s.beginTxnReadTimestamp(r.Context(), readTS, "s3 delete bucket: begin read timestamp") if err != nil { return errors.Wrap(err, "s3: allocate startTS for mutation") } + readTS = readTimestamp.Timestamp() + startTS := readTS readPin := s.pinReadTS(readTS) defer readPin.Release() @@ -747,7 +752,8 @@ func (s *S3Server) deleteBucket(w http.ResponseWriter, r *http.Request, bucket s // groups (kv/sharded_coordinator.go: dispatchDelPrefixBroadcast). // See AdminDeleteBucket's doc comment for the full // rationale. - _, err = s.coordinator.Dispatch(r.Context(), &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(r.Context()) + _, err = kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, Elems: []*kv.Elem[kv.OP]{{Op: kv.Del, Key: s3keys.BucketMetaKey(bucket)}}, @@ -853,10 +859,12 @@ func (s *S3Server) putBucketAcl(w http.ResponseWriter, r *http.Request, bucket s err := s.retryS3Mutation(r.Context(), func() error { readTS := s.readTS() - startTS, err := s.txnStartTS(r.Context(), readTS) + readTimestamp, err := s.beginTxnReadTimestamp(r.Context(), readTS, "s3 put bucket acl: begin read timestamp") if err != nil { return errors.Wrap(err, "s3: allocate startTS for mutation") } + readTS = readTimestamp.Timestamp() + startTS := readTS readPin := s.pinReadTS(readTS) defer readPin.Release() @@ -878,7 +886,8 @@ func (s *S3Server) putBucketAcl(w http.ResponseWriter, r *http.Request, bucket s if err != nil { return errors.WithStack(err) } - _, err = s.coordinator.Dispatch(r.Context(), &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(r.Context()) + _, err = kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, Elems: []*kv.Elem[kv.OP]{ @@ -1135,10 +1144,12 @@ func (s *S3Server) deleteObject(w http.ResponseWriter, r *http.Request, bucket s var generation uint64 err := s.retryS3Mutation(r.Context(), func() error { readTS := s.readTS() - startTS, err := s.txnStartTS(r.Context(), readTS) + readTimestamp, err := s.beginTxnReadTimestamp(r.Context(), readTS, "s3 delete object: begin read timestamp") if err != nil { return errors.Wrap(err, "s3: allocate startTS for mutation") } + readTS = readTimestamp.Timestamp() + startTS := readTS readPin := s.pinReadTS(readTS) defer readPin.Release() @@ -1165,7 +1176,8 @@ func (s *S3Server) deleteObject(w http.ResponseWriter, r *http.Request, bucket s cleanupManifest = nil return nil } - _, err = s.coordinator.Dispatch(r.Context(), &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(r.Context()) + _, err = kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, Elems: []*kv.Elem[kv.OP]{ @@ -1191,6 +1203,12 @@ func (s *S3Server) deleteObject(w http.ResponseWriter, r *http.Request, bucket s func (s *S3Server) createMultipartUpload(w http.ResponseWriter, r *http.Request, bucket string, objectKey string) { readTS := s.readTS() + readTimestamp, err := s.beginTxnReadTimestamp(r.Context(), readTS, "s3 create multipart upload: begin read timestamp") + if err != nil { + writeS3InternalError(w, err) + return + } + readTS = readTimestamp.Timestamp() readPin := s.pinReadTS(readTS) defer readPin.Release() @@ -1205,11 +1223,7 @@ func (s *S3Server) createMultipartUpload(w http.ResponseWriter, r *http.Request, } uploadID := newS3UploadID(s.clock()) - startTS, err := s.txnStartTS(r.Context(), readTS) - if err != nil { - writeS3InternalError(w, err) - return - } + startTS := readTS commitTS, err := s.nextTxnCommitTS(r.Context(), startTS) if err != nil { writeS3InternalError(w, err) @@ -1234,7 +1248,8 @@ func (s *S3Server) createMultipartUpload(w http.ResponseWriter, r *http.Request, writeS3InternalError(w, err) return } - if _, err := s.coordinator.Dispatch(r.Context(), &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(r.Context()) + if _, err := kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, CommitTS: commitTS, @@ -1326,10 +1341,12 @@ func (s *S3Server) abortMultipartUpload(w http.ResponseWriter, r *http.Request, var generation uint64 err := s.retryS3Mutation(r.Context(), func() error { readTS := s.readTS() - startTS, err := s.txnStartTS(r.Context(), readTS) + readTimestamp, err := s.beginTxnReadTimestamp(r.Context(), readTS, "s3 abort multipart upload: begin read timestamp") if err != nil { return errors.Wrap(err, "s3: allocate startTS for mutation") } + readTS = readTimestamp.Timestamp() + startTS := readTS readPin := s.pinReadTS(readTS) defer readPin.Release() @@ -1363,7 +1380,8 @@ func (s *S3Server) abortMultipartUpload(w http.ResponseWriter, r *http.Request, } // Transactional delete with startTS fencing — conflicts with concurrent Complete. - _, err = s.coordinator.Dispatch(r.Context(), &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(r.Context()) + _, err = kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, Elems: []*kv.Elem[kv.OP]{ @@ -2493,6 +2511,24 @@ func (s *S3Server) txnStartTS(ctx context.Context, readTS uint64) (uint64, error return ts, nil } +func (s *S3Server) beginTxnReadTimestamp(ctx context.Context, readTS uint64, label string) (kv.ReadTimestamp, error) { + if readTS == ^uint64(0) { + if alloc, ok := kv.TimestampAllocatorThrough(s.coordinator); ok { + if phaseD, phaseDOK := alloc.(kv.TSOPhaseDState); phaseDOK && (phaseD.PhaseDRequired() || phaseD.PhaseDActive()) { + readTS = 1 + readTimestamp, err := kv.BeginReadTimestampThrough(ctx, s.coordinator, readTS, label) + return readTimestamp, errors.WithStack(err) + } + } + } + startTS, err := s.txnStartTS(ctx, readTS) + if err != nil { + return kv.ReadTimestamp{}, err + } + readTimestamp, err := kv.BeginReadTimestampThrough(ctx, s.coordinator, startTS, label) + return readTimestamp, errors.WithStack(err) +} + // newS3UploadID generates an upload identifier for multipart uploads. // This is NOT a persistence timestamp — it is an opaque identifier // returned to the client and is only used as a lookup key for diff --git a/adapter/s3_admin.go b/adapter/s3_admin.go index 0d8c58d28..227e74ee0 100644 --- a/adapter/s3_admin.go +++ b/adapter/s3_admin.go @@ -259,10 +259,12 @@ func (s *S3Server) AdminCreateBucket(ctx context.Context, principal AdminPrincip // error path that the wrapping retry harness needs to see. func (s *S3Server) adminCreateBucketTxn(ctx context.Context, principal AdminPrincipal, name, acl string) (*AdminBucketSummary, error) { readTS := s.readTS() - startTS, err := s.txnStartTS(ctx, readTS) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, readTS, "s3 admin create bucket: begin read timestamp") if err != nil { return nil, errors.Wrap(err, "s3 admin: allocate startTS for adminCreateBucketTxn") } + readTS = readTimestamp.Timestamp() + startTS := readTS readPin := s.pinReadTS(readTS) defer readPin.Release() @@ -293,7 +295,8 @@ func (s *S3Server) adminCreateBucketTxn(ctx context.Context, principal AdminPrin if err != nil { return nil, errors.WithStack(err) } - _, err = s.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + _, err = kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, CommitTS: commitTS, @@ -327,10 +330,12 @@ func (s *S3Server) AdminPutBucketAcl(ctx context.Context, principal AdminPrincip err := s.retryS3Mutation(ctx, func() error { readTS := s.readTS() - startTS, err := s.txnStartTS(ctx, readTS) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, readTS, "s3 admin put bucket acl: begin read timestamp") if err != nil { return errors.Wrap(err, "s3 admin: allocate startTS for mutation") } + readTS = readTimestamp.Timestamp() + startTS := readTS readPin := s.pinReadTS(readTS) defer readPin.Release() @@ -438,10 +443,12 @@ func (s *S3Server) AdminDeleteBucket(ctx context.Context, principal AdminPrincip // allocation (PR #867 Phase 2a). func (s *S3Server) adminDeleteBucketTxnBody(ctx context.Context, name string, deletedGeneration *uint64) error { readTS := s.readTS() - startTS, err := s.txnStartTS(ctx, readTS) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, readTS, "s3 admin delete bucket: begin read timestamp") if err != nil { return errors.Wrap(err, "s3 admin: allocate startTS for adminDeleteBucket") } + readTS = readTimestamp.Timestamp() + startTS := readTS readPin := s.pinReadTS(readTS) defer readPin.Release() diff --git a/adapter/s3_admin_objects.go b/adapter/s3_admin_objects.go index d01db3a0a..4aebb4981 100644 --- a/adapter/s3_admin_objects.go +++ b/adapter/s3_admin_objects.go @@ -116,10 +116,12 @@ func (s *S3Server) AdminDeleteObject(ctx context.Context, principal AdminPrincip // silent-no-op semantics and AWS S3. func (s *S3Server) adminDeleteObjectTxn(ctx context.Context, bucket, key string) (*s3ObjectManifest, uint64, error) { readTS := s.readTS() - startTS, err := s.txnStartTS(ctx, readTS) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, readTS, "s3 admin delete object: begin read timestamp") if err != nil { return nil, 0, errors.Wrap(err, "s3 admin: allocate startTS for adminDeleteObjectTxn") } + readTS = readTimestamp.Timestamp() + startTS := readTS readPin := s.pinReadTS(readTS) defer readPin.Release() @@ -139,7 +141,8 @@ func (s *S3Server) adminDeleteObjectTxn(ctx context.Context, bucket, key string) if !found { return nil, 0, nil } - _, err = s.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + _, err = kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, Elems: []*kv.Elem[kv.OP]{ @@ -215,10 +218,12 @@ func (s *S3Server) AdminPutObject(ctx context.Context, principal AdminPrincipal, //nolint:cyclop,gocognit,nestif // see comment above func (s *S3Server) adminPutObjectStream(ctx context.Context, bucket, key string, body io.Reader, contentType string) (*s3ObjectManifest, uint64, error) { readTS := s.readTS() - startTS, err := s.txnStartTS(ctx, readTS) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, readTS, "s3 admin put object: begin read timestamp") if err != nil { return nil, 0, errors.Wrap(err, "s3 admin: allocate startTS for adminPutObjectStream") } + readTS = readTimestamp.Timestamp() + startTS := readTS readPin := s.pinReadTS(readTS) defer readPin.Release() @@ -343,7 +348,8 @@ func (s *S3Server) adminPutObjectStream(ctx context.Context, bucket, key string, s.cleanupManifestBlobs(ctx, bucket, meta.Generation, key, uploadedManifest()) return nil, 0, errors.WithStack(err) } - if _, err := s.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + if _, err := kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, CommitTS: commitTS, diff --git a/adapter/s3_hlc_fence_test.go b/adapter/s3_hlc_fence_test.go index 30f6b3b4b..303120f2f 100644 --- a/adapter/s3_hlc_fence_test.go +++ b/adapter/s3_hlc_fence_test.go @@ -2,10 +2,15 @@ package adapter import ( "context" + "net/http" + "net/http/httptest" + "strings" "testing" "time" + "github.com/bootjp/elastickv/internal/s3keys" "github.com/bootjp/elastickv/kv" + "github.com/bootjp/elastickv/store" "github.com/stretchr/testify/require" ) @@ -54,6 +59,137 @@ func TestS3TxnStartTSPassesThroughExplicitReadTS(t *testing.T) { require.Equal(t, uint64(42), ts) } +func TestS3BeginTxnReadTimestampPhaseDPreservesAppliedWatermark(t *testing.T) { + t.Parallel() + + allocator := &distributionTSOAllocator{base: 100, phaseD: true, phaseDFloor: 10} + coord := newDistributionCoordinatorStub(nil, true) + coord.allocator = allocator + srv := &S3Server{coordinator: coord} + + readTimestamp, err := srv.beginTxnReadTimestamp(context.Background(), 42, "test") + require.NoError(t, err) + require.Equal(t, uint64(42), readTimestamp.Timestamp()) + require.Zero(t, allocator.count, "an applied Phase-D read watermark must not allocate ahead of Raft apply") +} + +func TestS3BeginTxnReadTimestampPhaseDNormalizesEmptySnapshotSentinel(t *testing.T) { + t.Parallel() + + allocator := &distributionTSOAllocator{base: 100, phaseD: true, phaseDFloor: 10} + coord := newDistributionCoordinatorStub(nil, true) + coord.allocator = allocator + srv := &S3Server{coordinator: coord} + + readTimestamp, err := srv.beginTxnReadTimestamp(context.Background(), ^uint64(0), "test") + require.NoError(t, err) + require.Equal(t, uint64(1), readTimestamp.Timestamp()) + require.Zero(t, allocator.count, "an empty S3 snapshot must not allocate ahead of Raft apply") +} + +func TestS3CreateBucketPhaseDBindsReadVoucher(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + allocator := &distributionTSOAllocator{base: 100, phaseD: true, phaseDFloor: 10} + coord := newDistributionCoordinatorStub(st, true) + coord.allocator = allocator + server := NewS3Server(nil, "", st, coord, nil) + + req := httptest.NewRequest(http.MethodPut, "/voucher-bucket", nil) + rec := httptest.NewRecorder() + server.createBucket(rec, req, "voucher-bucket") + + require.Equalf(t, http.StatusOK, rec.Code, "body=%s", rec.Body.String()) + require.Equal(t, 1, coord.vouchCalls, "create bucket must carry the applied-read voucher into dispatch") + require.Equal(t, uint64(1), coord.lastStartTS) +} + +func TestS3AdminCreateBucketPhaseDBindsReadVoucher(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + allocator := &distributionTSOAllocator{base: 100, phaseD: true, phaseDFloor: 10} + coord := newDistributionCoordinatorStub(st, true) + coord.allocator = allocator + server := NewS3Server(nil, "", st, coord, nil) + + _, err := server.AdminCreateBucket(context.Background(), fullAdminBucketsPrincipal(), "admin-voucher-bucket", s3AclPrivate) + + require.NoError(t, err) + require.Equal(t, 1, coord.vouchCalls, "admin create bucket must carry the applied-read voucher into dispatch") + require.Equal(t, uint64(1), coord.lastStartTS) +} + +func TestS3AdminPutObjectPhaseDBindsReadVoucher(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + const generation = uint64(1) + bucketMeta, err := encodeS3BucketMeta(&s3BucketMeta{ + BucketName: "admin-voucher-bucket", + Generation: generation, + CreatedAtHLC: 1, + Region: s3DefaultRegion, + Owner: "AKIA_FULL", + Acl: s3AclPrivate, + }) + require.NoError(t, err) + require.NoError(t, st.PutAt(ctx, s3keys.BucketMetaKey("admin-voucher-bucket"), bucketMeta, 1, 0)) + + allocator := &distributionTSOAllocator{base: 100, phaseD: true, phaseDFloor: 10} + coord := newDistributionCoordinatorStub(st, true) + coord.allocator = allocator + server := NewS3Server(nil, "", st, coord, nil) + + err = server.AdminPutObject(ctx, fullAdminBucketsPrincipal(), "admin-voucher-bucket", "key.txt", strings.NewReader(""), "text/plain") + + require.NoError(t, err) + require.Equal(t, 1, coord.vouchCalls, "admin put object must carry the applied-read voucher into final dispatch") + require.Equal(t, uint64(1), coord.lastStartTS) +} + +func TestS3DeleteObjectPhaseDBindsReadVoucher(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + seedS3ObjectForReadVoucherTest(t, ctx, st, "voucher-bucket", "key.txt") + + allocator := &distributionTSOAllocator{base: 100, phaseD: true, phaseDFloor: 10} + coord := newDistributionCoordinatorStub(st, true) + coord.allocator = allocator + server := NewS3Server(nil, "", st, coord, nil) + + req := httptest.NewRequest(http.MethodDelete, "/voucher-bucket/key.txt", nil) + rec := httptest.NewRecorder() + server.deleteObject(rec, req, "voucher-bucket", "key.txt") + + require.Equalf(t, http.StatusNoContent, rec.Code, "body=%s", rec.Body.String()) + require.Equal(t, 1, coord.vouchCalls, "delete object must carry the applied-read voucher into dispatch") + require.Equal(t, uint64(1), coord.lastStartTS) +} + +func TestS3AdminDeleteObjectPhaseDBindsReadVoucher(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + seedS3ObjectForReadVoucherTest(t, ctx, st, "admin-voucher-bucket", "key.txt") + + allocator := &distributionTSOAllocator{base: 100, phaseD: true, phaseDFloor: 10} + coord := newDistributionCoordinatorStub(st, true) + coord.allocator = allocator + server := NewS3Server(nil, "", st, coord, nil) + + err := server.AdminDeleteObject(ctx, fullAdminBucketsPrincipal(), "admin-voucher-bucket", "key.txt") + + require.NoError(t, err) + require.Equal(t, 1, coord.vouchCalls, "admin delete object must carry the applied-read voucher into dispatch") + require.Equal(t, uint64(1), coord.lastStartTS) +} + // TestS3NextTxnCommitTSFailsClosedOnExpiredCeiling verifies that // nextTxnCommitTS surfaces ErrCeilingExpired through the // NextFenced() it calls after Observe(startTS). This is the @@ -81,3 +217,77 @@ func TestS3NextTxnCommitTSFailsClosedOnExpiredCeiling(t *testing.T) { require.ErrorIs(t, err, kv.ErrCeilingExpired, "s3 nextTxnCommitTS must propagate ErrCeilingExpired from NextFenced") } + +func TestS3CommitUploadPartRechecksUploadAtLatestAppliedWatermark(t *testing.T) { + st := store.NewMVCCStore() + const generation = uint64(1) + uploadMetaKey := s3keys.UploadMetaKey("bucket", generation, "object", "upload") + require.NoError(t, st.PutAt(context.Background(), uploadMetaKey, []byte("meta"), 10, 0)) + require.NoError(t, st.DeleteAt(context.Background(), uploadMetaKey, 20)) + coord := &recordingS3DispatchCoordinator{} + server := NewS3Server(nil, "", st, coord, nil) + + _, err := server.commitS3UploadPart(context.Background(), &s3UploadPartState{ + partNo: 1, + readTS: 10, + meta: &s3BucketMeta{Generation: generation}, + uploadMetaKey: uploadMetaKey, + }, s3ChunkUploadResult{}, "bucket", "object", "upload", s3UploadPartVersion{startTS: 10, commitTS: 30}) + require.ErrorContains(t, err, "upload not found") + require.Nil(t, coord.request, "an upload removed after startTS must not dispatch an orphan part") +} + +func TestS3CommitUploadPartIncludesUploadMetaInReadSet(t *testing.T) { + st := store.NewMVCCStore() + const generation = uint64(1) + uploadMetaKey := s3keys.UploadMetaKey("bucket", generation, "object", "upload") + require.NoError(t, st.PutAt(context.Background(), uploadMetaKey, []byte("meta"), 10, 0)) + coord := &recordingS3DispatchCoordinator{} + server := NewS3Server(nil, "", st, coord, nil) + + _, err := server.commitS3UploadPart(context.Background(), &s3UploadPartState{ + partNo: 1, + readTS: 10, + meta: &s3BucketMeta{Generation: generation}, + uploadMetaKey: uploadMetaKey, + }, s3ChunkUploadResult{}, "bucket", "object", "upload", s3UploadPartVersion{startTS: 10, commitTS: 30}) + require.NoError(t, err) + require.NotNil(t, coord.request) + require.Equal(t, [][]byte{uploadMetaKey}, coord.request.ReadKeys) +} + +func seedS3ObjectForReadVoucherTest(t *testing.T, ctx context.Context, st store.MVCCStore, bucket, key string) { + t.Helper() + + const generation = uint64(1) + bucketMeta, err := encodeS3BucketMeta(&s3BucketMeta{ + BucketName: bucket, + Generation: generation, + CreatedAtHLC: 1, + Region: s3DefaultRegion, + Owner: "AKIA_FULL", + Acl: s3AclPrivate, + }) + require.NoError(t, err) + require.NoError(t, st.PutAt(ctx, s3keys.BucketMetaKey(bucket), bucketMeta, 1, 0)) + + manifest, err := encodeS3ObjectManifest(&s3ObjectManifest{ + UploadID: "upload", + ETag: quoteS3ETag("etag"), + SizeBytes: 0, + LastModifiedHLC: 1, + ContentType: "text/plain", + }) + require.NoError(t, err) + require.NoError(t, st.PutAt(ctx, s3keys.ObjectManifestKey(bucket, generation, key), manifest, 1, 0)) +} + +type recordingS3DispatchCoordinator struct { + stubAdapterCoordinator + request *kv.OperationGroup[kv.OP] +} + +func (c *recordingS3DispatchCoordinator) Dispatch(_ context.Context, request *kv.OperationGroup[kv.OP]) (*kv.CoordinateResponse, error) { + c.request = request + return &kv.CoordinateResponse{}, nil +} diff --git a/adapter/s3_multipart_complete.go b/adapter/s3_multipart_complete.go index 7cded1e9a..3f9e08502 100644 --- a/adapter/s3_multipart_complete.go +++ b/adapter/s3_multipart_complete.go @@ -67,6 +67,11 @@ func validateS3MultipartCompletionParts(parts []s3CompleteMultipartUploadPart, b func (s *S3Server) loadS3MultipartCompletion(ctx context.Context, bucket, objectKey, uploadID string, request s3CompleteMultipartUploadRequest) (s3MultipartCompletion, error) { completion := s3MultipartCompletion{bucket: bucket, objectKey: objectKey, uploadID: uploadID} readTS := s.readTS() + readTimestamp, err := s.beginTxnReadTimestamp(ctx, readTS, "s3 complete multipart upload preparation: begin read timestamp") + if err != nil { + return completion, errors.WithStack(err) + } + readTS = readTimestamp.Timestamp() readPin := s.pinReadTS(readTS) defer readPin.Release() @@ -164,10 +169,12 @@ func (s *S3Server) commitS3MultipartCompletion(ctx context.Context, completion s func (s *S3Server) commitS3MultipartCompletionAttempt(ctx context.Context, completion s3MultipartCompletion) (*s3ObjectManifest, error) { readTS := s.readTS() - startTS, err := s.txnStartTS(ctx, readTS) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, readTS, "s3 complete multipart upload: begin read timestamp") if err != nil { return nil, errors.Wrap(err, "s3: allocate startTS for completeMultipartUpload retry") } + readTS = readTimestamp.Timestamp() + startTS := readTS readPin := s.pinReadTS(readTS) defer readPin.Release() @@ -200,7 +207,8 @@ func (s *S3Server) commitS3MultipartCompletionAttempt(ctx context.Context, compl if err != nil { return nil, errors.WithStack(err) } - _, err = s.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + _, err = kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: startTS, CommitTS: commitTS, diff --git a/adapter/s3_put_object.go b/adapter/s3_put_object.go index 3b53c2e83..1e65acae2 100644 --- a/adapter/s3_put_object.go +++ b/adapter/s3_put_object.go @@ -12,21 +12,24 @@ import ( ) type s3PutObjectState struct { - startTS uint64 - meta *s3BucketMeta - headKey []byte - previous *s3ObjectManifest - uploadID string - readPin *kv.ActiveTimestampToken + startTS uint64 + readTimestamp kv.ReadTimestamp + meta *s3BucketMeta + headKey []byte + previous *s3ObjectManifest + uploadID string + readPin *kv.ActiveTimestampToken } func (s *S3Server) prepareS3PutObject(ctx context.Context, request *http.Request, bucket, objectKey string) (*s3PutObjectState, error) { readTS := s.readTS() - startTS, err := s.txnStartTS(ctx, readTS) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, readTS, "s3 put object: begin read timestamp") if err != nil { return nil, errors.WithStack(err) } - state := &s3PutObjectState{startTS: startTS, readPin: s.pinReadTS(readTS)} + readTS = readTimestamp.Timestamp() + startTS := readTS + state := &s3PutObjectState{startTS: startTS, readTimestamp: readTimestamp, readPin: s.pinReadTS(readTS)} prepared := false defer func() { if !prepared { @@ -124,7 +127,8 @@ func (s *S3Server) commitS3PutObject(ctx context.Context, request *http.Request, &kv.Elem[kv.OP]{Op: kv.Put, Key: s3keys.BucketMetaKey(bucket), Value: bucketFence}, &kv.Elem[kv.OP]{Op: kv.Put, Key: state.headKey, Value: body}, ) - _, err = s.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ + dispatchCtx := state.readTimestamp.WithDispatchVoucher(ctx) + _, err = kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, &kv.OperationGroup[kv.OP]{ IsTxn: true, StartTS: state.startTS, CommitTS: commitTS, diff --git a/adapter/s3_upload_part.go b/adapter/s3_upload_part.go index 1a9e154b2..c547a5091 100644 --- a/adapter/s3_upload_part.go +++ b/adapter/s3_upload_part.go @@ -22,6 +22,7 @@ type s3UploadPartState struct { type s3UploadPartVersion struct { startTS uint64 + readTimestamp kv.ReadTimestamp commitTS uint64 chunkRefVersion uint64 } @@ -31,7 +32,11 @@ func (s *S3Server) prepareS3UploadPart(ctx context.Context, bucket, objectKey, u if err != nil { return nil, err } - state := &s3UploadPartState{partNo: partNo, readTS: s.readTS()} + readTimestamp, err := s.beginTxnReadTimestamp(ctx, s.readTS(), "s3 upload part preparation: begin read timestamp") + if err != nil { + return nil, errors.WithStack(err) + } + state := &s3UploadPartState{partNo: partNo, readTS: readTimestamp.Timestamp()} state.readPin = s.pinReadTS(state.readTS) prepared := false defer func() { @@ -110,11 +115,11 @@ func (s *S3Server) storeS3UploadPart(ctx context.Context, request *http.Request, } func (s *S3Server) allocateS3UploadPartVersionForMode(ctx context.Context, offloaded bool) (s3UploadPartVersion, error) { - startTS, commitTS, err := s.allocateS3UploadPartVersion(ctx) + readTimestamp, startTS, commitTS, err := s.allocateS3UploadPartVersion(ctx) if err != nil { return s3UploadPartVersion{}, err } - version := s3UploadPartVersion{startTS: startTS, commitTS: commitTS, chunkRefVersion: startTS} + version := s3UploadPartVersion{startTS: startTS, readTimestamp: readTimestamp, commitTS: commitTS, chunkRefVersion: startTS} if offloaded { // Reserve the initial commit timestamp as this attempt's immutable // chunkref namespace before delaying the descriptor commit timestamp. @@ -134,17 +139,19 @@ func (s *S3Server) finalizeS3UploadPartCommitTS(ctx context.Context, startTS, co return ts, errors.WithStack(err) } -func (s *S3Server) allocateS3UploadPartVersion(ctx context.Context) (uint64, uint64, error) { +func (s *S3Server) allocateS3UploadPartVersion(ctx context.Context) (kv.ReadTimestamp, uint64, uint64, error) { readTS := s.readTS() - startTS, err := s.txnStartTS(ctx, readTS) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, readTS, "s3 upload part: begin read timestamp") if err != nil { - return 0, 0, errors.WithStack(err) + return kv.ReadTimestamp{}, 0, 0, errors.WithStack(err) } + readTS = readTimestamp.Timestamp() + startTS := readTS commitTS, err := s.nextTxnCommitTS(ctx, startTS) if err != nil { - return 0, 0, errors.WithStack(err) + return kv.ReadTimestamp{}, 0, 0, errors.WithStack(err) } - return startTS, commitTS, nil + return readTimestamp, startTS, commitTS, nil } func (s *S3Server) commitS3UploadPart(ctx context.Context, state *s3UploadPartState, upload s3ChunkUploadResult, bucket, objectKey, uploadID string, version s3UploadPartVersion) (*s3PartDescriptor, error) { @@ -159,15 +166,19 @@ func (s *S3Server) commitS3UploadPart(ctx context.Context, state *s3UploadPartSt } partKey := s3keys.UploadPartKey(bucket, state.meta.Generation, objectKey, uploadID, state.partNo) previous := s.loadPreviousS3PartDescriptor(ctx, partKey, state.readTS) - if err := s.verifyS3UploadStillExists(ctx, state.uploadMetaKey, bucket, objectKey); err != nil { + if err := s.verifyS3UploadStillExists(ctx, state.uploadMetaKey, bucket, objectKey, s.readTS()); err != nil { return nil, err } elems := make([]*kv.Elem[kv.OP], 0, len(upload.ChunkRefElems)+1) elems = append(elems, upload.ChunkRefElems...) elems = append(elems, &kv.Elem[kv.OP]{Op: kv.Put, Key: partKey, Value: body}) - _, err = s.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{ - IsTxn: true, StartTS: version.startTS, CommitTS: version.commitTS, - Elems: elems, + dispatchCtx := version.readTimestamp.WithDispatchVoucher(ctx) + _, err = kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, &kv.OperationGroup[kv.OP]{ + IsTxn: true, + StartTS: version.startTS, + CommitTS: version.commitTS, + Elems: elems, + ReadKeys: [][]byte{state.uploadMetaKey}, }) if err != nil { return nil, errors.WithStack(err) @@ -187,8 +198,8 @@ func (s *S3Server) loadPreviousS3PartDescriptor(ctx context.Context, partKey []b return &descriptor } -func (s *S3Server) verifyS3UploadStillExists(ctx context.Context, uploadMetaKey []byte, bucket, objectKey string) error { - if _, err := s.store.GetAt(ctx, uploadMetaKey, s.readTS()); err != nil { +func (s *S3Server) verifyS3UploadStillExists(ctx context.Context, uploadMetaKey []byte, bucket, objectKey string, readTS uint64) error { + if _, err := s.store.GetAt(ctx, uploadMetaKey, readTS); err != nil { if errors.Is(err, store.ErrKeyNotFound) { return newS3ResponseError(http.StatusNotFound, "NoSuchUpload", "upload not found", bucket, objectKey) } diff --git a/adapter/sqs_catalog.go b/adapter/sqs_catalog.go index 6cb18d92e..435adb804 100644 --- a/adapter/sqs_catalog.go +++ b/adapter/sqs_catalog.go @@ -878,6 +878,11 @@ func (s *SQSServer) nextTxnReadTS(ctx context.Context) uint64 { return maxTS } +func (s *SQSServer) beginTxnReadTimestamp(ctx context.Context, label string) (kv.ReadTimestamp, error) { + readTimestamp, err := kv.BeginReadTimestampThrough(ctx, s.coordinator, s.nextTxnReadTS(ctx), label) + return readTimestamp, errors.WithStack(err) +} + func (s *SQSServer) loadQueueMetaAt(ctx context.Context, queueName string, ts uint64) (*sqsQueueMeta, bool, error) { b, err := s.store.GetAt(ctx, sqsQueueMetaKey(queueName), ts) if err != nil { @@ -983,11 +988,11 @@ func (s *SQSServer) createQueueWithRetry(ctx context.Context, requested *sqsQueu // with the requested attributes, false means the dispatch hit a retryable // conflict and should be retried after backoff. func (s *SQSServer) tryCreateQueueOnce(ctx context.Context, requested *sqsQueueMeta) (bool, error) { - readTS := s.nextTxnReadTS(ctx) - existing, exists, err := s.loadQueueMetaAt(ctx, requested.Name, readTS) + readTimestamp, existing, exists, err := s.createQueueSnapshot(ctx, requested.Name) if err != nil { return false, errors.WithStack(err) } + readTS := readTimestamp.Timestamp() if exists { if attributesEqual(existing, requested) { return true, nil @@ -1079,6 +1084,21 @@ func (s *SQSServer) tryCreateQueueOnce(ctx context.Context, requested *sqsQueueM return true, nil } +func (s *SQSServer) createQueueSnapshot( + ctx context.Context, + queueName string, +) (kv.ReadTimestamp, *sqsQueueMeta, bool, error) { + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs create queue: begin read timestamp") + if err != nil { + return kv.ReadTimestamp{}, nil, false, errors.WithStack(err) + } + existing, exists, err := s.loadQueueMetaAt(ctx, queueName, readTimestamp.Timestamp()) + if err != nil { + return kv.ReadTimestamp{}, nil, false, errors.WithStack(err) + } + return readTimestamp, existing, exists, nil +} + func (s *SQSServer) deleteQueue(w http.ResponseWriter, r *http.Request) { var in sqsDeleteQueueInput if err := decodeSQSJSONInput(r, &in); err != nil { @@ -1124,7 +1144,11 @@ func (s *SQSServer) deleteQueueWithRetry(ctx context.Context, queueName string) backoff := transactRetryInitialBackoff deadline := time.Now().Add(transactRetryMaxDuration) for range transactRetryMaxAttempts { - readTS := s.nextTxnReadTS(ctx) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs delete queue: begin read timestamp") + if err != nil { + return 0, errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() existing, exists, err := s.loadQueueMetaAt(ctx, queueName, readTS) if err != nil { return 0, errors.WithStack(err) @@ -1656,7 +1680,11 @@ func applyAndValidateSetAttributes(meta *sqsQueueMeta, attrs map[string]string) // successful Dispatch so post-commit request-path gauges are not removed by // the caller's cleanup. func (s *SQSServer) trySetQueueAttributesOnce(ctx context.Context, queueName string, attrs map[string]string) (bool, *sqsQueueThrottle, []string, uint64, bool, error) { - readTS := s.nextTxnReadTS(ctx) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs set queue attributes: begin read timestamp") + if err != nil { + return false, nil, nil, 0, false, errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() meta, exists, err := s.loadQueueMetaAt(ctx, queueName, readTS) if err != nil { return false, nil, nil, 0, false, errors.WithStack(err) diff --git a/adapter/sqs_fifo.go b/adapter/sqs_fifo.go index b4445adeb..d41b40ef0 100644 --- a/adapter/sqs_fifo.go +++ b/adapter/sqs_fifo.go @@ -174,8 +174,9 @@ func (s *SQSServer) sendFifoMessage( in sqsSendMessageInput, dedupID string, delay int64, - readTS uint64, + readTimestamp kv.ReadTimestamp, ) (map[string]string, bool, error) { + readTS := readTimestamp.Timestamp() // HT-FIFO: hash the MessageGroupId once at the entry point so // every key built in this transaction (data, vis, byage, dedup, // group-lock, sequence) lands in the same partition. partitionFor @@ -248,7 +249,8 @@ func (s *SQSServer) sendFifoMessage( {Op: kv.Put, Key: seqKey, Value: []byte(strconv.FormatUint(nextSeq, 10))}, }, } - if _, err := s.coordinator.Dispatch(ctx, req); err != nil { + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + if _, err := kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, req); err != nil { if isRetryableTransactWriteError(err) { return nil, true, nil } diff --git a/adapter/sqs_messages.go b/adapter/sqs_messages.go index 516d53d87..ca18825ec 100644 --- a/adapter/sqs_messages.go +++ b/adapter/sqs_messages.go @@ -631,7 +631,11 @@ func (s *SQSServer) sendMessageFifoLoop(w http.ResponseWriter, r *http.Request, // concurrent DeleteQueue / PurgeQueue could slip in between our read // and the write, storing a message under a dead generation. func (s *SQSServer) loadQueueMetaForSend(ctx context.Context, queueName string, body []byte) (*sqsQueueMeta, uint64, error) { - readTS := s.nextTxnReadTS(ctx) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs send message: begin read timestamp") + if err != nil { + return nil, 0, errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() meta, exists, err := s.loadQueueMetaAt(ctx, queueName, readTS) if err != nil { return nil, readTS, errors.WithStack(err) @@ -991,7 +995,12 @@ func (s *SQSServer) longPollReceive(ctx context.Context, queueName string, opts // elapses prevents false-empty returns under poison-message backlogs // or hot-FIFO-group fan-in. func (s *SQSServer) scanAndDeliverOnce(ctx context.Context, queueName string, opts sqsReceiveOptions) ([]map[string]any, error) { - readTS := s.nextTxnReadTS(ctx) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs receive message: begin read timestamp") + if err != nil { + return nil, errors.WithStack(err) + } + ctx = readTimestamp.WithDispatchVoucher(ctx) + readTS := readTimestamp.Timestamp() meta, exists, err := s.loadQueueMetaAt(ctx, queueName, readTS) if err != nil { return nil, errors.WithStack(err) @@ -1291,7 +1300,7 @@ func (s *SQSServer) expireMessage(ctx context.Context, queueName string, meta *s ReadKeys: readKeys, Elems: elems, } - if _, err := s.coordinator.Dispatch(ctx, req); err != nil { + if _, err := kv.DispatchWithReadTimestamp(ctx, s.coordinator, req); err != nil { if isRetryableTransactWriteError(err) { return nil } @@ -1440,7 +1449,7 @@ func (s *SQSServer) commitReceiveRotation(ctx context.Context, queueName string, if err != nil { return nil, false, err } - if _, err := s.coordinator.Dispatch(ctx, req); err != nil { + if _, err := kv.DispatchWithReadTimestamp(ctx, s.coordinator, req); err != nil { if isRetryableTransactWriteError(err) { return nil, true, nil } @@ -1613,18 +1622,20 @@ func (s *SQSServer) deleteMessageWithRetry(ctx context.Context, queueName string backoff := transactRetryInitialBackoff deadline := time.Now().Add(transactRetryMaxDuration) for range transactRetryMaxAttempts { - meta, rec, dataKey, readTS, outcome, err := s.loadMessageForDelete(ctx, queueName, handle) + meta, rec, dataKey, readTimestamp, outcome, err := s.loadMessageForDelete(ctx, queueName, handle) if err != nil { return err } if outcome == sqsDeleteNoOp { return nil } + readTS := readTimestamp.Timestamp() req, err := s.buildDeleteOps(ctx, queueName, meta, handle, rec, dataKey, readTS) if err != nil { return err } - if _, err := s.coordinator.Dispatch(ctx, req); err == nil { + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + if _, err := kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, req); err == nil { // Hot-partition observability (§11 PR 7): record the // successful delete on the partitioned commit branch // only. Legacy queues stay off the metric. @@ -1688,8 +1699,9 @@ const ( // outcome for AWS-compatible DeleteMessage semantics: structural errors // propagate; missing records and token mismatches on an otherwise-valid // queue return sqsDeleteNoOp; matching tokens return sqsDeleteProceed -// with the loaded record. The readTS it took the snapshot at is -// returned so the caller can pass it as StartTS on the OCC dispatch, +// with the loaded record. The ReadTimestamp it took the snapshot at is +// returned so the caller can pass it as StartTS on the OCC dispatch and +// bind any Phase-D applied-read voucher to that dispatch, // pinning the read-write conflict detection window. // // The caller-supplied QueueUrl is cross-checked against the handle's @@ -1698,37 +1710,41 @@ const ( // to a different (or recreated) queue and we reject it as a structural // error — silently succeeding would let misrouted deletes ack messages // that cannot possibly be deleted on this queue. -func (s *SQSServer) loadMessageForDelete(ctx context.Context, queueName string, handle *decodedReceiptHandle) (*sqsQueueMeta, *sqsMessageRecord, []byte, uint64, sqsDeleteOutcome, error) { - readTS := s.nextTxnReadTS(ctx) +func (s *SQSServer) loadMessageForDelete(ctx context.Context, queueName string, handle *decodedReceiptHandle) (*sqsQueueMeta, *sqsMessageRecord, []byte, kv.ReadTimestamp, sqsDeleteOutcome, error) { + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs delete message: begin read timestamp") + if err != nil { + return nil, nil, nil, kv.ReadTimestamp{}, sqsDeleteProceed, errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() meta, exists, err := s.loadQueueMetaAt(ctx, queueName, readTS) if err != nil { - return nil, nil, nil, readTS, sqsDeleteProceed, errors.WithStack(err) + return nil, nil, nil, readTimestamp, sqsDeleteProceed, errors.WithStack(err) } if !exists { - return nil, nil, nil, readTS, sqsDeleteProceed, newSQSAPIError(http.StatusBadRequest, sqsErrQueueDoesNotExist, "queue does not exist") + return nil, nil, nil, readTimestamp, sqsDeleteProceed, newSQSAPIError(http.StatusBadRequest, sqsErrQueueDoesNotExist, "queue does not exist") } if meta.Generation != handle.QueueGeneration { - return nil, nil, nil, readTS, sqsDeleteProceed, newSQSAPIError(http.StatusBadRequest, sqsErrReceiptHandleInvalid, "receipt handle does not belong to this queue") + return nil, nil, nil, readTimestamp, sqsDeleteProceed, newSQSAPIError(http.StatusBadRequest, sqsErrReceiptHandleInvalid, "receipt handle does not belong to this queue") } if err := validateReceiptHandleVersion(meta, handle); err != nil { - return nil, nil, nil, readTS, sqsDeleteProceed, newSQSAPIError(http.StatusBadRequest, sqsErrReceiptHandleInvalid, "receipt handle is not valid for this queue") + return nil, nil, nil, readTimestamp, sqsDeleteProceed, newSQSAPIError(http.StatusBadRequest, sqsErrReceiptHandleInvalid, "receipt handle is not valid for this queue") } dataKey := sqsMsgDataKeyDispatch(meta, queueName, handle.Partition, handle.QueueGeneration, handle.MessageIDHex) raw, err := s.store.GetAt(ctx, dataKey, readTS) if err != nil { if errors.Is(err, store.ErrKeyNotFound) { - return meta, nil, nil, readTS, sqsDeleteNoOp, nil + return meta, nil, nil, readTimestamp, sqsDeleteNoOp, nil } - return nil, nil, nil, readTS, sqsDeleteProceed, errors.WithStack(err) + return nil, nil, nil, readTimestamp, sqsDeleteProceed, errors.WithStack(err) } rec, err := decodeSQSMessageRecord(raw) if err != nil { - return nil, nil, nil, readTS, sqsDeleteProceed, errors.WithStack(err) + return nil, nil, nil, readTimestamp, sqsDeleteProceed, errors.WithStack(err) } if !bytes.Equal(rec.CurrentReceiptToken, handle.ReceiptToken) { - return meta, nil, nil, readTS, sqsDeleteNoOp, nil + return meta, nil, nil, readTimestamp, sqsDeleteNoOp, nil } - return meta, rec, dataKey, readTS, sqsDeleteProceed, nil + return meta, rec, dataKey, readTimestamp, sqsDeleteProceed, nil } func (s *SQSServer) changeMessageVisibility(w http.ResponseWriter, r *http.Request) { @@ -1782,10 +1798,11 @@ func (s *SQSServer) changeVisibilityWithRetry(ctx context.Context, queueName str backoff := transactRetryInitialBackoff deadline := time.Now().Add(transactRetryMaxDuration) for range transactRetryMaxAttempts { - meta, rec, dataKey, readTS, apiErr := s.loadAndVerifyMessage(ctx, queueName, handle) + meta, rec, dataKey, readTimestamp, apiErr := s.loadAndVerifyMessage(ctx, queueName, handle) if apiErr != nil { return apiErr } + readTS := readTimestamp.Timestamp() now := time.Now().UnixMilli() if rec.VisibleAtMillis <= now { return newSQSAPIError(http.StatusBadRequest, sqsErrMessageNotInflight, "message is not currently in flight") @@ -1813,7 +1830,8 @@ func (s *SQSServer) changeVisibilityWithRetry(ctx context.Context, queueName str {Op: kv.Put, Key: dataKey, Value: recordBytes}, }, } - if _, err := s.coordinator.Dispatch(ctx, req); err == nil { + dispatchCtx := readTimestamp.WithDispatchVoucher(ctx) + if _, err := kv.DispatchWithReadTimestamp(dispatchCtx, s.coordinator, req); err == nil { return nil } else if !isRetryableTransactWriteError(err) { return errors.WithStack(err) @@ -1842,9 +1860,10 @@ func (s *SQSServer) parseQueueAndReceipt(queueUrl, receiptHandle string) (string // loadAndVerifyMessage reads the data record for the given handle and // verifies that the receipt token matches the current one on record. -// Returns the record, its key, the snapshot timestamp the read ran at, -// or a typed SQS error. Callers use the snapshot as StartTS on the -// OCC dispatch so concurrent commits cannot slip past ReadKeys. +// Returns the record, its key, the ReadTimestamp the read ran at, or a +// typed SQS error. Callers use the snapshot as StartTS on the OCC +// dispatch and bind its Phase-D applied-read voucher so concurrent commits +// cannot slip past ReadKeys. // // The caller-supplied QueueUrl is cross-checked against the handle's // embedded queue_generation, mirroring loadMessageForDelete: an @@ -1852,37 +1871,41 @@ func (s *SQSServer) parseQueueAndReceipt(queueUrl, receiptHandle string) (string // cleans them up, so a handle from a deleted / recreated queue must // be rejected with ReceiptHandleIsInvalid instead of silently // mutating the orphan record. -func (s *SQSServer) loadAndVerifyMessage(ctx context.Context, queueName string, handle *decodedReceiptHandle) (*sqsQueueMeta, *sqsMessageRecord, []byte, uint64, error) { - readTS := s.nextTxnReadTS(ctx) +func (s *SQSServer) loadAndVerifyMessage(ctx context.Context, queueName string, handle *decodedReceiptHandle) (*sqsQueueMeta, *sqsMessageRecord, []byte, kv.ReadTimestamp, error) { + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs change message visibility: begin read timestamp") + if err != nil { + return nil, nil, nil, kv.ReadTimestamp{}, errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() meta, exists, err := s.loadQueueMetaAt(ctx, queueName, readTS) if err != nil { - return nil, nil, nil, readTS, errors.WithStack(err) + return nil, nil, nil, readTimestamp, errors.WithStack(err) } if !exists { - return nil, nil, nil, readTS, newSQSAPIError(http.StatusBadRequest, sqsErrQueueDoesNotExist, "queue does not exist") + return nil, nil, nil, readTimestamp, newSQSAPIError(http.StatusBadRequest, sqsErrQueueDoesNotExist, "queue does not exist") } if meta.Generation != handle.QueueGeneration { - return nil, nil, nil, readTS, newSQSAPIError(http.StatusBadRequest, sqsErrReceiptHandleInvalid, "receipt handle does not belong to this queue") + return nil, nil, nil, readTimestamp, newSQSAPIError(http.StatusBadRequest, sqsErrReceiptHandleInvalid, "receipt handle does not belong to this queue") } if err := validateReceiptHandleVersion(meta, handle); err != nil { - return nil, nil, nil, readTS, newSQSAPIError(http.StatusBadRequest, sqsErrReceiptHandleInvalid, "receipt handle is not valid for this queue") + return nil, nil, nil, readTimestamp, newSQSAPIError(http.StatusBadRequest, sqsErrReceiptHandleInvalid, "receipt handle is not valid for this queue") } dataKey := sqsMsgDataKeyDispatch(meta, queueName, handle.Partition, handle.QueueGeneration, handle.MessageIDHex) raw, err := s.store.GetAt(ctx, dataKey, readTS) if err != nil { if errors.Is(err, store.ErrKeyNotFound) { - return nil, nil, nil, readTS, newSQSAPIError(http.StatusBadRequest, sqsErrReceiptHandleInvalid, "message not found") + return nil, nil, nil, readTimestamp, newSQSAPIError(http.StatusBadRequest, sqsErrReceiptHandleInvalid, "message not found") } - return nil, nil, nil, readTS, errors.WithStack(err) + return nil, nil, nil, readTimestamp, errors.WithStack(err) } rec, err := decodeSQSMessageRecord(raw) if err != nil { - return nil, nil, nil, readTS, errors.WithStack(err) + return nil, nil, nil, readTimestamp, errors.WithStack(err) } if !bytes.Equal(rec.CurrentReceiptToken, handle.ReceiptToken) { - return nil, nil, nil, readTS, newSQSAPIError(http.StatusBadRequest, sqsErrInvalidReceiptHandle, "receipt handle token does not match") + return nil, nil, nil, readTimestamp, newSQSAPIError(http.StatusBadRequest, sqsErrInvalidReceiptHandle, "receipt handle token does not match") } - return meta, rec, dataKey, readTS, nil + return meta, rec, dataKey, readTimestamp, nil } // ------------------------ small helpers ------------------------ diff --git a/adapter/sqs_messages_batch.go b/adapter/sqs_messages_batch.go index ea2de434b..4ca825682 100644 --- a/adapter/sqs_messages_batch.go +++ b/adapter/sqs_messages_batch.go @@ -181,7 +181,11 @@ func (s *SQSServer) trySendMessageBatchOnce( entries []sqsSendMessageBatchEntryInput, identities []sqsSendIdentity, ) ([]sqsSendMessageBatchResultEntry, []sqsBatchResultErrorEntry, bool, error) { - readTS := s.nextTxnReadTS(ctx) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs send message batch: begin read timestamp") + if err != nil { + return nil, nil, false, errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() meta, exists, err := s.loadQueueMetaAt(ctx, queueName, readTS) if err != nil { return nil, nil, false, errors.WithStack(err) @@ -367,12 +371,16 @@ func (s *SQSServer) runFifoSendWithRetry( backoff := transactRetryInitialBackoff deadline := time.Now().Add(transactRetryMaxDuration) for range transactRetryMaxAttempts { - readTS := s.nextTxnReadTS(ctx) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs fifo send: begin read timestamp") + if err != nil { + return nil, errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() meta, dedupID, delay, err := s.resolveFreshFifoSnapshot(ctx, queueName, in, readTS) if err != nil { return nil, err } - resp, retry, err := s.sendFifoMessage(ctx, queueName, meta, in, dedupID, delay, readTS) + resp, retry, err := s.sendFifoMessage(ctx, queueName, meta, in, dedupID, delay, readTimestamp) if err != nil { return nil, err } diff --git a/adapter/sqs_purge.go b/adapter/sqs_purge.go index 10dc0a06d..f2e993ab8 100644 --- a/adapter/sqs_purge.go +++ b/adapter/sqs_purge.go @@ -91,7 +91,11 @@ func (s *SQSServer) purgeQueueWithRetry(ctx context.Context, queueName string) ( // pre-bump and post-bump generations so the caller can audit-log the // committed value without a second meta read. func (s *SQSServer) tryPurgeQueueOnce(ctx context.Context, queueName string) (bool, uint64, uint64, error) { - readTS := s.nextTxnReadTS(ctx) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs purge queue: begin read timestamp") + if err != nil { + return false, 0, 0, errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() meta, exists, err := s.loadQueueMetaAt(ctx, queueName, readTS) if err != nil { return false, 0, 0, errors.WithStack(err) diff --git a/adapter/sqs_reaper.go b/adapter/sqs_reaper.go index ae7f6c20d..5c5d43d06 100644 --- a/adapter/sqs_reaper.go +++ b/adapter/sqs_reaper.go @@ -73,8 +73,13 @@ func (s *SQSServer) reapAllQueues(ctx context.Context) error { if err := ctx.Err(); err != nil { return errors.WithStack(err) } - readTS := s.nextTxnReadTS(ctx) - meta, exists, err := s.loadQueueMetaAt(ctx, name, readTS) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs reaper queue: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } + queueCtx := readTimestamp.WithDispatchVoucher(ctx) + readTS := readTimestamp.Timestamp() + meta, exists, err := s.loadQueueMetaAt(queueCtx, name, readTS) if err != nil || !exists { // Even when meta is gone (DeleteQueue), prior-generation // orphans need reaping; reapTombstonedQueues (called @@ -82,10 +87,10 @@ func (s *SQSServer) reapAllQueues(ctx context.Context) error { // the queue if loading itself failed (transient). continue } - if err := s.reapQueue(ctx, name, meta, readTS); err != nil { + if err := s.reapQueue(queueCtx, name, meta, readTS); err != nil { slog.Warn("sqs reaper queue pass failed", "queue", name, "err", err) } - if err := s.reapExpiredDedup(ctx, name, meta, readTS); err != nil { + if err := s.reapExpiredDedup(queueCtx, name, meta, readTS); err != nil { slog.Warn("sqs dedup reaper pass failed", "queue", name, "err", err) } } @@ -108,8 +113,13 @@ func (s *SQSServer) reapTombstonedQueues(ctx context.Context) error { upper := prefixScanEnd(prefix) start := bytes.Clone(prefix) for { - readTS := s.nextTxnReadTS(ctx) - page, err := s.store.ScanAt(ctx, start, upper, sqsReaperPageLimit, readTS) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs tombstone reaper: begin read timestamp") + if err != nil { + return errors.WithStack(err) + } + pageCtx := readTimestamp.WithDispatchVoucher(ctx) + readTS := readTimestamp.Timestamp() + page, err := s.store.ScanAt(pageCtx, start, upper, sqsReaperPageLimit, readTS) if err != nil { return errors.WithStack(err) } @@ -129,7 +139,7 @@ func (s *SQSServer) reapTombstonedQueues(ctx context.Context) error { // non-canonical values to 1 so pre-PR-6a tombstones // retain their byte-identical legacy reaper path. partitionCount := decodeQueueTombstoneValue(kvp.Value) - s.reapTombstonedGeneration(ctx, queueName, gen, partitionCount, kvp.Key, readTS) + s.reapTombstonedGeneration(pageCtx, queueName, gen, partitionCount, kvp.Key, readTS) } if len(page) < sqsReaperPageLimit { return nil @@ -663,7 +673,7 @@ func (s *SQSServer) reapOneRecord(ctx context.Context, queueName string, meta *s if err != nil { return err } - if _, err := s.coordinator.Dispatch(ctx, req); err != nil { + if _, err := kv.DispatchWithReadTimestamp(ctx, s.coordinator, req); err != nil { if isRetryableTransactWriteError(err) { return nil } @@ -717,7 +727,7 @@ func (s *SQSServer) dispatchOrphanByAgeDrop(ctx context.Context, byAgeKey []byte {Op: kv.Del, Key: byAgeKey}, }, } - _, _ = s.coordinator.Dispatch(ctx, req) + _, _ = kv.DispatchWithReadTimestamp(ctx, s.coordinator, req) } // reapExpiredDedup walks every FIFO dedup record under the given @@ -872,7 +882,7 @@ func (s *SQSServer) dispatchDedupDelete(ctx context.Context, key []byte, readTS {Op: kv.Del, Key: key}, }, } - if _, err := s.coordinator.Dispatch(ctx, req); err != nil { + if _, err := kv.DispatchWithReadTimestamp(ctx, s.coordinator, req); err != nil { if isRetryableTransactWriteError(err) { return nil } diff --git a/adapter/sqs_redrive.go b/adapter/sqs_redrive.go index 3fb13b8d3..c87547a55 100644 --- a/adapter/sqs_redrive.go +++ b/adapter/sqs_redrive.go @@ -236,7 +236,7 @@ func (s *SQSServer) redriveCandidateToDLQ( if err != nil { return false, err } - if _, err := s.coordinator.Dispatch(ctx, req); err != nil { + if _, err := kv.DispatchWithReadTimestamp(ctx, s.coordinator, req); err != nil { if isRetryableTransactWriteError(err) { return true, nil } diff --git a/adapter/sqs_tags.go b/adapter/sqs_tags.go index 7635b89f0..eec8b74f2 100644 --- a/adapter/sqs_tags.go +++ b/adapter/sqs_tags.go @@ -164,7 +164,11 @@ func (s *SQSServer) tryMutateQueueTagsOnce( queueName string, mutate func(*sqsQueueMeta) error, ) (bool, error) { - readTS := s.nextTxnReadTS(ctx) + readTimestamp, err := s.beginTxnReadTimestamp(ctx, "sqs mutate queue tags: begin read timestamp") + if err != nil { + return false, errors.WithStack(err) + } + readTS := readTimestamp.Timestamp() meta, exists, err := s.loadQueueMetaAt(ctx, queueName, readTS) if err != nil { return false, errors.WithStack(err) diff --git a/distribution/catalog.go b/distribution/catalog.go index a27734c1f..df22a64c9 100644 --- a/distribution/catalog.go +++ b/distribution/catalog.go @@ -240,6 +240,16 @@ func (s *CatalogStore) Snapshot(ctx context.Context) (CatalogSnapshot, error) { return s.SnapshotAt(ctx, s.store.LastCommitTS()) } +// LatestCommitTS returns the local catalog store watermark without reading +// catalog data. Callers use it as the pre-Phase-D legacy timestamp input before +// selecting the transaction snapshot through the dedicated TSO. +func (s *CatalogStore) LatestCommitTS() uint64 { + if s == nil || s.store == nil { + return 0 + } + return s.store.LastCommitTS() +} + // SnapshotAt reads a consistent route catalog snapshot at a specific MVCC // timestamp. func (s *CatalogStore) SnapshotAt(ctx context.Context, ts uint64) (CatalogSnapshot, error) { diff --git a/distribution/catalog_test.go b/distribution/catalog_test.go index a3ffe852f..33b59e8c3 100644 --- a/distribution/catalog_test.go +++ b/distribution/catalog_test.go @@ -8,6 +8,7 @@ import ( "github.com/bootjp/elastickv/store" "github.com/cockroachdb/errors" + "github.com/stretchr/testify/require" ) func TestCatalogVersionCodecRoundTrip(t *testing.T) { @@ -345,6 +346,33 @@ func TestCatalogStoreSaveAndSnapshot(t *testing.T) { assertRouteEqual(t, saved.Routes[1], snapshot.Routes[1]) } +func TestCatalogStoreSnapshotAtPreservesRequestedVersion(t *testing.T) { + cs := NewCatalogStore(store.NewMVCCStore()) + ctx := context.Background() + first, err := cs.Save(ctx, 0, []RouteDescriptor{{ + RouteID: 1, + Start: []byte(""), + GroupID: 1, + State: RouteStateActive, + }}) + require.NoError(t, err) + firstTS := cs.LatestCommitTS() + _, err = cs.Save(ctx, first.Version, []RouteDescriptor{{ + RouteID: 2, + Start: []byte(""), + GroupID: 2, + State: RouteStateActive, + }}) + require.NoError(t, err) + + snapshot, err := cs.SnapshotAt(ctx, firstTS) + require.NoError(t, err) + require.Equal(t, first.Version, snapshot.Version) + require.Equal(t, firstTS, snapshot.ReadTS) + require.Equal(t, uint64(1), snapshot.Routes[0].RouteID) + require.GreaterOrEqual(t, cs.LatestCommitTS(), firstTS) +} + func TestCatalogStoreSaveAndSnapshotSortsRoutesByStart(t *testing.T) { cs := NewCatalogStore(store.NewMVCCStore()) ctx := context.Background() diff --git a/docs/design/2026_04_16_partial_centralized_tso.md b/docs/design/2026_04_16_partial_centralized_tso.md index 58ea65acb..d0247099c 100644 --- a/docs/design/2026_04_16_partial_centralized_tso.md +++ b/docs/design/2026_04_16_partial_centralized_tso.md @@ -1,12 +1,13 @@ # Centralized Timestamp Oracle (TSO) Design -- Status: Partial — M1 all-led-group HLC renewal, the M2 reserved-group - bootstrap bridge, the M3-M5 TSO allocator/batch cutover, and the minimal - TSO-only FSM are implemented; group-0 timestamp issuance, follower - redirect/admin exposure, and shadow validation remain open +- Status: Partial — M1-M7 are implemented, including the dedicated group-0 + FSM, leader-routed durable windows, strict term bootstrap, serialized shadow + migration, one-way rolling cutover, durable Phase-D retirement, and + cross-shard SSI timestamp validation. Runtime config reload and production + latency/alerting work remain open. - Author: bootjp - Date: 2026-04-16 -- Updated: 2026-07-23 +- Updated: 2026-07-24 --- @@ -41,25 +42,109 @@ 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. -9. `TSOStateMachine` implements a TSO-only Raft FSM that accepts HLC lease - entries, snapshots exactly the physical ceiling as eight big-endian bytes, - restores with malformed-input rejection, advances the HLC handoff floor - through the committed ceiling, and classifies HLC lease entries as - volatile-only for safe cold-start replay. This removes the requirement for - group 0 to carry the KV FSM's data store once startup wiring switches to - the dedicated FSM. + 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 + full ceiling/floor mirror entries as volatile-only so post-snapshot HLC + mirror raises are not lost on duplicate replay skip paths. +10. Runtime bootstrap instantiates `TSOStateMachine` for `groupID = 0` without + opening an MVCC store. Group 0 retains the normal wrap-aware proposer path + for Raft-envelope cutovers, while route validation keeps all user data out + of the control group. Restore accepts snapshots written by the earlier + compatibility `kvFSM`, imports their HLC ceiling, derives a safe allocation + floor, and drains the legacy MVCC payload for full CRC verification. The + group-0 `EncryptionAdmin` surface is capability-only: mutators are left + unwired and return `FailedPrecondition`. During upgrade replay, valid + encryption control entries committed by the earlier compatibility FSM are + decoded and deterministically rejected without halting the TSO apply loop; + malformed control entries still halt fail-closed. +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. +15. `--tsoPhaseDEnabled` commits a second one-way group-0 marker after cutover. + The marker captures the maximum pre-Phase-D timestamp floor. Its state and + floor are persisted in the V4 TSO snapshot. Before the marker applies, the + FSM continues writing V3 snapshots so older followers can install them + during the rolling upgrade. Malformed ordering or a changed replay floor + halts the TSO FSM fail-closed. The marker has its own versioned TSO envelope + with the same legacy fail-closed prefix, so old or misrouted data FSMs halt + and no reserved encryption byte can activate Phase D. +16. Once Phase D is durable, data groups no longer receive HLC ceiling renewal, + coordinator legacy HLC issuance is rejected, and shadow mode bypasses + legacy candidate generation/comparison. Group 0 remains renewed while it + is needed for the dedicated allocator's physical-ceiling fence. +17. Cross-shard transactions with a caller-supplied `StartTS` are accepted only + when the group-0 leader verifies `phase_d_floor < StartTS <= + allocation_floor`. Phase-D reservations are contiguous from the previous + committed floor, so that interval contains only group-0-issued timestamps; + a `min_timestamp` above the floor is rejected instead of opening an + unallocated gap. The upper bound is a committed TSO reservation floor and the + lower bound excludes every legacy or pre-Phase-D value. Follower-local state + is never authoritative for this check. +18. Adapter read-modify-write paths call `BeginReadTimestampThrough` before the + first read and pass an applied store/catalog watermark, never a freshly + allocated clock value. If Phase D is required but not yet locally active, + the read boundary first forces the marker and its post-marker allocation + window to commit, then discards that allocation. The adapter still uses the + exact applied watermark for every read and `StartTS`. When that watermark + predates the Phase-D floor, the read boundary registers a bounded, + process-local capability that identifies the audited applied snapshot. The + adapter binds that capability to the logical operation and reserves exactly + one one-use coordinator voucher immediately before each dispatch that + shares the snapshot. Unvouched external `StartTS` values remain subject to + group-0 validation and values at/below the floor fail closed. A zero, + latest-sentinel, unavailable activation, or otherwise unprovable watermark + also fails closed. Retries reload and revalidate the applied watermark; + retries that intentionally reuse an exact OCC write set reserve a fresh use + without widening the capability to another timestamp. Coordinator + decorators forward both the allocator provider and voucher operation so + startup and keyviz wrappers cannot silently fall back to an unvalidated + value. +19. `BatchAllocator` validates cached candidates once Phase D is required. A + candidate at/below the Phase-D floor invalidates the entire local window and + forces a refill above the marker before any timestamp is returned. Remaining: -1. Wire group 0 startup to `TSOStateMachine` instead of the compatibility KV - FSM and pin production timestamp issuance to the group-0 leader. -2. Add follower redirect/admin exposure for the dedicated TSO leader. -3. Add Phase B shadow-read validation before making dedicated TSO the only - production timestamp path. +1. Runtime config reload for the mode switch; current flags are startup-only. +2. Production benchmark, divergence metrics, and alert thresholds. ### 1.1 Original Limitation @@ -77,10 +162,10 @@ Node B: not in defaultGroup → ceiling never updated ❌ → may collide with the previous leader's committed window ``` -M1 fixes that per-node renewal gap by proposing to every group this node -currently leads. **Global timestamp monotonicity is still not guaranteed when -different coordinators on different nodes allocate timestamps for cross-group -work**; that is the dedicated TSO / single-oracle work left open by M2-M7. +M1 fixed that per-node renewal gap by proposing to every group this node +currently leads. It did not by itself guarantee global timestamp monotonicity +when different coordinators allocated timestamps for cross-group work; M2-M7 +add the dedicated TSO and retire that distributed issuance path. ### 1.2 Near-Term Workaround @@ -191,7 +276,10 @@ Node-3 ──┘ │ - `groupID = 0` is reserved for the TSO group; all user-data shards use `groupID >= 1`. -- The TSO FSM applies only HLC lease entries. It carries no key-value storage, +- The TSO FSM applies HLC lease entries and explicit allocation-floor entries. + Allocation floors use a versioned TSO envelope whose first byte remains in + the old data-FSM fail-closed range; the full magic prevents an encryption + opcode from being interpreted as TSO state. It carries no key-value storage, so compaction and snapshot overhead are negligible. - Membership mirrors the full cluster so any node can become TSO leader. @@ -199,62 +287,108 @@ Node-3 ──┘ │ ```go // TSOStateMachine implements raftengine.StateMachine. It is a minimal FSM -// that only tracks the HLC physical ceiling; its entire persisted state is a -// single int64 ceiling value. +// that tracks TSO-owned HLC physical ceiling and allocation floor state. The +// shared HLC is a volatile mirror only; snapshots are sourced from the FSM's +// own fields so data-group lease renewals cannot advance group-0 state outside +// group-0 consensus. // // Interface mapping (raftengine.StateMachine vs raw raft.FSM): // Apply([]byte) ← engine strips log envelope, passes raw payload // Snapshot() (Snapshot, _) ← returns raftengine.Snapshot (WriteTo/Close) // Restore(io.Reader) ← caller owns and closes the reader type TSOStateMachine struct { - hlc *HLC + hlc *HLC + ceilingMs atomic.Int64 + allocationFloor atomic.Uint64 } -// Apply processes an HLC lease entry. The engine strips the raft.Log +// Apply processes TSO ceiling/floor entries. The engine strips the raft.Log // envelope and passes only the raw command bytes, matching the // raftengine.StateMachine.Apply([]byte) signature. func (f *TSOStateMachine) Apply(data []byte) any { - if len(data) == 0 || data[0] != raftEncodeHLCLease { - return nil + if len(data) == 0 { + return haltErr(errors.Wrap(ErrTSOStateMachineInvalidEntry, + "empty entry")) } - return f.applyHLCLease(data[1:]) + switch { + case data[0] == raftEncodeHLCLease: + if len(data) != hlcLeaseEntryLen { + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "expected %d bytes, got %d", hlcLeaseEntryLen, len(data))) + } + ceiling := int64(binary.BigEndian.Uint64(data[1:])) + if ceiling <= 0 { + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "non-positive HLC lease ceiling %d", ceiling)) + } + f.applyLeaseCeiling(ceiling) + case bytes.HasPrefix(data, []byte(tsoAllocationFloorEnvelope)): + if len(data) != len(tsoAllocationFloorEnvelope)+8 { + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "invalid allocation-floor envelope length %d", len(data))) + } + floor := binary.BigEndian.Uint64(data[len(tsoAllocationFloorEnvelope):]) + if floor == 0 { + return haltErr(errors.Wrap(ErrTSOStateMachineInvalidEntry, + "zero TSO allocation floor")) + } + f.applyAllocationFloor(floor) + default: + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "unexpected tag 0x%02x", data[0])) + } + return nil } -// Snapshot serialises the ceiling as 8 big-endian bytes. +// Snapshot serialises ceiling and allocationFloor as two 8-byte big-endian +// fields. The source is TSO-applied state, not the shared HLC mirror. // Returns raftengine.Snapshot (WriteTo/Close) rather than raft.FSMSnapshot. func (f *TSOStateMachine) Snapshot() (raftengine.Snapshot, error) { - var ceiling int64 - if f.hlc != nil { - ceiling = f.hlc.PhysicalCeiling() - } - return &tsoSnapshot{ceiling: ceiling}, nil + return &tsoSnapshot{ + ceiling: f.ceilingMs.Load(), + allocationFloor: f.allocationFloor.Load(), + }, nil } -// Restore deserialises the ceiling and updates the shared HLC. +// Restore deserialises ceiling/floor state and mirrors it to the shared HLC. +// Current snapshots are 16 bytes; legacy 8-byte ceiling snapshots derive the +// allocation floor from the ceiling used by the previous FSM format. // The caller owns the reader and is responsible for closing it. func (f *TSOStateMachine) Restore(r io.Reader) error { - var buf [8]byte - if _, err := io.ReadFull(r, buf[:]); err != nil { + payload, err := io.ReadAll(io.LimitReader(r, 17)) + if err != nil { return err } - ceiling := int64(binary.BigEndian.Uint64(buf[:])) - if f.hlc != nil && ceiling > 0 { - f.hlc.SetPhysicalCeiling(ceiling) + var ceiling int64 + var floor uint64 + switch len(payload) { + case 8: + ceiling = int64(binary.BigEndian.Uint64(payload[:8])) + floor = tsoLeaseAllocationFloor(ceiling) + case 16: + ceiling = int64(binary.BigEndian.Uint64(payload[:8])) + floor = binary.BigEndian.Uint64(payload[8:]) + default: + return errors.Newf("expected 8 or 16 snapshot bytes, got %d", + len(payload)) } + f.restoreSnapshotState(ceiling, floor) return nil } -// tsoSnapshot implements raftengine.Snapshot. It serialises the HLC physical -// ceiling as 8 big-endian bytes — the entire persisted state of the TSO group. +// tsoSnapshot implements raftengine.Snapshot. It serialises TSO-owned ceiling +// and floor state as 16 bytes — the entire persisted state of the TSO group. // WriteTo/Close matches the raftengine.Snapshot interface (not Persist/Release // from raft.FSMSnapshot). type tsoSnapshot struct { - ceiling int64 + ceiling int64 + allocationFloor uint64 } func (s *tsoSnapshot) WriteTo(w io.Writer) (int64, error) { - var buf [8]byte + var buf [16]byte binary.BigEndian.PutUint64(buf[:], uint64(s.ceiling)) + binary.BigEndian.PutUint64(buf[8:], s.allocationFloor) n, err := w.Write(buf[:]) return int64(n), err } @@ -476,10 +610,11 @@ New TSO leader elected via Raft ├─ FSM.Restore() or Raft log replay │ → physicalCeiling restored to the last committed value │ - └─ HLC.NextBatchFenced() rejects the restored exhausted ceiling; group-0 - timestamp serving stays deferred until it can publish a fresh usable - lease window before allocation - → new leader does not reuse timestamps from the old leader's window ✅ + ├─ 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 ✅ ``` --- @@ -579,7 +714,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) ``` ┌──────────────────────────────────────────────────────────┐ @@ -588,50 +723,104 @@ 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 -- Remove per-shard ceiling proposals. -- Remove `legacyHLC` from `ShardedCoordinator`. -- Remove the shadow comparison code. +- Roll every member to a binary that understands the Phase-D entry and + `ValidateTimestamp` RPC. Run all members on the dedicated TSO path before + enabling `--tsoPhaseDEnabled`; an older member would correctly halt on the + unknown control entry, so mixed-version activation is prohibited. +- The first allocation with `--tsoPhaseDEnabled` commits cutover (if needed), + then the Phase-D marker and its pre-Phase-D floor, then a new allocation + window strictly above that floor. The switch is one-way and survives restart + through the TSO V4 snapshot. +- From that marker onward, group 0 reserves contiguous windows from the previous + committed allocation floor even if wall time or the local HLC mirror has + advanced. Requests carrying a `min_timestamp` above the committed floor fail + closed; callers must first validate or use an already issued timestamp rather + than forcing the allocator to bless a gap. +- `ShardedCoordinator` dynamically stops data-group ceiling proposals and + fails closed instead of issuing from its legacy HLC. It continues group-0 + renewal only. Shadow mode observes durable cutover and directly uses the + dedicated allocator without generating or comparing a legacy candidate. +- Every adapter read-modify-write attempt obtains its transaction snapshot + from an applied store/catalog watermark before reading. If Phase D is required + but not yet active locally, the read boundary first commits the marker and a + post-marker allocation window, then discards the allocated value; it is never + substituted for data-group apply progress. The same applied watermark is used + for direct MVCC reads, active-snapshot pinning, and `OperationGroup.StartTS`; + a retry reloads and revalidates the watermark and repeats all reads. An applied + watermark at/below the Phase-D floor is admitted only through a bounded, + process-local capability registered by that audited read boundary. Every + dispatch sharing the snapshot reserves and consumes its own one-use voucher; + the capability is timestamp-bound and cannot authorize another value. + Arbitrary or repeated unvouched caller values at/below the floor still fail + closed at group 0. Pre-Phase-D deployments retain the prior + committed-watermark behavior. +- After Phase D, the coordinator asks the current group-0 leader to validate a + caller-supplied cross-shard `StartTS` before allocating `CommitTS` or proposing + to any data group. The allocator's configured `PhaseDRequired` state activates + this check before the local group-0 replica has applied the marker. Values + at/below the marker floor, beyond the committed TSO allocation floor, inside a + rejected allocation gap, zero values, inactive Phase-D state, missing + validation support, and unavailable leadership all fail closed. Existing + single-shard caller-supplied `StartTS` remains compatible because it does not + establish a cross-shard SSI snapshot. ### 7.5 Monotonicity Invariant Across Phases -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. --- @@ -640,12 +829,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 and warmed by the HLC renewal bridge; minimal `TSOStateMachine` is implemented; TSO-leader-only timestamp issuance remains 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 — shipped | Commit the durable Phase-D floor marker, preserve V3 snapshots until activation, retire data-shard HLC renewal and legacy/shadow issuance after cutover, activate before read validation while preserving exact applied snapshots through timestamp-bound per-dispatch vouchers, invalidate pre-Phase-D batch windows, and validate unvouched caller-supplied cross-shard SSI timestamps at the group-0 leader from activation onward. | Low | --- @@ -661,10 +850,19 @@ 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? - -5. **Backward compatibility:** Existing snapshots and Raft logs store ceiling - values inline in shard FSMs. Migration to a dedicated TSO FSM requires a - coordinated snapshot + compaction pass. +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` + recognizes legacy `kvFSM` snapshot headers, imports the ceiling, derives the + allocation floor, and drains the old store payload before Raft resumes log + replay. Headerless legacy store payloads that exceed every TSO snapshot + length are also drained instead of being parsed as raw TSO state. Valid + historical encryption control entries are decoded and rejected as obsolete + ordinary apply responses, so replay advances without mutating TSO state; new + group-0 encryption mutators are disabled at the RPC boundary, while the + group-0 engine remains wired as a leader view so read-only sidecar recovery + cannot be served from a stale follower. Runtime wiring therefore + does not require deleting group-0 state. diff --git a/kv/coordinator.go b/kv/coordinator.go index 15d5268e1..e7be835fb 100644 --- a/kv/coordinator.go +++ b/kv/coordinator.go @@ -78,6 +78,15 @@ func WithTSOAllocator(alloc TimestampAllocator) CoordinatorOption { } } +// TimestampAllocator exposes the configured allocator to coordinator +// decorators without widening the Coordinator interface. +func (c *Coordinate) TimestampAllocator() TimestampAllocator { + if c == nil { + return nil + } + return c.tsAllocator +} + // LeaseReadObserver records lease-read fast-path vs slow-path outcomes // without coupling kv to a concrete monitoring backend. It is called once // per LeaseRead invocation that actually evaluates the lease (the initial @@ -241,6 +250,17 @@ type Coordinate struct { } var _ Coordinator = (*Coordinate)(nil) +var _ AppliedReadTimestampVoucher = (*Coordinate)(nil) +var _ AppliedReadTimestampVoucherRevoker = (*Coordinate)(nil) + +// VouchAppliedReadTimestamp is a no-op for the single-group coordinator. The +// sharded coordinator consumes vouchers before cross-group StartTS validation. +func (c *Coordinate) VouchAppliedReadTimestamp(uint64, AppliedReadTimestampVoucherRef) error { + return nil +} + +// RevokeAppliedReadTimestamp is a no-op for the single-group coordinator. +func (c *Coordinate) RevokeAppliedReadTimestamp(uint64, AppliedReadTimestampVoucherRef) {} type Coordinator interface { Dispatch(ctx context.Context, reqs *OperationGroup[OP]) (*CoordinateResponse, error) diff --git a/kv/keyviz_label.go b/kv/keyviz_label.go index 80d727638..3bcabab95 100644 --- a/kv/keyviz_label.go +++ b/kv/keyviz_label.go @@ -90,6 +90,11 @@ func (c keyVizLabeledCoordinator) RaftLeaderForKey(key []byte) string { func (c keyVizLabeledCoordinator) Clock() *HLC { return c.inner.Clock() } +func (c keyVizLabeledCoordinator) TimestampAllocator() TimestampAllocator { + alloc, _ := TimestampAllocatorThrough(c.inner) + return alloc +} + func (c keyVizLabeledCoordinator) Next(ctx context.Context) (uint64, error) { ctx = contextWithKeyVizLabel(ctx, c.label) if alloc, ok := c.inner.(TimestampAllocator); ok { @@ -116,6 +121,20 @@ func (c keyVizLabeledCoordinator) RecoverHLCLease(ctx context.Context) error { return errors.WithStack(errHLCLeaseRecoveryUnavailable) } +func (c keyVizLabeledCoordinator) VouchAppliedReadTimestamp(timestamp uint64, ref AppliedReadTimestampVoucherRef) error { + voucher, ok := c.inner.(AppliedReadTimestampVoucher) + if !ok { + return errors.WithStack(ErrTSOProtocolUnsupported) + } + return errors.WithStack(voucher.VouchAppliedReadTimestamp(timestamp, ref)) +} + +func (c keyVizLabeledCoordinator) RevokeAppliedReadTimestamp(timestamp uint64, ref AppliedReadTimestampVoucherRef) { + if revoker, ok := c.inner.(AppliedReadTimestampVoucherRevoker); ok { + revoker.RevokeAppliedReadTimestamp(timestamp, ref) + } +} + func (c keyVizLabeledCoordinator) LeaseRead(ctx context.Context) (uint64, error) { if lr, ok := c.inner.(LeaseReadableCoordinator); ok { idx, err := lr.LeaseRead(ctx) diff --git a/kv/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/lease_warmup_test.go b/kv/lease_warmup_test.go index e3ca74f79..e433f6c66 100644 --- a/kv/lease_warmup_test.go +++ b/kv/lease_warmup_test.go @@ -295,6 +295,33 @@ func TestShardedCoordinator_RenewHLCLeases_ProposesToEveryLedGroup(t *testing.T) "the non-default group lease must be warmed by all-group renewal") } +func TestShardedCoordinator_RenewHLCLeases_PhaseDOnlyRenewsTimestampGroup(t *testing.T) { + t.Parallel() + eng0 := newShardedLeaseEngine(50) + eng1 := newShardedLeaseEngine(100) + eng2 := newShardedLeaseEngine(200) + distEngine := distribution.NewEngine() + distEngine.UpdateRoute([]byte("a"), []byte("m"), 1) + distEngine.UpdateRoute([]byte("m"), nil, 2) + phaseD := NewTSOStateMachine(NewHLC()) + require.Nil(t, phaseD.Apply(marshalTSOCutover())) + require.Nil(t, phaseD.Apply(marshalTSOPhaseD(0))) + coord := NewShardedCoordinator(distEngine, map[uint64]*ShardGroup{ + 0: {Engine: eng0}, + 1: {Engine: eng1}, + 2: {Engine: eng2}, + }, 1, NewHLC(), nil). + WithTimestampGroup(0). + WithTSOCutoverState(phaseD) + + done := coord.renewHLCLeases(context.Background()) + requireRenewalDone(t, done) + + require.Equal(t, int32(1), eng0.proposeCalls.Load()) + require.Equal(t, int32(0), eng1.proposeCalls.Load()) + require.Equal(t, int32(0), eng2.proposeCalls.Load()) +} + func TestShardedCoordinator_RecoverHLCLease_ProposesToEveryLedGroup(t *testing.T) { t.Parallel() clock := NewHLC() diff --git a/kv/shard_store.go b/kv/shard_store.go index fc3276de9..932cf54a7 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,142 @@ 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) + } + readCtx, cancel := context.WithTimeout(nonNilTSOContext(ctx), proxyForwardTimeout) + defer cancel() + if _, err := linearizableReadEngineCtx(readCtx, 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 715086523..fc36425e4 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. @@ -360,7 +364,8 @@ const ( // surfaces unchanged so the client (or a wrapping retry harness in // the adapter) sees the failure rather than the coordinator spinning // on a persistent route shift. - composed1RetryAttempts = 1 + composed1RetryAttempts = 1 + maxAppliedReadTimestampVouches = 4096 ) // ShardedCoordinator routes operations to shard-specific raft groups. @@ -380,6 +385,14 @@ type ShardedCoordinator struct { // behavior where any locally-led shard group can issue TSO timestamps. timestampGroup uint64 timestampGroupConfigured bool + // tsoCutoverState is consensus-owned group-0 migration state. Phase D uses + // it to retire data-shard HLC renewal and reject legacy cross-shard startTS. + tsoCutoverState interface { + CutoverActive() bool + PhaseDActive() bool + } + appliedReadVoucherMu sync.Mutex + appliedReadVouchers map[appliedReadVoucherKey]uint64 // allShardGroupIDs, when configured, is the explicit set of data groups // that whole-keyspace operations must visit. It lets callers keep // non-data groups (for example a reserved timestamp group) in c.groups @@ -421,6 +434,11 @@ type ShardedCoordinator struct { registrationGate *RegistrationGate } +type appliedReadVoucherKey struct { + timestamp uint64 + ref AppliedReadTimestampVoucherRef +} + // RegistrationGate carries the Stage 7a §4.1 registration-before- // first-write barrier into the coordinator. main.go owns the // encryption StateCache and the registration goroutine and supplies @@ -464,6 +482,82 @@ func (c *ShardedCoordinator) WithTSOAllocator(alloc TimestampAllocator) *Sharded return c } +// TimestampAllocator exposes the configured allocator to coordinator +// decorators without widening the Coordinator interface. +func (c *ShardedCoordinator) TimestampAllocator() TimestampAllocator { + if c == nil { + return nil + } + return c.tsAllocator +} + +// VouchAppliedReadTimestamp records one use of an audited adapter watermark. +// The bounded map prevents abandoned requests from growing process memory. +func (c *ShardedCoordinator) VouchAppliedReadTimestamp(timestamp uint64, ref AppliedReadTimestampVoucherRef) error { + if c == nil || timestamp == 0 || timestamp == ^uint64(0) || ref.id == 0 { + return errors.WithStack(ErrTSOTimestampInvalid) + } + key := appliedReadVoucherKey{timestamp: timestamp, ref: ref} + c.appliedReadVoucherMu.Lock() + defer c.appliedReadVoucherMu.Unlock() + if c.appliedReadVouchers == nil { + c.appliedReadVouchers = make(map[appliedReadVoucherKey]uint64) + } + if uses, ok := c.appliedReadVouchers[key]; ok { + c.appliedReadVouchers[key] = uses + 1 + return nil + } + if len(c.appliedReadVouchers) >= maxAppliedReadTimestampVouches { + return errors.WithStack(ErrTSOReadVoucherLimit) + } + c.appliedReadVouchers[key] = 1 + return nil +} + +// RevokeAppliedReadTimestamp removes one prepared voucher that did not reach +// ShardedCoordinator dispatch validation, for example because an outer +// coordinator decorator rejected the dispatch first. +func (c *ShardedCoordinator) RevokeAppliedReadTimestamp(timestamp uint64, ref AppliedReadTimestampVoucherRef) { + if c == nil || timestamp == 0 || timestamp == ^uint64(0) || ref.id == 0 { + return + } + key := appliedReadVoucherKey{timestamp: timestamp, ref: ref} + c.appliedReadVoucherMu.Lock() + defer c.appliedReadVoucherMu.Unlock() + uses, ok := c.appliedReadVouchers[key] + if !ok { + return + } + if uses <= 1 { + delete(c.appliedReadVouchers, key) + return + } + c.appliedReadVouchers[key] = uses - 1 +} + +func (c *ShardedCoordinator) consumeAppliedReadTimestampVoucher(ctx context.Context, timestamp uint64) bool { + if c == nil { + return false + } + ref, ok := appliedReadTimestampVoucherRefFromContext(ctx, timestamp) + if !ok { + return false + } + key := appliedReadVoucherKey{timestamp: timestamp, ref: ref} + c.appliedReadVoucherMu.Lock() + defer c.appliedReadVoucherMu.Unlock() + uses, ok := c.appliedReadVouchers[key] + if !ok { + return false + } + if uses <= 1 { + delete(c.appliedReadVouchers, key) + } else { + c.appliedReadVouchers[key] = uses - 1 + } + return true +} + // WithTimestampGroup pins timestamp issuance leadership to one Raft group. // Callers should only enable this once a data-shard leader can redirect // timestamp allocation to that group; otherwise data leaders would stop being @@ -474,6 +568,17 @@ func (c *ShardedCoordinator) WithTimestampGroup(groupID uint64) *ShardedCoordina return c } +// WithTSOCutoverState wires the durable group-0 migration state. The state is +// read dynamically so a marker applied after startup changes renewal and +// validation behavior without a process-local mode race. +func (c *ShardedCoordinator) WithTSOCutoverState(state interface { + CutoverActive() bool + PhaseDActive() bool +}) *ShardedCoordinator { + c.tsoCutoverState = state + return c +} + // WithAllShardGroups restricts whole-keyspace operations to the supplied data // groups. When unset, the coordinator preserves the legacy behaviour and uses // every group it owns. @@ -885,7 +990,17 @@ func (c *ShardedCoordinator) dispatchTxnWithComposed1Retry(ctx context.Context, c.maybeAutoPinObservedRouteVersion(reqs, callerSuppliedStartTS) for attempt := 0; attempt <= composed1RetryAttempts; attempt++ { - resp, err := c.dispatchTxn(ctx, reqs.StartTS, reqs.CommitTS, reqs.PrevCommitTS, reqs.Elems, reqs.ReadKeys, reqs.ObservedRouteVersion, reqs.KeyVizLabel) + resp, err := c.dispatchTxn( + ctx, + reqs.StartTS, + reqs.CommitTS, + reqs.PrevCommitTS, + reqs.Elems, + reqs.ReadKeys, + reqs.ObservedRouteVersion, + reqs.KeyVizLabel, + callerSuppliedStartTS, + ) if err == nil { return resp, nil } @@ -1114,7 +1229,17 @@ func (c *ShardedCoordinator) broadcastToAllGroups(ctx context.Context, requests return &CoordinateResponse{CommitIndex: maxIndex.Load()}, nil } -func (c *ShardedCoordinator) dispatchTxn(ctx context.Context, startTS uint64, commitTS uint64, prevCommitTS uint64, elems []*Elem[OP], readKeys [][]byte, observedRouteVersion uint64, label keyviz.Label) (*CoordinateResponse, error) { +func (c *ShardedCoordinator) dispatchTxn( + ctx context.Context, + startTS uint64, + commitTS uint64, + prevCommitTS uint64, + elems []*Elem[OP], + readKeys [][]byte, + observedRouteVersion uint64, + label keyviz.Label, + callerSuppliedStartTS bool, +) (*CoordinateResponse, error) { if len(readKeys) > maxReadKeys { return nil, errors.WithStack(ErrInvalidRequest) } @@ -1126,16 +1251,17 @@ func (c *ShardedCoordinator) dispatchTxn(ctx context.Context, startTS uint64, co if len(primaryKey) == 0 { return nil, errors.WithStack(ErrTxnPrimaryKeyRequired) } - - commitTS, err = c.resolveTxnCommitTS(ctx, startTS, commitTS) - if err != nil { + singleShard := len(gids) == 1 && c.allReadKeysInShard(readKeys, gids[0]) + if err := c.validateCallerSuppliedTxnStart(ctx, startTS, singleShard, callerSuppliedStartTS); err != nil { return nil, err } - if err := ValidateElemCommitTSPatches(elems, commitTS); err != nil { + + commitTS, err = c.prepareTxnCommitTimestamp(ctx, startTS, commitTS, elems) + if err != nil { return nil, err } - if len(gids) == 1 && c.allReadKeysInShard(readKeys, gids[0]) { + if singleShard { // Fast path: all mutations and read keys are in a single shard. // Use the one-phase path without allocating a grouped-read-keys map. // If any read key belongs to a different shard the 2PC path is required @@ -1149,9 +1275,47 @@ func (c *ShardedCoordinator) dispatchTxn(ctx context.Context, startTS uint64, co return c.dispatchMultiShardTxn(ctx, startTS, commitTS, prevCommitTS, primaryKey, grouped, gids, readKeys, observedRouteVersion) } +func (c *ShardedCoordinator) validateCallerSuppliedTxnStart(ctx context.Context, startTS uint64, singleShard, callerSupplied bool) error { + if !callerSupplied { + return nil + } + if c.consumeAppliedReadTimestampVoucher(ctx, startTS) || singleShard { + return nil + } + return c.validateCrossShardReadTimestamp(ctx, startTS) +} + +func (c *ShardedCoordinator) prepareTxnCommitTimestamp(ctx context.Context, startTS, commitTS uint64, elems []*Elem[OP]) (uint64, error) { + resolved, err := c.resolveTxnCommitTS(ctx, startTS, commitTS) + if err != nil { + return 0, err + } + if err := ValidateElemCommitTSPatches(elems, resolved); err != nil { + return 0, err + } + return resolved, nil +} + +func (c *ShardedCoordinator) validateCrossShardReadTimestamp( + ctx context.Context, + startTS uint64, +) error { + validator, _, required, err := phaseDTimestampValidator(c.tsAllocator) + if err != nil { + return errors.Wrap(err, "cross-shard read timestamp validator is unavailable") + } + if !required { + return nil + } + if err := validator.ValidateDurableTimestamp(ctx, startTS); err != nil { + return errors.Wrap(err, "validate cross-shard read/start timestamp") + } + return nil +} + // dispatchMultiShardTxn runs the 2PC path. Extracted from dispatchTxn to keep -// that function under the cyclop budget after the prevCommitTS reject (codex -// P2 round-10) was added; the multi-shard branch already carries five linear +// that function under the cyclop budget after the prevCommitTS reject was +// added; the multi-shard branch already carries five linear // error checks (groupReadKeys, prewrite, commitPrimary, abortCleanup, // commitSecondaries) that pushed the parent over the 10-edge limit. func (c *ShardedCoordinator) dispatchMultiShardTxn(ctx context.Context, startTS, commitTS, prevCommitTS uint64, primaryKey []byte, grouped map[uint64][]*pb.Mutation, gids []uint64, readKeys [][]byte, observedRouteVersion uint64) (*CoordinateResponse, error) { @@ -1504,6 +1668,10 @@ func (c *ShardedCoordinator) allocateTimestampAfter(ctx context.Context, label s } return nextTimestampFromAllocator(ctx, c.tsAllocator, label) } + if c.tsoCutoverState != nil && c.tsoCutoverState.PhaseDActive() { + return 0, errors.Wrap(ErrTSOAllocatorRequired, + "legacy HLC issuance is disabled by durable TSO phase D") + } if c.clock == nil { return 0, errors.Wrap(ErrTSOClockNil, label) } @@ -2380,10 +2548,7 @@ func (c *ShardedCoordinator) renewHLCLeases(ctx context.Context) <-chan struct{} done := make(chan struct{}) var wg sync.WaitGroup for gid, group := range c.groups { - if group == nil || group.Engine == nil { - continue - } - if !shardGroupAcceptingWrites(group) { + if !c.shouldRenewHLCGroup(gid, group) { continue } if !c.startHLCLeaseRenewal(gid) { @@ -2405,6 +2570,14 @@ func (c *ShardedCoordinator) renewHLCLeases(ctx context.Context) <-chan struct{} return done } +func (c *ShardedCoordinator) shouldRenewHLCGroup(gid uint64, group *ShardGroup) bool { + if c.tsoCutoverState != nil && c.tsoCutoverState.PhaseDActive() && + (!c.timestampGroupConfigured || gid != c.timestampGroup) { + return false + } + return shardGroupAcceptingWrites(group) +} + func (c *ShardedCoordinator) startHLCLeaseRenewal(gid uint64) bool { c.hlcRenewalMu.Lock() defer c.hlcRenewalMu.Unlock() diff --git a/kv/sharded_coordinator_txn_test.go b/kv/sharded_coordinator_txn_test.go index 42b65f328..a58ae92c8 100644 --- a/kv/sharded_coordinator_txn_test.go +++ b/kv/sharded_coordinator_txn_test.go @@ -203,6 +203,444 @@ func TestShardedCoordinatorDispatchTxn_CrossShardPhasesAndCommitIndex(t *testing require.Zero(t, binary.BigEndian.Uint64(value2[4:12])) } +func TestShardedCoordinatorDispatchTxn_PhaseDRejectsInvalidCallerStartTSBeforeProposal(t *testing.T) { + t.Parallel() + coord, g1Txn, g2Txn, alloc := newPhaseDCrossShardCoordinator(t, ErrTSOTimestampInvalid) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + IsTxn: true, + StartTS: 10, + Elems: []*Elem[OP]{ + {Op: Put, Key: []byte("b"), Value: []byte("v1")}, + {Op: Put, Key: []byte("x"), Value: []byte("v2")}, + }, + }) + require.ErrorIs(t, err, ErrTSOTimestampInvalid) + require.Equal(t, uint64(1), alloc.validateCalls.Load()) + require.Empty(t, g1Txn.requests) + require.Empty(t, g2Txn.requests) +} + +func TestShardedCoordinatorDispatchTxn_PhaseDAcceptsValidatedCallerStartTS(t *testing.T) { + t.Parallel() + coord, g1Txn, g2Txn, alloc := newPhaseDCrossShardCoordinator(t, nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + IsTxn: true, + StartTS: 100, + Elems: []*Elem[OP]{ + {Op: Put, Key: []byte("b"), Value: []byte("v1")}, + {Op: Put, Key: []byte("x"), Value: []byte("v2")}, + }, + }) + require.NoError(t, err) + require.Equal(t, uint64(1), alloc.validateCalls.Load()) + require.Len(t, g1Txn.requests, 2) + require.Len(t, g2Txn.requests, 2) + require.Equal(t, uint64(100), g1Txn.requests[0].Ts) + require.Equal(t, uint64(100), g2Txn.requests[0].Ts) +} + +func TestShardedCoordinatorDispatchTxn_PhaseDAcceptsBoundAppliedWatermark(t *testing.T) { + t.Parallel() + prePhaseDErr := errors.Join(ErrTSOTimestampInvalid, ErrTSOTimestampPrePhaseD) + coord, g1Txn, g2Txn, alloc := newPhaseDCrossShardCoordinator(t, prePhaseDErr) + + readTS, err := BeginReadTimestampThrough(context.Background(), coord, 10, "vouch applied watermark") + require.NoError(t, err) + require.Equal(t, uint64(10), readTS.Timestamp()) + require.Equal(t, uint64(1), alloc.validateCalls.Load()) + + request := func() *OperationGroup[OP] { + return &OperationGroup[OP]{ + IsTxn: true, + StartTS: readTS.Timestamp(), + Elems: []*Elem[OP]{ + {Op: Put, Key: []byte("b"), Value: []byte("v1")}, + {Op: Put, Key: []byte("x"), Value: []byte("v2")}, + }, + } + } + + _, err = coord.Dispatch(context.Background(), request()) + require.ErrorIs(t, err, ErrTSOTimestampPrePhaseD, "a numeric timestamp without the bound capability must not steal a voucher") + require.Equal(t, uint64(2), alloc.validateCalls.Load()) + require.Empty(t, g1Txn.requests) + require.Empty(t, g2Txn.requests) + + ctx := readTS.WithDispatchVoucher(context.Background()) + _, err = DispatchWithReadTimestamp(ctx, coord, request()) + require.NoError(t, err) + require.Equal(t, uint64(2), alloc.validateCalls.Load(), "bound voucher must bypass numeric Phase-D validation") + require.Len(t, g1Txn.requests, 2) + require.Len(t, g2Txn.requests, 2) + + _, err = coord.Dispatch(context.Background(), request()) + require.ErrorIs(t, err, ErrTSOTimestampPrePhaseD) + require.Equal(t, uint64(3), alloc.validateCalls.Load(), "voucher must be bound and single-use") +} + +func TestDispatchWithReadTimestampVouchesEveryBoundDispatch(t *testing.T) { + t.Parallel() + prePhaseDErr := errors.Join(ErrTSOTimestampInvalid, ErrTSOTimestampPrePhaseD) + coord, g1Txn, g2Txn, alloc := newPhaseDCrossShardCoordinator(t, prePhaseDErr) + + readTimestamp, err := BeginReadTimestampThrough(context.Background(), coord, 10, "vouch reused applied watermark") + require.NoError(t, err) + ctx := readTimestamp.WithDispatchVoucher(context.Background()) + request := func(startTS uint64) *OperationGroup[OP] { + return &OperationGroup[OP]{ + IsTxn: true, + StartTS: startTS, + Elems: []*Elem[OP]{ + {Op: Put, Key: []byte("b"), Value: []byte("v1")}, + {Op: Put, Key: []byte("x"), Value: []byte("v2")}, + }, + } + } + + for range 2 { + _, err = DispatchWithReadTimestamp(ctx, coord, request(readTimestamp.Timestamp())) + require.NoError(t, err) + } + require.Equal(t, uint64(1), alloc.validateCalls.Load(), "each dispatch must consume a reserved voucher") + require.Len(t, g1Txn.requests, 4) + require.Len(t, g2Txn.requests, 4) + + _, err = DispatchWithReadTimestamp(ctx, coord, request(readTimestamp.Timestamp()+1)) + require.ErrorIs(t, err, ErrTSOTimestampInvalid, "the bound capability must not authorize another timestamp") + + _, err = coord.Dispatch(context.Background(), request(readTimestamp.Timestamp())) + require.ErrorIs(t, err, ErrTSOTimestampPrePhaseD, "no unused voucher may remain after the bound dispatches") + require.Equal(t, uint64(2), alloc.validateCalls.Load()) +} + +func TestDispatchWithReadTimestampUsesDistinctRefsForOverlappingDispatches(t *testing.T) { + t.Parallel() + + coord := newOverlappingReadVoucherCoordinator() + readTimestamp := ReadTimestamp{ + timestamp: 10, + voucher: newAppliedReadDispatchVoucher(), + } + ctx := readTimestamp.WithDispatchVoucher(context.Background()) + request := func() *OperationGroup[OP] { + return &OperationGroup[OP]{ + IsTxn: true, + StartTS: readTimestamp.Timestamp(), + Elems: []*Elem[OP]{ + {Op: Put, Key: []byte("b"), Value: []byte("v1")}, + {Op: Put, Key: []byte("x"), Value: []byte("v2")}, + }, + } + } + + firstErr := make(chan error, 1) + go func() { + _, err := DispatchWithReadTimestamp(ctx, coord, request()) + firstErr <- err + close(coord.firstDone) + }() + requireChannelClosed(t, coord.firstEntered) + + secondErr := make(chan error, 1) + go func() { + _, err := DispatchWithReadTimestamp(ctx, coord, request()) + secondErr <- err + }() + requireChannelClosed(t, coord.secondEntered) + + close(coord.releaseFirst) + require.NoError(t, <-firstErr) + require.NoError(t, <-secondErr) +} + +func TestDispatchWithReadTimestampRevokesVoucherWhenOuterGateRejects(t *testing.T) { + t.Parallel() + prePhaseDErr := errors.Join(ErrTSOTimestampInvalid, ErrTSOTimestampPrePhaseD) + coord, _, _, _ := newPhaseDCrossShardCoordinator(t, prePhaseDErr) + gateErr := errors.New("startup gate rejected dispatch") + gated := phaseDGateCoordinator{inner: coord, err: gateErr} + + readTimestamp, err := BeginReadTimestampThrough(context.Background(), gated, 10, "vouch gated applied watermark") + require.NoError(t, err) + _, err = DispatchWithReadTimestamp( + readTimestamp.WithDispatchVoucher(context.Background()), + gated, + &OperationGroup[OP]{ + IsTxn: true, + StartTS: readTimestamp.Timestamp(), + Elems: []*Elem[OP]{ + {Op: Put, Key: []byte("b"), Value: []byte("v1")}, + {Op: Put, Key: []byte("x"), Value: []byte("v2")}, + }, + }, + ) + require.ErrorIs(t, err, gateErr) + + coord.appliedReadVoucherMu.Lock() + defer coord.appliedReadVoucherMu.Unlock() + require.Empty(t, coord.appliedReadVouchers) +} + +func TestReadTimestampVoucherBindingShadowsParentCapability(t *testing.T) { + prePhaseDErr := errors.Join(ErrTSOTimestampInvalid, ErrTSOTimestampPrePhaseD) + coord, _, _, alloc := newPhaseDCrossShardCoordinator(t, prePhaseDErr) + + oldRead, err := BeginReadTimestampThrough(context.Background(), coord, 10, "vouch old applied watermark") + require.NoError(t, err) + ctx := oldRead.WithDispatchVoucher(context.Background()) + + alloc.validateErr = nil + currentRead, err := BeginReadTimestampThrough(ctx, coord, 100, "validate current durable watermark") + require.NoError(t, err) + ctx = currentRead.WithDispatchVoucher(ctx) + + _, err = DispatchWithReadTimestamp(ctx, coord, &OperationGroup[OP]{ + IsTxn: true, + StartTS: currentRead.Timestamp(), + Elems: []*Elem[OP]{ + {Op: Put, Key: []byte("b"), Value: []byte("v1")}, + {Op: Put, Key: []byte("x"), Value: []byte("v2")}, + }, + }) + require.NoError(t, err, "the current timestamp must shadow the parent capability") +} + +func TestShardedCoordinatorDispatchTxn_PhaseDPreservesSingleShardCallerStartTS(t *testing.T) { + t.Parallel() + coord, g1Txn, _, alloc := newPhaseDCrossShardCoordinator(t, ErrTSOTimestampInvalid) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + IsTxn: true, + StartTS: 10, + Elems: []*Elem[OP]{ + {Op: Put, Key: []byte("b"), Value: []byte("v1")}, + }, + }) + require.NoError(t, err) + require.Zero(t, alloc.validateCalls.Load()) + require.Len(t, g1Txn.requests, 1) +} + +func TestShardedCoordinatorDispatchTxn_PrePhaseDPreservesCrossShardCallerStartTS(t *testing.T) { + t.Parallel() + coord, g1Txn, g2Txn, alloc := newPhaseDCrossShardCoordinator(t, ErrTSOTimestampInvalid) + coord.WithTSOCutoverState(NewTSOStateMachine(NewHLC())) + alloc.phaseDActive = false + alloc.phaseDRequired = false + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + IsTxn: true, + StartTS: 10, + Elems: []*Elem[OP]{ + {Op: Put, Key: []byte("b"), Value: []byte("v1")}, + {Op: Put, Key: []byte("x"), Value: []byte("v2")}, + }, + }) + require.NoError(t, err) + require.Zero(t, alloc.validateCalls.Load()) + require.Len(t, g1Txn.requests, 2) + require.Len(t, g2Txn.requests, 2) +} + +func TestShardedCoordinatorDispatchTxn_PhaseDActivationValidatesBeforeLocalMarker(t *testing.T) { + t.Parallel() + coord, g1Txn, g2Txn, alloc := newPhaseDCrossShardCoordinator(t, ErrTSOTimestampInvalid) + coord.WithTSOCutoverState(NewTSOStateMachine(NewHLC())) + alloc.phaseDActive = false + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + IsTxn: true, + StartTS: 10, + Elems: []*Elem[OP]{ + {Op: Put, Key: []byte("b"), Value: []byte("v1")}, + {Op: Put, Key: []byte("x"), Value: []byte("v2")}, + }, + }) + require.ErrorIs(t, err, ErrTSOTimestampInvalid) + require.Equal(t, uint64(1), alloc.validateCalls.Load()) + require.Empty(t, g1Txn.requests) + require.Empty(t, g2Txn.requests) +} + +func newPhaseDCrossShardCoordinator( + t *testing.T, + validateErr error, +) (*ShardedCoordinator, *recordingTransactional, *recordingTransactional, *phaseDTestAllocator) { + t.Helper() + engine := distribution.NewEngine() + engine.UpdateRoute([]byte("a"), []byte("m"), 1) + engine.UpdateRoute([]byte("m"), nil, 2) + g1Txn := &recordingTransactional{} + g2Txn := &recordingTransactional{} + alloc := &phaseDTestAllocator{ + next: 200, + phaseDActive: true, + phaseDRequired: true, + validateErr: validateErr, + } + state := NewTSOStateMachine(NewHLC()) + require.Nil(t, state.Apply(marshalTSOCutover())) + require.Nil(t, state.Apply(marshalTSOPhaseD(0))) + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + }, 1, NewHLC(), nil). + WithTSOAllocator(alloc). + WithTSOCutoverState(state) + return coord, g1Txn, g2Txn, alloc +} + +type phaseDGateCoordinator struct { + inner *ShardedCoordinator + err error +} + +type overlappingReadVoucherCoordinator struct { + mu sync.Mutex + clock *HLC + vouchers map[appliedReadVoucherKey]uint64 + dispatchCalls int + firstEntered chan struct{} + secondEntered chan struct{} + releaseFirst chan struct{} + firstDone chan struct{} +} + +func newOverlappingReadVoucherCoordinator() *overlappingReadVoucherCoordinator { + return &overlappingReadVoucherCoordinator{ + clock: NewHLC(), + vouchers: make(map[appliedReadVoucherKey]uint64), + firstEntered: make(chan struct{}), + secondEntered: make(chan struct{}), + releaseFirst: make(chan struct{}), + firstDone: make(chan struct{}), + } +} + +func (c *overlappingReadVoucherCoordinator) Dispatch(ctx context.Context, req *OperationGroup[OP]) (*CoordinateResponse, error) { + ref, ok := appliedReadTimestampVoucherRefFromContext(ctx, req.StartTS) + if !ok { + return nil, errors.Join(ErrTSOTimestampInvalid, ErrTSOTimestampPrePhaseD) + } + c.mu.Lock() + c.dispatchCalls++ + call := c.dispatchCalls + c.mu.Unlock() + switch call { + case 1: + close(c.firstEntered) + <-c.releaseFirst + case 2: + close(c.secondEntered) + <-c.firstDone + default: + return nil, errors.New("unexpected dispatch") + } + if !c.consumeAppliedReadTimestampVoucher(req.StartTS, ref) { + return nil, errors.Join(ErrTSOTimestampInvalid, ErrTSOTimestampPrePhaseD) + } + return &CoordinateResponse{}, nil +} + +func (c *overlappingReadVoucherCoordinator) VouchAppliedReadTimestamp(timestamp uint64, ref AppliedReadTimestampVoucherRef) error { + c.mu.Lock() + defer c.mu.Unlock() + c.vouchers[appliedReadVoucherKey{timestamp: timestamp, ref: ref}]++ + return nil +} + +func (c *overlappingReadVoucherCoordinator) RevokeAppliedReadTimestamp(timestamp uint64, ref AppliedReadTimestampVoucherRef) { + c.mu.Lock() + defer c.mu.Unlock() + key := appliedReadVoucherKey{timestamp: timestamp, ref: ref} + uses := c.vouchers[key] + if uses <= 1 { + delete(c.vouchers, key) + return + } + c.vouchers[key] = uses - 1 +} + +func (c *overlappingReadVoucherCoordinator) consumeAppliedReadTimestampVoucher(timestamp uint64, ref AppliedReadTimestampVoucherRef) bool { + c.mu.Lock() + defer c.mu.Unlock() + key := appliedReadVoucherKey{timestamp: timestamp, ref: ref} + uses := c.vouchers[key] + if uses == 0 { + return false + } + if uses == 1 { + delete(c.vouchers, key) + return true + } + c.vouchers[key] = uses - 1 + return true +} + +func (c *overlappingReadVoucherCoordinator) IsLeader() bool { return true } + +func (c *overlappingReadVoucherCoordinator) VerifyLeader(context.Context) error { return nil } + +func (c *overlappingReadVoucherCoordinator) LinearizableRead(context.Context) (uint64, error) { + return 0, nil +} + +func (c *overlappingReadVoucherCoordinator) RaftLeader() string { return "" } + +func (c *overlappingReadVoucherCoordinator) IsLeaderForKey([]byte) bool { return true } + +func (c *overlappingReadVoucherCoordinator) VerifyLeaderForKey(context.Context, []byte) error { + return nil +} + +func (c *overlappingReadVoucherCoordinator) RaftLeaderForKey([]byte) string { return "" } + +func (c *overlappingReadVoucherCoordinator) Clock() *HLC { return c.clock } + +func (c phaseDGateCoordinator) Dispatch(context.Context, *OperationGroup[OP]) (*CoordinateResponse, error) { + return nil, c.err +} + +func (c phaseDGateCoordinator) IsLeader() bool { return c.inner.IsLeader() } + +func (c phaseDGateCoordinator) VerifyLeader(ctx context.Context) error { + return c.inner.VerifyLeader(ctx) +} + +func (c phaseDGateCoordinator) LinearizableRead(ctx context.Context) (uint64, error) { + return c.inner.LinearizableRead(ctx) +} + +func (c phaseDGateCoordinator) RaftLeader() string { return c.inner.RaftLeader() } + +func (c phaseDGateCoordinator) IsLeaderForKey(key []byte) bool { + return c.inner.IsLeaderForKey(key) +} + +func (c phaseDGateCoordinator) VerifyLeaderForKey(ctx context.Context, key []byte) error { + return c.inner.VerifyLeaderForKey(ctx, key) +} + +func (c phaseDGateCoordinator) RaftLeaderForKey(key []byte) string { + return c.inner.RaftLeaderForKey(key) +} + +func (c phaseDGateCoordinator) Clock() *HLC { return c.inner.Clock() } + +func (c phaseDGateCoordinator) TimestampAllocator() TimestampAllocator { + return c.inner.TimestampAllocator() +} + +func (c phaseDGateCoordinator) VouchAppliedReadTimestamp(timestamp uint64, ref AppliedReadTimestampVoucherRef) error { + return c.inner.VouchAppliedReadTimestamp(timestamp, ref) +} + +func (c phaseDGateCoordinator) RevokeAppliedReadTimestamp(timestamp uint64, ref AppliedReadTimestampVoucherRef) { + c.inner.RevokeAppliedReadTimestamp(timestamp, ref) +} + func TestShardedCoordinatorDispatchTxn_SingleShardUsesOnePhase(t *testing.T) { t.Parallel() diff --git a/kv/tso.go b/kv/tso.go index 48e800c33..317873e82 100644 --- a/kv/tso.go +++ b/kv/tso.go @@ -11,11 +11,19 @@ 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") - ErrTSOClockNil = errors.New("tso: coordinator clock is nil") - ErrInvalidTSOBatchSize = errors.New("tso: invalid batch size") + ErrTSOAllocatorRequired = errors.New("tso: allocator is required") + ErrTSOCoordinatorNil = errors.New("tso: coordinator is required") + ErrTSOClockNil = errors.New("tso: coordinator clock is nil") + ErrInvalidTSOBatchSize = errors.New("tso: invalid batch size") + ErrTSOPhaseDInactive = errors.New("tso: phase D is not active") + ErrTSOTimestampInvalid = errors.New("tso: timestamp is not a durable phase-D allocation") + ErrTSOTimestampPrePhaseD = errors.New("tso: timestamp predates phase D") + ErrTSOReadVoucherLimit = errors.New("tso: applied read timestamp voucher limit reached") errHLCLeaseRecoveryUnavailable = errors.New("hlc lease recovery unavailable") errHLCLeaseRecoveryBlocked = errors.New("hlc lease recovery blocked") @@ -37,10 +45,159 @@ 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) } +// DurableTimestampValidator verifies that a timestamp belongs to the durable +// post-Phase-D allocation range owned by the dedicated TSO group. +type DurableTimestampValidator interface { + ValidateDurableTimestamp(context.Context, uint64) error +} + +// TSOPhaseDState exposes the one-way Phase-D state to coordinators and adapter +// migration helpers without coupling them to TSOStateMachine. +type TSOPhaseDState interface { + PhaseDActive() bool + PhaseDRequired() bool +} + +// AppliedReadTimestampVoucher records an adapter-provided applied watermark. +// It is a process-local capability used only to distinguish audited adapter +// snapshots from arbitrary caller-supplied StartTS values during Phase D. +type AppliedReadTimestampVoucher interface { + VouchAppliedReadTimestamp(uint64, AppliedReadTimestampVoucherRef) error +} + +// AppliedReadTimestampVoucherRevoker removes an unused voucher registration. +// Decorators that forward VouchAppliedReadTimestamp should forward revocation +// too so pre-dispatch gates cannot leak inner-coordinator voucher entries. +type AppliedReadTimestampVoucherRevoker interface { + RevokeAppliedReadTimestamp(uint64, AppliedReadTimestampVoucherRef) +} + +// AppliedReadTimestampVoucherRef is an opaque process-local dispatch +// capability. Only this package can mint a non-zero ref. +type AppliedReadTimestampVoucherRef struct { + id uint64 +} + +// ReadTimestamp is the adapter-side result of beginning a transaction snapshot. +// When it represents an applied pre-Phase-D watermark, it also carries a +// process-local capability that can reserve exactly one coordinator voucher per +// DispatchWithReadTimestamp call. The capability cannot be constructed outside +// this package because both the timestamp and voucher state are private. +type ReadTimestamp struct { + timestamp uint64 + voucher *appliedReadDispatchVoucher +} + +func (t ReadTimestamp) Timestamp() uint64 { + return t.timestamp +} + +func (t ReadTimestamp) voucherRef() (AppliedReadTimestampVoucherRef, bool) { + if t.voucher == nil || t.voucher.ref.id == 0 { + return AppliedReadTimestampVoucherRef{}, false + } + return t.voucher.ref, true +} + +var nextAppliedReadDispatchVoucherID atomic.Uint64 + +type appliedReadDispatchVoucher struct { + mu sync.Mutex + prepared uint64 + ref AppliedReadTimestampVoucherRef +} + +type appliedReadDispatchVoucherContextKey struct{} + +// WithDispatchVoucher binds this read timestamp's process-local capability to +// ctx. A timestamp without a voucher is still bound so it shadows any parent +// capability instead of accidentally inheriting authority for an older read. +func (t ReadTimestamp) WithDispatchVoucher(ctx context.Context) context.Context { + return context.WithValue(nonNilTSOContext(ctx), appliedReadDispatchVoucherContextKey{}, t) +} + +func appliedReadTimestampVoucherRefFromContext(ctx context.Context, timestamp uint64) (AppliedReadTimestampVoucherRef, bool) { + if ctx == nil { + return AppliedReadTimestampVoucherRef{}, false + } + readTimestamp, ok := ctx.Value(appliedReadDispatchVoucherContextKey{}).(ReadTimestamp) + if !ok || readTimestamp.timestamp != timestamp { + return AppliedReadTimestampVoucherRef{}, false + } + return readTimestamp.voucherRef() +} + +// DispatchWithReadTimestamp dispatches an OCC operation under the applied-read +// capability bound by ReadTimestamp.WithDispatchVoucher. Each dispatch reserves +// one token tied to that bound capability immediately before dispatching, so a +// same-valued StartTS from another request cannot consume it. +func DispatchWithReadTimestamp( + ctx context.Context, + coord Coordinator, + reqs *OperationGroup[OP], +) (*CoordinateResponse, error) { + if ctx == nil { + resp, err := coord.Dispatch(ctx, reqs) + return resp, errors.WithStack(err) + } + readTimestamp, ok := ctx.Value(appliedReadDispatchVoucherContextKey{}).(ReadTimestamp) + if !ok || readTimestamp.voucher == nil { + resp, err := coord.Dispatch(ctx, reqs) + return resp, errors.WithStack(err) + } + if reqs == nil || reqs.StartTS != readTimestamp.timestamp { + return nil, errors.WithStack(ErrTSOTimestampInvalid) + } + preparedReadTimestamp, revoke, err := readTimestamp.voucher.prepare(coord, readTimestamp.timestamp) + if err != nil { + return nil, err + } + defer revoke() + dispatchCtx := preparedReadTimestamp.WithDispatchVoucher(ctx) + resp, err := coord.Dispatch(dispatchCtx, reqs) + return resp, errors.WithStack(err) +} + +func newAppliedReadDispatchVoucher() *appliedReadDispatchVoucher { + id := nextAppliedReadDispatchVoucherID.Add(1) + if id == 0 { + id = nextAppliedReadDispatchVoucherID.Add(1) + } + return &appliedReadDispatchVoucher{ref: AppliedReadTimestampVoucherRef{id: id}} +} + +func (v *appliedReadDispatchVoucher) prepare(coord Coordinator, timestamp uint64) (ReadTimestamp, func(), error) { + v.mu.Lock() + defer v.mu.Unlock() + voucher, ok := coord.(AppliedReadTimestampVoucher) + if !ok { + return ReadTimestamp{}, nil, errors.WithStack(ErrTSOProtocolUnsupported) + } + preparedVoucher := newAppliedReadDispatchVoucher() + if err := voucher.VouchAppliedReadTimestamp(timestamp, preparedVoucher.ref); err != nil { + return ReadTimestamp{}, nil, errors.WithStack(err) + } + v.prepared++ + revoke := func() {} + if revoker, ok := coord.(AppliedReadTimestampVoucherRevoker); ok { + ref := preparedVoucher.ref + revoke = func() { + revoker.RevokeAppliedReadTimestamp(timestamp, ref) + } + } + return ReadTimestamp{timestamp: timestamp, voucher: preparedVoucher}, revoke, nil +} + type tsoBatchAfterAllocator interface { NextBatchAfter(ctx context.Context, n int, min uint64) (uint64, error) } @@ -95,7 +252,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 { @@ -113,6 +276,97 @@ func coordinatorTimestampAllocator(coord Coordinator) (TimestampAllocator, bool) } } +func coordinatorTimestampAllocator(coord Coordinator) (TimestampAllocator, bool) { + return TimestampAllocatorThrough(coord) +} + +// BeginReadTimestampThrough preserves the caller's applied-snapshot watermark. +// Once Phase D is requested, this boundary activates it before validation. An +// applied pre-D watermark receives a bounded one-use coordinator voucher; +// arbitrary caller timestamps remain subject to group-0 numeric validation. +// The returned timestamp must be used for every read and OperationGroup.StartTS. +func BeginReadTimestampThrough( + ctx context.Context, + coord Coordinator, + legacyTimestamp uint64, + label string, +) (ReadTimestamp, error) { + alloc, ok := coordinatorTimestampAllocator(coord) + if !ok { + return ReadTimestamp{timestamp: legacyTimestamp}, nil + } + validator, phaseD, phaseDRequired, err := phaseDTimestampValidator(alloc) + if err != nil { + return ReadTimestamp{}, errors.Wrap(err, label) + } + if !phaseDRequired { + return ReadTimestamp{timestamp: legacyTimestamp}, nil + } + if legacyTimestamp == 0 || legacyTimestamp == ^uint64(0) { + return ReadTimestamp{}, errors.Wrap(ErrTSOTimestampInvalid, label) + } + if err := activatePhaseDForRead(ctx, alloc, phaseD, label); err != nil { + return ReadTimestamp{}, err + } + vouched, err := validateAppliedReadTimestamp(ctx, coord, validator, legacyTimestamp, label) + if err != nil { + return ReadTimestamp{}, err + } + readTimestamp := ReadTimestamp{timestamp: legacyTimestamp} + if vouched { + readTimestamp.voucher = newAppliedReadDispatchVoucher() + } + return readTimestamp, nil +} + +func activatePhaseDForRead( + ctx context.Context, + alloc TimestampAllocator, + phaseD TSOPhaseDState, + label string, +) error { + if phaseD.PhaseDActive() { + return nil + } + // Reserve and discard one post-D timestamp to commit the marker. The + // caller's applied watermark remains the read snapshot; using this fresh + // allocation for reads could run ahead of data-group apply. + _, err := nextTimestampFromAllocator(nonNilTSOContext(ctx), alloc, label+": activate phase D") + return err +} + +func validateAppliedReadTimestamp( + ctx context.Context, + coord Coordinator, + validator DurableTimestampValidator, + timestamp uint64, + label string, +) (bool, error) { + err := validator.ValidateDurableTimestamp(nonNilTSOContext(ctx), timestamp) + if err == nil { + return false, nil + } + if !errors.Is(err, ErrTSOTimestampPrePhaseD) { + return false, errors.Wrap(err, label) + } + if _, ok := coord.(AppliedReadTimestampVoucher); !ok { + return false, errors.Wrap(ErrTSOProtocolUnsupported, label+": applied read voucher unavailable") + } + return true, nil +} + +func phaseDTimestampValidator(alloc TimestampAllocator) (DurableTimestampValidator, TSOPhaseDState, bool, error) { + phaseD, ok := alloc.(TSOPhaseDState) + if !ok || (!phaseD.PhaseDRequired() && !phaseD.PhaseDActive()) { + return nil, nil, false, nil + } + validator, ok := alloc.(DurableTimestampValidator) + if !ok { + return nil, phaseD, false, ErrTSOProtocolUnsupported + } + return validator, phaseD, true, nil +} + func nextTimestampFromAllocator(ctx context.Context, alloc TimestampAllocator, label string) (uint64, error) { ts, err := alloc.Next(ctx) if err != nil { @@ -292,8 +546,8 @@ func (a *LocalTSOAllocator) NextBatchAfter(ctx context.Context, n int, min uint6 } func (a *LocalTSOAllocator) nextBatchAfter(ctx context.Context, n int, min uint64) (uint64, error) { - if n <= 0 { - return 0, errors.WithStack(ErrInvalidTSOBatchSize) + if err := validateTSOMinimumWindow(min, n); err != nil { + return 0, err } if err := a.waitLeader(ctx); err != nil { return 0, err @@ -360,10 +614,11 @@ type windowSnapshot struct { // TSOAllocator. The hot path is lock-free: callers claim a slot with atomic Add // on the currently published window. type BatchAllocator struct { - tso TSOAllocator - batchSize int - win atomic.Pointer[windowSnapshot] - epoch atomic.Uint64 + tso TSOAllocator + batchSize int + win atomic.Pointer[windowSnapshot] + epoch atomic.Uint64 + phaseDSeen atomic.Bool mu sync.Mutex refillDone chan struct{} @@ -398,12 +653,38 @@ func (b *BatchAllocator) NextAfter(ctx context.Context, min uint64) (uint64, err return b.nextAfter(ctx, min) } +func (b *BatchAllocator) ValidateDurableTimestamp(ctx context.Context, timestamp uint64) error { + validator, ok := b.tso.(DurableTimestampValidator) + if !ok { + return errors.WithStack(ErrTSOProtocolUnsupported) + } + return errors.WithStack(validator.ValidateDurableTimestamp(ctx, timestamp)) +} + +func (b *BatchAllocator) PhaseDActive() bool { + state, ok := b.tso.(TSOPhaseDState) + return ok && state.PhaseDActive() +} + +func (b *BatchAllocator) PhaseDRequired() bool { + state, ok := b.tso.(TSOPhaseDState) + return ok && state.PhaseDRequired() +} + func (b *BatchAllocator) nextAfter(ctx context.Context, min uint64) (uint64, error) { for { if err := ctxErr(ctx); err != nil { return 0, err } + phaseDAtStart, invalidated := b.ensurePhaseDTransition() + if invalidated { + continue + } if ts, ok := b.tryWindowAfter(min); ok { + phaseDAfterClaim, invalidated := b.ensurePhaseDTransition() + if invalidated || phaseDAfterClaim != phaseDAtStart { + continue + } return ts, nil } if err := b.refill(ctx, min); err != nil { @@ -412,6 +693,21 @@ func (b *BatchAllocator) nextAfter(ctx context.Context, min uint64) (uint64, err } } +func (b *BatchAllocator) ensurePhaseDTransition() (required, invalidated bool) { + if !b.PhaseDRequired() { + return false, false + } + if b.phaseDSeen.Load() { + return true, false + } + // Publish phaseDSeen only after the old epoch is invalidated. Concurrent + // callers may invalidate redundantly, but none can observe the transition as + // complete while a pre-Phase-D window is still current. + b.Invalidate() + b.phaseDSeen.CompareAndSwap(false, true) + return true, true +} + func (b *BatchAllocator) tryWindowAfter(min uint64) (uint64, bool) { w := b.win.Load() if w == nil { diff --git a/kv/tso_floor_test.go b/kv/tso_floor_test.go new file mode 100644 index 000000000..88732c08e --- /dev/null +++ b/kv/tso_floor_test.go @@ -0,0 +1,261 @@ +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 TestShardStoreGroupCommittedTimestampFloorBoundsLocalReadIndex(t *testing.T) { + t.Parallel() + + local := store.NewMVCCStore() + require.NoError(t, local.PutAt(context.Background(), []byte("local"), []byte("v"), 88, 0)) + deadlineSeen := false + engine := &recordingTSOEngine{ + state: raftengine.StateLeader, + readIndex: func(ctx context.Context) (uint64, error) { + _, deadlineSeen = ctx.Deadline() + return 0, nil + }, + } + shards := NewShardStore(nil, map[uint64]*ShardGroup{ + 1: {Engine: engine, Store: local}, + }) + t.Cleanup(func() { require.NoError(t, shards.Close()) }) + + floor, err := shards.GroupCommittedTimestampFloor(context.Background(), 1) + require.NoError(t, err) + require.Equal(t, uint64(88), floor) + require.True(t, deadlineSeen, "local TSO floor ReadIndex must be bounded even when caller has no deadline") +} + +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 e721bb740..986904096 100644 --- a/kv/tso_fsm.go +++ b/kv/tso_fsm.go @@ -1,27 +1,57 @@ package kv import ( + "bufio" + "bytes" "encoding/binary" "io" "sync/atomic" + "github.com/bootjp/elastickv/internal/encryption/fsmwire" "github.com/bootjp/elastickv/internal/raftengine" "github.com/cockroachdb/errors" ) -const tsoSnapshotLen = 8 -const maxHLCPhysicalMillis = int64((uint64(1) << hlcPhysicalBits) - 1) - var _ raftengine.StateMachine = (*TSOStateMachine)(nil) -var _ raftengine.Snapshot = (*tsoSnapshot)(nil) +var _ raftengine.Snapshot = (*tsoFSMSnapshot)(nil) var _ raftengine.VolatileEntryClassifier = (*TSOStateMachine)(nil) -// TSOStateMachine is the minimal FSM for the dedicated timestamp-oracle Raft -// group. It tracks only the Raft-agreed HLC physical ceiling; no KV state, -// route catalog, or sidecar state is attached to group 0. +var ( + ErrTSOStateMachineInvalidEntry = errors.New("tso fsm: invalid entry") + ErrTSOLegacyEncryptionEntryRejected = errors.New("tso fsm: legacy encryption entry rejected") +) + +const ( + // tsoAllocationFloorEnvelope starts in kvFSM's fail-closed encryption + // range so old data-group FSMs halt on a misrouted entry. The remaining + // magic and version bytes keep it distinct from every encryption opcode. + tsoAllocationFloorEnvelope = "\x07TSOF\x01" + // 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" + // tsoPhaseDEnvelope retains the rolling-upgrade halt prefix and gives the + // irreversible Phase-D marker its own exact wire identity. + tsoPhaseDEnvelope = "\x07TSOD\x01" + tsoSnapshotV1Len = hlcLeasePayloadLen + tsoSnapshotV2Len = hlcLeasePayloadLen * 2 + tsoSnapshotV3Len = tsoSnapshotV2Len + 1 + tsoSnapshotV4Len = tsoSnapshotV3Len + 1 + hlcLeasePayloadLen + + maxHLCPhysicalMillis = int64((uint64(1) << hlcPhysicalBits) - 1) +) + +// TSOStateMachine is the minimal state machine for the dedicated timestamp +// group. It accepts HLC lease-renewal entries plus explicit allocation-floor +// entries. The HLC is only a volatile mirror; snapshots are sourced from the +// TSO FSM's own applied state so unrelated shard-group lease renewals cannot +// advance group-0 state outside the group-0 log. type TSOStateMachine struct { - hlc *HLC - ceilingMs atomic.Int64 + hlc *HLC + ceilingMs atomic.Int64 + allocationFloor atomic.Uint64 + cutoverActive atomic.Bool + phaseDActive atomic.Bool + phaseDFloor atomic.Uint64 } // NewTSOStateMachine constructs the dedicated TSO FSM over the shared HLC. @@ -30,119 +60,574 @@ func NewTSOStateMachine(hlc *HLC) *TSOStateMachine { } func (f *TSOStateMachine) Apply(data []byte) any { - if len(data) == 0 || data[0] != raftEncodeHLCLease { - return nil + if len(data) == 0 { + return haltErr(errors.Wrap(ErrTSOStateMachineInvalidEntry, "empty entry")) + } + switch { + case data[0] == raftEncodeHLCLease: + 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 bytes.HasPrefix(data, []byte(tsoPhaseDEnvelope)): + return f.applyPhaseDEntry(data) + case data[0] >= fsmwire.OpEncryptionMin && data[0] <= fsmwire.OpEncryptionMax: + return rejectLegacyTSOEncryptionEntry(data) + default: + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, "unexpected tag 0x%02x", data[0])) } - return f.applyHLCLease(data[1:]) } -func (f *TSOStateMachine) applyHLCLease(data []byte) any { - if len(data) != hlcLeasePayloadLen { - return errors.Newf("tso fsm: hlc lease: expected %d bytes, got %d", hlcLeasePayloadLen, len(data)) //nolint:wrapcheck // creating new error, nothing to wrap +// rejectLegacyTSOEncryptionEntry lets an upgraded group 0 advance past a +// control entry committed while it still ran kvFSM. New group-0 listeners do +// not expose encryption mutators, so these entries can only be historical. +// Validate the old wire shape before returning an ordinary apply response: +// malformed control bytes still halt, while a valid obsolete entry is +// deterministically rejected without mutating TSO state or halting replay. +func rejectLegacyTSOEncryptionEntry(data []byte) any { + var err error + switch data[0] { + case fsmwire.OpRegistration: + _, err = fsmwire.DecodeRegistration(data[1:]) + case fsmwire.OpBootstrap: + _, err = fsmwire.DecodeBootstrap(data[1:]) + case fsmwire.OpRotation: + _, err = fsmwire.DecodeRotation(data[1:]) + default: + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "reserved encryption entry tag 0x%02x", data[0])) } - ceilingMs, err := decodeTSOCeiling(binary.BigEndian.Uint64(data)) if err != nil { - return err + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "malformed legacy encryption entry tag 0x%02x: %v", data[0], err)) + } + return errors.Wrapf(ErrTSOLegacyEncryptionEntryRejected, + "tag 0x%02x is not part of dedicated TSO state", data[0]) +} + +func (f *TSOStateMachine) applyLeaseEntry(data []byte) any { + if len(data) != hlcLeaseEntryLen { + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, "expected HLC lease entry length %d, got %d", hlcLeaseEntryLen, len(data))) + } + ceilingMs, err := decodeTSOCeiling(binary.BigEndian.Uint64(data[1:]), "HLC lease") + if err != nil { + return haltErr(err) + } + if ceilingMs <= 0 { + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, "non-positive HLC lease ceiling %d", ceilingMs)) + } + if f != nil { + f.applyLeaseCeiling(ceilingMs) } - f.applyTSOCeiling(ceilingMs, false) return nil } +func (f *TSOStateMachine) applyAllocationFloorEntry(data []byte) any { + expectedLen := len(tsoAllocationFloorEnvelope) + hlcLeasePayloadLen + if len(data) != expectedLen { + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, "expected TSO allocation floor entry length %d, got %d", expectedLen, len(data))) + } + floor := binary.BigEndian.Uint64(data[len(tsoAllocationFloorEnvelope):]) + if floor == 0 { + return haltErr(errors.Wrap(ErrTSOStateMachineInvalidEntry, "zero TSO allocation floor")) + } + if f != nil { + f.applyAllocationFloor(floor) + } + return nil +} + +func (f *TSOStateMachine) 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 +} + +func (f *TSOStateMachine) applyPhaseDEntry(data []byte) any { + expectedLen := len(tsoPhaseDEnvelope) + hlcLeasePayloadLen + if len(data) != expectedLen { + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "expected TSO phase-D entry length %d, got %d", expectedLen, len(data))) + } + if f == nil { + return nil + } + if !f.cutoverActive.Load() { + return haltErr(errors.Wrap(ErrTSOStateMachineInvalidEntry, + "TSO phase-D marker requires the durable cutover marker")) + } + floor := binary.BigEndian.Uint64(data[len(tsoPhaseDEnvelope):]) + if f.phaseDActive.Load() && f.phaseDFloor.Load() != floor { + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "TSO phase-D floor changed from %d to %d", f.phaseDFloor.Load(), floor)) + } + if !f.phaseDActive.Load() && floor < f.allocationFloor.Load() { + return haltErr(errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "TSO phase-D floor %d is below allocation floor %d", floor, f.allocationFloor.Load())) + } + f.phaseDFloor.Store(floor) + f.phaseDActive.Store(true) + return nil +} + +// AllocationFloor returns the highest timestamp window end applied by the +// dedicated TSO group. It is consensus-owned state, unlike HLC.Current(). +func (f *TSOStateMachine) AllocationFloor() uint64 { + 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() +} + +// PhaseDActive reports whether the compatibility window has been durably +// closed. Once active, data-shard HLC renewal and caller-supplied cross-shard +// timestamps may no longer use legacy issuance semantics. +func (f *TSOStateMachine) PhaseDActive() bool { + return f != nil && f.phaseDActive.Load() +} + +// PhaseDFloor is the highest allocation floor that existed when Phase D was +// activated. Only timestamps reserved strictly above it are valid M7 durable +// read/start allocations. +func (f *TSOStateMachine) PhaseDFloor() uint64 { + if f == nil { + return 0 + } + return f.phaseDFloor.Load() +} + func (f *TSOStateMachine) Snapshot() (raftengine.Snapshot, error) { - return &tsoSnapshot{ceilingMs: f.committedCeiling()}, nil + var ceilingMs int64 + var allocationFloor uint64 + var cutoverActive bool + var phaseDActive bool + var phaseDFloor uint64 + if f != nil { + ceilingMs = f.ceilingMs.Load() + allocationFloor = f.allocationFloor.Load() + cutoverActive = f.cutoverActive.Load() + phaseDActive = f.phaseDActive.Load() + phaseDFloor = f.phaseDFloor.Load() + } + return &tsoFSMSnapshot{ + ceilingMs: ceilingMs, + allocationFloor: allocationFloor, + cutoverActive: cutoverActive, + phaseDActive: phaseDActive, + phaseDFloor: phaseDFloor, + }, nil } func (f *TSOStateMachine) Restore(r io.Reader) error { - var buf [tsoSnapshotLen]byte - if _, err := io.ReadFull(r, buf[:]); err != nil { - return errors.Wrap(err, "tso fsm: restore snapshot") - } - var extra [1]byte - n, err := r.Read(extra[:]) - if err != nil && !errors.Is(err, io.EOF) { - return errors.Wrap(err, "tso fsm: restore snapshot") + if r == nil { + return errors.New("tso fsm snapshot: reader is nil") } - if n != 0 { - return errors.New("tso fsm: restore snapshot: trailing bytes") //nolint:wrapcheck // creating new error, nothing to wrap + br := tsoSnapshotReader(r) + if legacy, err := restoreLegacyKVFSMSnapshot(f, br); legacy || err != nil { + return err } - ceilingMs, err := decodeTSOCeiling(binary.BigEndian.Uint64(buf[:])) + ceilingMs, allocationFloor, cutoverActive, phaseDActive, phaseDFloor, err := readTSOSnapshotState(br) if err != nil { - return errors.Wrap(err, "tso fsm: restore snapshot") + return err + } + if f != nil { + f.restoreSnapshotState(ceilingMs, allocationFloor, cutoverActive, phaseDActive, phaseDFloor) } - f.applyTSOCeiling(ceilingMs, true) return nil } -func decodeTSOCeiling(raw uint64) (int64, error) { +func tsoSnapshotReader(r io.Reader) *bufio.Reader { + if br, ok := r.(*bufio.Reader); ok { + return br + } + return bufio.NewReader(r) +} + +func readTSOSnapshotState(br *bufio.Reader) (int64, uint64, bool, bool, uint64, error) { + payload, err := io.ReadAll(io.LimitReader(br, tsoSnapshotV4Len+1)) + if err != nil { + return 0, 0, false, false, 0, errors.Wrap(err, "restore tso fsm snapshot") + } + ceilingMs, allocationFloor, cutoverActive, phaseDActive, phaseDFloor, legacySnapshot, err := decodeTSOSnapshotPayload(payload) + if err != nil { + return 0, 0, false, false, 0, err + } + if ceilingMs < 0 { + return 0, 0, false, false, 0, errors.Wrapf(ErrTSOStateMachineInvalidEntry, "tso fsm snapshot: negative ceiling %d", ceilingMs) + } + if legacySnapshot && ceilingMs > 0 { + allocationFloor = tsoLeaseAllocationFloor(ceilingMs) + } + return ceilingMs, allocationFloor, cutoverActive, phaseDActive, phaseDFloor, nil +} + +func decodeTSOSnapshotPayload(payload []byte) (int64, uint64, bool, bool, uint64, bool, error) { + var ceilingMs int64 + var allocationFloor uint64 + var cutoverActive bool + var phaseDActive bool + var phaseDFloor uint64 + var legacySnapshot bool + switch len(payload) { + case tsoSnapshotV1Len: + legacySnapshot = true + var err error + ceilingMs, err = decodeTSOCeiling(binary.BigEndian.Uint64(payload[:hlcLeasePayloadLen]), "legacy snapshot") + if err != nil { + return 0, 0, false, false, 0, false, err + } + case tsoSnapshotV2Len: + var err error + ceilingMs, err = decodeTSOCeiling(binary.BigEndian.Uint64(payload[:hlcLeasePayloadLen]), "snapshot") + if err != nil { + return 0, 0, false, false, 0, false, err + } + allocationFloor = binary.BigEndian.Uint64(payload[hlcLeasePayloadLen:]) + case tsoSnapshotV3Len: + var err error + ceilingMs, err = decodeTSOCeiling(binary.BigEndian.Uint64(payload[:hlcLeasePayloadLen]), "snapshot") + if err != nil { + return 0, 0, false, false, 0, false, err + } + allocationFloor = binary.BigEndian.Uint64(payload[hlcLeasePayloadLen:tsoSnapshotV2Len]) + cutoverActive, err = decodeTSOCutoverByte(payload[tsoSnapshotV2Len]) + if err != nil { + return 0, 0, false, false, 0, false, err + } + case tsoSnapshotV4Len: + return decodeTSOSnapshotV4(payload) + default: + return 0, 0, false, false, 0, false, errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "tso fsm snapshot: expected %d, %d, %d, or %d bytes, got %d", + tsoSnapshotV1Len, tsoSnapshotV2Len, tsoSnapshotV3Len, tsoSnapshotV4Len, len(payload)) + } + return ceilingMs, allocationFloor, cutoverActive, phaseDActive, phaseDFloor, legacySnapshot, nil +} + +func decodeTSOSnapshotV4(payload []byte) (int64, uint64, bool, bool, uint64, bool, error) { + ceilingMs, err := decodeTSOCeiling(binary.BigEndian.Uint64(payload[:hlcLeasePayloadLen]), "snapshot") + if err != nil { + return 0, 0, false, false, 0, false, err + } + allocationFloor := binary.BigEndian.Uint64(payload[hlcLeasePayloadLen:tsoSnapshotV2Len]) + cutoverActive, err := decodeTSOCutoverByte(payload[tsoSnapshotV2Len]) + if err != nil { + return 0, 0, false, false, 0, false, err + } + phaseDActive, err := decodeTSOBooleanByte("phase-D", payload[tsoSnapshotV3Len]) + if err != nil { + return 0, 0, false, false, 0, false, err + } + phaseDFloor := binary.BigEndian.Uint64(payload[tsoSnapshotV3Len+1:]) + if phaseDActive && !cutoverActive { + return 0, 0, false, false, 0, false, errors.Wrap(ErrTSOStateMachineInvalidEntry, + "tso fsm snapshot: phase-D active without cutover") + } + if !phaseDActive && phaseDFloor != 0 { + return 0, 0, false, false, 0, false, errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "tso fsm snapshot: inactive phase-D has floor %d", phaseDFloor) + } + return ceilingMs, allocationFloor, cutoverActive, phaseDActive, phaseDFloor, false, nil +} + +func decodeTSOCeiling(raw uint64, field string) (int64, error) { if raw > uint64(maxHLCPhysicalMillis) { - return 0, errors.Newf("tso fsm: hlc lease: physical ceiling %d exceeds max %d", raw, maxHLCPhysicalMillis) //nolint:wrapcheck // creating new error, nothing to wrap + return 0, errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "tso fsm: %s physical ceiling %d exceeds max %d", field, raw, maxHLCPhysicalMillis) } - return int64(raw), nil //nolint:gosec // raw is bounded by maxHLCPhysicalMillis. + return int64(raw), nil // #nosec G115 -- raw is bounded by maxHLCPhysicalMillis above. } -func (f *TSOStateMachine) applyTSOCeiling(ceilingMs int64, restore bool) { - if ceilingMs <= 0 { +func decodeTSOCutoverByte(value byte) (bool, error) { + return decodeTSOBooleanByte("cutover", value) +} + +func decodeTSOBooleanByte(name string, value byte) (bool, error) { + switch value { + case 0: + return false, nil + case 1: + return true, nil + default: + return false, errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "tso fsm snapshot: invalid %s byte %d", name, value) + } +} + +// restoreLegacyKVFSMSnapshot migrates snapshots produced while reserved group +// 0 still used kvFSM as a compatibility bridge. Only the HLC header is TSO +// state; the empty MVCC payload is drained so the raft engine can verify the +// complete snapshot CRC. Non-kvFSM snapshots are left untouched in br. +func restoreLegacyKVFSMSnapshot(f *TSOStateMachine, br *bufio.Reader) (bool, error) { + legacy, err := hasLegacyKVFSMSnapshotHeader(br) + if err != nil { + return false, err + } + if !legacy { + return restoreHeaderlessLegacyKVFSMSnapshot(br) + } + ceiling, _, err := ReadSnapshotHeader(br) + if err != nil { + return true, errors.Wrap(err, "tso fsm snapshot: read legacy kv fsm header") + } + if _, err := io.Copy(io.Discard, br); err != nil { + return true, errors.Wrap(err, "tso fsm snapshot: drain legacy kv fsm payload") + } + ceilingMs, err := decodeTSOCeiling(ceiling, "legacy kv fsm snapshot") + if err != nil { + return true, err + } + if f == nil || ceilingMs == 0 { + return true, nil + } + f.restoreSnapshotState(ceilingMs, tsoLeaseAllocationFloor(ceilingMs), false, false, 0) + return true, nil +} + +func restoreHeaderlessLegacyKVFSMSnapshot(br *bufio.Reader) (bool, error) { + headerless, err := hasHeaderlessLegacyKVFSMSnapshotPayload(br) + if err != nil || !headerless { + return headerless, err + } + if _, err := io.Copy(io.Discard, br); err != nil { + return true, errors.Wrap(err, "tso fsm snapshot: drain headerless legacy kv fsm payload") + } + return true, nil +} + +func hasLegacyKVFSMSnapshotHeader(br *bufio.Reader) (bool, error) { + peeked, err := br.Peek(len(hlcSnapshotMagic)) + if err != nil { + if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) { + return false, nil + } + return false, errors.Wrap(err, "tso fsm snapshot: peek legacy header") + } + switch { + case isV1Magic(peeked): + return true, nil + case isV2Magic(peeked): + return true, nil + case isUnknownEKVTHLC(peeked): + return true, nil + case isLegacyKVFSMStoreSnapshot(peeked): + return true, nil + default: + return false, nil + } +} + +func hasHeaderlessLegacyKVFSMSnapshotPayload(br *bufio.Reader) (bool, error) { + peeked, err := br.Peek(tsoSnapshotV4Len + 1) + if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, io.ErrUnexpectedEOF) { + return false, errors.Wrap(err, "tso fsm snapshot: peek headerless legacy payload") + } + return len(peeked) > tsoSnapshotV4Len, nil +} + +func isLegacyKVFSMStoreSnapshot(peeked []byte) bool { + for _, magic := range legacyKVFSMStoreSnapshotMagics() { + if bytes.Equal(peeked, magic) { + return true + } + } + return false +} + +func legacyKVFSMStoreSnapshotMagics() [][]byte { + return [][]byte{ + []byte("EKVMVCC2"), + []byte("EKVPBBL1"), + []byte("EKVSSTI1"), + } +} + +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)) || + len(payload) == len(tsoPhaseDEnvelope)+hlcLeasePayloadLen && + bytes.HasPrefix(payload, []byte(tsoPhaseDEnvelope)) +} + +func (f *TSOStateMachine) applyLeaseCeiling(ceilingMs int64) { + if f == nil || ceilingMs <= 0 { return } - prevCeiling, advanced := f.advanceCommittedCeiling(ceilingMs) - if !advanced && !restore { + storeMaxInt64(&f.ceilingMs, ceilingMs) + if f.hlc != nil { + f.hlc.SetPhysicalCeiling(f.ceilingMs.Load()) + } +} + +func (f *TSOStateMachine) applyAllocationFloor(floor uint64) { + if f == nil || floor == 0 { return } + storeMaxUint64(&f.allocationFloor, floor) if f.hlc != nil { - floorCeiling := prevCeiling - if restore { - floorCeiling = ceilingMs + f.hlc.Observe(f.allocationFloor.Load()) + } +} + +func (f *TSOStateMachine) restoreSnapshotState( + ceilingMs int64, + allocationFloor uint64, + cutoverActive bool, + phaseDActive bool, + phaseDFloor uint64, +) { + if f == nil { + return + } + if ceilingMs > 0 { + storeMaxInt64(&f.ceilingMs, ceilingMs) + } + if allocationFloor > 0 { + storeMaxUint64(&f.allocationFloor, allocationFloor) + } + if cutoverActive { + f.cutoverActive.Store(true) + } + if phaseDActive { + f.phaseDFloor.Store(phaseDFloor) + f.phaseDActive.Store(true) + } + if f.hlc != nil { + if currentCeiling := f.ceilingMs.Load(); currentCeiling > 0 { + f.hlc.SetPhysicalCeiling(currentCeiling) + } + if currentFloor := f.allocationFloor.Load(); currentFloor > 0 { + f.hlc.Observe(currentFloor) } - if floorCeiling > 0 { - f.hlc.Observe(tsoCeilingMaxTimestamp(floorCeiling)) + } +} + +func storeMaxInt64(value *atomic.Int64, candidate int64) { + for { + current := value.Load() + if candidate <= current { + return + } + if value.CompareAndSwap(current, candidate) { + return } - f.hlc.SetPhysicalCeiling(ceilingMs) } } -func (f *TSOStateMachine) advanceCommittedCeiling(ceilingMs int64) (int64, bool) { +func storeMaxUint64(value *atomic.Uint64, candidate uint64) { for { - prev := f.ceilingMs.Load() - if ceilingMs <= prev { - return prev, false + current := value.Load() + if candidate <= current { + return } - if f.ceilingMs.CompareAndSwap(prev, ceilingMs) { - return prev, true + if value.CompareAndSwap(current, candidate) { + return } } } -func (f *TSOStateMachine) committedCeiling() int64 { - return f.ceilingMs.Load() +func tsoLeaseAllocationFloor(ceilingMs int64) uint64 { + return (nonNegativeUint64(ceilingMs) << hlcLogicalBits) | hlcLogicalMask } -func tsoCeilingMaxTimestamp(ceilingMs int64) uint64 { - return (uint64(ceilingMs) << hlcLogicalBits) | hlcLogicalMask //nolint:gosec // ceilingMs is validated positive before conversion. +func marshalTSOAllocationFloor(floor uint64) []byte { + out := make([]byte, len(tsoAllocationFloorEnvelope)+hlcLeasePayloadLen) + copy(out, tsoAllocationFloorEnvelope) + binary.BigEndian.PutUint64(out[len(tsoAllocationFloorEnvelope):], floor) + return out } -// IsVolatileOnlyPayload classifies HLC lease entries for the cold-start replay -// gate. Re-applying them is monotonic and reconstructs the in-memory ceiling. -func (f *TSOStateMachine) IsVolatileOnlyPayload(payload []byte) bool { - return len(payload) > 0 && payload[0] == raftEncodeHLCLease +func marshalTSOCutover() []byte { + return []byte(tsoCutoverEnvelope) +} + +func marshalTSOPhaseD(floor uint64) []byte { + out := make([]byte, len(tsoPhaseDEnvelope)+hlcLeasePayloadLen) + copy(out, tsoPhaseDEnvelope) + binary.BigEndian.PutUint64(out[len(tsoPhaseDEnvelope):], floor) + return out } -type tsoSnapshot struct { - ceilingMs int64 +type tsoFSMSnapshot struct { + ceilingMs int64 + allocationFloor uint64 + cutoverActive bool + phaseDActive bool + phaseDFloor uint64 } -func (s *tsoSnapshot) WriteTo(w io.Writer) (int64, error) { - var buf [tsoSnapshotLen]byte - binary.BigEndian.PutUint64(buf[:], uint64(s.ceilingMs)) //nolint:gosec // ceilingMs is a Unix ms timestamp encoded as uint64. - n, err := w.Write(buf[:]) +func (s *tsoFSMSnapshot) WriteTo(w io.Writer) (int64, error) { + if w == nil { + return 0, errors.New("tso fsm snapshot: writer is nil") + } + var ceilingMs int64 + var allocationFloor uint64 + var cutoverActive bool + var phaseDActive bool + var phaseDFloor uint64 + if s != nil { + ceilingMs = s.ceilingMs + allocationFloor = s.allocationFloor + cutoverActive = s.cutoverActive + phaseDActive = s.phaseDActive + phaseDFloor = s.phaseDFloor + } + snapshotLen := tsoSnapshotV3Len + if phaseDActive { + snapshotLen = tsoSnapshotV4Len + } + buf := make([]byte, snapshotLen) + if err := encodeTSOCeiling(buf, ceilingMs); err != nil { + return 0, err + } + binary.BigEndian.PutUint64(buf[hlcLeasePayloadLen:tsoSnapshotV2Len], allocationFloor) + if cutoverActive { + buf[tsoSnapshotV2Len] = 1 + } + if phaseDActive { + buf[tsoSnapshotV3Len] = 1 + binary.BigEndian.PutUint64(buf[tsoSnapshotV3Len+1:], phaseDFloor) + } + n, err := w.Write(buf) if err != nil { - return int64(n), errors.WithStack(err) + return int64(n), errors.Wrap(err, "write tso fsm snapshot") } - if n != tsoSnapshotLen { + if n != len(buf) { return int64(n), errors.WithStack(io.ErrShortWrite) } - return tsoSnapshotLen, nil + return int64(n), nil +} + +func encodeTSOCeiling(dst []byte, ceilingMs int64) error { + if len(dst) < hlcLeasePayloadLen { + return errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "tso fsm snapshot: ceiling buffer too short: %d", len(dst)) + } + if ceilingMs < 0 || ceilingMs > maxHLCPhysicalMillis { + return errors.Wrapf(ErrTSOStateMachineInvalidEntry, + "tso fsm snapshot: invalid ceiling %d", ceilingMs) + } + binary.BigEndian.PutUint64(dst, uint64(ceilingMs)) // #nosec G115 -- ceilingMs is validated non-negative above. + return nil } -func (s *tsoSnapshot) Close() error { +func (s *tsoFSMSnapshot) Close() error { return nil } diff --git a/kv/tso_fsm_test.go b/kv/tso_fsm_test.go index 5c3d65151..269397994 100644 --- a/kv/tso_fsm_test.go +++ b/kv/tso_fsm_test.go @@ -1,149 +1,347 @@ package kv import ( + "bufio" "bytes" "encoding/binary" "io" "testing" "time" + "github.com/bootjp/elastickv/internal/encryption/fsmwire" + "github.com/bootjp/elastickv/store" "github.com/stretchr/testify/require" ) func TestTSOStateMachineApplyHLCLeaseUpdatesCeiling(t *testing.T) { t.Parallel() - const ceilingMs = int64(9_999_999_999_999) - clock := NewHLC() - fsm := NewTSOStateMachine(clock) + const ceilingMs = int64(1_700_000_123_456) + hlc := NewHLC() + fsm := NewTSOStateMachine(hlc) - require.Nil(t, fsm.Apply(marshalHLCLeaseRenew(ceilingMs))) - require.Equal(t, ceilingMs, clock.PhysicalCeiling()) - require.Zero(t, clock.Current()) - require.Equal(t, ceilingMs, fsm.committedCeiling()) + result := fsm.Apply(marshalHLCLeaseRenew(ceilingMs)) + require.Nil(t, result) + require.Equal(t, ceilingMs, hlc.PhysicalCeiling()) + require.Zero(t, hlc.Current()) + require.Equal(t, ceilingMs, fsm.ceilingMs.Load()) } -func TestTSOStateMachineApplyHLCLeaseKeepsRenewedWindowAllocatable(t *testing.T) { +func TestTSOStateMachineApplyRejectsOutOfRangeHLCLease(t *testing.T) { t.Parallel() - firstCeilingMs := time.Now().Add(time.Hour).UnixMilli() - secondCeilingMs := firstCeilingMs + 100 clock := NewHLC() + clock.Observe(123) + clock.SetPhysicalCeiling(456) fsm := NewTSOStateMachine(clock) - require.Nil(t, fsm.Apply(marshalHLCLeaseRenew(firstCeilingMs))) - first, err := clock.NextBatchFenced(1) - require.NoError(t, err) - require.EqualValues(t, firstCeilingMs, first>>hlcLogicalBits) + err := requireTSOHaltError(t, fsm.Apply(marshalRawHLCLeaseRenew(uint64(maxHLCPhysicalMillis)+1))) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) + require.Equal(t, uint64(123), clock.Current()) + require.Equal(t, int64(456), clock.PhysicalCeiling()) + require.Zero(t, fsm.ceilingMs.Load()) +} + +func TestTSOStateMachineApplyAllocationFloorAdvancesHLC(t *testing.T) { + t.Parallel() + + ceilingMs := time.Now().Add(time.Hour).UnixMilli() + hlc := NewHLC() + fsm := NewTSOStateMachine(hlc) + + require.Nil(t, fsm.Apply(marshalHLCLeaseRenew(ceilingMs))) + floor := tsoLeaseAllocationFloor(ceilingMs - 1) + require.Zero(t, hlc.Current()) - require.Nil(t, fsm.Apply(marshalHLCLeaseRenew(secondCeilingMs))) - require.Equal(t, tsoCeilingMaxTimestamp(firstCeilingMs), clock.Current()) - second, err := clock.NextBatchFenced(1) + require.Nil(t, fsm.Apply(marshalTSOAllocationFloor(floor))) + require.Equal(t, floor, hlc.Current()) + + base, err := hlc.NextBatchFenced(1) require.NoError(t, err) - require.EqualValues(t, secondCeilingMs, second>>hlcLogicalBits) - require.Equal(t, uint64(0), clock.NextFencedRejections()) + require.Greater(t, base, floor) } -func TestTSOStateMachineApplyRejectsOutOfRangeHLCLease(t *testing.T) { +func TestTSOStateMachineRejectsNonLeaseEntry(t *testing.T) { t.Parallel() - clock := NewHLC() - clock.Observe(123) - clock.SetPhysicalCeiling(456) - fsm := NewTSOStateMachine(clock) + payload := make([]byte, hlcLeaseEntryLen) + payload[0] = raftEncodeSingle - resp := fsm.Apply(marshalRawHLCLeaseRenew(uint64(maxHLCPhysicalMillis) + 1)) - requireApplyError(t, resp) - require.Equal(t, uint64(123), clock.Current()) - require.Equal(t, int64(456), clock.PhysicalCeiling()) - require.Zero(t, fsm.committedCeiling()) + err := requireTSOHaltError(t, NewTSOStateMachine(NewHLC()).Apply(payload)) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) } -func TestTSOStateMachineApplyIgnoresNonLeasePayload(t *testing.T) { +func TestTSOStateMachineRejectsLegacyEncryptionEntriesWithoutHalting(t *testing.T) { t.Parallel() - clock := NewHLC() - fsm := NewTSOStateMachine(clock) + registration := fsmwire.RegistrationPayload{DEKID: 1, FullNodeID: 2, LocalEpoch: 3} + tests := []struct { + name string + opcode byte + payload []byte + }{ + { + name: "registration", + opcode: fsmwire.OpRegistration, + payload: fsmwire.EncodeRegistration(registration), + }, + { + name: "bootstrap", + opcode: fsmwire.OpBootstrap, + payload: fsmwire.EncodeBootstrap(fsmwire.BootstrapPayload{ + StorageDEKID: 1, + WrappedStorage: []byte("storage"), + RaftDEKID: 2, + WrappedRaft: []byte("raft"), + BatchRegistry: []fsmwire.RegistrationPayload{registration}, + }), + }, + { + name: "rotation", + opcode: fsmwire.OpRotation, + payload: fsmwire.EncodeRotation(fsmwire.RotationPayload{ + SubTag: fsmwire.RotateSubRotateDEK, + DEKID: 4, + Purpose: fsmwire.PurposeStorage, + Wrapped: []byte("rotated"), + ProposerRegistration: registration, + }), + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + entry := append([]byte{tc.opcode}, tc.payload...) + result := NewTSOStateMachine(NewHLC()).Apply(entry) + err, ok := result.(error) + require.Truef(t, ok, "legacy control response type = %T, want ordinary error", result) + require.ErrorIs(t, err, ErrTSOLegacyEncryptionEntryRejected) + _, halts := result.(interface{ HaltApply() error }) + require.False(t, halts, "valid legacy control entry must not halt replay") + }) + } +} + +func TestTSOStateMachineHaltsMalformedLegacyEncryptionEntry(t *testing.T) { + t.Parallel() + + err := requireTSOHaltError(t, NewTSOStateMachine(NewHLC()).Apply([]byte{fsmwire.OpRegistration})) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) +} + +func TestTSOStateMachineRejectsMalformedLease(t *testing.T) { + t.Parallel() + + for _, payload := range [][]byte{ + {}, + {raftEncodeHLCLease}, + append([]byte{raftEncodeHLCLease}, make([]byte, hlcLeasePayloadLen+1)...), + } { + err := requireTSOHaltError(t, NewTSOStateMachine(NewHLC()).Apply(payload)) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) + } +} + +func TestTSOStateMachineRejectsNonPositiveLeaseCeiling(t *testing.T) { + t.Parallel() + + for _, ceilingMs := range []int64{0, -1} { + err := requireTSOHaltError(t, NewTSOStateMachine(NewHLC()).Apply(marshalHLCLeaseRenew(ceilingMs))) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) + } +} + +func TestTSOStateMachineRejectsMalformedAllocationFloor(t *testing.T) { + t.Parallel() + + for _, payload := range [][]byte{ + []byte(tsoAllocationFloorEnvelope), + append([]byte(tsoAllocationFloorEnvelope), make([]byte, hlcLeasePayloadLen+1)...), + marshalTSOAllocationFloor(0), + } { + err := requireTSOHaltError(t, NewTSOStateMachine(NewHLC()).Apply(payload)) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) + } +} + +func TestTSOStateMachineRejectsBareEncryptionReservedAllocationFloor(t *testing.T) { + t.Parallel() - require.Nil(t, fsm.Apply([]byte{raftEncodeSingle})) - require.Zero(t, clock.PhysicalCeiling()) + payload := make([]byte, hlcLeaseEntryLen) + payload[0] = fsmwire.OpEncryptionMax + binary.BigEndian.PutUint64(payload[1:], 42) + err := requireTSOHaltError(t, NewTSOStateMachine(NewHLC()).Apply(payload)) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) } -func TestTSOStateMachineRejectsMalformedHLCLease(t *testing.T) { +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") - requireApplyError(t, fsm.Apply([]byte{raftEncodeHLCLease})) - requireApplyError(t, fsm.Apply(append([]byte{raftEncodeHLCLease}, make([]byte, hlcLeasePayloadLen+1)...))) + err := requireTSOHaltError(t, fsm.Apply(append([]byte(tsoCutoverEnvelope), 1))) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) } -func TestTSOStateMachineSnapshotRestoreRoundTrip(t *testing.T) { +func TestTSOStateMachineAppliesOneWayPhaseDMarker(t *testing.T) { t.Parallel() - ceilingMs := time.Now().Add(time.Hour).UnixMilli() - sourceClock := NewHLC() - source := NewTSOStateMachine(sourceClock) - require.Nil(t, source.Apply(marshalHLCLeaseRenew(ceilingMs))) + const floor = uint64(1234) + fsm := NewTSOStateMachine(NewHLC()) - snap, err := source.Snapshot() + err := requireTSOHaltError(t, fsm.Apply(marshalTSOPhaseD(floor))) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) + require.False(t, fsm.PhaseDActive()) + + require.Nil(t, fsm.Apply(marshalTSOCutover())) + require.Nil(t, fsm.Apply(marshalTSOPhaseD(floor))) + require.True(t, fsm.PhaseDActive()) + require.Equal(t, floor, fsm.PhaseDFloor()) + require.Nil(t, fsm.Apply(marshalTSOPhaseD(floor)), "phase-D replay must be idempotent") + + err = requireTSOHaltError(t, fsm.Apply(marshalTSOPhaseD(floor+1))) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) + require.ErrorContains(t, err, "floor changed") + err = requireTSOHaltError(t, fsm.Apply([]byte(tsoPhaseDEnvelope))) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) +} + +func TestTSOStateMachineRejectsPhaseDFloorBelowExistingAllocation(t *testing.T) { + t.Parallel() + + fsm := NewTSOStateMachine(NewHLC()) + require.Nil(t, fsm.Apply(marshalTSOAllocationFloor(100))) + require.Nil(t, fsm.Apply(marshalTSOCutover())) + err := requireTSOHaltError(t, fsm.Apply(marshalTSOPhaseD(99))) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) + require.ErrorContains(t, err, "below allocation floor") + require.False(t, fsm.PhaseDActive()) +} + +func TestTSOStateMachineRejectsInvalidCutoverSnapshotByte(t *testing.T) { + t.Parallel() + + 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 TestTSOStateMachineRejectsInvalidPhaseDSnapshot(t *testing.T) { + t.Parallel() + + payload := make([]byte, tsoSnapshotV4Len) + payload[tsoSnapshotV2Len] = 1 + payload[tsoSnapshotV3Len] = 2 + err := NewTSOStateMachine(NewHLC()).Restore(bytes.NewReader(payload)) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) + require.ErrorContains(t, err, "invalid phase-D byte") + + payload[tsoSnapshotV2Len] = 0 + payload[tsoSnapshotV3Len] = 1 + err = NewTSOStateMachine(NewHLC()).Restore(bytes.NewReader(payload)) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) + require.ErrorContains(t, err, "phase-D active without cutover") + + payload[tsoSnapshotV3Len] = 0 + binary.BigEndian.PutUint64(payload[tsoSnapshotV3Len+1:], 1) + err = NewTSOStateMachine(NewHLC()).Restore(bytes.NewReader(payload)) + require.ErrorIs(t, err, ErrTSOStateMachineInvalidEntry) + require.ErrorContains(t, err, "inactive phase-D has floor") +} + +func TestTSOStateMachineNilHLCDoesNotPanic(t *testing.T) { + t.Parallel() + + require.Nil(t, NewTSOStateMachine(nil).Apply(marshalHLCLeaseRenew(1_700_000_123_456))) + require.Nil(t, NewTSOStateMachine(nil).Apply(marshalTSOAllocationFloor(1))) +} + +func TestTSOStateMachineSnapshotWithNilHLCWritesZeroState(t *testing.T) { + t.Parallel() + + fsm := NewTSOStateMachine(nil) + snap, err := fsm.Snapshot() require.NoError(t, err) defer func() { require.NoError(t, snap.Close()) }() var buf bytes.Buffer n, err := snap.WriteTo(&buf) require.NoError(t, err) - require.EqualValues(t, 8, n) - require.Len(t, buf.Bytes(), 8) - - restoredClock := NewHLC() - restored := NewTSOStateMachine(restoredClock) - require.NoError(t, restored.Restore(bytes.NewReader(buf.Bytes()))) - require.Equal(t, ceilingMs, restoredClock.PhysicalCeiling()) - require.Equal(t, tsoCeilingMaxTimestamp(ceilingMs), restoredClock.Current()) - require.Equal(t, ceilingMs, restored.committedCeiling()) - _, err = restoredClock.NextBatchFenced(1) - require.ErrorIs(t, err, ErrCeilingExpired) + require.EqualValues(t, tsoSnapshotV3Len, n) + require.Equal(t, make([]byte, tsoSnapshotV3Len), buf.Bytes()) } -func TestTSOStateMachineSnapshotUsesCommittedCeiling(t *testing.T) { +func TestTSOStateMachineSnapshotRestoreRoundTrip(t *testing.T) { t.Parallel() - committedCeilingMs := time.Now().Add(time.Hour).UnixMilli() - unrelatedCeilingMs := committedCeilingMs + int64(time.Hour/time.Millisecond) - sourceClock := NewHLC() - source := NewTSOStateMachine(sourceClock) - - require.Nil(t, source.Apply(marshalHLCLeaseRenew(committedCeilingMs))) - sourceClock.SetPhysicalCeiling(unrelatedCeilingMs) + const ceilingMs = int64(1_700_000_654_321) + sourceHLC := NewHLC() + source := NewTSOStateMachine(sourceHLC) + require.Nil(t, source.Apply(marshalHLCLeaseRenew(ceilingMs))) + floor := tsoLeaseAllocationFloor(ceilingMs) + require.Nil(t, source.Apply(marshalTSOAllocationFloor(floor))) + require.Nil(t, source.Apply(marshalTSOCutover())) + require.Nil(t, source.Apply(marshalTSOPhaseD(floor))) + postPhaseDFloor := floor + 10 + require.Nil(t, source.Apply(marshalTSOAllocationFloor(postPhaseDFloor))) snap, err := source.Snapshot() require.NoError(t, err) defer func() { require.NoError(t, snap.Close()) }() var buf bytes.Buffer - _, err = snap.WriteTo(&buf) + n, err := snap.WriteTo(&buf) require.NoError(t, err) - - restoredClock := NewHLC() - restored := NewTSOStateMachine(restoredClock) - require.NoError(t, restored.Restore(bytes.NewReader(buf.Bytes()))) - require.Equal(t, committedCeilingMs, restoredClock.PhysicalCeiling()) - require.Equal(t, tsoCeilingMaxTimestamp(committedCeilingMs), restoredClock.Current()) + require.EqualValues(t, tsoSnapshotV4Len, n) + require.Len(t, buf.Bytes(), tsoSnapshotV4Len) + + targetHLC := NewHLC() + target := NewTSOStateMachine(targetHLC) + require.NoError(t, target.Restore(bytes.NewReader(buf.Bytes()))) + require.Equal(t, ceilingMs, targetHLC.PhysicalCeiling()) + require.Equal(t, postPhaseDFloor, targetHLC.Current()) + require.True(t, target.CutoverActive()) + require.True(t, target.PhaseDActive()) + require.Equal(t, floor, target.PhaseDFloor()) } -func TestTSOStateMachineSnapshotWithNilHLCWritesZero(t *testing.T) { +func TestTSOStateMachineSnapshotUsesTSOOwnedCeiling(t *testing.T) { t.Parallel() - fsm := NewTSOStateMachine(nil) - snap, err := fsm.Snapshot() + const ( + tsoCeiling = int64(1_000) + unrelatedCeiling = int64(2_000) + ) + sourceHLC := NewHLC() + source := NewTSOStateMachine(sourceHLC) + require.Nil(t, source.Apply(marshalHLCLeaseRenew(tsoCeiling))) + sourceHLC.SetPhysicalCeiling(unrelatedCeiling) + + snap, err := source.Snapshot() require.NoError(t, err) defer func() { require.NoError(t, snap.Close()) }() var buf bytes.Buffer _, err = snap.WriteTo(&buf) require.NoError(t, err) - require.Equal(t, make([]byte, 8), buf.Bytes()) + + targetHLC := NewHLC() + require.NoError(t, NewTSOStateMachine(targetHLC).Restore(bytes.NewReader(buf.Bytes()))) + require.Equal(t, tsoCeiling, targetHLC.PhysicalCeiling()) + require.Zero(t, targetHLC.Current()) +} + +func TestTSOStateMachineRestoreRejectsTruncatedSnapshot(t *testing.T) { + t.Parallel() + + err := NewTSOStateMachine(NewHLC()).Restore(bytes.NewReader([]byte{0x01, 0x02})) + require.Error(t, err) } func TestTSOStateMachineRestoreRejectsOutOfRangeSnapshot(t *testing.T) { @@ -154,48 +352,197 @@ func TestTSOStateMachineRestoreRejectsOutOfRangeSnapshot(t *testing.T) { clock.SetPhysicalCeiling(456) fsm := NewTSOStateMachine(clock) - var buf [tsoSnapshotLen]byte + var buf [tsoSnapshotV1Len]byte binary.BigEndian.PutUint64(buf[:], uint64(maxHLCPhysicalMillis)+1) require.Error(t, fsm.Restore(bytes.NewReader(buf[:]))) require.Equal(t, uint64(123), clock.Current()) require.Equal(t, int64(456), clock.PhysicalCeiling()) - require.Zero(t, fsm.committedCeiling()) + require.Zero(t, fsm.ceilingMs.Load()) +} + +func TestTSOStateMachineLegacySnapshotProbeRejectsEveryShortMagicPrefix(t *testing.T) { + t.Parallel() + + for size := range len(hlcSnapshotMagic) { + payload := append([]byte(nil), hlcSnapshotMagic[:size]...) + legacy, err := hasLegacyKVFSMSnapshotHeader(bufio.NewReader(bytes.NewReader(payload))) + require.NoError(t, err) + require.False(t, legacy) + } +} + +func TestTSOStateMachineRestoreLegacySnapshotDerivesAllocationFloor(t *testing.T) { + t.Parallel() + + const ceilingMs = int64(1_700_000_654_321) + var buf [hlcLeasePayloadLen]byte + binary.BigEndian.PutUint64(buf[:], uint64(ceilingMs)) + + hlc := NewHLC() + require.NoError(t, NewTSOStateMachine(hlc).Restore(bytes.NewReader(buf[:]))) + require.Equal(t, ceilingMs, hlc.PhysicalCeiling()) + require.Equal(t, tsoLeaseAllocationFloor(ceilingMs), hlc.Current()) } -func TestTSOStateMachineRestoreRejectsMalformedSnapshot(t *testing.T) { +func TestTSOStateMachineRestoresV3SnapshotWithoutPhaseD(t *testing.T) { t.Parallel() + const ( + ceilingMs = int64(1_700_000_654_321) + floor = uint64(4567) + ) + payload := make([]byte, tsoSnapshotV3Len) + binary.BigEndian.PutUint64(payload[:hlcLeasePayloadLen], uint64(ceilingMs)) + binary.BigEndian.PutUint64(payload[hlcLeasePayloadLen:tsoSnapshotV2Len], floor) + payload[tsoSnapshotV2Len] = 1 + fsm := NewTSOStateMachine(NewHLC()) + require.NoError(t, fsm.Restore(bytes.NewReader(payload))) + require.Equal(t, floor, fsm.AllocationFloor()) + require.True(t, fsm.CutoverActive()) + require.False(t, fsm.PhaseDActive()) + require.Zero(t, fsm.PhaseDFloor()) +} + +func TestTSOStateMachineRestoreKeepsMonotonicCeiling(t *testing.T) { + t.Parallel() + + const ( + higherCeiling = int64(2_000) + lowerCeiling = int64(1_000) + ) + hlc := NewHLC() + fsm := NewTSOStateMachine(hlc) + require.Nil(t, fsm.Apply(marshalHLCLeaseRenew(higherCeiling))) + require.Nil(t, fsm.Apply(marshalTSOAllocationFloor(tsoLeaseAllocationFloor(higherCeiling)))) + + var buf [tsoSnapshotV2Len]byte + binary.BigEndian.PutUint64(buf[:], uint64(lowerCeiling)) + + require.NoError(t, fsm.Restore(bytes.NewReader(buf[:]))) + require.Equal(t, higherCeiling, hlc.PhysicalCeiling()) + require.Equal(t, tsoLeaseAllocationFloor(higherCeiling), hlc.Current()) + + snap, err := fsm.Snapshot() + require.NoError(t, err) + defer func() { require.NoError(t, snap.Close()) }() + + var snapBuf bytes.Buffer + n, err := snap.WriteTo(&snapBuf) + require.NoError(t, err) + require.EqualValues(t, tsoSnapshotV3Len, n) + require.Equal(t, uint64(higherCeiling), binary.BigEndian.Uint64(snapBuf.Bytes()[:hlcLeasePayloadLen])) + require.Equal(t, tsoLeaseAllocationFloor(higherCeiling), binary.BigEndian.Uint64(snapBuf.Bytes()[hlcLeasePayloadLen:tsoSnapshotV2Len])) +} + +func TestTSOStateMachineRestoresLegacyKVFSMSnapshot(t *testing.T) { + const ceilingMs = int64(1_700_000_123_456) + tests := []struct { + name string + opts []FSMOption + }{ + {name: "v1"}, + {name: "v2", opts: []FSMOption{WithCutoverSource(&staticCutoverSource{v: 42})}}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + legacyClock := NewHLC() + legacyClock.SetPhysicalCeiling(ceilingMs) + legacyFSM := NewKvFSMWithHLC(store.NewMVCCStore(), legacyClock, tc.opts...) + legacySnapshot, err := legacyFSM.Snapshot() + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, legacySnapshot.Close()) }) + + var legacy bytes.Buffer + _, err = legacySnapshot.WriteTo(&legacy) + require.NoError(t, err) + + targetClock := NewHLC() + target := NewTSOStateMachine(targetClock) + require.NoError(t, target.Restore(bytes.NewReader(legacy.Bytes()))) + require.Equal(t, ceilingMs, targetClock.PhysicalCeiling()) + require.Equal(t, tsoLeaseAllocationFloor(ceilingMs), targetClock.Current()) + + got, err := target.Snapshot() + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, got.Close()) }) + var payload bytes.Buffer + _, err = got.WriteTo(&payload) + require.NoError(t, err) + require.Len(t, payload.Bytes(), tsoSnapshotV3Len) + require.Equal(t, uint64(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]) + }) + } +} + +func TestTSOStateMachineDrainsHeaderlessLegacyKVFSMSnapshot(t *testing.T) { + t.Parallel() + + legacySnapshot, err := store.NewMVCCStore().Snapshot() + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, legacySnapshot.Close()) }) - require.Error(t, fsm.Restore(bytes.NewReader(nil))) - require.Error(t, fsm.Restore(bytes.NewReader(make([]byte, 9)))) + var legacy bytes.Buffer + _, err = legacySnapshot.WriteTo(&legacy) + require.NoError(t, err) + + targetClock := NewHLC() + target := NewTSOStateMachine(targetClock) + require.NoError(t, target.Restore(bytes.NewReader(legacy.Bytes()))) + require.Zero(t, targetClock.PhysicalCeiling()) + require.Zero(t, targetClock.Current()) +} + +func TestTSOStateMachineDrainsLongHeaderlessLegacyKVFSMSnapshot(t *testing.T) { + t.Parallel() + + legacy := bytes.Repeat([]byte("legacy-kvfsm-headerless-body"), 2) + require.Greater(t, len(legacy), tsoSnapshotV4Len) + + targetClock := NewHLC() + target := NewTSOStateMachine(targetClock) + require.NoError(t, target.Restore(bytes.NewReader(legacy))) + require.Zero(t, targetClock.PhysicalCeiling()) + require.Zero(t, targetClock.Current()) } func TestTSOStateMachineSnapshotWriteWrapsShortWrite(t *testing.T) { t.Parallel() - snap := &tsoSnapshot{ceilingMs: 1} + snap := &tsoFSMSnapshot{ceilingMs: 1} n, err := snap.WriteTo(shortTSOWriter{}) - require.EqualValues(t, tsoSnapshotLen-1, n) + require.EqualValues(t, tsoSnapshotV3Len-1, n) require.ErrorIs(t, err, io.ErrShortWrite) } -func TestTSOStateMachineClassifiesHLCLeaseAsVolatileOnly(t *testing.T) { +func TestTSOStateMachineClassifiesOnlyFullLeaseEntriesAsVolatile(t *testing.T) { t.Parallel() fsm := NewTSOStateMachine(NewHLC()) - - require.True(t, fsm.IsVolatileOnlyPayload(marshalHLCLeaseRenew(1))) + require.True(t, fsm.IsVolatileOnlyPayload(marshalHLCLeaseRenew(1_700_000_123_456))) + require.True(t, fsm.IsVolatileOnlyPayload(marshalTSOAllocationFloor(1))) + require.True(t, fsm.IsVolatileOnlyPayload(marshalTSOCutover())) + require.True(t, fsm.IsVolatileOnlyPayload(marshalTSOPhaseD(1))) + require.False(t, fsm.IsVolatileOnlyPayload([]byte{raftEncodeHLCLease})) + require.False(t, fsm.IsVolatileOnlyPayload([]byte(tsoAllocationFloorEnvelope))) + require.False(t, fsm.IsVolatileOnlyPayload(append([]byte(tsoCutoverEnvelope), 1))) + require.False(t, fsm.IsVolatileOnlyPayload([]byte(tsoPhaseDEnvelope))) require.False(t, fsm.IsVolatileOnlyPayload([]byte{raftEncodeSingle})) - require.False(t, fsm.IsVolatileOnlyPayload(nil)) } -func requireApplyError(t *testing.T, got any) { +func requireTSOHaltError(t *testing.T, result any) error { t.Helper() - err, ok := got.(error) - require.True(t, ok) + if _, ok := result.(error); ok { + t.Fatalf("expected HaltApply response, got plain error %T", result) + } + halt, ok := result.(interface{ HaltApply() error }) + require.Truef(t, ok, "expected HaltApply response, got %T", result) + err := halt.HaltApply() require.Error(t, err) + return err } func marshalRawHLCLeaseRenew(raw uint64) []byte { diff --git a/kv/tso_raft.go b/kv/tso_raft.go new file mode 100644 index 000000000..ca3f37558 --- /dev/null +++ b/kv/tso_raft.go @@ -0,0 +1,1005 @@ +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 + PhaseDActive bool + PhaseDFloor uint64 +} + +// TSOReservationAllocator is the migration-aware extension exposed by the +// dedicated group leader. Ordinary callers continue to use TSOAllocator. +type TSOReservationAllocator interface { + ReserveBatchAfter(context.Context, int, uint64, bool, bool) (TSOReservation, error) +} + +type TSOShadowReservationAllocator interface { + ValidateShadowTimestamp(context.Context, uint64) (TSOReservation, error) +} + +type tsoCutoverState interface { + CutoverActive() bool +} + +// RaftTSOAllocator reserves timestamp windows on the dedicated TSO leader. +// A window is returned only after its inclusive end has committed to Raft. +// Failed proposals may leak a local window, but they can never expose an +// 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, 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, + activatePhaseD bool, +) (TSOReservation, error) { + var empty TSOReservation + if err := validateTSOMinimumWindow(min, n); err != nil { + return empty, err + } + ctx = nonNilTSOContext(ctx) + a.mu.Lock() + defer a.mu.Unlock() + return a.reserveBatchAfterLocked(ctx, n, min, activateCutover, activatePhaseD) +} + +func (a *RaftTSOAllocator) reserveBatchAfterLocked( + ctx context.Context, + n int, + min uint64, + activateCutover bool, + activatePhaseD 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, commitPhaseD, err := a.prepareLeaderTermReservation(ctx, engine, term, activateCutover, activatePhaseD) + if err != nil { + return empty, err + } + base, end, err := a.reservePreparedLocalWindow(n, min, previousFloor, commitPhaseD) + if err != nil { + return empty, err + } + reservationFloor := previousFloor + if commitPhaseD { + reservationFloor, err = a.commitPhaseDReservation(ctx, engine, term, base) + if err != nil { + return empty, err + } + } + 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: reservationFloor, + CutoverActive: a.state.CutoverActive(), + PhaseDActive: a.state.PhaseDActive(), + PhaseDFloor: a.state.PhaseDFloor(), + }, nil +} + +func (a *RaftTSOAllocator) reservePreparedLocalWindow( + n int, + min uint64, + previousFloor uint64, + commitPhaseD bool, +) (uint64, uint64, error) { + phaseDContiguous := commitPhaseD || a.state.PhaseDActive() + if phaseDContiguous && min > previousFloor { + return 0, 0, errors.Wrapf(ErrTSOTimestampInvalid, + "phase-D TSO minimum %d exceeds contiguous allocation floor %d", min, previousFloor) + } + return a.reserveLocalWindow(n, max(min, previousFloor), phaseDContiguous) +} + +func (a *RaftTSOAllocator) reserveLocalWindow(n int, minimum uint64, phaseDContiguous bool) (uint64, uint64, error) { + if !phaseDContiguous { + if minimum > 0 { + a.clock.Observe(minimum) + } + base, err := a.clock.NextBatchFenced(n) + if err != nil { + return 0, 0, errors.Wrap(err, "tso reserve local batch") + } + end := base + positiveIntToUint64(n) - 1 + return base, end, nil + } + return a.reserveContiguousPhaseDWindow(n, minimum) +} + +func (a *RaftTSOAllocator) reserveContiguousPhaseDWindow(n int, floor uint64) (uint64, uint64, error) { + if err := validateTSOMinimumWindow(floor, n); err != nil { + return 0, 0, err + } + base := floor + 1 + end := base + positiveIntToUint64(n) - 1 + if err := a.validateContiguousPhaseDWindow(end); err != nil { + return 0, 0, err + } + a.clock.Observe(end) + return base, end, nil +} + +func (a *RaftTSOAllocator) validateContiguousPhaseDWindow(end uint64) error { + ceiling := a.clock.PhysicalCeiling() + if ceiling <= 0 { + return nil + } + if time.Now().UnixMilli() >= ceiling { + a.clock.nextFencedRejections.Add(1) + return errors.Wrap(ErrCeilingExpired, "tso reserve contiguous phase-D window") + } + if end>>hlcLogicalBits > uint64(ceiling) { + a.clock.nextFencedRejections.Add(1) + return errors.Wrap(ErrCeilingExpired, "tso contiguous phase-D window exceeds ceiling") + } + return nil +} + +func (a *RaftTSOAllocator) commitPhaseDReservation( + ctx context.Context, + engine raftengine.Engine, + term uint64, + base uint64, +) (uint64, error) { + if base == 0 { + return 0, errors.Wrap(ErrTxnCommitTSRequired, "tso phase-D reservation base is zero") + } + reservationFloor := base - 1 + if err := a.commitPhaseD(ctx, engine, reservationFloor); err != nil { + return 0, err + } + if err := verifyTSOLeaderTerm(ctx, engine, term, true); err != nil { + return 0, err + } + return reservationFloor, nil +} + +func (a *RaftTSOAllocator) prepareLeaderTermReservation( + ctx context.Context, + engine raftengine.Engine, + term uint64, + activateCutover bool, + activatePhaseD bool, +) (uint64, bool, error) { + // The ordinary per-term cache protects all reservations served by this TSO + // leader term. Phase-D activation is a one-way cutover boundary, so it must + // fence legacy data-group commits that may have landed after this term's + // first non-Phase-D reservation. + forceFreshFloor := activatePhaseD && !a.state.PhaseDActive() + termFloor, err := a.termCommitFloor(ctx, term, forceFreshFloor) + if err != nil { + return 0, false, err + } + previousFloor := max(a.state.AllocationFloor(), termFloor, a.state.PhaseDFloor()) + commitPhaseD, err := a.activateDurableMarkers(ctx, engine, activateCutover, activatePhaseD) + if err != nil { + return 0, false, err + } + if err := verifyTSOLeaderTerm(ctx, engine, term, true); err != nil { + return 0, false, err + } + return previousFloor, commitPhaseD, nil +} + +func (a *RaftTSOAllocator) activateDurableMarkers( + ctx context.Context, + engine raftengine.Engine, + activateCutover bool, + activatePhaseD bool, +) (bool, error) { + if activateCutover && !a.state.CutoverActive() { + if err := a.commitCutover(ctx, engine); err != nil { + return false, err + } + } + if !activatePhaseD || a.state.PhaseDActive() { + return false, nil + } + if !a.state.CutoverActive() { + return false, errors.Wrap(ErrTSOPhaseDInactive, "phase D activation requires durable cutover") + } + return true, nil +} + +func (a *RaftTSOAllocator) termCommitFloor(ctx context.Context, term uint64, forceFresh bool) (uint64, error) { + if a.initializedTerm == term && !forceFresh { + return a.state.AllocationFloor(), nil + } + floor, err := a.floorProvider.GlobalCommittedTimestampFloor(ctx) + 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 (a *RaftTSOAllocator) commitPhaseD(ctx context.Context, engine raftengine.Engine, floor uint64) error { + if _, err := a.group.Proposer().Propose(ctx, marshalTSOPhaseD(floor)); err != nil { + wrapped := errors.Wrap(err, "tso commit phase-D marker") + if tsoProposalLostLeadership(engine, err) { + return stderrors.Join(ErrTSONotLeader, wrapped) + } + return wrapped + } + if !a.state.PhaseDActive() || a.state.PhaseDFloor() != floor { + return errors.New("tso phase-D proposal committed without applied marker") + } + return nil +} + +func (a *RaftTSOAllocator) ValidateDurableTimestamp(ctx context.Context, timestamp uint64) error { + ctx = nonNilTSOContext(ctx) + a.mu.Lock() + defer a.mu.Unlock() + if _, err := a.verifiedLeader(ctx); err != nil { + return err + } + if !a.state.PhaseDActive() { + return errors.WithStack(ErrTSOPhaseDInactive) + } + floor := a.state.PhaseDFloor() + end := a.state.AllocationFloor() + if timestamp == 0 || timestamp > end { + return errors.Wrapf(ErrTSOTimestampInvalid, + "timestamp=%d phase_d_floor=%d allocation_floor=%d", timestamp, floor, end) + } + if timestamp <= floor { + return errors.Wrapf(stderrors.Join(ErrTSOTimestampInvalid, ErrTSOTimestampPrePhaseD), + "timestamp=%d phase_d_floor=%d allocation_floor=%d", timestamp, floor, end) + } + return nil +} + +func (a *RaftTSOAllocator) PhaseDActive() bool { + return a != nil && a.state != nil && a.state.PhaseDActive() +} + +func (a *RaftTSOAllocator) PhaseDRequired() bool { + return a.PhaseDActive() +} + +func (a *RaftTSOAllocator) PhaseDFloor() uint64 { + if a == nil || a.state == nil { + return 0 + } + return a.state.PhaseDFloor() +} + +func (a *RaftTSOAllocator) AllocationFloor() uint64 { + if a == nil || a.state == nil { + return 0 + } + return a.state.AllocationFloor() +} + +func tsoProposalLostLeadership(engine raftengine.Engine, err error) bool { + return !isLeaderEngine(engine) || errors.Is(err, raftengine.ErrNotLeader) || isTransientLeaderError(err) +} + +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, bool) (TSOReservation, error) + +type tsoRemoteValidation func(context.Context, string, uint64) error + +// LeaderRoutedTSOAllocator serves local requests on the TSO leader and sends +// follower requests to the leader address published by the group-0 engine. +// It re-resolves that address after transient errors so a leadership change +// does not pin a BatchAllocator refill to a stale endpoint. +type LeaderRoutedTSOAllocator struct { + local TSOAllocator + leader raftengine.LeaderView + connCache GRPCConnCache + remoteRequest tsoRemoteRequest + retryBudget time.Duration + retryInterval time.Duration + activate bool + activatePhaseD bool + clock *HLC + remoteValidate tsoRemoteValidation +} + +type LeaderRoutedTSOAllocatorOption func(*LeaderRoutedTSOAllocator) + +func WithTSOCutoverActivation() LeaderRoutedTSOAllocatorOption { + return func(a *LeaderRoutedTSOAllocator) { a.activate = true } +} + +func WithTSOPhaseDActivation() LeaderRoutedTSOAllocatorOption { + return func(a *LeaderRoutedTSOAllocator) { + a.activate = true + a.activatePhaseD = true + } +} + +func WithTSORoutedClock(clock *HLC) LeaderRoutedTSOAllocatorOption { + return func(a *LeaderRoutedTSOAllocator) { a.clock = clock } +} + +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 + a.remoteValidate = a.requestRemoteValidation + 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.activatePhaseD, 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, false, true) + if err == nil { + a.observeReservation(reservation) + } + return reservation, err +} + +func (a *LeaderRoutedTSOAllocator) nextReservation( + ctx context.Context, + n int, + min uint64, + activate bool, + activatePhaseD bool, + requireReservation bool, +) (TSOReservation, error) { + var empty TSOReservation + if err := validateTSOMinimumWindow(min, 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, activatePhaseD, 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, + activatePhaseD bool, + requireReservation bool, +) (TSOReservation, error) { + var empty TSOReservation + if a.local.IsLeader() { + return a.tryLocalBatch(ctx, n, min, activate, activatePhaseD, requireReservation) + } + addr := leaderAddrFromEngine(a.leader) + if addr == "" { + return empty, errors.WithStack(ErrLeaderNotFound) + } + reservation, err := a.remoteRequest(ctx, addr, n, min, activate, activatePhaseD) + 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, + activatePhaseD 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, activatePhaseD) + 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, + activatePhaseD 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, + ActivatePhaseD: activatePhaseD, + }) + 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") + } + if activatePhaseD && !resp.GetPhaseDActive() { + return empty, errors.Wrap(ErrTSOProtocolUnsupported, + "leader did not confirm durable TSO phase D") + } + return TSOReservation{ + Base: resp.GetTimestamp(), + Count: n, + PreviousAllocationFloor: resp.GetPreviousAllocationFloor(), + CutoverActive: resp.GetCutoverActive(), + PhaseDActive: resp.GetPhaseDActive(), + PhaseDFloor: resp.GetPhaseDFloor(), + }, nil +} + +func (a *LeaderRoutedTSOAllocator) ValidateDurableTimestamp(ctx context.Context, timestamp uint64) error { + if timestamp == 0 { + return errors.WithStack(ErrTSOTimestampInvalid) + } + ctx = nonNilTSOContext(ctx) + deadline := time.Now().Add(a.retryBudget) + ctx, cancel := context.WithDeadline(ctx, deadline) + defer cancel() + var lastErr error + for { + if a.local.IsLeader() { + validator, ok := a.local.(DurableTimestampValidator) + if !ok { + return errors.WithStack(ErrTSOProtocolUnsupported) + } + lastErr = errors.WithStack(validator.ValidateDurableTimestamp(ctx, timestamp)) + } else { + addr := leaderAddrFromEngine(a.leader) + if addr == "" { + lastErr = errors.WithStack(ErrLeaderNotFound) + } else { + lastErr = a.remoteValidate(ctx, addr, timestamp) + } + } + if lastErr == nil { + return nil + } + if !isTransientTSORouteError(lastErr) { + return lastErr + } + if err := waitTSORouteRetry(ctx, a.retryInterval); err != nil { + return errors.Wrap(stderrors.Join(ctx.Err(), lastErr), "tso durable timestamp validation") + } + } +} + +func (a *LeaderRoutedTSOAllocator) requestRemoteValidation(ctx context.Context, addr string, timestamp uint64) error { + conn, err := a.connCache.ConnFor(addr) + if err != nil { + return errors.Wrap(err, "tso dial validation leader") + } + resp, err := pb.NewDistributionClient(conn).ValidateTimestamp(ctx, &pb.ValidateTimestampRequest{Timestamp: timestamp}) + if err != nil { + code := status.Code(err) + if code == codes.OutOfRange { + return errors.Wrap(stderrors.Join(ErrTSOTimestampInvalid, ErrTSOTimestampPrePhaseD), err.Error()) + } + if code == codes.InvalidArgument { + return errors.Wrap(ErrTSOTimestampInvalid, err.Error()) + } + return errors.Wrap(err, "tso validate timestamp at leader") + } + if !resp.GetValid() || !resp.GetPhaseDActive() { + return errors.Wrapf(ErrTSOTimestampInvalid, + "leader validation valid=%t phase_d_active=%t", resp.GetValid(), resp.GetPhaseDActive()) + } + return nil +} + +func (a *LeaderRoutedTSOAllocator) PhaseDActive() bool { + state, ok := a.local.(TSOPhaseDState) + return ok && state.PhaseDActive() +} + +func (a *LeaderRoutedTSOAllocator) PhaseDRequired() bool { + return a != nil && (a.activatePhaseD || a.PhaseDActive()) +} + +func (a *LeaderRoutedTSOAllocator) observeReservation(reservation TSOReservation) { + if a == nil || a.clock == nil || reservation.Base == 0 || reservation.Count <= 0 { + return + } + 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 validateTSOMinimumWindow(min uint64, n int) error { + if err := validateTSOBatchSize(n); err != nil { + return err + } + if min > ^uint64(0)-uint64(n) { //nolint:gosec // n is already positive and bounded. + return errors.Wrapf(ErrTxnCommitTSRequired, + "tso minimum %d cannot fit a %d-timestamp window", min, n) + } + 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 + cutover tsoCutoverState + log *slog.Logger +} + +type ShadowTimestampAllocatorOption func(*ShadowTimestampAllocator) + +func WithTSOShadowCutoverState(state tsoCutoverState) ShadowTimestampAllocatorOption { + return func(a *ShadowTimestampAllocator) { a.cutover = state } +} + +func NewShadowTimestampAllocator( + legacy *HLC, + shadow TSOShadowReservationAllocator, + logger *slog.Logger, + opts ...ShadowTimestampAllocatorOption, +) (*ShadowTimestampAllocator, error) { + if legacy == nil { + return nil, errors.WithStack(ErrTSOClockNil) + } + if shadow == nil { + return nil, errors.WithStack(ErrTSOAllocatorRequired) + } + if logger == nil { + logger = slog.Default() + } + a := &ShadowTimestampAllocator{ + legacy: legacy, + shadow: shadow, + log: logger, + } + for _, opt := range opts { + if opt != nil { + opt(a) + } + } + return a, 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) + if a.cutover != nil && a.cutover.CutoverActive() { + return a.nextDedicatedAfter(ctx, min) + } + 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) nextDedicatedAfter(ctx context.Context, min uint64) (uint64, error) { + alloc, ok := a.shadow.(TimestampAllocator) + if !ok { + return 0, errors.WithStack(ErrTSOProtocolUnsupported) + } + return nextTimestampAfterFromAllocator(ctx, alloc, min, "tso post-cutover shadow bypass") +} + +func (a *ShadowTimestampAllocator) ValidateDurableTimestamp(ctx context.Context, timestamp uint64) error { + validator, ok := a.shadow.(DurableTimestampValidator) + if !ok { + return errors.WithStack(ErrTSOProtocolUnsupported) + } + return errors.WithStack(validator.ValidateDurableTimestamp(ctx, timestamp)) +} + +func (a *ShadowTimestampAllocator) PhaseDActive() bool { + state, ok := a.shadow.(TSOPhaseDState) + return ok && state.PhaseDActive() +} + +func (a *ShadowTimestampAllocator) PhaseDRequired() bool { + state, ok := a.shadow.(TSOPhaseDState) + return ok && state.PhaseDRequired() +} + +func (a *ShadowTimestampAllocator) Close() error { + return nil +} diff --git a/kv/tso_raft_test.go b/kv/tso_raft_test.go new file mode 100644 index 000000000..422953aec --- /dev/null +++ b/kv/tso_raft_test.go @@ -0,0 +1,874 @@ +package kv + +import ( + "bytes" + "context" + "encoding/binary" + "io" + "log/slog" + "sync" + "sync/atomic" + "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 TestRaftTSOAllocatorRejectsTermChangeAfterPhaseDMarker(t *testing.T) { + clock := NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) + fsm := NewTSOStateMachine(clock) + engine := &recordingTSOEngine{ + state: raftengine.StateLeader, + leader: raftengine.LeaderInfo{Address: "self"}, + term: 1, + apply: applyTSOTestFSM(fsm), + } + var changeTerm sync.Once + engine.afterPropose = func(payload []byte) { + if bytes.HasPrefix(payload, []byte(tsoPhaseDEnvelope)) { + changeTerm.Do(func() { engine.setTerm(2) }) + } + } + provider := &recordingTSOFloorProvider{floor: 100} + alloc, err := NewRaftTSOAllocator( + &ShardGroup{Engine: engine, TSOState: fsm}, + clock, + WithTSOCutoverFloorProvider(provider), + ) + require.NoError(t, err) + + _, err = alloc.ReserveBatchAfter(context.Background(), testTSOBatchSize, 0, true, true) + require.ErrorIs(t, err, ErrTSONotLeader) + require.True(t, fsm.PhaseDActive()) + require.Zero(t, fsm.AllocationFloor()) + + payloads := engine.proposedPayloads() + require.Len(t, payloads, 2) + require.Equal(t, []byte(tsoCutoverEnvelope), payloads[0]) + require.True(t, bytes.HasPrefix(payloads[1], []byte(tsoPhaseDEnvelope))) + require.Greater(t, fsm.PhaseDFloor(), uint64(0)) + require.Equal(t, marshalTSOPhaseD(fsm.PhaseDFloor()), payloads[1]) +} + +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, false) + 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 TestRaftTSOAllocatorCommitsPhaseDBeforeWindowAndValidatesRange(t *testing.T) { + clock := NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) + fsm := NewTSOStateMachine(clock) + engine := &recordingTSOEngine{ + state: raftengine.StateLeader, + leader: raftengine.LeaderInfo{Address: "self"}, + term: 1, + apply: applyTSOTestFSM(fsm), + } + provider := &recordingTSOFloorProvider{floor: 100} + alloc, err := NewRaftTSOAllocator( + &ShardGroup{Engine: engine, TSOState: fsm}, + clock, + WithTSOCutoverFloorProvider(provider), + ) + require.NoError(t, err) + + reservation, err := alloc.ReserveBatchAfter(context.Background(), testTSOBatchSize, 0, true, true) + require.NoError(t, err) + require.True(t, reservation.CutoverActive) + require.True(t, reservation.PhaseDActive) + require.Equal(t, reservation.PreviousAllocationFloor, reservation.PhaseDFloor) + require.Equal(t, reservation.Base-1, reservation.PhaseDFloor) + require.Greater(t, reservation.Base, reservation.PhaseDFloor) + + payloads := engine.proposedPayloads() + require.Len(t, payloads, 3) + require.Equal(t, []byte(tsoCutoverEnvelope), payloads[0]) + require.Equal(t, marshalTSOPhaseD(reservation.Base-1), payloads[1]) + require.True(t, bytes.HasPrefix(payloads[2], []byte(tsoAllocationFloorEnvelope))) + + end := reservation.Base + positiveIntToUint64(reservation.Count) - 1 + require.NoError(t, alloc.ValidateDurableTimestamp(context.Background(), reservation.Base)) + require.NoError(t, alloc.ValidateDurableTimestamp(context.Background(), end)) + err = alloc.ValidateDurableTimestamp(context.Background(), reservation.PhaseDFloor) + require.ErrorIs(t, err, ErrTSOTimestampInvalid) + require.ErrorIs(t, alloc.ValidateDurableTimestamp(context.Background(), 1), ErrTSOTimestampInvalid) + require.ErrorIs(t, alloc.ValidateDurableTimestamp(context.Background(), end+1), ErrTSOTimestampInvalid) +} + +func TestRaftTSOAllocatorPhaseDReservationsStayContiguousAcrossWallClockJump(t *testing.T) { + clock := NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(2 * testTSOFutureCeiling).UnixMilli()) + fsm := NewTSOStateMachine(clock) + engine := &recordingTSOEngine{ + state: raftengine.StateLeader, + leader: raftengine.LeaderInfo{Address: "self"}, + term: 1, + apply: applyTSOTestFSM(fsm), + } + provider := &recordingTSOFloorProvider{} + provider.setFloor(uint64(time.Now().UnixMilli()) << hlcLogicalBits) //nolint:gosec // test floor is positive. + alloc, err := NewRaftTSOAllocator( + &ShardGroup{Engine: engine, TSOState: fsm}, + clock, + WithTSOCutoverFloorProvider(provider), + ) + require.NoError(t, err) + + first, err := alloc.ReserveBatchAfter(context.Background(), testTSOBatchSize, 0, true, true) + require.NoError(t, err) + firstEnd := first.Base + positiveIntToUint64(first.Count) - 1 + require.Equal(t, firstEnd, fsm.AllocationFloor()) + require.Equal(t, first.Base-1, fsm.PhaseDFloor()) + + clock.Observe(uint64(time.Now().Add(testTSOFutureCeiling).UnixMilli()) << hlcLogicalBits) //nolint:gosec // test HLC timestamp is positive. + second, err := alloc.ReserveBatchAfter(context.Background(), testTSOBatchSize, 0, false, false) + require.NoError(t, err) + secondEnd := second.Base + positiveIntToUint64(second.Count) - 1 + require.Equal(t, firstEnd+1, second.Base) + require.NoError(t, alloc.ValidateDurableTimestamp(context.Background(), second.Base)) + require.NoError(t, alloc.ValidateDurableTimestamp(context.Background(), secondEnd)) + require.ErrorIs(t, alloc.ValidateDurableTimestamp(context.Background(), secondEnd+1), ErrTSOTimestampInvalid) + + _, err = alloc.ReserveBatchAfter(context.Background(), 1, secondEnd+10, false, false) + require.ErrorIs(t, err, ErrTSOTimestampInvalid) + require.Equal(t, secondEnd, fsm.AllocationFloor()) +} + +func TestRaftTSOAllocatorRejectsNearOverflowMinimumBeforeObserve(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) + + _, err = alloc.ReserveBatchAfter(context.Background(), 2, ^uint64(0)-1, false, false) + require.ErrorIs(t, err, ErrTxnCommitTSRequired) + require.Zero(t, clock.Current()) + require.Zero(t, fsm.AllocationFloor()) + require.Empty(t, engine.proposedPayloads()) +} + +func TestRaftTSOAllocatorResamplesCommitFloorWhenActivatingPhaseDInInitializedTerm(t *testing.T) { + clock := NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) + fsm := NewTSOStateMachine(clock) + engine := &recordingTSOEngine{ + state: raftengine.StateLeader, + leader: raftengine.LeaderInfo{Address: "self"}, + term: 1, + apply: applyTSOTestFSM(fsm), + } + provider := &recordingTSOFloorProvider{} + alloc, err := NewRaftTSOAllocator( + &ShardGroup{Engine: engine, TSOState: fsm}, + clock, + WithTSOCutoverFloorProvider(provider), + ) + require.NoError(t, err) + + _, err = alloc.Next(context.Background()) + require.NoError(t, err) + require.Equal(t, 1, provider.callCount()) + + laterDataFloor := fsm.AllocationFloor() + 1_000 + provider.setFloor(laterDataFloor) + reservation, err := alloc.ReserveBatchAfter(context.Background(), testTSOBatchSize, 0, true, true) + require.NoError(t, err) + require.Equal(t, 2, provider.callCount(), "Phase-D activation must resample the data commit floor even within an initialized term") + require.GreaterOrEqual(t, reservation.PreviousAllocationFloor, laterDataFloor) + require.Equal(t, reservation.Base-1, reservation.PreviousAllocationFloor) + require.Equal(t, reservation.Base-1, reservation.PhaseDFloor) + require.Greater(t, reservation.Base, laterDataFloor) +} + +func TestRaftTSOAllocatorFencesAboveRestoredPhaseDFloor(t *testing.T) { + clock := NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) + fsm := NewTSOStateMachine(clock) + phaseDFloor := uint64(time.Now().Add(30*time.Minute).UnixMilli()) << hlcLogicalBits //nolint:gosec // future test HLC. + require.Nil(t, fsm.Apply(marshalTSOCutover())) + require.Nil(t, fsm.Apply(marshalTSOPhaseD(phaseDFloor))) + engine := &recordingTSOEngine{ + state: raftengine.StateLeader, + leader: raftengine.LeaderInfo{Address: "self"}, + term: 1, + apply: applyTSOTestFSM(fsm), + } + alloc, err := newTestRaftTSOAllocator(&ShardGroup{Engine: engine, TSOState: fsm}, clock) + require.NoError(t, err) + + reservation, err := alloc.ReserveBatchAfter(context.Background(), 1, 0, false, false) + require.NoError(t, err) + require.Equal(t, phaseDFloor, reservation.PreviousAllocationFloor) + require.Greater(t, reservation.Base, phaseDFloor) + require.NoError(t, alloc.ValidateDurableTimestamp(context.Background(), reservation.Base)) + err = alloc.ValidateDurableTimestamp(context.Background(), phaseDFloor) + require.ErrorIs(t, err, ErrTSOTimestampInvalid) + require.ErrorIs(t, err, ErrTSOTimestampPrePhaseD) +} + +func TestLeaderRoutedTSOAllocatorReResolvesAfterStaleLeader(t *testing.T) { + local := &fakeTSOAllocator{leader: false, nextBase: testTSOInitialBase} + engine := &recordingTSOEngine{state: raftengine.StateFollower, leader: raftengine.LeaderInfo{Address: "old"}} + 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, activatePhaseD bool) (TSOReservation, error) { + addresses = append(addresses, addr) + require.Equal(t, testTSOBatchSize, n) + require.Equal(t, uint64(testTSOInitialBase), min) + require.False(t, activate) + require.False(t, activatePhaseD) + if addr == "old" { + engine.setLeaderAddress("new") + return TSOReservation{}, status.Error(codes.FailedPrecondition, "tso: not leader") + } + 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, 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, 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, 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 TestLeaderRoutedTSOAllocatorRetriesValidationAfterLocalLeadershipLoss(t *testing.T) { + local := &leadershipLosingValidationTSO{fakeTSOAllocator: &fakeTSOAllocator{leader: true}} + engine := &recordingTSOEngine{ + state: raftengine.StateFollower, + leader: raftengine.LeaderInfo{Address: "new-leader"}, + } + alloc, err := NewLeaderRoutedTSOAllocator(local, engine) + require.NoError(t, err) + alloc.retryBudget = time.Second + alloc.retryInterval = time.Millisecond + var remoteCalls atomic.Uint64 + alloc.remoteValidate = func(_ context.Context, addr string, timestamp uint64) error { + remoteCalls.Add(1) + require.Equal(t, "new-leader", addr) + require.Equal(t, uint64(42), timestamp) + return nil + } + + require.NoError(t, alloc.ValidateDurableTimestamp(context.Background(), 42)) + require.Equal(t, uint64(1), local.validateCalls.Load()) + require.Equal(t, uint64(1), remoteCalls.Load()) +} + +func TestShadowTimestampAllocatorReturnsLegacyAndAdvancesTSO(t *testing.T) { + legacy := NewHLC() + legacy.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) + 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 TestShadowTimestampAllocatorBypassesLegacyAfterObservedCutover(t *testing.T) { + legacy := NewHLC() + fsm := NewTSOStateMachine(NewHLC()) + require.Nil(t, fsm.Apply(marshalTSOCutover())) + dedicated := &dedicatedShadowAllocator{next: testTSOInitialBase} + alloc, err := NewShadowTimestampAllocator( + legacy, + dedicated, + slog.New(slog.NewTextHandler(io.Discard, nil)), + WithTSOShadowCutoverState(fsm), + ) + require.NoError(t, err) + + issued, err := alloc.NextAfter(context.Background(), testTSOInitialBase-1) + require.NoError(t, err) + require.Equal(t, uint64(testTSOInitialBase), issued) + require.Zero(t, legacy.Current(), "post-cutover issuance must not sample legacy HLC") + require.Zero(t, dedicated.shadowCalls, "post-cutover issuance must not run shadow comparison") +} + +func TestShadowAndCutoverAllocatorsSerializeMigrationOnGroupZero(t *testing.T) { + tsoClock := NewHLC() + tsoClock.SetPhysicalCeiling(time.Now().Add(testTSOFutureCeiling).UnixMilli()) + 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 +} + +type dedicatedShadowAllocator struct { + next uint64 + shadowCalls int +} + +func (a *dedicatedShadowAllocator) Next(context.Context) (uint64, error) { + return a.next, nil +} + +func (a *dedicatedShadowAllocator) ValidateShadowTimestamp(context.Context, uint64) (TSOReservation, error) { + a.shadowCalls++ + return TSOReservation{}, errors.New("shadow comparison must not run after cutover") +} + +func (a *recordingShadowReservationAllocator) ValidateShadowTimestamp(_ context.Context, min uint64) (TSOReservation, error) { + a.mu.Lock() + defer a.mu.Unlock() + 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) + readIndex func(context.Context) (uint64, error) + term uint64 +} + +type leadershipLosingValidationTSO struct { + *fakeTSOAllocator + validateCalls atomic.Uint64 +} + +func (a *leadershipLosingValidationTSO) ValidateDurableTimestamp(context.Context, uint64) error { + a.validateCalls.Add(1) + a.leader = false + return errors.WithStack(ErrTSONotLeader) +} + +func (e *recordingTSOEngine) Propose(_ context.Context, payload []byte) (*raftengine.ProposalResult, error) { + e.mu.Lock() + if e.proposeErr != nil { + 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(ctx context.Context) (uint64, error) { + if e.readIndex != nil { + return e.readIndex(ctx) + } + 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 +} + +func (p *recordingTSOFloorProvider) setFloor(floor uint64) { + p.mu.Lock() + defer p.mu.Unlock() + p.floor = floor +} diff --git a/kv/tso_test.go b/kv/tso_test.go index 3c5f8f13f..99a2cf26f 100644 --- a/kv/tso_test.go +++ b/kv/tso_test.go @@ -330,6 +330,82 @@ func TestNextTimestampAfterThroughUsesAllocatorFloor(t *testing.T) { require.EqualValues(t, testTSOInitialBase+6, got) } +func TestBeginReadTimestampThroughPreservesLegacyBeforePhaseD(t *testing.T) { + alloc := &phaseDTestAllocator{next: testTSOInitialBase} + coord := &Coordinate{tsAllocator: alloc} + + readTS, err := BeginReadTimestampThrough(context.Background(), coord, 42, "test read timestamp") + require.NoError(t, err) + require.Equal(t, uint64(42), readTS.Timestamp()) + require.Zero(t, alloc.nextCalls.Load()) + require.Zero(t, alloc.validateCalls.Load()) +} + +func TestBeginReadTimestampThroughValidatesAppliedWatermarkDuringPhaseD(t *testing.T) { + alloc := &phaseDTestAllocator{next: testTSOInitialBase, phaseDActive: true, phaseDRequired: true} + coord := &Coordinate{tsAllocator: alloc} + + readTS, err := BeginReadTimestampThrough(context.Background(), coord, 42, "test read timestamp") + require.NoError(t, err) + require.Equal(t, uint64(42), readTS.Timestamp()) + require.Zero(t, alloc.nextCalls.Load()) + require.Equal(t, uint64(1), alloc.validateCalls.Load()) + require.Equal(t, uint64(42), alloc.validated.Load()) +} + +func TestBeginReadTimestampThroughFailsClosedOnValidationError(t *testing.T) { + alloc := &phaseDTestAllocator{ + next: testTSOInitialBase, + phaseDActive: true, + phaseDRequired: true, + validateErr: ErrTSOTimestampInvalid, + } + coord := &Coordinate{tsAllocator: alloc} + + _, err := BeginReadTimestampThrough(context.Background(), coord, 42, "test read timestamp") + require.ErrorIs(t, err, ErrTSOTimestampInvalid) + require.Zero(t, alloc.nextCalls.Load()) + require.Equal(t, uint64(1), alloc.validateCalls.Load()) +} + +func TestBeginReadTimestampThroughActivatesPhaseDBeforeValidation(t *testing.T) { + alloc := &phaseDTestAllocator{next: testTSOInitialBase, phaseDRequired: true} + coord := &Coordinate{tsAllocator: alloc} + + readTS, err := BeginReadTimestampThrough(context.Background(), coord, 42, "test phase D activation") + require.NoError(t, err) + require.Equal(t, uint64(42), readTS.Timestamp()) + require.Equal(t, uint64(1), alloc.nextCalls.Load()) + require.Equal(t, uint64(1), alloc.validateCalls.Load()) +} + +func TestBatchAllocatorDropsPrePhaseDWindowAfterActivation(t *testing.T) { + raw := &phaseDWindowTSO{nextBase: testTSOInitialBase, leader: true} + alloc, err := NewBatchAllocator(raw, testTSOBatchSize) + require.NoError(t, err) + + first, err := alloc.Next(context.Background()) + require.NoError(t, err) + require.Equal(t, uint64(testTSOInitialBase), first) + raw.phaseD = true + raw.floor = testTSOInitialBase + testTSOBatchSize - 1 + + readTS, err := alloc.Next(context.Background()) + require.NoError(t, err) + require.Equal(t, uint64(testTSOInitialBase+testTSOBatchSize), readTS) + require.Equal(t, uint64(2), raw.calls.Load()) +} + +func TestBeginReadTimestampThroughFindsAllocatorBehindKeyVizDecorator(t *testing.T) { + alloc := &phaseDTestAllocator{next: testTSOInitialBase, phaseDActive: true, phaseDRequired: true} + coord := WithKeyVizLabel(&Coordinate{tsAllocator: alloc}, keyviz.LabelRedis) + + readTS, err := BeginReadTimestampThrough(context.Background(), coord, 42, "test decorated read timestamp") + require.NoError(t, err) + require.Equal(t, uint64(42), readTS.Timestamp()) + require.Equal(t, uint64(1), alloc.validateCalls.Load()) +} + func TestKeyVizLabeledCoordinatorRecoversExpiredHLCCeiling(t *testing.T) { t.Parallel() clock := NewHLC() @@ -506,6 +582,71 @@ type fakeTSOAllocator struct { leader bool } +type phaseDTestAllocator struct { + next uint64 + phaseDActive bool + phaseDRequired bool + validateErr error + nextCalls atomic.Uint64 + validateCalls atomic.Uint64 + validated atomic.Uint64 +} + +func (a *phaseDTestAllocator) Next(context.Context) (uint64, error) { + a.nextCalls.Add(1) + return a.next, nil +} + +func (a *phaseDTestAllocator) NextAfter(_ context.Context, min uint64) (uint64, error) { + a.nextCalls.Add(1) + if a.next <= min { + return min + 1, nil + } + return a.next, nil +} + +func (a *phaseDTestAllocator) ValidateDurableTimestamp(_ context.Context, timestamp uint64) error { + a.validateCalls.Add(1) + a.validated.Store(timestamp) + return a.validateErr +} + +func (a *phaseDTestAllocator) PhaseDActive() bool { return a.phaseDActive } +func (a *phaseDTestAllocator) PhaseDRequired() bool { return a.phaseDRequired } + +type phaseDWindowTSO struct { + nextBase uint64 + floor uint64 + phaseD bool + leader bool + calls atomic.Uint64 +} + +func (a *phaseDWindowTSO) Next(ctx context.Context) (uint64, error) { + return a.NextBatch(ctx, 1) +} + +func (a *phaseDWindowTSO) NextBatch(_ context.Context, n int) (uint64, error) { + a.calls.Add(1) + base := a.nextBase + a.nextBase += uint64(n) //nolint:gosec // positive test batch size. + return base, nil +} + +func (a *phaseDWindowTSO) IsLeader() bool { return a.leader } + +func (a *phaseDWindowTSO) RunLeaseRenewal(ctx context.Context) { <-ctx.Done() } + +func (a *phaseDWindowTSO) ValidateDurableTimestamp(_ context.Context, timestamp uint64) error { + if timestamp <= a.floor { + return ErrTSOTimestampInvalid + } + return nil +} + +func (a *phaseDWindowTSO) PhaseDActive() bool { return a.phaseD } +func (a *phaseDWindowTSO) PhaseDRequired() bool { return a.phaseD } + func (f *fakeTSOAllocator) Next(ctx context.Context) (uint64, error) { return f.NextBatch(ctx, 1) } diff --git a/main.go b/main.go index 94a4cda03..05b542c83 100644 --- a/main.go +++ b/main.go @@ -124,8 +124,10 @@ 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") + tsoPhaseDEnabled = flag.Bool("tsoPhaseDEnabled", false, "Close the centralized-TSO compatibility window after every group-0 member understands the M7 marker") + tsoBatchSize = flag.Int("tsoBatchSize", defaultTSOBatchSize, "Timestamp batch size used by TSO cutover and shadow validation") leaderBalance = flag.Bool("leaderBalance", false, "Enable automatic count-based Raft-group leader balancing on the default-group leader") leaderBalanceInterval = flag.Duration("leaderBalanceInterval", defaultLeaderBalanceInterval, "Interval between leader-balance scheduler evaluations") leaderBalanceGroupCooldown = flag.Duration("leaderBalanceGroupCooldown", defaultLeaderBalanceGroupCooldown, "Minimum time before the scheduler can move the same raft group again") @@ -525,6 +527,12 @@ func run() error { WithKeyVizLabelsEnabled(*keyvizLabelsEnabled). WithAllShardGroups(dataGroupIDs(cfg.groups)...). WithPartitionResolver(buildSQSPartitionResolver(cfg.sqsFifoPartitionMap)) + 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 // TransferLeadership when it acquires (or already holds) @@ -578,6 +586,7 @@ func run() error { cfg.engine, distCatalog, adapter.WithDistributionCoordinator(coordinate), + adapter.WithDistributionTimestampAllocator(tsoWiring.serverAllocator), adapter.WithDistributionActiveTimestampTracker(readTracker), adapter.WithCatalogWatchLeaderCheck(func() bool { engine := catalogRuntime.snapshotEngine() @@ -631,6 +640,10 @@ func run() error { leaderBalanceConfigFromFlags(*raftId, cfg.defaultGroup, cfg.sqsFifoPartitionMap, metricsRegistry.Registerer()), ) + return waitRunGroup(eg) +} + +func waitRunGroup(eg *errgroup.Group) error { if err := eg.Wait(); err != nil { return errors.Wrapf(err, "failed to serve") } @@ -1415,106 +1428,149 @@ func buildShardGroups( // breaking the shared-cache invariant 6D-6c-1 relies on. encWiring = encWiring.withDefaultedCache() multi = effectiveMultiDataDirs(groups, multi) + builder := shardGroupBuilder{ + raftID: raftID, + raftDir: raftDir, + multi: multi, + bootstrap: bootstrap, + bootstrapCfg: bootstrapCfg, + factory: factory, + proposalObserverForGroup: proposalObserverForGroup, + clock: clock, + kekWrapper: kekWrapper, + keystore: keystore, + sidecarPath: sidecarPath, + encWiring: encWiring, + routeEngine: routeEngine, + applyObservers: applyObservers, + } runtimes := make([]*raftGroupRuntime, 0, len(groups)) shardGroups := make(map[uint64]*kv.ShardGroup, len(groups)) for _, g := range groups { - dir := groupDataDir(raftDir, raftID, g.id, multi) - if err := os.MkdirAll(dir, dirPerm); err != nil { - return nil, nil, errors.Wrapf(err, "failed to create fsm store dir for group %d", g.id) - } - st, err := store.NewPebbleStore(filepath.Join(dir, "fsm.db"), encWiring.pebbleOptions()...) - if err != nil { - return nil, nil, errors.Wrapf(err, "failed to open pebble fsm store for group %d", g.id) - } - // Each shard FSM shares the same HLC so any shard's lease renewal advances - // the global physicalCeiling. The logical counter remains in-memory only. - // - // §6.3 EncryptionApplier wiring (Stage 6A). Constructs an - // Applier backed by the FSM's Pebble store and threads it - // into the FSM dispatch via kv.WithEncryption. The applier's - // ApplyBootstrap / ApplyRotation paths return ErrKEKNotConfigured - // until Stage 6B threads in the KEK plumbing; only - // ApplyRegistration is fully functional in 6A, but the gRPC - // mutator gate (registerEncryptionAdminServer, Stage 5D - // posture) keeps the operator surface inert until 6B re-enables - // it gated on (--encryption-enabled AND KEKConfigured()). - reg, err := store.WriterRegistryFor(st) + runtime, sg, err := builder.build(g) if err != nil { - for _, rt := range runtimes { - rt.Close() - } - _ = st.Close() - return nil, nil, errors.Wrapf(err, "failed to construct writer registry for group %d", g.id) - } - // Stage 6B-2: thread WithKEK + WithKeystore + WithSidecarPath - // into the Applier when all three are supplied. Without - // any of them the applier stays in the Stage 6A posture - // (ApplyBootstrap / ApplyRotation return ErrKEKNotConfigured) - // which is exactly the desired apply-time fail-closed - // behaviour when an operator has not opted in to encryption. - // - // Stage 6D-6c-1: WithStateCache threads the process-shared - // StateCache so an encryption apply landing on this shard's - // FSM updates the atomics every shard's storage layer reads. - // Stage 7b' §3.1: WithLocalEpoch threads this process load's - // pinned storage write-path w.epoch into the Applier so - // applyRotateDEK (PurposeStorage) can write each node's own - // highest-emitted local_epoch into Keys[newDEK].LocalEpoch - // rather than the proposer's 0. Stage 6E does the same for - // raft rotations through WithRaftLocalEpoch using the raft - // DEK's separate per-process epoch. - applierOpts := encryptionApplierOptionsFor(kekWrapper, keystore, sidecarPath, encWiring) - applier, err := encryption.NewApplier(reg, applierOpts...) - if err != nil { - for _, rt := range runtimes { - rt.Close() - } - _ = st.Close() - return nil, nil, errors.Wrapf(err, "failed to construct encryption applier for group %d", g.id) - } - // Composed-1 M2 plumbing: wire the shared route catalog - // engine and this shard's owning group ID so M3's - // verifyComposed1 apply-time gate can resolve the - // observed-version owner-of-key without further plumbing - // work. At M2 the FSM stores both but does not consult them; - // see docs/design/2026_05_29_implemented_composed1_cross_group_commit_guard.md - // §M2. - sm := kv.NewKvFSMWithHLC(st, clock, fsmOptionsForGroup(applier, routeEngine, g.id, encWiring, applyObservers...)...) - groupBootstrap, groupBootstrapServers, groupBootstrapSeed := bootstrapSettingsForGroup(bootstrapCfg, g.id, bootstrap) - runtime, err := buildRuntimeForGroup( - raftID, g, raftDir, multi, groupBootstrap, - groupBootstrapServers, groupBootstrapSeed, - st, sm, factory, *raftJoinAsLearner) - if err != nil { - for _, rt := range runtimes { - rt.Close() - } - _ = st.Close() - return nil, nil, errors.Wrapf(err, "failed to start raft group %d", g.id) + closeRaftGroupRuntimes(runtimes) + return nil, nil, err } runtimes = append(runtimes, runtime) - // Stage 6E-2c: route every shard group's TransactionManager - // through NewLeaderProxyForShardGroup so the proposer chain - // consults sg.raftPayloadWrap on every Propose / ProposeAdmin. - // The cell is nil at startup (the Stage 3 default — payloads - // pass through cleartext); Stage 6E-2d's EnableRaftEnvelope - // handler will publish the active wrap closure the instant - // the cutover entry commits. Going through this constructor - // for every group avoids a future "I forgot to wire wrap on - // this code path" regression — if a new shard group bypasses - // this helper, the wrap install would silently no-op on - // proposals for that group. - sg := &kv.ShardGroup{ - Engine: runtime.engine, - Store: st, - } - sg.Txn = kv.NewLeaderProxyForShardGroup(sg, kv.WithProposalObserver(observerForGroup(proposalObserverForGroup, g.id))) shardGroups[g.id] = sg encWiring.attachRaftEnvelopeGroup(g.id, sg) } return runtimes, shardGroups, nil } +type shardGroupBuilder struct { + raftID string + raftDir string + multi bool + bootstrap bool + bootstrapCfg raftBootstrapConfig + factory raftengine.Factory + proposalObserverForGroup func(uint64) kv.ProposalObserver + clock *kv.HLC + kekWrapper kek.Wrapper + keystore *encryption.Keystore + sidecarPath string + encWiring encryptionWriteWiring + routeEngine *distribution.Engine + applyObservers []kv.ApplyObserver +} + +func (b shardGroupBuilder) build(group groupSpec) (*raftGroupRuntime, *kv.ShardGroup, error) { + groupBootstrap, groupBootstrapServers, groupBootstrapSeed := bootstrapSettingsForGroup(b.bootstrapCfg, group.id, b.bootstrap) + observer := observerForGroup(b.proposalObserverForGroup, group.id) + if group.id == dedicatedTSORaftGroupID { + runtime, sg, err := buildDedicatedTSOGroup( + b.raftID, group, b.raftDir, b.multi, groupBootstrap, + groupBootstrapServers, groupBootstrapSeed, + b.factory, b.clock, observer, + ) + return runtime, sg, errors.Wrap(err, "failed to start dedicated TSO group") + } + return b.buildDataGroup(group, groupBootstrap, groupBootstrapServers, groupBootstrapSeed, observer) +} + +func (b shardGroupBuilder) buildDataGroup( + group groupSpec, + bootstrap bool, + bootstrapServers []raftengine.Server, + bootstrapSeed []raftengine.Server, + proposalObserver kv.ProposalObserver, +) (*raftGroupRuntime, *kv.ShardGroup, error) { + dir := groupDataDir(b.raftDir, b.raftID, group.id, b.multi) + if err := os.MkdirAll(dir, dirPerm); err != nil { + return nil, nil, errors.Wrapf(err, "failed to create fsm store dir for group %d", group.id) + } + st, err := store.NewPebbleStore(filepath.Join(dir, "fsm.db"), b.encWiring.pebbleOptions()...) + if err != nil { + return nil, nil, errors.Wrapf(err, "failed to open pebble fsm store for group %d", group.id) + } + reg, err := store.WriterRegistryFor(st) + if err != nil { + _ = st.Close() + return nil, nil, errors.Wrapf(err, "failed to construct writer registry for group %d", group.id) + } + applierOpts := encryptionApplierOptionsFor(b.kekWrapper, b.keystore, b.sidecarPath, b.encWiring) + applier, err := encryption.NewApplier(reg, applierOpts...) + if err != nil { + _ = st.Close() + return nil, nil, errors.Wrapf(err, "failed to construct encryption applier for group %d", group.id) + } + sm := kv.NewKvFSMWithHLC(st, b.clock, fsmOptionsForGroup(applier, b.routeEngine, group.id, b.encWiring, b.applyObservers...)...) + runtime, err := buildRuntimeForGroup( + b.raftID, group, b.raftDir, b.multi, bootstrap, + bootstrapServers, bootstrapSeed, + st, sm, b.factory, *raftJoinAsLearner, + ) + if err != nil { + _ = st.Close() + return nil, nil, errors.Wrapf(err, "failed to start raft group %d", group.id) + } + // Every group goes through the wrap-aware proposer so HLC renewals and + // transactional proposals observe raft-envelope cutovers uniformly. + sg := &kv.ShardGroup{Engine: runtime.engine, Store: st} + sg.Txn = kv.NewLeaderProxyForShardGroup(sg, kv.WithProposalObserver(proposalObserver)) + return runtime, sg, nil +} + +func closeRaftGroupRuntimes(runtimes []*raftGroupRuntime) { + for _, runtime := range runtimes { + runtime.Close() + } +} + +// buildDedicatedTSOGroup constructs group 0 with the minimal TSO state machine. +// It intentionally does not open an MVCC store: shard routing validation keeps +// user data out of group 0, while the state machine persists 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, + baseDir string, + multi bool, + bootstrap bool, + bootstrapServers []raftengine.Server, + bootstrapSeed []raftengine.Server, + factory raftengine.Factory, + clock *kv.HLC, + proposalObserver kv.ProposalObserver, +) (*raftGroupRuntime, *kv.ShardGroup, error) { + sm := kv.NewTSOStateMachine(clock) + runtime, err := buildRuntimeForGroup( + raftID, group, baseDir, multi, bootstrap, + bootstrapServers, bootstrapSeed, + nil, sm, factory, *raftJoinAsLearner, + ) + if err != nil { + return nil, nil, err + } + sg := &kv.ShardGroup{Engine: runtime.engine, TSOState: sm} + sg.Txn = kv.NewLeaderProxyForShardGroup(sg, kv.WithProposalObserver(proposalObserver)) + return runtime, sg, nil +} + func bootstrapSettingsForGroup( cfg raftBootstrapConfig, groupID uint64, @@ -2039,22 +2095,155 @@ 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") + } + if *tsoPhaseDEnabled && !*tsoEnabled { + return wiring, errors.New("--tsoPhaseDEnabled requires --tsoEnabled") + } + + 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 *tsoPhaseDEnabled { + return wiring, errors.Wrap(kv.ErrTSOGroupRequired, "configure tso phase D") + } + 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 + coordinate.WithTimestampGroup(dedicatedTSORaftGroupID) + coordinate.WithTSOCutoverState(tsoGroup.TSOState) + if !dedicatedTSOAllocatorRequired(tsoGroup.TSOState) { + return wiring, nil + } + + routedOpts := dedicatedTSORoutingOptions(coordinate.Clock()) + routed, err := kv.NewLeaderRoutedTSOAllocator(local, tsoGroup.Engine, routedOpts...) + if err != nil { + return wiring, errors.Wrap(err, "configure tso leader routing") + } + wiring.routedAllocator = routed + return installDedicatedCoordinatorAllocator(coordinate, tsoGroup.TSOState, wiring, routed) +} + +func dedicatedTSOAllocatorRequired(state *kv.TSOStateMachine) bool { + return *tsoEnabled || *tsoShadowEnabled || (state != nil && state.CutoverActive()) +} + +func dedicatedTSORoutingOptions(clock *kv.HLC) []kv.LeaderRoutedTSOAllocatorOption { + opts := []kv.LeaderRoutedTSOAllocatorOption{kv.WithTSORoutedClock(clock)} + if *tsoEnabled { + opts = append(opts, kv.WithTSOCutoverActivation()) + } + if *tsoPhaseDEnabled { + opts = append(opts, kv.WithTSOPhaseDActivation()) + } + return opts +} + +func installDedicatedCoordinatorAllocator( + coordinate *kv.ShardedCoordinator, + state *kv.TSOStateMachine, + wiring coordinatorTSOWiring, + routed *kv.LeaderRoutedTSOAllocator, +) (coordinatorTSOWiring, error) { + if *tsoShadowEnabled && (state == nil || !state.CutoverActive()) { + shadow, err := kv.NewShadowTimestampAllocator( + coordinate.Clock(), + routed, + slog.Default(), + kv.WithTSOShadowCutoverState(state), + ) + if err != nil { + _ = wiring.Close() + return coordinatorTSOWiring{}, errors.Wrap(err, "configure tso shadow allocator") + } + wiring.shadowAllocator = shadow + coordinate.WithTSOAllocator(shadow) + return wiring, nil + } + batch, err := kv.NewBatchAllocator(routed, *tsoBatchSize) + if err != nil { + _ = wiring.Close() + return coordinatorTSOWiring{}, errors.Wrap(err, "configure tso batch allocator") + } + coordinate.WithTSOAllocator(batch) + return wiring, nil } type hlcLeaseRenewalBlocker interface { @@ -2204,8 +2393,11 @@ var _ kv.LeaseReadableCoordinator = (*startupGatedCoordinator)(nil) var _ kv.AllGroupsLeaseReadableCoordinator = (*startupGatedCoordinator)(nil) var _ kv.GroupRoutableCoordinator = (*startupGatedCoordinator)(nil) var _ kv.RaftMembershipCoordinator = (*startupGatedCoordinator)(nil) +var _ kv.TimestampAllocatorProvider = (*startupGatedCoordinator)(nil) var _ kv.TimestampAllocator = (*startupGatedCoordinator)(nil) var _ kv.TimestampAfterAllocator = (*startupGatedCoordinator)(nil) +var _ kv.AppliedReadTimestampVoucher = (*startupGatedCoordinator)(nil) +var _ kv.AppliedReadTimestampVoucherRevoker = (*startupGatedCoordinator)(nil) func (c startupGatedCoordinator) Dispatch(ctx context.Context, reqs *kv.OperationGroup[kv.OP]) (*kv.CoordinateResponse, error) { if c.gate != nil && c.gate.blocked() { @@ -2246,7 +2438,51 @@ func (c startupGatedCoordinator) Clock() *kv.HLC { return c.inner.Clock() } +func (c startupGatedCoordinator) TimestampAllocator() kv.TimestampAllocator { + return c +} + +func (c startupGatedCoordinator) innerTimestampAllocator() (kv.TimestampAllocator, bool) { + if c.inner == nil { + return nil, false + } + return kv.TimestampAllocatorThrough(c.inner) +} + +func (c startupGatedCoordinator) PhaseDActive() bool { + alloc, ok := c.innerTimestampAllocator() + if !ok { + return false + } + phaseD, ok := alloc.(kv.TSOPhaseDState) + return ok && phaseD.PhaseDActive() +} + +func (c startupGatedCoordinator) PhaseDRequired() bool { + alloc, ok := c.innerTimestampAllocator() + if !ok { + return false + } + phaseD, ok := alloc.(kv.TSOPhaseDState) + return ok && phaseD.PhaseDRequired() +} + +func (c startupGatedCoordinator) ValidateDurableTimestamp(ctx context.Context, timestamp uint64) error { + alloc, ok := c.innerTimestampAllocator() + if !ok { + return errors.WithStack(kv.ErrTSOProtocolUnsupported) + } + validator, ok := alloc.(kv.DurableTimestampValidator) + if !ok { + return errors.WithStack(kv.ErrTSOProtocolUnsupported) + } + return errors.WithStack(validator.ValidateDurableTimestamp(ctx, timestamp)) +} + func (c startupGatedCoordinator) Next(ctx context.Context) (uint64, error) { + if c.gate != nil && c.gate.blocked() { + return 0, status.Error(codes.Unavailable, "startup rotation has not completed") //nolint:wrapcheck // Preserve the gRPC status for adapters. + } if alloc, ok := c.inner.(kv.TimestampAllocator); ok { ts, err := alloc.Next(ctx) return ts, errors.WithStack(err) @@ -2256,6 +2492,9 @@ func (c startupGatedCoordinator) Next(ctx context.Context) (uint64, error) { } func (c startupGatedCoordinator) NextAfter(ctx context.Context, min uint64) (uint64, error) { + if c.gate != nil && c.gate.blocked() { + return 0, status.Error(codes.Unavailable, "startup rotation has not completed") //nolint:wrapcheck // Preserve the gRPC status for adapters. + } if alloc, ok := c.inner.(kv.TimestampAfterAllocator); ok { ts, err := alloc.NextAfter(ctx, min) return ts, errors.WithStack(err) @@ -2274,6 +2513,20 @@ func (c startupGatedCoordinator) RecoverHLCLease(ctx context.Context) error { return errors.New("startup-gated coordinator: hlc lease recovery unavailable") } +func (c startupGatedCoordinator) VouchAppliedReadTimestamp(timestamp uint64, ref kv.AppliedReadTimestampVoucherRef) error { + voucher, ok := c.inner.(kv.AppliedReadTimestampVoucher) + if !ok { + return errors.WithStack(kv.ErrTSOProtocolUnsupported) + } + return errors.WithStack(voucher.VouchAppliedReadTimestamp(timestamp, ref)) +} + +func (c startupGatedCoordinator) RevokeAppliedReadTimestamp(timestamp uint64, ref kv.AppliedReadTimestampVoucherRef) { + if revoker, ok := c.inner.(kv.AppliedReadTimestampVoucherRevoker); ok { + revoker.RevokeAppliedReadTimestamp(timestamp, ref) + } +} + func (c startupGatedCoordinator) LeaseRead(ctx context.Context) (uint64, error) { return kv.LeaseReadThrough(c.inner, ctx) //nolint:wrapcheck // Pass through coordinator errors unchanged. } @@ -2362,6 +2615,7 @@ func startupRotationGatedMethod(fullMethod string) bool { switch fullMethod { case pb.Internal_Forward_FullMethodName, pb.AdminForward_Forward_FullMethodName, + pb.Distribution_GetTimestamp_FullMethodName, pb.Distribution_SplitRange_FullMethodName, pb.RaftAdmin_AddVoter_FullMethodName, pb.RaftAdmin_AddLearner_FullMethodName, @@ -2814,7 +3068,7 @@ func registerAdminServerIfPresent(gs grpc.ServiceRegistrar, adminServer *adapter } 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 @@ -3073,9 +3327,6 @@ func setupDistributionRuntimeDependencies( raftID string, sidecarPath string, ) (*distribution.CatalogStore, *raftGroupRuntime, *raftGroupRuntime, error) { - if err := configureCoordinatorTSO(coordinate); err != nil { - return nil, nil, nil, err - } catalog, err := setupDistributionAndRegistration( runCtx, eg, runtimes, engine, coordinate, defaultGroup, encWiring, raftID, sidecarPath) if err != nil { diff --git a/main_encryption_admin.go b/main_encryption_admin.go index c6ffc54d5..eeeac6d5c 100644 --- a/main_encryption_admin.go +++ b/main_encryption_admin.go @@ -79,6 +79,20 @@ type encryptionAdminEngine interface { raftengine.LeaderView } +// encryptionAdminWiringForGroup keeps the read-only capability service on the +// dedicated TSO listener while withholding every mutating proposal path. The +// TSO FSM has no encryption applier or writer registry, so allowing a mutator +// to propose there would commit an entry that cannot update dedicated TSO +// state. The engine remains available as a LeaderView so ResyncSidecar keeps +// its quorum-confirmed leader-only freshness contract. Data groups retain the +// normal flag-driven mutator wiring. +func encryptionAdminWiringForGroup(groupID uint64, enableMutators bool, engine encryptionAdminEngine) (bool, encryptionAdminEngine) { + if groupID == dedicatedTSORaftGroupID { + return false, engine + } + return enableMutators, engine +} + func encryptionAdminWriterRegistry(sidecarPath string, st store.MVCCStore, groupID uint64) (encryption.WriterRegistryStore, error) { if sidecarPath == "" { return nil, nil @@ -137,20 +151,18 @@ func encryptionAdminDefaultRuntimeOptions(runtimes []*raftGroupRuntime, shardGro // number — Codex r1 P1 on PR #760 caught the original wiring // passing the shard id by mistake. // -// Stage 6B-2: the mutator wiring (Proposer + LeaderView) is now -// gated on the supplied enableMutators boolean, which the caller -// in main.go computes as (--encryption-enabled AND --kekFile -// non-empty AND engine non-nil). When enableMutators is false, -// Proposer + LeaderView stay unwired and EncryptionAdminServer's -// BootstrapEncryption / RotateDEK / RegisterEncryptionWriter -// short-circuit at the gRPC boundary with FailedPrecondition — -// identical to the Stage 5D posture. When enableMutators is -// true, both options are wired and the mutators reach the §6.3 -// applier through the supplied engine. In multi-group production -// wiring the supplied engine is the default-group engine for every -// per-group listener, because writer-registry rows and the sidecar -// applied-index contract are default-group control-plane state. -// Callers may append a +// Stage 6B-2 gates the mutating Proposer on enableMutators, which +// the caller in main.go computes as (--encryption-enabled AND +// --kekFile non-empty AND engine non-nil). The LeaderView remains +// wired whenever an engine is present because read-only recovery +// RPCs also require quorum-confirmed leader freshness. In multi-group +// production wiring the supplied engine is the default-group engine +// for every per-group listener, because writer-registry rows and the +// sidecar applied-index contract are default-group control-plane state. +// When enableMutators is false, BootstrapEncryption / RotateDEK / +// RegisterEncryptionWriter short-circuit at the gRPC boundary with +// FailedPrecondition. When true, the proposer reaches the §6.3 +// applier through the supplied engine. Callers may append a // wrap-aware post-cutover proposer option so non-cutover admin // entries wrap after EnableRaftEnvelope while the cutover marker // itself remains on the raw engine path. @@ -171,9 +183,9 @@ func encryptionAdminDefaultRuntimeOptions(runtimes []*raftGroupRuntime, shardGro // practice. // // Validate() panics on a misconfiguration that wires a proposer -// without a LeaderView. The wiring below pairs both options -// together (both wired iff enableMutators) so the invariant -// holds by construction. +// without a LeaderView. The wiring below always pairs a mutating +// proposer with the engine's LeaderView, while permitting a +// LeaderView-only read surface. // // sidecarPath controls only the read-only capability surface: // when empty, GetCapability reports encryption_capable=false @@ -187,6 +199,9 @@ func newEncryptionAdminServer(fullNodeID uint64, sidecarPath string, enableMutat if sidecarPath != "" { opts = append(opts, adapter.WithEncryptionAdminSidecarPath(sidecarPath)) } + if engine != nil { + opts = append(opts, adapter.WithEncryptionAdminLeaderView(engine)) + } if latestAppliedIndex != nil { opts = append(opts, adapter.WithEncryptionAdminLatestAppliedIndex(latestAppliedIndex)) } @@ -199,11 +214,10 @@ func newEncryptionAdminServer(fullNodeID uint64, sidecarPath string, enableMutat if enableMutators && engine != nil { opts = append(opts, adapter.WithEncryptionAdminProposer(engine), - adapter.WithEncryptionAdminLeaderView(engine), ) // The §4 capability fan-out is only meaningful when the // cutover mutator is reachable (same enableMutators gate as - // the Proposer/LeaderView). Without it the cutover RPC + // the Proposer). Without it the cutover RPC // refuses with the §4 "fan-out not wired" FailedPrecondition. if capabilityFanout != nil { opts = append(opts, adapter.WithEncryptionAdminCapabilityFanout(capabilityFanout)) diff --git a/main_encryption_admin_test.go b/main_encryption_admin_test.go index 5cb9bfd12..2c9f45bcc 100644 --- a/main_encryption_admin_test.go +++ b/main_encryption_admin_test.go @@ -45,6 +45,10 @@ func (stubEncryptionAdminEngine) LinearizableRead(context.Context) (uint64, erro return 0, nil } +type followerEncryptionAdminEngine struct{ stubEncryptionAdminEngine } + +func (followerEncryptionAdminEngine) State() raftengine.State { return raftengine.StateFollower } + // TestEncryptionAdminFullNodeID_DistinctPerRaftId pins the // PR760 r1 Codex P1 regression: full_node_id must be derived from // --raftId (per-node-stable) and NOT from the Raft group id @@ -248,6 +252,42 @@ func TestEncryptionAdmin_MutatingRPCEnabledWhenGateOn(t *testing.T) { } } +func TestEncryptionAdmin_DedicatedTSOGroupAlwaysRefusesMutators(t *testing.T) { + want := stubEncryptionAdminEngine{} + enabled, engine := encryptionAdminWiringForGroup( + dedicatedTSORaftGroupID, + true, + want, + ) + if enabled || engine == nil { + t.Fatalf("dedicated TSO mutator wiring = (%v, %T), want (false, leader view)", enabled, engine) + } + err := callBootstrapAgainstNewServer(t, enabled, engine) + if got := status.Code(err); got != codes.FailedPrecondition { + t.Fatalf("BootstrapEncryption status=%v, want FailedPrecondition; err=%v", got, err) + } +} + +func TestEncryptionAdmin_DedicatedTSOGroupResyncRejectsFollower(t *testing.T) { + enabled, engine := encryptionAdminWiringForGroup( + dedicatedTSORaftGroupID, + true, + followerEncryptionAdminEngine{}, + ) + err := callResyncAgainstNewServer(t, enabled, engine) + if got := status.Code(err); got != codes.FailedPrecondition { + t.Fatalf("ResyncSidecar status=%v, want FailedPrecondition; err=%v", got, err) + } +} + +func TestEncryptionAdmin_DataGroupPreservesMutatorGate(t *testing.T) { + want := stubEncryptionAdminEngine{} + enabled, engine := encryptionAdminWiringForGroup(1, true, want) + if !enabled || engine == nil { + t.Fatalf("data-group mutator wiring = (%v, %T), want (true, non-nil)", enabled, engine) + } +} + func callBootstrapAgainstNewServer(t *testing.T, enableMutators bool, engine encryptionAdminEngine) error { t.Helper() listener := bufconn.Listen(1024 * 1024) @@ -276,6 +316,28 @@ func callBootstrapAgainstNewServer(t *testing.T, enableMutators bool, engine enc return err } +func callResyncAgainstNewServer(t *testing.T, enableMutators bool, engine encryptionAdminEngine) error { + t.Helper() + listener := bufconn.Listen(1024 * 1024) + gs := grpc.NewServer() + registerEncryptionAdminServer(gs, newEncryptionAdminServer(1, "", enableMutators, engine, nil, nil, nil)) + go func() { _ = gs.Serve(listener) }() + t.Cleanup(gs.Stop) + conn, err := grpc.NewClient( + "passthrough:///bufnet", + grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { + return listener.Dial() + }), + grpc.WithTransportCredentials(insecure.NewCredentials()), + ) + if err != nil { + t.Fatalf("grpc.NewClient: %v", err) + } + t.Cleanup(func() { _ = conn.Close() }) + _, err = pb.NewEncryptionAdminClient(conn).ResyncSidecar(context.Background(), &pb.ResyncSidecarRequest{}) + return err +} + func TestRegisterEncryptionAdminServer_Registers(t *testing.T) { gs := grpc.NewServer() registerEncryptionAdminServer(gs, newEncryptionAdminServer(1, "", false, nil, nil, nil, nil)) diff --git a/main_encryption_rotate_on_startup_test.go b/main_encryption_rotate_on_startup_test.go index 48b4adc9d..81b8ac677 100644 --- a/main_encryption_rotate_on_startup_test.go +++ b/main_encryption_rotate_on_startup_test.go @@ -441,6 +441,7 @@ func TestStartupPublicKVGate_BlocksMutatorsUntilReady(t *testing.T) { pb.TransactionalKV_Get_FullMethodName, pb.Internal_Forward_FullMethodName, pb.AdminForward_Forward_FullMethodName, + pb.Distribution_GetTimestamp_FullMethodName, pb.Distribution_SplitRange_FullMethodName, pb.RaftAdmin_AddVoter_FullMethodName, pb.RaftAdmin_AddLearner_FullMethodName, @@ -627,20 +628,45 @@ func TestStartupGatedCoordinator_ForwardsTimestampCapabilities(t *testing.T) { t.Parallel() sentinel := errors.New("recover failed") inner := &stubStartupCoordinator{clock: kv.NewHLC(), recoverErr: sentinel} + blocked := true + gate := &startupPublicKVGate{blockMutator: func() bool { return blocked }} + gate.markReady() coord := startupGatedCoordinator{ inner: inner, - gate: &startupPublicKVGate{blockMutator: func() bool { return true }}, + gate: gate, + } + + assertStartupGatedTimestampBlocked(t, coord) + blocked = false + assertStartupGatedTimestampForwarding(t, coord, inner, sentinel) +} + +func assertStartupGatedTimestampBlocked(t *testing.T, coord startupGatedCoordinator) { + t.Helper() + if _, err := coord.Next(context.Background()); status.Code(err) != codes.Unavailable { + t.Fatalf("blocked Next() err=%v, want Unavailable", err) } + if _, err := coord.NextAfter(context.Background(), 200); status.Code(err) != codes.Unavailable { + t.Fatalf("blocked NextAfter() err=%v, want Unavailable", err) + } +} +func assertStartupGatedTimestampForwarding( + t *testing.T, + coord startupGatedCoordinator, + inner *stubStartupCoordinator, + recoverErr error, +) { + t.Helper() ts, err := coord.Next(context.Background()) if err != nil || ts != 101 { - t.Fatalf("Next() = (%d, %v), want (101, nil)", ts, err) + t.Fatalf("Next() after unblock = (%d, %v), want (101, nil)", ts, err) } next, err := coord.NextAfter(context.Background(), 200) if err != nil || next != 201 { - t.Fatalf("NextAfter() = (%d, %v), want (201, nil)", next, err) + t.Fatalf("NextAfter() after unblock = (%d, %v), want (201, nil)", next, err) } - if err := coord.RecoverHLCLease(context.Background()); !errors.Is(err, sentinel) { + if err := coord.RecoverHLCLease(context.Background()); !errors.Is(err, recoverErr) { t.Fatalf("RecoverHLCLease() err=%v, want sentinel", err) } if got := inner.nextCalls.Load(); got != 1 { diff --git a/main_tso_routing_test.go b/main_tso_routing_test.go new file mode 100644 index 000000000..6e3b60828 --- /dev/null +++ b/main_tso_routing_test.go @@ -0,0 +1,288 @@ +package main + +import ( + "context" + "fmt" + "sync/atomic" + "testing" + "time" + + "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/internal/raftengine" + "github.com/bootjp/elastickv/keyviz" + "github.com/bootjp/elastickv/kv" + "github.com/stretchr/testify/require" +) + +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 TestConfigureCoordinatorTSOPhaseDWorksThroughAdapterDecorators(t *testing.T) { + setTSOModeFlags(t, true, false) + *tsoPhaseDEnabled = true + clock := kv.NewHLC() + clock.SetPhysicalCeiling(time.Now().Add(time.Minute).UnixMilli()) + fsm := kv.NewTSOStateMachine(clock) + engine := &mainTSOEngine{state: raftengine.StateLeader, tsoState: fsm} + groups := map[uint64]*kv.ShardGroup{ + dedicatedTSORaftGroupID: {Engine: engine, TSOState: fsm}, + } + coord := newMainTSOCoordinator(clock, groups) + + wiring, err := configureCoordinatorTSO(coord, groups, mainTSOFloorProvider{floor: 42}) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, wiring.Close()) }) + decorated := kv.WithKeyVizLabel(startupGatedCoordinator{inner: coord}, keyviz.LabelRedis) + + readTS, err := kv.BeginReadTimestampThrough(context.Background(), decorated, 42, "test decorated phase D") + require.NoError(t, err) + require.NotZero(t, readTS.Timestamp()) + require.True(t, fsm.CutoverActive()) + require.True(t, fsm.PhaseDActive()) + require.Equal(t, uint64(3), engine.proposals.Load(), + "cutover and phase-D markers must commit before the timestamp window") +} + +func TestConfigureCoordinatorTSOPhaseDRequiresCutoverMode(t *testing.T) { + setTSOModeFlags(t, false, false) + *tsoPhaseDEnabled = true + coord := newMainTSOCoordinator(kv.NewHLC(), nil) + + _, err := configureCoordinatorTSO(coord, nil) + require.ErrorContains(t, err, "requires --tsoEnabled") +} + +func TestConfigureCoordinatorTSOShadowReturnsLegacyTimestamp(t *testing.T) { + setTSOModeFlags(t, false, true) + clock := kv.NewHLC() + 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.IsType(t, startupGatedCoordinator{}, got) + _, err := got.Next(context.Background()) + require.Error(t, err) + require.Len(t, internalTimestampOptions(gated), 1) + + legacy := startupGatedCoordinator{inner: newMainTSOCoordinator(kv.NewHLC(), nil)} + require.Len(t, internalTimestampOptions(legacy), 1) +} + +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, false) + 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 + oldPhaseD := *tsoPhaseDEnabled + oldBatchSize := *tsoBatchSize + *tsoEnabled = enabled + *tsoShadowEnabled = shadow + *tsoPhaseDEnabled = false + *tsoBatchSize = 8 + t.Cleanup(func() { + *tsoEnabled = oldEnabled + *tsoShadowEnabled = oldShadow + *tsoPhaseDEnabled = oldPhaseD + *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 { + floor uint64 +} + +type mainTimestampAllocator struct { + next uint64 +} + +func (a *mainTimestampAllocator) Next(context.Context) (uint64, error) { + return a.next, nil +} + +func (p mainTSOFloorProvider) GlobalCommittedTimestampFloor(context.Context) (uint64, error) { + return p.floor, 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/multiraft_runtime_test.go b/multiraft_runtime_test.go index 11670c9e9..c7b7da1e9 100644 --- a/multiraft_runtime_test.go +++ b/multiraft_runtime_test.go @@ -7,6 +7,7 @@ import ( "testing" "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/internal/encryption/fsmwire" "github.com/bootjp/elastickv/kv" "github.com/bootjp/elastickv/store" "github.com/stretchr/testify/require" @@ -74,14 +75,138 @@ func TestBuildShardGroupsWithDedicatedTSOPreservesSingleDataGroupDir(t *testing. }) require.Contains(t, shardGroups, uint64(dedicatedTSORaftGroupID)) require.Contains(t, shardGroups, uint64(1)) + require.Nil(t, shardGroups[dedicatedTSORaftGroupID].Store) + require.IsType(t, &kv.TSOStateMachine{}, runtimes[0].stateMachine) + require.Nil(t, runtimes[0].store) require.DirExists(t, filepath.Join(legacyDir, "fsm.db")) - require.DirExists(t, filepath.Join(legacyDir, "group-0", "fsm.db")) + require.DirExists(t, filepath.Join(legacyDir, "group-0")) + require.NoDirExists(t, filepath.Join(legacyDir, "group-0", "fsm.db")) require.NoDirExists(t, filepath.Join(legacyDir, "group-1"), "group 1 should stay in legacy dir") got, err := os.ReadFile(legacyMarker) require.NoError(t, err) require.Equal(t, []byte("keep"), got) } +func TestBuildShardGroupsWithDedicatedTSOPreservesLegacyGroupZeroStore(t *testing.T) { + baseDir := t.TempDir() + legacyStoreDir := filepath.Join(baseDir, "n1", "group-0", "fsm.db") + require.NoError(t, os.MkdirAll(legacyStoreDir, raftDirPerm)) + legacyMarker := filepath.Join(legacyStoreDir, "legacy.marker") + require.NoError(t, os.WriteFile(legacyMarker, []byte("preserve"), 0o600)) + + factory, err := newRaftFactory(raftEngineEtcd, nil) + require.NoError(t, err) + runtimes, shardGroups, err := buildShardGroups( + "n1", + baseDir, + []groupSpec{{id: dedicatedTSORaftGroupID, address: "127.0.0.1:17002"}}, + true, + true, + raftBootstrapConfig{}, + factory, + nil, + kv.NewHLC(), + nil, + nil, + "", + encryptionWriteWiring{}, + nil, + ) + require.NoError(t, err) + t.Cleanup(func() { closeRaftGroupRuntimes(runtimes) }) + require.Nil(t, shardGroups[dedicatedTSORaftGroupID].Store) + require.Nil(t, runtimes[0].store) + + got, err := os.ReadFile(legacyMarker) + require.NoError(t, err) + require.Equal(t, []byte("preserve"), got) +} + +func TestBuildShardGroupsClosesEarlierStoresWhenLaterGroupFails(t *testing.T) { + baseDir := t.TempDir() + badGroupDir := groupDataDir(baseDir, "n1", 2, true) + require.NoError(t, os.MkdirAll(badGroupDir, raftDirPerm)) + require.NoError(t, os.WriteFile( + filepath.Join(badGroupDir, raftEngineMarkerFile), + []byte("unsupported\n"), + raftEngineMarkerPerm, + )) + + factory, err := newRaftFactory(raftEngineEtcd, nil) + require.NoError(t, err) + runtimes, shardGroups, err := buildShardGroups( + "n1", + baseDir, + []groupSpec{ + {id: 1, address: "127.0.0.1:17011"}, + {id: 2, address: "127.0.0.1:17012"}, + }, + true, + true, + raftBootstrapConfig{}, + factory, + nil, + kv.NewHLC(), + nil, + nil, + "", + encryptionWriteWiring{}, + nil, + ) + require.ErrorIs(t, err, ErrUnsupportedRaftEngine) + require.Nil(t, runtimes) + require.Nil(t, shardGroups) + + // Reopening the first group's Pebble directory proves the error path ran + // raftGroupRuntime.Close, which owns both engine and store cleanup. + reopened, err := store.NewPebbleStore(filepath.Join( + groupDataDir(baseDir, "n1", 1, true), + "fsm.db", + )) + require.NoError(t, err) + require.NoError(t, reopened.Close()) +} + +func TestDedicatedTSOGroupContinuesAfterLegacyEncryptionEntry(t *testing.T) { + baseDir := t.TempDir() + factory, err := newRaftFactory(raftEngineEtcd, nil) + require.NoError(t, err) + clock := kv.NewHLC() + runtimes, _, err := buildShardGroups( + "n1", + baseDir, + []groupSpec{{id: dedicatedTSORaftGroupID, address: "127.0.0.1:17020"}}, + true, + true, + raftBootstrapConfig{}, + factory, + nil, + clock, + nil, + nil, + "", + encryptionWriteWiring{}, + nil, + ) + require.NoError(t, err) + t.Cleanup(func() { closeRaftGroupRuntimes(runtimes) }) + + legacyEntry := append([]byte{fsmwire.OpRegistration}, fsmwire.EncodeRegistration( + fsmwire.RegistrationPayload{DEKID: 1, FullNodeID: 2, LocalEpoch: 3}, + )...) + result, err := runtimes[0].engine.ProposeAdmin(context.Background(), legacyEntry) + require.NoError(t, err) + legacyErr, ok := result.Response.(error) + require.Truef(t, ok, "legacy response type = %T, want error", result.Response) + require.ErrorIs(t, legacyErr, kv.ErrTSOLegacyEncryptionEntryRejected) + + const ceilingMs = int64(1_700_000_123_456) + coordinator := kv.NewCoordinatorWithEngine(nil, runtimes[0].engine) + t.Cleanup(func() { require.NoError(t, coordinator.Close()) }) + require.NoError(t, coordinator.ProposeHLCLease(context.Background(), ceilingMs)) + require.Equal(t, ceilingMs, clock.PhysicalCeiling()) +} + func TestParseRaftEngineType(t *testing.T) { t.Run("default", func(t *testing.T) { engineType, err := parseRaftEngineType("") diff --git a/proto/distribution.pb.go b/proto/distribution.pb.go index 14a3db51f..63fb3f639 100644 --- a/proto/distribution.pb.go +++ b/proto/distribution.pb.go @@ -406,9 +406,22 @@ 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"` + // ActivatePhaseD closes the legacy compatibility window after every member + // understands the M7 FSM entry. It requires ActivateCutover and is one-way. + ActivatePhaseD bool `protobuf:"varint,4,opt,name=activate_phase_d,json=activatePhaseD,proto3" json:"activate_phase_d,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *GetTimestampRequest) Reset() { @@ -441,9 +454,53 @@ 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 +} + +func (x *GetTimestampRequest) GetActivatePhaseD() bool { + if x != nil { + return x.ActivatePhaseD + } + return false +} + type GetTimestampResponse struct { - state protoimpl.MessageState `protogen:"open.v1"` - Timestamp uint64 `protobuf:"varint,1,opt,name=timestamp,proto3" json:"timestamp,omitempty"` + 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"` + // PhaseDActive reports that legacy per-shard issuance has been retired. + PhaseDActive bool `protobuf:"varint,6,opt,name=phase_d_active,json=phaseDActive,proto3" json:"phase_d_active,omitempty"` + // PhaseDFloor is the allocation floor captured by the Phase-D marker. + PhaseDFloor uint64 `protobuf:"varint,7,opt,name=phase_d_floor,json=phaseDFloor,proto3" json:"phase_d_floor,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -485,6 +542,160 @@ 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 +} + +func (x *GetTimestampResponse) GetPhaseDActive() bool { + if x != nil { + return x.PhaseDActive + } + return false +} + +func (x *GetTimestampResponse) GetPhaseDFloor() uint64 { + if x != nil { + return x.PhaseDFloor + } + return 0 +} + +type ValidateTimestampRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + Timestamp uint64 `protobuf:"varint,1,opt,name=timestamp,proto3" json:"timestamp,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ValidateTimestampRequest) Reset() { + *x = ValidateTimestampRequest{} + mi := &file_distribution_proto_msgTypes[4] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ValidateTimestampRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ValidateTimestampRequest) ProtoMessage() {} + +func (x *ValidateTimestampRequest) ProtoReflect() protoreflect.Message { + mi := &file_distribution_proto_msgTypes[4] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ValidateTimestampRequest.ProtoReflect.Descriptor instead. +func (*ValidateTimestampRequest) Descriptor() ([]byte, []int) { + return file_distribution_proto_rawDescGZIP(), []int{4} +} + +func (x *ValidateTimestampRequest) GetTimestamp() uint64 { + if x != nil { + return x.Timestamp + } + return 0 +} + +type ValidateTimestampResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Valid bool `protobuf:"varint,1,opt,name=valid,proto3" json:"valid,omitempty"` + PhaseDActive bool `protobuf:"varint,2,opt,name=phase_d_active,json=phaseDActive,proto3" json:"phase_d_active,omitempty"` + PhaseDFloor uint64 `protobuf:"varint,3,opt,name=phase_d_floor,json=phaseDFloor,proto3" json:"phase_d_floor,omitempty"` + AllocationFloor uint64 `protobuf:"varint,4,opt,name=allocation_floor,json=allocationFloor,proto3" json:"allocation_floor,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ValidateTimestampResponse) Reset() { + *x = ValidateTimestampResponse{} + mi := &file_distribution_proto_msgTypes[5] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ValidateTimestampResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ValidateTimestampResponse) ProtoMessage() {} + +func (x *ValidateTimestampResponse) ProtoReflect() protoreflect.Message { + mi := &file_distribution_proto_msgTypes[5] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ValidateTimestampResponse.ProtoReflect.Descriptor instead. +func (*ValidateTimestampResponse) Descriptor() ([]byte, []int) { + return file_distribution_proto_rawDescGZIP(), []int{5} +} + +func (x *ValidateTimestampResponse) GetValid() bool { + if x != nil { + return x.Valid + } + return false +} + +func (x *ValidateTimestampResponse) GetPhaseDActive() bool { + if x != nil { + return x.PhaseDActive + } + return false +} + +func (x *ValidateTimestampResponse) GetPhaseDFloor() uint64 { + if x != nil { + return x.PhaseDFloor + } + return 0 +} + +func (x *ValidateTimestampResponse) GetAllocationFloor() uint64 { + if x != nil { + return x.AllocationFloor + } + return 0 +} + type RouteDescriptor struct { state protoimpl.MessageState `protogen:"open.v1"` RouteId uint64 `protobuf:"varint,1,opt,name=route_id,json=routeId,proto3" json:"route_id,omitempty"` @@ -500,7 +711,7 @@ type RouteDescriptor struct { func (x *RouteDescriptor) Reset() { *x = RouteDescriptor{} - mi := &file_distribution_proto_msgTypes[4] + mi := &file_distribution_proto_msgTypes[6] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -512,7 +723,7 @@ func (x *RouteDescriptor) String() string { func (*RouteDescriptor) ProtoMessage() {} func (x *RouteDescriptor) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[4] + mi := &file_distribution_proto_msgTypes[6] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -525,7 +736,7 @@ func (x *RouteDescriptor) ProtoReflect() protoreflect.Message { // Deprecated: Use RouteDescriptor.ProtoReflect.Descriptor instead. func (*RouteDescriptor) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{4} + return file_distribution_proto_rawDescGZIP(), []int{6} } func (x *RouteDescriptor) GetRouteId() uint64 { @@ -593,7 +804,7 @@ type SplitJobBracketProgress struct { func (x *SplitJobBracketProgress) Reset() { *x = SplitJobBracketProgress{} - mi := &file_distribution_proto_msgTypes[5] + mi := &file_distribution_proto_msgTypes[7] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -605,7 +816,7 @@ func (x *SplitJobBracketProgress) String() string { func (*SplitJobBracketProgress) ProtoMessage() {} func (x *SplitJobBracketProgress) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[5] + mi := &file_distribution_proto_msgTypes[7] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -618,7 +829,7 @@ func (x *SplitJobBracketProgress) ProtoReflect() protoreflect.Message { // Deprecated: Use SplitJobBracketProgress.ProtoReflect.Descriptor instead. func (*SplitJobBracketProgress) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{5} + return file_distribution_proto_rawDescGZIP(), []int{7} } func (x *SplitJobBracketProgress) GetBracketId() uint64 { @@ -718,7 +929,7 @@ type SplitJob struct { func (x *SplitJob) Reset() { *x = SplitJob{} - mi := &file_distribution_proto_msgTypes[6] + mi := &file_distribution_proto_msgTypes[8] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -730,7 +941,7 @@ func (x *SplitJob) String() string { func (*SplitJob) ProtoMessage() {} func (x *SplitJob) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[6] + mi := &file_distribution_proto_msgTypes[8] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -743,7 +954,7 @@ func (x *SplitJob) ProtoReflect() protoreflect.Message { // Deprecated: Use SplitJob.ProtoReflect.Descriptor instead. func (*SplitJob) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{6} + return file_distribution_proto_rawDescGZIP(), []int{8} } func (x *SplitJob) GetJobId() uint64 { @@ -985,7 +1196,7 @@ type ListRoutesRequest struct { func (x *ListRoutesRequest) Reset() { *x = ListRoutesRequest{} - mi := &file_distribution_proto_msgTypes[7] + mi := &file_distribution_proto_msgTypes[9] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -997,7 +1208,7 @@ func (x *ListRoutesRequest) String() string { func (*ListRoutesRequest) ProtoMessage() {} func (x *ListRoutesRequest) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[7] + mi := &file_distribution_proto_msgTypes[9] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1010,7 +1221,7 @@ func (x *ListRoutesRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use ListRoutesRequest.ProtoReflect.Descriptor instead. func (*ListRoutesRequest) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{7} + return file_distribution_proto_rawDescGZIP(), []int{9} } type ListRoutesResponse struct { @@ -1023,7 +1234,7 @@ type ListRoutesResponse struct { func (x *ListRoutesResponse) Reset() { *x = ListRoutesResponse{} - mi := &file_distribution_proto_msgTypes[8] + mi := &file_distribution_proto_msgTypes[10] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1035,7 +1246,7 @@ func (x *ListRoutesResponse) String() string { func (*ListRoutesResponse) ProtoMessage() {} func (x *ListRoutesResponse) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[8] + mi := &file_distribution_proto_msgTypes[10] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1048,7 +1259,7 @@ func (x *ListRoutesResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use ListRoutesResponse.ProtoReflect.Descriptor instead. func (*ListRoutesResponse) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{8} + return file_distribution_proto_rawDescGZIP(), []int{10} } func (x *ListRoutesResponse) GetCatalogVersion() uint64 { @@ -1076,7 +1287,7 @@ type SplitRangeRequest struct { func (x *SplitRangeRequest) Reset() { *x = SplitRangeRequest{} - mi := &file_distribution_proto_msgTypes[9] + mi := &file_distribution_proto_msgTypes[11] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1088,7 +1299,7 @@ func (x *SplitRangeRequest) String() string { func (*SplitRangeRequest) ProtoMessage() {} func (x *SplitRangeRequest) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[9] + mi := &file_distribution_proto_msgTypes[11] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1101,7 +1312,7 @@ func (x *SplitRangeRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use SplitRangeRequest.ProtoReflect.Descriptor instead. func (*SplitRangeRequest) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{9} + return file_distribution_proto_rawDescGZIP(), []int{11} } func (x *SplitRangeRequest) GetExpectedCatalogVersion() uint64 { @@ -1136,7 +1347,7 @@ type SplitRangeResponse struct { func (x *SplitRangeResponse) Reset() { *x = SplitRangeResponse{} - mi := &file_distribution_proto_msgTypes[10] + mi := &file_distribution_proto_msgTypes[12] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1148,7 +1359,7 @@ func (x *SplitRangeResponse) String() string { func (*SplitRangeResponse) ProtoMessage() {} func (x *SplitRangeResponse) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[10] + mi := &file_distribution_proto_msgTypes[12] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1161,7 +1372,7 @@ func (x *SplitRangeResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use SplitRangeResponse.ProtoReflect.Descriptor instead. func (*SplitRangeResponse) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{10} + return file_distribution_proto_rawDescGZIP(), []int{12} } func (x *SplitRangeResponse) GetCatalogVersion() uint64 { @@ -1193,7 +1404,7 @@ type CatalogCapabilitiesRequest struct { func (x *CatalogCapabilitiesRequest) Reset() { *x = CatalogCapabilitiesRequest{} - mi := &file_distribution_proto_msgTypes[11] + mi := &file_distribution_proto_msgTypes[13] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1205,7 +1416,7 @@ func (x *CatalogCapabilitiesRequest) String() string { func (*CatalogCapabilitiesRequest) ProtoMessage() {} func (x *CatalogCapabilitiesRequest) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[11] + mi := &file_distribution_proto_msgTypes[13] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1218,7 +1429,7 @@ func (x *CatalogCapabilitiesRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use CatalogCapabilitiesRequest.ProtoReflect.Descriptor instead. func (*CatalogCapabilitiesRequest) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{11} + return file_distribution_proto_rawDescGZIP(), []int{13} } type CatalogCapabilitiesResponse struct { @@ -1233,7 +1444,7 @@ type CatalogCapabilitiesResponse struct { func (x *CatalogCapabilitiesResponse) Reset() { *x = CatalogCapabilitiesResponse{} - mi := &file_distribution_proto_msgTypes[12] + mi := &file_distribution_proto_msgTypes[14] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1245,7 +1456,7 @@ func (x *CatalogCapabilitiesResponse) String() string { func (*CatalogCapabilitiesResponse) ProtoMessage() {} func (x *CatalogCapabilitiesResponse) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[12] + mi := &file_distribution_proto_msgTypes[14] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1258,7 +1469,7 @@ func (x *CatalogCapabilitiesResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use CatalogCapabilitiesResponse.ProtoReflect.Descriptor instead. func (*CatalogCapabilitiesResponse) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{12} + return file_distribution_proto_rawDescGZIP(), []int{14} } func (x *CatalogCapabilitiesResponse) GetSupportedProtocolVersions() []uint32 { @@ -1300,7 +1511,7 @@ type CatalogWatchRequest struct { func (x *CatalogWatchRequest) Reset() { *x = CatalogWatchRequest{} - mi := &file_distribution_proto_msgTypes[13] + mi := &file_distribution_proto_msgTypes[15] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1312,7 +1523,7 @@ func (x *CatalogWatchRequest) String() string { func (*CatalogWatchRequest) ProtoMessage() {} func (x *CatalogWatchRequest) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[13] + mi := &file_distribution_proto_msgTypes[15] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1325,7 +1536,7 @@ func (x *CatalogWatchRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use CatalogWatchRequest.ProtoReflect.Descriptor instead. func (*CatalogWatchRequest) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{13} + return file_distribution_proto_rawDescGZIP(), []int{15} } func (x *CatalogWatchRequest) GetProtocolVersion() uint32 { @@ -1360,7 +1571,7 @@ type CatalogDeltaMutation struct { func (x *CatalogDeltaMutation) Reset() { *x = CatalogDeltaMutation{} - mi := &file_distribution_proto_msgTypes[14] + mi := &file_distribution_proto_msgTypes[16] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1372,7 +1583,7 @@ func (x *CatalogDeltaMutation) String() string { func (*CatalogDeltaMutation) ProtoMessage() {} func (x *CatalogDeltaMutation) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[14] + mi := &file_distribution_proto_msgTypes[16] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1385,7 +1596,7 @@ func (x *CatalogDeltaMutation) ProtoReflect() protoreflect.Message { // Deprecated: Use CatalogDeltaMutation.ProtoReflect.Descriptor instead. func (*CatalogDeltaMutation) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{14} + return file_distribution_proto_rawDescGZIP(), []int{16} } func (x *CatalogDeltaMutation) GetOp() CatalogDeltaMutationOp { @@ -1420,7 +1631,7 @@ type CatalogDeltaRecord struct { func (x *CatalogDeltaRecord) Reset() { *x = CatalogDeltaRecord{} - mi := &file_distribution_proto_msgTypes[15] + mi := &file_distribution_proto_msgTypes[17] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1432,7 +1643,7 @@ func (x *CatalogDeltaRecord) String() string { func (*CatalogDeltaRecord) ProtoMessage() {} func (x *CatalogDeltaRecord) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[15] + mi := &file_distribution_proto_msgTypes[17] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1445,7 +1656,7 @@ func (x *CatalogDeltaRecord) ProtoReflect() protoreflect.Message { // Deprecated: Use CatalogDeltaRecord.ProtoReflect.Descriptor instead. func (*CatalogDeltaRecord) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{15} + return file_distribution_proto_rawDescGZIP(), []int{17} } func (x *CatalogDeltaRecord) GetPreviousVersion() uint64 { @@ -1479,7 +1690,7 @@ type CatalogSnapshotReset struct { func (x *CatalogSnapshotReset) Reset() { *x = CatalogSnapshotReset{} - mi := &file_distribution_proto_msgTypes[16] + mi := &file_distribution_proto_msgTypes[18] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1491,7 +1702,7 @@ func (x *CatalogSnapshotReset) String() string { func (*CatalogSnapshotReset) ProtoMessage() {} func (x *CatalogSnapshotReset) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[16] + mi := &file_distribution_proto_msgTypes[18] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1504,7 +1715,7 @@ func (x *CatalogSnapshotReset) ProtoReflect() protoreflect.Message { // Deprecated: Use CatalogSnapshotReset.ProtoReflect.Descriptor instead. func (*CatalogSnapshotReset) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{16} + return file_distribution_proto_rawDescGZIP(), []int{18} } func (x *CatalogSnapshotReset) GetVersion() uint64 { @@ -1534,7 +1745,7 @@ type CatalogWatchEvent struct { func (x *CatalogWatchEvent) Reset() { *x = CatalogWatchEvent{} - mi := &file_distribution_proto_msgTypes[17] + mi := &file_distribution_proto_msgTypes[19] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1546,7 +1757,7 @@ func (x *CatalogWatchEvent) String() string { func (*CatalogWatchEvent) ProtoMessage() {} func (x *CatalogWatchEvent) ProtoReflect() protoreflect.Message { - mi := &file_distribution_proto_msgTypes[17] + mi := &file_distribution_proto_msgTypes[19] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1559,7 +1770,7 @@ func (x *CatalogWatchEvent) ProtoReflect() protoreflect.Message { // Deprecated: Use CatalogWatchEvent.ProtoReflect.Descriptor instead. func (*CatalogWatchEvent) Descriptor() ([]byte, []int) { - return file_distribution_proto_rawDescGZIP(), []int{17} + return file_distribution_proto_rawDescGZIP(), []int{19} } func (x *CatalogWatchEvent) GetPayload() isCatalogWatchEvent_Payload { @@ -1613,10 +1824,27 @@ const file_distribution_proto_rawDesc = "" + "\x10GetRouteResponse\x12\x14\n" + "\x05start\x18\x01 \x01(\fR\x05start\x12\x10\n" + "\x03end\x18\x02 \x01(\fR\x03end\x12\"\n" + - "\rraft_group_id\x18\x03 \x01(\x04R\vraftGroupId\"\x15\n" + - "\x13GetTimestampRequest\"4\n" + + "\rraft_group_id\x18\x03 \x01(\x04R\vraftGroupId\"\xa5\x01\n" + + "\x13GetTimestampRequest\x12\x14\n" + + "\x05count\x18\x01 \x01(\rR\x05count\x12#\n" + + "\rmin_timestamp\x18\x02 \x01(\x04R\fminTimestamp\x12)\n" + + "\x10activate_cutover\x18\x03 \x01(\bR\x0factivateCutover\x12(\n" + + "\x10activate_phase_d\x18\x04 \x01(\bR\x0eactivatePhaseD\"\xb4\x02\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\x12$\n" + + "\x0ephase_d_active\x18\x06 \x01(\bR\fphaseDActive\x12\"\n" + + "\rphase_d_floor\x18\a \x01(\x04R\vphaseDFloor\"8\n" + + "\x18ValidateTimestampRequest\x12\x1c\n" + + "\ttimestamp\x18\x01 \x01(\x04R\ttimestamp\"\xa6\x01\n" + + "\x19ValidateTimestampResponse\x12\x14\n" + + "\x05valid\x18\x01 \x01(\bR\x05valid\x12$\n" + + "\x0ephase_d_active\x18\x02 \x01(\bR\fphaseDActive\x12\"\n" + + "\rphase_d_floor\x18\x03 \x01(\x04R\vphaseDFloor\x12)\n" + + "\x10allocation_floor\x18\x04 \x01(\x04R\x0fallocationFloor\"\xe5\x01\n" + "\x0fRouteDescriptor\x12\x19\n" + "\broute_id\x18\x01 \x01(\x04R\arouteId\x12\x14\n" + "\x05start\x18\x02 \x01(\fR\x05start\x12\x10\n" + @@ -1744,10 +1972,11 @@ const file_distribution_proto_rawDesc = "" + "\x16CatalogDeltaMutationOp\x12)\n" + "%CATALOG_DELTA_MUTATION_OP_UNSPECIFIED\x10\x00\x12$\n" + " CATALOG_DELTA_MUTATION_OP_UPSERT\x10\x01\x12$\n" + - " CATALOG_DELTA_MUTATION_OP_DELETE\x10\x022\x87\x03\n" + + " CATALOG_DELTA_MUTATION_OP_DELETE\x10\x022\xd5\x03\n" + "\fDistribution\x121\n" + "\bGetRoute\x12\x10.GetRouteRequest\x1a\x11.GetRouteResponse\"\x00\x12=\n" + - "\fGetTimestamp\x12\x14.GetTimestampRequest\x1a\x15.GetTimestampResponse\"\x00\x127\n" + + "\fGetTimestamp\x12\x14.GetTimestampRequest\x1a\x15.GetTimestampResponse\"\x00\x12L\n" + + "\x11ValidateTimestamp\x12\x19.ValidateTimestampRequest\x1a\x1a.ValidateTimestampResponse\"\x00\x127\n" + "\n" + "ListRoutes\x12\x12.ListRoutesRequest\x1a\x13.ListRoutesResponse\"\x00\x127\n" + "\n" + @@ -1768,7 +1997,7 @@ func file_distribution_proto_rawDescGZIP() []byte { } var file_distribution_proto_enumTypes = make([]protoimpl.EnumInfo, 5) -var file_distribution_proto_msgTypes = make([]protoimpl.MessageInfo, 18) +var file_distribution_proto_msgTypes = make([]protoimpl.MessageInfo, 20) var file_distribution_proto_goTypes = []any{ (RouteState)(0), // 0: RouteState (SplitJobPhase)(0), // 1: SplitJobPhase @@ -1779,20 +2008,22 @@ var file_distribution_proto_goTypes = []any{ (*GetRouteResponse)(nil), // 6: GetRouteResponse (*GetTimestampRequest)(nil), // 7: GetTimestampRequest (*GetTimestampResponse)(nil), // 8: GetTimestampResponse - (*RouteDescriptor)(nil), // 9: RouteDescriptor - (*SplitJobBracketProgress)(nil), // 10: SplitJobBracketProgress - (*SplitJob)(nil), // 11: SplitJob - (*ListRoutesRequest)(nil), // 12: ListRoutesRequest - (*ListRoutesResponse)(nil), // 13: ListRoutesResponse - (*SplitRangeRequest)(nil), // 14: SplitRangeRequest - (*SplitRangeResponse)(nil), // 15: SplitRangeResponse - (*CatalogCapabilitiesRequest)(nil), // 16: CatalogCapabilitiesRequest - (*CatalogCapabilitiesResponse)(nil), // 17: CatalogCapabilitiesResponse - (*CatalogWatchRequest)(nil), // 18: CatalogWatchRequest - (*CatalogDeltaMutation)(nil), // 19: CatalogDeltaMutation - (*CatalogDeltaRecord)(nil), // 20: CatalogDeltaRecord - (*CatalogSnapshotReset)(nil), // 21: CatalogSnapshotReset - (*CatalogWatchEvent)(nil), // 22: CatalogWatchEvent + (*ValidateTimestampRequest)(nil), // 9: ValidateTimestampRequest + (*ValidateTimestampResponse)(nil), // 10: ValidateTimestampResponse + (*RouteDescriptor)(nil), // 11: RouteDescriptor + (*SplitJobBracketProgress)(nil), // 12: SplitJobBracketProgress + (*SplitJob)(nil), // 13: SplitJob + (*ListRoutesRequest)(nil), // 14: ListRoutesRequest + (*ListRoutesResponse)(nil), // 15: ListRoutesResponse + (*SplitRangeRequest)(nil), // 16: SplitRangeRequest + (*SplitRangeResponse)(nil), // 17: SplitRangeResponse + (*CatalogCapabilitiesRequest)(nil), // 18: CatalogCapabilitiesRequest + (*CatalogCapabilitiesResponse)(nil), // 19: CatalogCapabilitiesResponse + (*CatalogWatchRequest)(nil), // 20: CatalogWatchRequest + (*CatalogDeltaMutation)(nil), // 21: CatalogDeltaMutation + (*CatalogDeltaRecord)(nil), // 22: CatalogDeltaRecord + (*CatalogSnapshotReset)(nil), // 23: CatalogSnapshotReset + (*CatalogWatchEvent)(nil), // 24: CatalogWatchEvent } var file_distribution_proto_depIdxs = []int32{ 0, // 0: RouteDescriptor.state:type_name -> RouteState @@ -1802,30 +2033,32 @@ var file_distribution_proto_depIdxs = []int32{ 1, // 4: SplitJob.abandon_from_phase:type_name -> SplitJobPhase 2, // 5: SplitJob.cutover_read_fence_state:type_name -> SplitJobBarrierState 2, // 6: SplitJob.target_staged_readiness_state:type_name -> SplitJobBarrierState - 10, // 7: SplitJob.bracket_progress:type_name -> SplitJobBracketProgress - 9, // 8: ListRoutesResponse.routes:type_name -> RouteDescriptor - 9, // 9: SplitRangeResponse.left:type_name -> RouteDescriptor - 9, // 10: SplitRangeResponse.right:type_name -> RouteDescriptor + 12, // 7: SplitJob.bracket_progress:type_name -> SplitJobBracketProgress + 11, // 8: ListRoutesResponse.routes:type_name -> RouteDescriptor + 11, // 9: SplitRangeResponse.left:type_name -> RouteDescriptor + 11, // 10: SplitRangeResponse.right:type_name -> RouteDescriptor 4, // 11: CatalogDeltaMutation.op:type_name -> CatalogDeltaMutationOp - 9, // 12: CatalogDeltaMutation.route:type_name -> RouteDescriptor - 19, // 13: CatalogDeltaRecord.mutations:type_name -> CatalogDeltaMutation - 9, // 14: CatalogSnapshotReset.routes:type_name -> RouteDescriptor - 21, // 15: CatalogWatchEvent.snapshot:type_name -> CatalogSnapshotReset - 20, // 16: CatalogWatchEvent.delta:type_name -> CatalogDeltaRecord + 11, // 12: CatalogDeltaMutation.route:type_name -> RouteDescriptor + 21, // 13: CatalogDeltaRecord.mutations:type_name -> CatalogDeltaMutation + 11, // 14: CatalogSnapshotReset.routes:type_name -> RouteDescriptor + 23, // 15: CatalogWatchEvent.snapshot:type_name -> CatalogSnapshotReset + 22, // 16: CatalogWatchEvent.delta:type_name -> CatalogDeltaRecord 5, // 17: Distribution.GetRoute:input_type -> GetRouteRequest 7, // 18: Distribution.GetTimestamp:input_type -> GetTimestampRequest - 12, // 19: Distribution.ListRoutes:input_type -> ListRoutesRequest - 14, // 20: Distribution.SplitRange:input_type -> SplitRangeRequest - 16, // 21: Distribution.GetCatalogCapabilities:input_type -> CatalogCapabilitiesRequest - 18, // 22: Distribution.WatchCatalog:input_type -> CatalogWatchRequest - 6, // 23: Distribution.GetRoute:output_type -> GetRouteResponse - 8, // 24: Distribution.GetTimestamp:output_type -> GetTimestampResponse - 13, // 25: Distribution.ListRoutes:output_type -> ListRoutesResponse - 15, // 26: Distribution.SplitRange:output_type -> SplitRangeResponse - 17, // 27: Distribution.GetCatalogCapabilities:output_type -> CatalogCapabilitiesResponse - 22, // 28: Distribution.WatchCatalog:output_type -> CatalogWatchEvent - 23, // [23:29] is the sub-list for method output_type - 17, // [17:23] is the sub-list for method input_type + 9, // 19: Distribution.ValidateTimestamp:input_type -> ValidateTimestampRequest + 14, // 20: Distribution.ListRoutes:input_type -> ListRoutesRequest + 16, // 21: Distribution.SplitRange:input_type -> SplitRangeRequest + 18, // 22: Distribution.GetCatalogCapabilities:input_type -> CatalogCapabilitiesRequest + 20, // 23: Distribution.WatchCatalog:input_type -> CatalogWatchRequest + 6, // 24: Distribution.GetRoute:output_type -> GetRouteResponse + 8, // 25: Distribution.GetTimestamp:output_type -> GetTimestampResponse + 10, // 26: Distribution.ValidateTimestamp:output_type -> ValidateTimestampResponse + 15, // 27: Distribution.ListRoutes:output_type -> ListRoutesResponse + 17, // 28: Distribution.SplitRange:output_type -> SplitRangeResponse + 19, // 29: Distribution.GetCatalogCapabilities:output_type -> CatalogCapabilitiesResponse + 24, // 30: Distribution.WatchCatalog:output_type -> CatalogWatchEvent + 24, // [24:31] is the sub-list for method output_type + 17, // [17:24] is the sub-list for method input_type 17, // [17:17] is the sub-list for extension type_name 17, // [17:17] is the sub-list for extension extendee 0, // [0:17] is the sub-list for field type_name @@ -1836,7 +2069,7 @@ func file_distribution_proto_init() { if File_distribution_proto != nil { return } - file_distribution_proto_msgTypes[17].OneofWrappers = []any{ + file_distribution_proto_msgTypes[19].OneofWrappers = []any{ (*CatalogWatchEvent_Snapshot)(nil), (*CatalogWatchEvent_Delta)(nil), } @@ -1846,7 +2079,7 @@ func file_distribution_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_distribution_proto_rawDesc), len(file_distribution_proto_rawDesc)), NumEnums: 5, - NumMessages: 18, + NumMessages: 20, NumExtensions: 0, NumServices: 1, }, diff --git a/proto/distribution.proto b/proto/distribution.proto index a71a64565..b6a00494f 100644 --- a/proto/distribution.proto +++ b/proto/distribution.proto @@ -5,6 +5,7 @@ option go_package = "github.com/bootjp/elastickv/proto"; service Distribution { rpc GetRoute (GetRouteRequest) returns (GetRouteResponse) {} rpc GetTimestamp (GetTimestampRequest) returns (GetTimestampResponse) {} + rpc ValidateTimestamp (ValidateTimestampRequest) returns (ValidateTimestampResponse) {} rpc ListRoutes (ListRoutesRequest) returns (ListRoutesResponse) {} rpc SplitRange (SplitRangeRequest) returns (SplitRangeResponse) {} rpc GetCatalogCapabilities (CatalogCapabilitiesRequest) returns (CatalogCapabilitiesResponse) {} @@ -23,10 +24,51 @@ 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; + // ActivatePhaseD closes the legacy compatibility window after every member + // understands the M7 FSM entry. It requires ActivateCutover and is one-way. + bool activate_phase_d = 4; +} message GetTimestampResponse { 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; + // PhaseDActive reports that legacy per-shard issuance has been retired. + bool phase_d_active = 6; + // PhaseDFloor is the allocation floor captured by the Phase-D marker. + uint64 phase_d_floor = 7; +} + +message ValidateTimestampRequest { + uint64 timestamp = 1; +} + +message ValidateTimestampResponse { + bool valid = 1; + bool phase_d_active = 2; + uint64 phase_d_floor = 3; + uint64 allocation_floor = 4; } enum RouteState { diff --git a/proto/distribution_grpc.pb.go b/proto/distribution_grpc.pb.go index 610f7d6f2..c274d3506 100644 --- a/proto/distribution_grpc.pb.go +++ b/proto/distribution_grpc.pb.go @@ -21,6 +21,7 @@ const _ = grpc.SupportPackageIsVersion9 const ( Distribution_GetRoute_FullMethodName = "/Distribution/GetRoute" Distribution_GetTimestamp_FullMethodName = "/Distribution/GetTimestamp" + Distribution_ValidateTimestamp_FullMethodName = "/Distribution/ValidateTimestamp" Distribution_ListRoutes_FullMethodName = "/Distribution/ListRoutes" Distribution_SplitRange_FullMethodName = "/Distribution/SplitRange" Distribution_GetCatalogCapabilities_FullMethodName = "/Distribution/GetCatalogCapabilities" @@ -33,6 +34,7 @@ const ( type DistributionClient interface { GetRoute(ctx context.Context, in *GetRouteRequest, opts ...grpc.CallOption) (*GetRouteResponse, error) GetTimestamp(ctx context.Context, in *GetTimestampRequest, opts ...grpc.CallOption) (*GetTimestampResponse, error) + ValidateTimestamp(ctx context.Context, in *ValidateTimestampRequest, opts ...grpc.CallOption) (*ValidateTimestampResponse, error) ListRoutes(ctx context.Context, in *ListRoutesRequest, opts ...grpc.CallOption) (*ListRoutesResponse, error) SplitRange(ctx context.Context, in *SplitRangeRequest, opts ...grpc.CallOption) (*SplitRangeResponse, error) GetCatalogCapabilities(ctx context.Context, in *CatalogCapabilitiesRequest, opts ...grpc.CallOption) (*CatalogCapabilitiesResponse, error) @@ -67,6 +69,16 @@ func (c *distributionClient) GetTimestamp(ctx context.Context, in *GetTimestampR return out, nil } +func (c *distributionClient) ValidateTimestamp(ctx context.Context, in *ValidateTimestampRequest, opts ...grpc.CallOption) (*ValidateTimestampResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(ValidateTimestampResponse) + err := c.cc.Invoke(ctx, Distribution_ValidateTimestamp_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + func (c *distributionClient) ListRoutes(ctx context.Context, in *ListRoutesRequest, opts ...grpc.CallOption) (*ListRoutesResponse, error) { cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) out := new(ListRoutesResponse) @@ -122,6 +134,7 @@ type Distribution_WatchCatalogClient = grpc.ServerStreamingClient[CatalogWatchEv type DistributionServer interface { GetRoute(context.Context, *GetRouteRequest) (*GetRouteResponse, error) GetTimestamp(context.Context, *GetTimestampRequest) (*GetTimestampResponse, error) + ValidateTimestamp(context.Context, *ValidateTimestampRequest) (*ValidateTimestampResponse, error) ListRoutes(context.Context, *ListRoutesRequest) (*ListRoutesResponse, error) SplitRange(context.Context, *SplitRangeRequest) (*SplitRangeResponse, error) GetCatalogCapabilities(context.Context, *CatalogCapabilitiesRequest) (*CatalogCapabilitiesResponse, error) @@ -142,6 +155,9 @@ func (UnimplementedDistributionServer) GetRoute(context.Context, *GetRouteReques func (UnimplementedDistributionServer) GetTimestamp(context.Context, *GetTimestampRequest) (*GetTimestampResponse, error) { return nil, status.Error(codes.Unimplemented, "method GetTimestamp not implemented") } +func (UnimplementedDistributionServer) ValidateTimestamp(context.Context, *ValidateTimestampRequest) (*ValidateTimestampResponse, error) { + return nil, status.Error(codes.Unimplemented, "method ValidateTimestamp not implemented") +} func (UnimplementedDistributionServer) ListRoutes(context.Context, *ListRoutesRequest) (*ListRoutesResponse, error) { return nil, status.Error(codes.Unimplemented, "method ListRoutes not implemented") } @@ -211,6 +227,24 @@ func _Distribution_GetTimestamp_Handler(srv interface{}, ctx context.Context, de return interceptor(ctx, in, info, handler) } +func _Distribution_ValidateTimestamp_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(ValidateTimestampRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(DistributionServer).ValidateTimestamp(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: Distribution_ValidateTimestamp_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(DistributionServer).ValidateTimestamp(ctx, req.(*ValidateTimestampRequest)) + } + return interceptor(ctx, in, info, handler) +} + func _Distribution_ListRoutes_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { in := new(ListRoutesRequest) if err := dec(in); err != nil { @@ -291,6 +325,10 @@ var Distribution_ServiceDesc = grpc.ServiceDesc{ MethodName: "GetTimestamp", Handler: _Distribution_GetTimestamp_Handler, }, + { + MethodName: "ValidateTimestamp", + Handler: _Distribution_ValidateTimestamp_Handler, + }, { MethodName: "ListRoutes", Handler: _Distribution_ListRoutes_Handler, diff --git a/proto/service.pb.go b/proto/service.pb.go index 450c4e5de..4f1e8aa8d 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"` @@ -2506,12 +2530,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 c1f5a003e..6634f5eee 100644 --- a/proto/service.proto +++ b/proto/service.proto @@ -69,11 +69,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 {