diff --git a/docs/design/2026_04_16_partial_centralized_tso.md b/docs/design/2026_04_16_partial_centralized_tso.md index 1e5a9a55f..631ae4dbf 100644 --- a/docs/design/2026_04_16_partial_centralized_tso.md +++ b/docs/design/2026_04_16_partial_centralized_tso.md @@ -1,11 +1,11 @@ # Centralized Timestamp Oracle (TSO) Design -- Status: Partial — M1 all-led-group HLC renewal and the M3-M5 default-group - TSO allocator/batch cutover are implemented; dedicated TSO Raft group wiring - remains open +- Status: Partial — M1 all-led-group HLC renewal, the M2 reserved-group + bootstrap bridge, and the M3-M5 TSO allocator/batch cutover are implemented; + the minimal TSO-only FSM and follower redirect/admin exposure remain open - Author: bootjp - Date: 2026-04-16 -- Updated: 2026-07-07 +- Updated: 2026-07-11 --- @@ -33,12 +33,25 @@ Implemented: wiring a `LocalTSOAllocator` through `BatchAllocator` so production coordinators can cut over to the default-group-backed TSO path without changing the legacy HLC default. +8. `groupID = 0` is reserved for the dedicated TSO group in runtime config: + shard ranges cannot route user data to it, the default data group skips it, + and SQS FIFO partition maps cannot route queue partitions to it. Existing + `--raftGroups` / `--raftGroupPeers` bootstrap the group as a compatibility + 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. Remaining: -1. Reserve and bootstrap a dedicated TSO Raft group. +1. Replace the compatibility bridge FSM with the minimal TSO-only FSM that + persists only the HLC ceiling. 2. Add follower redirect/admin exposure for the dedicated TSO leader. -3. Wire deployment/bootstrap config and mixed-version safety around the extra group. +3. Add Phase B shadow-read validation before making dedicated TSO the only + production timestamp path. ### 1.1 Original Limitation @@ -616,11 +629,11 @@ 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 | Phase A dual-write bridge: also propose ceiling to TSO group (Section 7.1) | Deferred until dedicated group | +| 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 | | 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 | Dedicated TSO Raft group (`groupID = 0`) with `TSOStateMachine` | Low | +| M6 — partial | Dedicated TSO Raft group (`groupID = 0`) is reserved/bootstrap-capable and warmed by the HLC renewal bridge; TSO-leader-only timestamp issuance and the minimal `TSOStateMachine` remain open | Low | | M7 | Phase D legacy cleanup + cross-shard SSI read-timestamp validation via TSO | Low | --- diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index b8b67dba5..b6d829d75 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -374,8 +374,19 @@ type ShardedCoordinator struct { // coordinator-owned persistence timestamp. Nil preserves the legacy shared // HLC path. tsAllocator TimestampAllocator - store store.MVCCStore - log *slog.Logger + // timestampGroup pins IsTimestampLeader to a dedicated Raft group when + // timestampGroupConfigured is true. Nil/false preserves the M3 bridge + // behavior where any locally-led shard group can issue TSO timestamps. + timestampGroup uint64 + timestampGroupConfigured bool + // 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 + // without letting Scan fences or DEL_PREFIX broadcasts touch them. + allShardGroupIDs []uint64 + allShardGroupsConfigured bool + store store.MVCCStore + log *slog.Logger // deregisterLeaseCbs removes the per-shard leader-loss callbacks // registered at construction. See Coordinate.Close for the // rationale. @@ -451,6 +462,27 @@ func (c *ShardedCoordinator) WithTSOAllocator(alloc TimestampAllocator) *Sharded return c } +// 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 +// able to commit writes when they do not also lead the timestamp group. +func (c *ShardedCoordinator) WithTimestampGroup(groupID uint64) *ShardedCoordinator { + c.timestampGroup = groupID + c.timestampGroupConfigured = true + 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. +func (c *ShardedCoordinator) WithAllShardGroups(groupIDs ...uint64) *ShardedCoordinator { + c.allShardGroupIDs = append([]uint64(nil), groupIDs...) + slices.Sort(c.allShardGroupIDs) + c.allShardGroupIDs = slices.Compact(c.allShardGroupIDs) + c.allShardGroupsConfigured = true + return c +} + // awaitRegistration implements the §4.1 first-write barrier (7a §3.2). // It returns nil (ungated) unless a registration is genuinely pending // AND this request is a mutating write that would land encrypted, in @@ -1007,9 +1039,10 @@ func validateDelPrefixOnly(elems []*Elem[OP]) error { } // dispatchDelPrefixBroadcast validates and broadcasts DEL_PREFIX operations -// to every shard group. Each element becomes a separate pb.Request (the FSM's -// extractDelPrefix processes only the first DEL_PREFIX mutation per request). -// All requests are batched into a single Commit call per shard group. +// to every configured all-shard data group. Each element becomes a separate +// pb.Request (the FSM's extractDelPrefix processes only the first DEL_PREFIX +// mutation per request). All requests are batched into a single Commit call +// per shard group. func (c *ShardedCoordinator) dispatchDelPrefixBroadcast(ctx context.Context, isTxn bool, elems []*Elem[OP]) (*CoordinateResponse, error) { if isTxn { return nil, errors.Wrap(ErrInvalidRequest, "DEL_PREFIX not supported in transactions") @@ -1035,8 +1068,8 @@ func (c *ShardedCoordinator) dispatchDelPrefixBroadcast(ctx context.Context, isT return c.broadcastToAllGroups(ctx, requests) } -// broadcastToAllGroups sends the same set of requests to every shard group in -// parallel and returns the maximum commit index. +// broadcastToAllGroups sends the same set of requests to every configured +// all-shard data group in parallel and returns the maximum commit index. func (c *ShardedCoordinator) broadcastToAllGroups(ctx context.Context, requests []*pb.Request) (*CoordinateResponse, error) { var ( maxIndex atomic.Uint64 @@ -1044,7 +1077,11 @@ func (c *ShardedCoordinator) broadcastToAllGroups(ctx context.Context, requests errMu sync.Mutex wg sync.WaitGroup ) - for _, g := range c.groups { + groups, err := c.allShardGroups() + if err != nil { + return nil, err + } + for _, g := range groups { wg.Add(1) go func(g *ShardGroup) { defer wg.Done() @@ -1583,14 +1620,44 @@ func (c *ShardedCoordinator) IsTimestampLeader() bool { if c == nil { return false } - for _, g := range c.groups { - if isLeaderEngine(engineForGroup(g)) { + if c.timestampGroupConfigured { + return isLeaderEngine(engineForGroup(c.groups[c.timestampGroup])) + } + for _, groupID := range c.timestampBridgeCandidateGroupIDs() { + if isLeaderEngine(engineForGroup(c.groups[groupID])) { return true } } return false } +func (c *ShardedCoordinator) timestampBridgeCandidateGroupIDs() []uint64 { + if c == nil { + return nil + } + if c.allShardGroupsConfigured { + ids := make([]uint64, 0, len(c.allShardGroupIDs)) + for _, groupID := range c.allShardGroupIDs { + if groupID == 0 { + continue + } + if _, ok := c.groups[groupID]; ok { + ids = append(ids, groupID) + } + } + return ids + } + ids := make([]uint64, 0, len(c.groups)) + for groupID := range c.groups { + if groupID == 0 { + continue + } + ids = append(ids, groupID) + } + slices.Sort(ids) + return ids +} + func (c *ShardedCoordinator) invalidateTimestampWindow() { if c == nil { return @@ -1677,20 +1744,21 @@ func (c *ShardedCoordinator) LeaseReadForKey(ctx context.Context, key []byte) (u return groupLeaseRead(ctx, g, c.leaseObserver) } -// LeaseReadAllGroups establishes the lease freshness bound on every shard -// group this coordinator owns. Multi-shard reads (Scan, GSI/whole-table -// Query) visit all intersecting routes across all groups (see -// ShardStore.ScanAt), so fencing only the default group would let those -// reads sample a snapshot on a non-default group without the freshness -// bound. It fails closed on the first group that cannot confirm its lease, -// since a partially-fenced read is exactly the stale read this guards -// against. Group iteration order is unspecified; correctness does not -// depend on it because every group must succeed. +// LeaseReadAllGroups establishes the lease freshness bound on every configured +// all-shard data group. Multi-shard reads (Scan, GSI/whole-table Query) visit +// all intersecting data routes across all groups (see ShardStore.ScanAt), so +// fencing only the default group would let those reads sample a snapshot on a +// non-default group without the freshness bound. It fails closed on the first +// group that cannot confirm its lease, since a partially-fenced read is exactly +// the stale read this guards against. Group iteration order is deterministic +// but correctness does not depend on it because every configured group must +// succeed. func (c *ShardedCoordinator) LeaseReadAllGroups(ctx context.Context) error { - if len(c.groups) == 0 { - return errors.WithStack(ErrLeaderNotFound) + groups, err := c.allShardGroups() + if err != nil { + return err } - for _, g := range c.groups { + for _, g := range groups { if _, err := groupLeaseRead(ctx, g, c.leaseObserver); err != nil { return errors.WithStack(err) } @@ -1698,6 +1766,48 @@ func (c *ShardedCoordinator) LeaseReadAllGroups(ctx context.Context) error { return nil } +func (c *ShardedCoordinator) allShardGroups() ([]*ShardGroup, error) { + if c == nil || len(c.groups) == 0 { + return nil, errors.WithStack(ErrLeaderNotFound) + } + if c.allShardGroupsConfigured { + return c.configuredAllShardGroups() + } + return c.ownedAllShardGroups() +} + +func (c *ShardedCoordinator) configuredAllShardGroups() ([]*ShardGroup, error) { + if len(c.allShardGroupIDs) == 0 { + return nil, errors.WithStack(ErrLeaderNotFound) + } + out := make([]*ShardGroup, 0, len(c.allShardGroupIDs)) + for _, gid := range c.allShardGroupIDs { + g, ok := c.groups[gid] + if !ok || g == nil { + return nil, errors.WithStack(ErrLeaderNotFound) + } + out = append(out, g) + } + return out, nil +} + +func (c *ShardedCoordinator) ownedAllShardGroups() ([]*ShardGroup, error) { + gids := make([]uint64, 0, len(c.groups)) + for gid := range c.groups { + gids = append(gids, gid) + } + slices.Sort(gids) + out := make([]*ShardGroup, 0, len(gids)) + for _, gid := range gids { + g := c.groups[gid] + if g == nil { + return nil, errors.WithStack(ErrLeaderNotFound) + } + out = append(out, g) + } + return out, nil +} + // observeLeaseRead forwards a hit / miss signal to observer when it // is non-nil. Kept as a package-level helper so both ShardedCoordinator // and any future sharded caller share one nil-safe entrypoint. diff --git a/kv/sharded_coordinator_leader_test.go b/kv/sharded_coordinator_leader_test.go index 10ec2de3f..0fda49e77 100644 --- a/kv/sharded_coordinator_leader_test.go +++ b/kv/sharded_coordinator_leader_test.go @@ -66,6 +66,20 @@ func TestShardedCoordinatorLeaseReadAllGroups_FencesEveryLeader(t *testing.T) { require.NoError(t, coord.LeaseReadAllGroups(context.Background())) } +func TestShardedCoordinatorLeaseReadAllGroups_UsesConfiguredAllShardGroups(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + engine.UpdateRoute([]byte(""), nil, 1) + + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 0: {Engine: &stubFollowerEngine{}}, + 1: {Engine: stubLeaderEngine{}}, + }, 1, NewHLC(), nil).WithAllShardGroups(1) + + require.NoError(t, coord.LeaseReadAllGroups(context.Background())) +} + // TestShardedCoordinatorLeaseReadAllGroups_FailsClosedOnUnreadableGroup // asserts LeaseReadAllGroups fails closed when ANY group cannot confirm its // lease: a partially-fenced multi-shard read is exactly the stale read the diff --git a/kv/sharded_coordinator_txn_test.go b/kv/sharded_coordinator_txn_test.go index 964f2713f..d0ea680c4 100644 --- a/kv/sharded_coordinator_txn_test.go +++ b/kv/sharded_coordinator_txn_test.go @@ -81,6 +81,32 @@ func TestShardedCoordinatorDispatchTxn_RejectsMissingPrimaryKey(t *testing.T) { require.ErrorIs(t, err, ErrTxnPrimaryKeyRequired) } +func TestShardedCoordinatorDelPrefixBroadcast_UsesConfiguredAllShardGroups(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + engine.UpdateRoute([]byte(""), nil, 1) + + tsoTxn := &recordingTransactional{} + dataTxn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 0: {Txn: tsoTxn}, + 1: {Txn: dataTxn}, + }, 1, NewHLC(), nil).WithAllShardGroups(1) + + resp, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{ + {Op: DelPrefix, Key: []byte("tenant/")}, + }, + }) + require.NoError(t, err) + require.NotNil(t, resp) + require.Empty(t, tsoTxn.requests) + require.Len(t, dataTxn.requests, 1) + require.Equal(t, pb.Op_DEL_PREFIX, dataTxn.requests[0].Mutations[0].Op) + require.Equal(t, []byte("tenant/"), dataTxn.requests[0].Mutations[0].Key) +} + func TestShardedCoordinatorDispatchTxn_CrossShardPhasesAndCommitIndex(t *testing.T) { t.Parallel() diff --git a/kv/tso_test.go b/kv/tso_test.go index 8daa8651d..ef4e58510 100644 --- a/kv/tso_test.go +++ b/kv/tso_test.go @@ -207,6 +207,58 @@ func TestShardedCoordinatorReportsAnyShardAsTimestampLeader(t *testing.T) { require.True(t, coord.IsTimestampLeader()) } +func TestShardedCoordinatorDedicatedGroupDoesNotBlockDataLeaderTimestampBridge(t *testing.T) { + engine := distribution.NewEngine() + engine.UpdateRoute([]byte(""), nil, 1) + + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 0: {Engine: &stubFollowerEngine{}}, + 1: {Engine: stubLeaderEngine{}}, + }, 1, NewHLC(), nil) + + require.True(t, coord.IsTimestampLeader()) +} + +func TestShardedCoordinatorTimestampBridgeIgnoresReservedGroupLeader(t *testing.T) { + engine := distribution.NewEngine() + engine.UpdateRoute([]byte(""), nil, 1) + + newCoord := func() *ShardedCoordinator { + return NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 0: {Engine: stubLeaderEngine{}}, + 1: {Engine: &stubFollowerEngine{}}, + }, 1, NewHLC(), nil) + } + + t.Run("configured data groups", func(t *testing.T) { + require.False(t, newCoord().WithAllShardGroups(1).IsTimestampLeader()) + }) + + t.Run("legacy fallback", func(t *testing.T) { + require.False(t, newCoord().IsTimestampLeader()) + }) +} + +func TestShardedCoordinatorUsesDedicatedTimestampGroupLeader(t *testing.T) { + engine := distribution.NewEngine() + engine.UpdateRoute([]byte("a"), []byte("m"), 1) + engine.UpdateRoute([]byte("m"), nil, 2) + + dataLeaderOnly := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 0: {Engine: &stubFollowerEngine{}}, + 1: {Engine: stubLeaderEngine{}}, + 2: {Engine: &stubFollowerEngine{}}, + }, 1, NewHLC(), nil).WithTimestampGroup(0) + require.False(t, dataLeaderOnly.IsTimestampLeader()) + + tsoLeader := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 0: {Engine: stubLeaderEngine{}}, + 1: {Engine: &stubFollowerEngine{}}, + 2: {Engine: &stubFollowerEngine{}}, + }, 1, NewHLC(), nil).WithTimestampGroup(0) + require.True(t, tsoLeader.IsTimestampLeader()) +} + func TestNextTimestampAfterThroughFallbackWithoutCoordinatorClock(t *testing.T) { got, err := NextTimestampThrough(context.Background(), &Coordinate{}, "test") require.NoError(t, err) diff --git a/main.go b/main.go index 530dbf4b7..d63287c98 100644 --- a/main.go +++ b/main.go @@ -452,6 +452,7 @@ func run() error { WithLeaseReadObserver(metricsRegistry.LeaseReadObserver()). WithSampler(keyVizSamplerForCoordinator(sampler)). WithKeyVizLabelsEnabled(*keyvizLabelsEnabled). + WithAllShardGroups(dataGroupIDs(cfg.groups)...). WithPartitionResolver(buildSQSPartitionResolver(cfg.sqsFifoPartitionMap)) if err := configureCoordinatorTSO(coordinate); err != nil { return err @@ -713,7 +714,7 @@ func parseRuntimeConfig(myAddr, redisAddr, s3Addr, dynamoAddr, sqsAddr, raftGrou leaderDynamo: leaderDynamo, leaderSQS: leaderSQS, sqsFifoPartitionMap: sqsFifoPartitionMap, - multi: len(groups) > 1, + multi: dataGroupsNeedMultiDirs(groups), }, nil } @@ -821,6 +822,9 @@ func buildSQSFifoPartitionMap(groups []groupSpec, raw string) (map[string]sqsFif if len(parsed) == 0 { return parsed, nil } + if err := validateSQSFifoPartitionMapNoDedicatedTSOGroup(parsed); err != nil { + return nil, errors.Wrapf(err, "invalid sqs fifo partition map") + } groupIDs := make(map[string]struct{}, len(groups)) for _, g := range groups { groupIDs[strconv.FormatUint(g.id, 10)] = struct{}{} @@ -831,6 +835,20 @@ func buildSQSFifoPartitionMap(groups []groupSpec, raw string) (map[string]sqsFif return parsed, nil } +func validateSQSFifoPartitionMapNoDedicatedTSOGroup(partitionMap map[string]sqsFifoQueueRouting) error { + for _, queue := range slices.Sorted(maps.Keys(partitionMap)) { + routing := partitionMap[queue] + for partition, group := range routing.groups { + if group == strconv.FormatUint(dedicatedTSORaftGroupID, 10) { + return errors.Wrapf(ErrInvalidSQSFifoPartitionMapEntry, + "queue %q partition %d: group %q is reserved for TSO", + queue, partition, group) + } + } + } + return nil +} + func buildLeaderDynamo(groups []groupSpec, dynamoAddr string, raftDynamoMap string) (map[string]string, error) { return buildLeaderAddrMap(groups, dynamoAddr, raftDynamoMap, parseRaftDynamoMap) } @@ -1035,6 +1053,7 @@ func buildShardGroups( // NewApplier install a private per-applier cache, silently // breaking the shared-cache invariant 6D-6c-1 relies on. encWiring = encWiring.withDefaultedCache() + multi = effectiveMultiDataDirs(groups, multi) runtimes := make([]*raftGroupRuntime, 0, len(groups)) shardGroups := make(map[uint64]*kv.ShardGroup, len(groups)) for _, g := range groups { @@ -1576,6 +1595,8 @@ func configureCoordinatorTSO(coordinate *kv.ShardedCoordinator) error { if !*tsoEnabled { 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) if err != nil { return errors.Wrap(err, "configure tso allocator") diff --git a/main_encryption_startup_guard.go b/main_encryption_startup_guard.go index 9074e01a6..13c62aeb3 100644 --- a/main_encryption_startup_guard.go +++ b/main_encryption_startup_guard.go @@ -245,6 +245,7 @@ func checkEnvelopeCutoverDivergenceBeforeNonceBump( if !encryptionEnabled || sidecarPath == "" { return nil } + multi = effectiveMultiDataDirs(groups, multi) sc, err := readExistingSidecarForStartupGuard(sidecarPath) if err != nil { return err @@ -368,6 +369,7 @@ func checkEncryptionMembershipStartupGuardsBeforeEngine(in encryptionMembershipS if !in.encryptionEnabled || in.sidecarPath == "" { return nil } + in.multi = effectiveMultiDataDirs(in.groups, in.multi) sc, err := readExistingSidecarForStartupGuard(in.sidecarPath) if err != nil || sc == nil { return err diff --git a/multiraft_runtime.go b/multiraft_runtime.go index 6c41fac80..c46daacfe 100644 --- a/multiraft_runtime.go +++ b/multiraft_runtime.go @@ -109,7 +109,7 @@ func (r *raftGroupRuntime) registerGRPC(server grpc.ServiceRegistrar) { const raftDirPerm = 0o755 func groupDataDir(baseDir, raftID string, groupID uint64, multi bool) string { - if !multi { + if !multi && groupID != 0 { return filepath.Join(baseDir, raftID) } return filepath.Join(baseDir, raftID, fmt.Sprintf("group-%d", groupID)) diff --git a/multiraft_runtime_test.go b/multiraft_runtime_test.go index 7e20233b3..3de1ccc38 100644 --- a/multiraft_runtime_test.go +++ b/multiraft_runtime_test.go @@ -20,11 +20,67 @@ func TestGroupDataDir(t *testing.T) { require.Equal(t, filepath.Join(base, raftID), groupDataDir(base, raftID, 1, false)) }) + t.Run("reserved TSO group stays isolated in single data-group mode", func(t *testing.T) { + require.Equal(t, filepath.Join(base, raftID, "group-0"), groupDataDir(base, raftID, 0, false)) + }) + t.Run("multi", func(t *testing.T) { require.Equal(t, filepath.Join(base, raftID, "group-2"), groupDataDir(base, raftID, 2, true)) }) } +func TestBuildShardGroupsWithDedicatedTSOPreservesSingleDataGroupDir(t *testing.T) { + baseDir := t.TempDir() + raftID := "n1" + legacyDir := filepath.Join(baseDir, raftID) + require.NoError(t, os.MkdirAll(legacyDir, raftDirPerm)) + legacyMarker := filepath.Join(legacyDir, "legacy.marker") + require.NoError(t, os.WriteFile(legacyMarker, []byte("keep"), 0o600)) + + groups := []groupSpec{ + {id: dedicatedTSORaftGroupID, address: "127.0.0.1:17000"}, + {id: 1, address: "127.0.0.1:17001"}, + } + engine := distribution.NewEngine() + engine.UpdateRoute([]byte(""), nil, 1) + factory, err := newRaftFactory(raftEngineEtcd, nil) + require.NoError(t, err) + + runtimes, shardGroups, err := buildShardGroups( + raftID, + baseDir, + groups, + true, + true, + raftBootstrapConfig{}, + factory, + nil, + kv.NewHLC(), + nil, + nil, + "", + encryptionWriteWiring{}, + engine, + ) + require.NoError(t, err) + t.Cleanup(func() { + for _, rt := range runtimes { + rt.Close() + } + }) + require.Contains(t, shardGroups, uint64(dedicatedTSORaftGroupID)) + require.Contains(t, shardGroups, uint64(1)) + require.DirExists(t, filepath.Join(legacyDir, "fsm.db")) + require.DirExists(t, filepath.Join(legacyDir, "group-0", "fsm.db")) + if _, err := os.Stat(filepath.Join(legacyDir, "group-1")); !os.IsNotExist(err) { + require.NoError(t, err) + t.Fatalf("group 1 should stay in legacy dir, but group-1 dir exists") + } + got, err := os.ReadFile(legacyMarker) + require.NoError(t, err) + require.Equal(t, []byte("keep"), got) +} + func TestParseRaftEngineType(t *testing.T) { t.Run("default", func(t *testing.T) { engineType, err := parseRaftEngineType("") diff --git a/shard_config.go b/shard_config.go index 7fbc50206..771dd9a42 100644 --- a/shard_config.go +++ b/shard_config.go @@ -23,7 +23,10 @@ type rangeSpec struct { groupID uint64 } -const splitParts = 2 +const ( + splitParts = 2 + dedicatedTSORaftGroupID = 0 +) var ( ErrAddressRequired = errors.New("address is required") @@ -39,6 +42,7 @@ var ( ErrInvalidSQSFifoPartitionMapEntry = errors.New("invalid sqsFifoPartitionMap entry") ErrInvalidRaftBootstrapMembersEntry = errors.New("invalid raftBootstrapMembers entry") ErrInvalidRaftGroupPeersEntry = errors.New("invalid raftGroupPeers entry") + ErrShardRangeReservedTSOGroup = errors.New("shard range references reserved TSO group") ) // sqsFifoPartitionMaxPartitions caps the per-queue partition count so @@ -479,6 +483,9 @@ func parseRaftGroupPeer(entry, part string) (raftengine.Server, error) { func defaultGroupID(groups []groupSpec) uint64 { min := uint64(0) for _, g := range groups { + if g.id == dedicatedTSORaftGroupID { + continue + } if min == 0 || g.id < min { min = g.id } @@ -489,12 +496,35 @@ func defaultGroupID(groups []groupSpec) uint64 { return min } +func dataGroupIDs(groups []groupSpec) []uint64 { + ids := make([]uint64, 0, len(groups)) + for _, g := range groups { + if g.id == dedicatedTSORaftGroupID { + continue + } + ids = append(ids, g.id) + } + slices.Sort(ids) + return slices.Compact(ids) +} + +func dataGroupsNeedMultiDirs(groups []groupSpec) bool { + return len(dataGroupIDs(groups)) > 1 +} + +func effectiveMultiDataDirs(groups []groupSpec, requested bool) bool { + return requested && dataGroupsNeedMultiDirs(groups) +} + func validateShardRanges(ranges []rangeSpec, groups []groupSpec) error { ids := map[uint64]struct{}{} for _, g := range groups { ids[g.id] = struct{}{} } for _, r := range ranges { + if r.groupID == dedicatedTSORaftGroupID { + return errors.Wrapf(ErrShardRangeReservedTSOGroup, "group %d", r.groupID) + } if _, ok := ids[r.groupID]; !ok { return errors.WithStack(errors.Newf("shard range references unknown group %d", r.groupID)) } diff --git a/shard_config_test.go b/shard_config_test.go index dd89ff5ec..c24ebe0fa 100644 --- a/shard_config_test.go +++ b/shard_config_test.go @@ -305,6 +305,30 @@ func TestValidateSQSFifoPartitionMap(t *testing.T) { }) } +func TestBuildSQSFifoPartitionMapRejectsDedicatedTSOGroup(t *testing.T) { + t.Parallel() + + _, err := buildSQSFifoPartitionMap( + []groupSpec{{id: 0}, {id: 1}}, + "orders.fifo:2=1,0", + ) + require.ErrorIs(t, err, ErrInvalidSQSFifoPartitionMapEntry) + require.Contains(t, err.Error(), "reserved for TSO") + require.Contains(t, err.Error(), "partition 1") +} + +func TestBuildSQSFifoPartitionMapRejectsDedicatedTSOGroupDeterministically(t *testing.T) { + t.Parallel() + + _, err := buildSQSFifoPartitionMap( + []groupSpec{{id: 0}, {id: 1}}, + "zeta.fifo:1=0;alpha.fifo:1=0", + ) + require.ErrorIs(t, err, ErrInvalidSQSFifoPartitionMapEntry) + require.Contains(t, err.Error(), "alpha.fifo") + require.NotContains(t, err.Error(), "zeta.fifo") +} + func TestParseRaftBootstrapMembers(t *testing.T) { t.Run("parses members", func(t *testing.T) { members, err := parseRaftBootstrapMembers("n1=10.0.0.11:50051, n2=10.0.0.12:50051") @@ -368,10 +392,43 @@ func TestParseRaftGroupPeers(t *testing.T) { func TestDefaultGroupID(t *testing.T) { require.Equal(t, uint64(1), defaultGroupID(nil)) require.Equal(t, uint64(2), defaultGroupID([]groupSpec{{id: 3}, {id: 2}})) + require.Equal(t, uint64(2), defaultGroupID([]groupSpec{{id: 0}, {id: 2}})) + require.Equal(t, uint64(1), defaultGroupID([]groupSpec{{id: 0}})) +} + +func TestDataGroupIDsExcludeDedicatedTSOGroup(t *testing.T) { + require.Equal(t, []uint64{1, 2}, dataGroupIDs([]groupSpec{{id: 0}, {id: 2}, {id: 1}})) + require.Empty(t, dataGroupIDs([]groupSpec{{id: 0}})) +} + +func TestParseRuntimeConfigAllowsDedicatedTSOGroupWithoutDataRoute(t *testing.T) { + cfg, err := parseRuntimeConfig( + "127.0.0.1:50051", + "", "", "", "", + "0=127.0.0.1:50050,2=127.0.0.1:50052", + "", "", "", "", "", "", + ) + require.NoError(t, err) + require.Equal(t, uint64(2), cfg.defaultGroup) + route, ok := cfg.engine.GetRoute([]byte("user-key")) + require.True(t, ok) + require.Equal(t, uint64(2), route.GroupID) + require.False(t, cfg.multi) +} + +func TestParseRuntimeConfigMultipleDataGroupsEnableMultiDataDirs(t *testing.T) { + cfg, err := parseRuntimeConfig( + "127.0.0.1:50051", + "", "", "", "", + "0=127.0.0.1:50050,1=127.0.0.1:50051,2=127.0.0.1:50052", + "", "", "", "", "", "", + ) + require.NoError(t, err) + require.True(t, cfg.multi) } func TestValidateShardRanges(t *testing.T) { - groups := []groupSpec{{id: 1}, {id: 2}} + groups := []groupSpec{{id: 0}, {id: 1}, {id: 2}} t.Run("valid", func(t *testing.T) { err := validateShardRanges([]rangeSpec{{groupID: 1}}, groups) @@ -382,4 +439,9 @@ func TestValidateShardRanges(t *testing.T) { err := validateShardRanges([]rangeSpec{{groupID: 3}}, groups) require.Error(t, err) }) + + t.Run("reserved tso group", func(t *testing.T) { + err := validateShardRanges([]rangeSpec{{groupID: 0}}, groups) + require.ErrorIs(t, err, ErrShardRangeReservedTSOGroup) + }) }