diff --git a/pkg/keyspace/keyspace.go b/pkg/keyspace/keyspace.go index 71536d9eb07..4ea04e4677b 100644 --- a/pkg/keyspace/keyspace.go +++ b/pkg/keyspace/keyspace.go @@ -1450,7 +1450,7 @@ func (manager *Manager) PatrolKeyspaceAssignment(startKeyspaceID, endKeyspaceID return err } manager.kgm.Lock() - manager.kgm.groups[endpoint.StringUserKind(defaultKeyspaceGroup.UserKind)].Put(defaultKeyspaceGroup) + manager.kgm.putKeyspaceGroupLocked(defaultKeyspaceGroup) manager.kgm.Unlock() // If all keyspaces in the current batch are assigned, update the next start ID. manager.nextPatrolStartID = nextStartID diff --git a/pkg/keyspace/tso_keyspace_group.go b/pkg/keyspace/tso_keyspace_group.go index 08bec239d85..42d3dbe75dd 100644 --- a/pkg/keyspace/tso_keyspace_group.go +++ b/pkg/keyspace/tso_keyspace_group.go @@ -15,6 +15,7 @@ package keyspace import ( + "bytes" "context" "encoding/json" "slices" @@ -41,6 +42,7 @@ import ( "github.com/tikv/pd/pkg/slice" "github.com/tikv/pd/pkg/storage/endpoint" "github.com/tikv/pd/pkg/storage/kv" + "github.com/tikv/pd/pkg/tso" "github.com/tikv/pd/pkg/utils/etcdutil" "github.com/tikv/pd/pkg/utils/keypath" "github.com/tikv/pd/pkg/utils/logutil" @@ -49,10 +51,10 @@ import ( ) const ( - defaultBalancerPolicy = balancer.PolicyRoundRobin - allocNodesToKeyspaceGroupsInterval = 1 * time.Second - allocNodesTimeout = 1 * time.Second - allocNodesInterval = 10 * time.Millisecond + defaultBalancerPolicy = balancer.PolicyRoundRobin + allocNodesTimeout = 1 * time.Second + allocNodesInterval = 10 * time.Millisecond + keyspaceGroupReconcileRetryInterval = 1 * time.Second // defaultKeyspaceCountSplitThreshold is the keyspace count threshold for auto-splitting // a keyspace group. When a group's keyspace count exceeds this value, a new group will be split automatically. defaultKeyspaceCountSplitThreshold = 40000 @@ -88,6 +90,13 @@ type GroupManager struct { serviceRegistryMap map[string]string // tsoNodesWatcher is the watcher for the registered tso servers. tsoNodesWatcher *etcdutil.LoopWatcher + // reconcileCh coalesces changes that may require keyspace group node allocation. + reconcileCh chan struct{} + reconcileCancel context.CancelFunc + // groupWatcherTerm fences watch events and reconciliation work from older API leader terms. + groupWatcherTerm uint64 + // groupMemberships is a lightweight cache used only to find groups that may need reconciliation. + groupMemberships map[uint32][]endpoint.KeyspaceGroupMember } // NewKeyspaceGroupManager creates a Manager of keyspace group related data. @@ -109,6 +118,7 @@ func NewKeyspaceGroupManager( client: client, nodesBalancer: balancer.GenByPolicy[string](defaultBalancerPolicy), serviceRegistryMap: make(map[string]string), + groupMemberships: make(map[uint32][]endpoint.KeyspaceGroupMember), } // If the etcd client is not nil, start the watch loop for the registered tso servers. @@ -133,33 +143,119 @@ func (m *GroupManager) Bootstrap(ctx context.Context) error { Keyspaces: keyspaces, } - m.Lock() - defer m.Unlock() - // Ignore the error if default keyspace group already exists in the storage (e.g. PD restart/recover). + m.Lock() err := m.saveKeyspaceGroups([]*endpoint.KeyspaceGroup{defaultKeyspaceGroup}, false) + m.Unlock() if err != nil && err != errs.ErrKeyspaceGroupExists { return err } - // Load all the keyspace groups from the storage and add to the respective userKind groups. + if m.client == nil { + if err := m.reloadKeyspaceGroups(); err != nil { + return err + } + } else { + termCtx, term, reconcileCh := m.beginKeyspaceGroupReconcileTerm(ctx) + if err := m.initKeyspaceGroupsWatcher(termCtx, term, reconcileCh); err != nil { + // Keep the previous bootstrap behavior as a fallback if the watcher's initial load fails. + log.Warn("failed to initialize keyspace group watcher, fall back to loading from storage", zap.Error(err)) + if err := m.reloadKeyspaceGroups(); err != nil { + m.cancelKeyspaceGroupReconcileTerm(term) + return err + } + } + + // It will only alloc node when the group manager is on API leader. + m.wg.Add(1) + go m.allocNodesToAllKeyspaceGroups(termCtx, term, reconcileCh) + notifyKeyspaceGroupReconcile(reconcileCh) + m.wg.Add(1) + go m.patrolKeyspaceGroupSizeForAutoSplit(termCtx) + } + return nil +} + +func (m *GroupManager) beginKeyspaceGroupReconcileTerm(leaderCtx context.Context) (context.Context, uint64, chan struct{}) { + termCtx, cancel := context.WithCancel(m.ctx) + stopLeaderCancel := context.AfterFunc(leaderCtx, cancel) + termCancel := func() { + stopLeaderCancel() + cancel() + } + + m.Lock() + if m.reconcileCancel != nil { + m.reconcileCancel() + } + m.groupWatcherTerm++ + term := m.groupWatcherTerm + reconcileCh := make(chan struct{}, 1) + m.reconcileCh = reconcileCh + m.reconcileCancel = termCancel + m.resetKeyspaceGroupCacheLocked() + m.Unlock() + return termCtx, term, reconcileCh +} + +func (m *GroupManager) cancelKeyspaceGroupReconcileTerm(term uint64) { + m.Lock() + defer m.Unlock() + if term != m.groupWatcherTerm { + return + } + if m.reconcileCancel != nil { + m.reconcileCancel() + m.reconcileCancel = nil + } + if m.reconcileCh != nil { + m.reconcileCh = nil + } +} + +func (m *GroupManager) reloadKeyspaceGroups() error { + m.Lock() + defer m.Unlock() groups, err := m.store.LoadKeyspaceGroups(constant.DefaultKeyspaceGroupID, 0) if err != nil { return err } + m.resetKeyspaceGroupCacheLocked() for _, group := range groups { - userKind := endpoint.StringUserKind(group.UserKind) - m.groups[userKind].Put(group) + m.putKeyspaceGroupLocked(group) } + return nil +} - // It will only alloc node when the group manager is on API leader. - if m.client != nil { - m.wg.Add(1) - go m.allocNodesToAllKeyspaceGroups(ctx) - m.wg.Add(1) - go m.patrolKeyspaceGroupSizeForAutoSplit(ctx) +func (m *GroupManager) resetKeyspaceGroupCacheLocked() { + for userKind := range endpoint.UserKindCount { + m.groups[userKind] = newIndexedHeap(int(mcs.MaxKeyspaceGroupCountInUse)) } - return nil + m.groupMemberships = make(map[uint32][]endpoint.KeyspaceGroupMember) +} + +func (m *GroupManager) putKeyspaceGroupLocked(group *endpoint.KeyspaceGroup) { + previousMembers, exists := m.groupMemberships[group.ID] + membershipChanged := !exists || !slices.Equal(previousMembers, group.Members) + for _, groupsByKind := range m.groups { + groupsByKind.Remove(group.ID) + } + userKind := endpoint.StringUserKind(group.UserKind) + m.groups[userKind].Put(group) + m.groupMemberships[group.ID] = slices.Clone(group.Members) + if membershipChanged && m.groupWatcherTerm != 0 && m.reconcileCh != nil && + m.keyspaceGroupMembersNeedReconcile(group.Members) { + notifyKeyspaceGroupReconcile(m.reconcileCh) + } +} + +// removeKeyspaceGroupLocked removes a group from both caches. +// The caller must hold m.Lock. +func (m *GroupManager) removeKeyspaceGroupLocked(groupID uint32) { + for _, groupsByKind := range m.groups { + groupsByKind.Remove(groupID) + } + delete(m.groupMemberships, groupID) } // Close closes the manager. @@ -168,14 +264,17 @@ func (m *GroupManager) Close() { m.wg.Wait() } -func (m *GroupManager) allocNodesToAllKeyspaceGroups(ctx context.Context) { +func (m *GroupManager) allocNodesToAllKeyspaceGroups(ctx context.Context, term uint64, reconcileCh <-chan struct{}) { defer logutil.LogPanic() defer m.wg.Done() - ticker := time.NewTicker(allocNodesToKeyspaceGroupsInterval) - failpoint.Inject("acceleratedAllocNodes", func() { - ticker.Reset(time.Millisecond * 100) - }) - defer ticker.Stop() + defer m.cancelKeyspaceGroupReconcileTerm(term) + var retryCh <-chan time.Time + updateRetry := func(needRetry bool) { + retryCh = nil + if needRetry { + retryCh = time.After(keyspaceGroupReconcileRetryInterval) + } + } log.Info("start to alloc nodes to all keyspace groups") for { select { @@ -190,42 +289,87 @@ func (m *GroupManager) allocNodesToAllKeyspaceGroups(ctx context.Context) { // When the API leader is changed, we should stop to alloc nodes to all keyspace groups. log.Info("the raftcluster is closed, stop to alloc nodes to all keyspace groups") return - case <-ticker.C: - if m.GetNodesCount() == 0 { - continue - } + case <-reconcileCh: + updateRetry(m.reconcileKeyspaceGroupsForTerm(ctx, term)) + case <-retryCh: + updateRetry(m.reconcileKeyspaceGroupsForTerm(ctx, term)) } - groups, err := m.store.LoadKeyspaceGroups(constant.DefaultKeyspaceGroupID, 0) - if err != nil { - log.Error("failed to load all keyspace groups", zap.Error(err)) - continue + } +} + +func (m *GroupManager) notifyKeyspaceGroupReconcile() { + m.RLock() + reconcileCh := m.reconcileCh + m.RUnlock() + notifyKeyspaceGroupReconcile(reconcileCh) +} + +func notifyKeyspaceGroupReconcile(reconcileCh chan<- struct{}) { + select { + case reconcileCh <- struct{}{}: + default: + } +} + +// reconcileKeyspaceGroupsForTerm checks the cached group memberships and allocates nodes only +// for groups that may be under-replicated. It returns whether a transient error needs a retry. +func (m *GroupManager) reconcileKeyspaceGroupsForTerm(ctx context.Context, term uint64) (needRetry bool) { + if !m.isKeyspaceGroupReconcileTermActive(ctx, term) { + return false + } + if m.GetNodesCount() == 0 { + return false + } + + type groupMembership struct { + id uint32 + members []endpoint.KeyspaceGroupMember + } + m.RLock() + groups := make([]groupMembership, 0, len(m.groupMemberships)) + for groupID, members := range m.groupMemberships { + groups = append(groups, groupMembership{id: groupID, members: slices.Clone(members)}) + } + m.RUnlock() + + for _, group := range groups { + if !m.isKeyspaceGroupReconcileTermActive(ctx, term) { + return false } - // if the default keyspace is not initialized, we should wait for the default keyspace to be initialized. - if len(groups) == 0 { + existMembers := make(map[string]struct{}, len(group.members)) + for _, member := range group.members { + if exist, addr := m.IsExistNode(member.Address); exist { + existMembers[addr] = struct{}{} + } + } + numExistMembers := len(existMembers) + if numExistMembers != 0 && numExistMembers == len(group.members) && numExistMembers == m.GetNodesCount() { continue } - for _, group := range groups { - existMembers := make(map[string]struct{}) - for _, member := range group.Members { - if exist, addr := m.IsExistNode(member.Address); exist { - existMembers[addr] = struct{}{} - } - } - numExistMembers := len(existMembers) - if numExistMembers != 0 && numExistMembers == len(group.Members) && numExistMembers == m.GetNodesCount() { + if numExistMembers < mcs.DefaultKeyspaceGroupReplicaCount { + nodes, err := m.reconcileNodesForKeyspaceGroup(ctx, group.id) + if err != nil { + log.Warn("failed to alloc nodes for keyspace group", zap.Uint32("keyspace-group-id", group.id), zap.Error(err)) + needRetry = true continue } - if numExistMembers < mcs.DefaultKeyspaceGroupReplicaCount { - nodes, err := m.AllocNodesForKeyspaceGroup(group.ID, existMembers, mcs.DefaultKeyspaceGroupReplicaCount) - if err != nil { - log.Error("failed to alloc nodes for keyspace group", zap.Uint32("keyspace-group-id", group.ID), zap.Error(err)) - continue - } - log.Info("alloc nodes for keyspace group", zap.Uint32("keyspace-group-id", group.ID), zap.Any("nodes", nodes)) - group.Members = nodes + if nodes != nil { + log.Info("alloc nodes for keyspace group", zap.Uint32("keyspace-group-id", group.id), zap.Any("nodes", nodes)) } } } + return needRetry +} + +func (m *GroupManager) isKeyspaceGroupReconcileTermActive(ctx context.Context, term uint64) bool { + select { + case <-ctx.Done(): + return false + default: + } + m.RLock() + defer m.RUnlock() + return term == m.groupWatcherTerm && m.reconcileCh != nil } // patrolKeyspaceGroupSizeForAutoSplit periodically checks all tso keyspace groups. @@ -386,9 +530,185 @@ func (m *GroupManager) initTSONodesWatcher(client *clientv3.Client) { func([]*clientv3.Event) error { return nil }, putFn, deleteFn, - func([]*clientv3.Event) error { return nil }, + func([]*clientv3.Event) error { + m.notifyKeyspaceGroupReconcile() + return nil + }, + true, /* withPrefix */ + ) +} + +func (m *GroupManager) initKeyspaceGroupsWatcher(ctx context.Context, term uint64, reconcileCh chan<- struct{}) error { + initialLoad := true + initialLoadDone := make(chan struct{}) + needReconcile := false + preEventsFn := func(events []*clientv3.Event) error { + // LoopWatcher retries a failed load, and calls postEventsFn after every attempt. + // Replace the full cache on each initial load attempt so a successful retry + // cannot retain entries from a previous partial attempt. + if len(events) != 0 { + return nil + } + select { + case <-initialLoadDone: + return nil + default: + } + m.Lock() + if term == m.groupWatcherTerm { + m.resetKeyspaceGroupCacheLocked() + } + m.Unlock() + initialLoad = true + return nil + } + putFn := func(kv *mvccpb.KeyValue) error { + members, changed, err := m.applyKeyspaceGroupMembership(term, kv, initialLoad) + if err == nil && !initialLoad && changed && m.keyspaceGroupMembersNeedReconcile(members) { + needReconcile = true + } + return err + } + deleteFn := func(kv *mvccpb.KeyValue) error { + return m.applyKeyspaceGroupMembershipDelete(term, kv) + } + postEventsFn := func([]*clientv3.Event) error { + if initialLoad { + initialLoad = false + needReconcile = true + } + if needReconcile { + notifyKeyspaceGroupReconcile(reconcileCh) + needReconcile = false + } + return nil + } + watcher := etcdutil.NewLoopWatcher( + ctx, + &m.wg, + m.client, + "keyspace-group-watcher", + keypath.KeyspaceGroupIDPrefix(), + preEventsFn, + putFn, + deleteFn, + postEventsFn, true, /* withPrefix */ ) + watcher.StartWatchLoop() + err := watcher.WaitLoad() + close(initialLoadDone) + return err +} + +func (m *GroupManager) applyKeyspaceGroupMembership( + term uint64, kv *mvccpb.KeyValue, updateGroupCache bool, +) ([]endpoint.KeyspaceGroupMember, bool, error) { + var ( + groupID uint32 + members []endpoint.KeyspaceGroupMember + group *endpoint.KeyspaceGroup + ) + if updateGroupCache { + group = &endpoint.KeyspaceGroup{} + if err := json.Unmarshal(kv.Value, group); err != nil { + return nil, false, err + } + groupID, members = group.ID, group.Members + } else { + var err error + groupID, members, err = decodeKeyspaceGroupMembership(kv.Value) + if err != nil { + return nil, false, err + } + } + m.Lock() + defer m.Unlock() + if term != m.groupWatcherTerm { + return nil, false, nil + } + previous, exists := m.groupMemberships[groupID] + changed := !exists || !slices.Equal(previous, members) + if updateGroupCache { + m.putKeyspaceGroupLocked(group) + } else if changed { + m.groupMemberships[groupID] = slices.Clone(members) + } + return members, changed, nil +} + +// decodeKeyspaceGroupMembership stops after ID and Members. With the persisted field order, +// it does not scan or allocate the potentially million-entry Keyspaces array. +func decodeKeyspaceGroupMembership(value []byte) (uint32, []endpoint.KeyspaceGroupMember, error) { + decoder := json.NewDecoder(bytes.NewReader(value)) + if _, err := decoder.Token(); err != nil { + return 0, nil, err + } + var ( + groupID uint32 + members []endpoint.KeyspaceGroupMember + gotID bool + gotMembers bool + ) + for decoder.More() { + token, err := decoder.Token() + if err != nil { + return 0, nil, err + } + field, ok := token.(string) + if !ok { + return 0, nil, errors.New("invalid keyspace group field") + } + switch field { + case "id": + err = decoder.Decode(&groupID) + gotID = err == nil + case "members": + err = decoder.Decode(&members) + gotMembers = err == nil + default: + var ignored json.RawMessage + err = decoder.Decode(&ignored) + } + if err != nil { + return 0, nil, err + } + if gotID && gotMembers { + return groupID, members, nil + } + } + return 0, nil, errors.New("keyspace group membership is incomplete") +} + +func (m *GroupManager) applyKeyspaceGroupMembershipDelete(term uint64, kv *mvccpb.KeyValue) error { + groupID, err := tso.ExtractKeyspaceGroupIDFromPath(keypath.GetCompiledKeyspaceGroupIDRegexp(), string(kv.Key)) + if err != nil { + return err + } + m.Lock() + defer m.Unlock() + if term != m.groupWatcherTerm { + return nil + } + delete(m.groupMemberships, groupID) + return nil +} + +func (m *GroupManager) keyspaceGroupMembersNeedReconcile(members []endpoint.KeyspaceGroupMember) bool { + if m.GetNodesCount() == 0 { + return false + } + existMembers := make(map[string]struct{}, len(members)) + for _, member := range members { + if exist, addr := m.IsExistNode(member.Address); exist { + existMembers[addr] = struct{}{} + } + } + numExistMembers := len(existMembers) + if numExistMembers != 0 && numExistMembers == len(members) && numExistMembers == m.GetNodesCount() { + return false + } + return numExistMembers < mcs.DefaultKeyspaceGroupReplicaCount } // CreateKeyspaceGroups creates keyspace groups. @@ -400,8 +720,7 @@ func (m *GroupManager) CreateKeyspaceGroups(keyspaceGroups []*endpoint.KeyspaceG } for _, keyspaceGroup := range keyspaceGroups { - userKind := endpoint.StringUserKind(keyspaceGroup.UserKind) - m.groups[userKind].Put(keyspaceGroup) + m.putKeyspaceGroupLocked(keyspaceGroup) } return nil @@ -465,10 +784,9 @@ func (m *GroupManager) DeleteKeyspaceGroupByID(id uint32) (*endpoint.KeyspaceGro return nil, err } - userKind := endpoint.StringUserKind(kg.UserKind) // TODO: move out the keyspace to another group // we don't need the keyspace group as the return value - m.groups[userKind].Remove(id) + m.removeKeyspaceGroupLocked(id) return kg, nil } @@ -625,9 +943,8 @@ func (m *GroupManager) RemoveKeyspacesFromGroup(groupID uint32, km *Manager, key return nil, err } - // Update the cache - userKind := endpoint.StringUserKind(kg.UserKind) - m.groups[userKind].Put(kg) + // Update the cache. + m.putKeyspaceGroupLocked(kg) return kg, nil } @@ -687,7 +1004,7 @@ func (m *GroupManager) updateKeyspaceForGroupLocked(userKind endpoint.UserKind, if err := m.saveKeyspaceGroups([]*endpoint.KeyspaceGroup{kg}, true); err != nil { return err } - m.groups[userKind].Put(kg) + m.putKeyspaceGroupLocked(kg) } return nil } @@ -751,11 +1068,11 @@ func (m *GroupManager) UpdateKeyspaceGroup(oldGroupID, newGroupID string, oldUse } if updateOld { - m.groups[oldUserKind].Put(oldKG) + m.putKeyspaceGroupLocked(oldKG) } if updateNew { - m.groups[newUserKind].Put(newKG) + m.putKeyspaceGroupLocked(newKG) } return nil @@ -840,8 +1157,8 @@ func (m *GroupManager) SplitKeyspaceGroupByID( return err } // Update the keyspace group cache. - m.groups[endpoint.StringUserKind(splitSourceKg.UserKind)].Put(splitSourceKg) - m.groups[endpoint.StringUserKind(splitTargetKg.UserKind)].Put(splitTargetKg) + m.putKeyspaceGroupLocked(splitSourceKg) + m.putKeyspaceGroupLocked(splitTargetKg) return nil } @@ -967,8 +1284,8 @@ func (m *GroupManager) FinishSplitKeyspaceByID(splitTargetID uint32) error { return err } // Update the keyspace group cache. - m.groups[endpoint.StringUserKind(splitTargetKg.UserKind)].Put(splitTargetKg) - m.groups[endpoint.StringUserKind(splitSourceKg.UserKind)].Put(splitSourceKg) + m.putKeyspaceGroupLocked(splitTargetKg) + m.putKeyspaceGroupLocked(splitSourceKg) log.Info("finish split keyspace group", zap.Uint32("split-source-id", splitSourceKg.ID), zap.Uint32("split-target-id", splitTargetID)) return nil } @@ -983,16 +1300,33 @@ func (m *GroupManager) GetNodesCount() int { // AllocNodesForKeyspaceGroup allocates nodes for the keyspace group. func (m *GroupManager) AllocNodesForKeyspaceGroup(id uint32, existMembers map[string]struct{}, desiredReplicaCount int) ([]endpoint.KeyspaceGroupMember, error) { + return m.allocNodesForKeyspaceGroupWithOptions(m.ctx, id, existMembers, desiredReplicaCount, false) +} + +// reconcileNodesForKeyspaceGroup revalidates the latest group membership in storage before +// allocating nodes. The cache is only a hint and must never be used as the source of truth for writes. +func (m *GroupManager) reconcileNodesForKeyspaceGroup(ctx context.Context, id uint32) ([]endpoint.KeyspaceGroupMember, error) { + return m.allocNodesForKeyspaceGroupWithOptions(ctx, id, nil, mcs.DefaultKeyspaceGroupReplicaCount, true) +} + +func (m *GroupManager) allocNodesForKeyspaceGroupWithOptions( + operationCtx context.Context, + id uint32, + existMembers map[string]struct{}, + desiredReplicaCount int, + revalidateMembers bool, +) ([]endpoint.KeyspaceGroupMember, error) { m.Lock() defer m.Unlock() - ctx, cancel := context.WithTimeout(m.ctx, allocNodesTimeout) + ctx, cancel := context.WithTimeout(operationCtx, allocNodesTimeout) defer cancel() ticker := time.NewTicker(allocNodesInterval) defer ticker.Stop() var kg *endpoint.KeyspaceGroup nodes := make([]endpoint.KeyspaceGroupMember, 0, desiredReplicaCount) - err := m.store.RunInTxn(m.ctx, func(txn kv.Txn) error { + changed := false + err := m.store.RunInTxn(operationCtx, func(txn kv.Txn) error { var err error kg, err = m.store.LoadKeyspaceGroup(txn, id) if err != nil { @@ -1007,6 +1341,19 @@ func (m *GroupManager) AllocNodesForKeyspaceGroup(id uint32, existMembers map[st if kg.IsMerging() { return errs.ErrKeyspaceGroupInMerging.FastGenByArgs(id) } + if revalidateMembers { + existMembers = make(map[string]struct{}, len(kg.Members)) + for _, member := range kg.Members { + if exist, addr := m.IsExistNode(member.Address); exist { + existMembers[addr] = struct{}{} + } + } + numExistMembers := len(existMembers) + if (numExistMembers != 0 && numExistMembers == len(kg.Members) && numExistMembers == m.GetNodesCount()) || + numExistMembers >= desiredReplicaCount { + return nil + } + } for addr := range existMembers { nodes = append(nodes, endpoint.KeyspaceGroupMember{ @@ -1041,15 +1388,21 @@ func (m *GroupManager) AllocNodesForKeyspaceGroup(id uint32, existMembers map[st }) } kg.Members = nodes + changed = true return m.store.SaveKeyspaceGroup(txn, kg) }) if err != nil { return nil, err } - m.groups[endpoint.StringUserKind(kg.UserKind)].Put(kg) - log.Info("alloc nodes for keyspace group", - zap.Uint32("keyspace-group-id", id), - zap.Reflect("nodes", nodes)) + if !changed { + return nil, nil + } + m.putKeyspaceGroupLocked(kg) + if !revalidateMembers { + log.Info("alloc nodes for keyspace group", + zap.Uint32("keyspace-group-id", id), + zap.Reflect("nodes", nodes)) + } return nodes, nil } @@ -1086,7 +1439,7 @@ func (m *GroupManager) SetNodesForKeyspaceGroup(id uint32, nodes []string) error if err != nil { return err } - m.groups[endpoint.StringUserKind(kg.UserKind)].Put(kg) + m.putKeyspaceGroupLocked(kg) return nil } @@ -1128,7 +1481,7 @@ func (m *GroupManager) SetPriorityForKeyspaceGroup(id uint32, node string, prior if err != nil { return err } - m.groups[endpoint.StringUserKind(kg.UserKind)].Put(kg) + m.putKeyspaceGroupLocked(kg) log.Info("set priority for keyspace group", zap.Uint32("keyspace-group-id", id), zap.String("node", node), @@ -1228,10 +1581,9 @@ func (m *GroupManager) MergeKeyspaceGroups(mergeTargetID uint32, mergeList []uin return err } // Update the keyspace group cache. - m.groups[endpoint.StringUserKind(mergeTargetKg.UserKind)].Put(mergeTargetKg) + m.putKeyspaceGroupLocked(mergeTargetKg) for _, kgID := range mergeList { - kg := groups[kgID] - m.groups[endpoint.StringUserKind(kg.UserKind)].Remove(kgID) + m.removeKeyspaceGroupLocked(kgID) } return nil } @@ -1274,7 +1626,7 @@ func (m *GroupManager) FinishMergeKeyspaceByID(mergeTargetID uint32) error { return err } // Update the keyspace group cache. - m.groups[endpoint.StringUserKind(mergeTargetKg.UserKind)].Put(mergeTargetKg) + m.putKeyspaceGroupLocked(mergeTargetKg) log.Info("finish merge keyspace group", zap.Uint32("merge-target-id", mergeTargetKg.ID), zap.Reflect("merge-list", mergeList)) @@ -1308,11 +1660,11 @@ func (m *GroupManager) MergeAllIntoDefaultKeyspaceGroup() error { groupsToMerge = make([]uint32, 0, maxBatchSize) ) for idx, group := range groups.GetAll() { - if group.ID == constant.DefaultKeyspaceGroupID { - continue + if group.ID != constant.DefaultKeyspaceGroupID { + groupsToMerge = append(groupsToMerge, group.ID) } - groupsToMerge = append(groupsToMerge, group.ID) - if len(groupsToMerge) < maxBatchSize && idx < mergeNum-1 { + if len(groupsToMerge) == 0 || + (len(groupsToMerge) < maxBatchSize && idx < mergeNum-1) { continue } log.Info("merge keyspace groups into the default one", diff --git a/pkg/keyspace/tso_keyspace_group_test.go b/pkg/keyspace/tso_keyspace_group_test.go index 489dbbf7efd..858a2fa59c5 100644 --- a/pkg/keyspace/tso_keyspace_group_test.go +++ b/pkg/keyspace/tso_keyspace_group_test.go @@ -16,16 +16,21 @@ package keyspace import ( "context" + "encoding/json" "errors" "fmt" + "runtime" + "sync/atomic" "testing" "time" "github.com/stretchr/testify/require" "github.com/stretchr/testify/suite" + "go.etcd.io/etcd/api/v3/mvccpb" "github.com/tikv/pd/pkg/errs" "github.com/tikv/pd/pkg/keyspace/constant" + "github.com/tikv/pd/pkg/mcs/discovery" mcs "github.com/tikv/pd/pkg/mcs/utils/constant" "github.com/tikv/pd/pkg/mock/mockcluster" "github.com/tikv/pd/pkg/mock/mockconfig" @@ -33,16 +38,59 @@ import ( "github.com/tikv/pd/pkg/storage/endpoint" "github.com/tikv/pd/pkg/storage/kv" "github.com/tikv/pd/pkg/utils/etcdutil" + "github.com/tikv/pd/pkg/utils/keypath" + "github.com/tikv/pd/pkg/utils/testutil" "github.com/tikv/pd/pkg/versioninfo/kerneltype" ) var errSaveKeyspaceGroup = errors.New("save keyspace group error") +var errLoadKeyspaceGroup = errors.New("load keyspace group error") type errorKeyspaceGroupStorage struct { *endpoint.StorageEndpoint failOnSaveID uint32 } +type countingKeyspaceGroupStorage struct { + endpoint.KeyspaceGroupStorage + loadGroupsCount atomic.Int64 + loadGroupCount atomic.Int64 + saveGroupCount atomic.Int64 +} + +type flakyKeyspaceGroupStorage struct { + endpoint.KeyspaceGroupStorage + loadFailures atomic.Int64 + loadFailureCh chan struct{} +} + +func (s *flakyKeyspaceGroupStorage) LoadKeyspaceGroup(txn kv.Txn, id uint32) (*endpoint.KeyspaceGroup, error) { + if s.loadFailures.Load() > 0 { + s.loadFailures.Add(-1) + select { + case s.loadFailureCh <- struct{}{}: + default: + } + return nil, errLoadKeyspaceGroup + } + return s.KeyspaceGroupStorage.LoadKeyspaceGroup(txn, id) +} + +func (s *countingKeyspaceGroupStorage) LoadKeyspaceGroups(startID uint32, limit int) ([]*endpoint.KeyspaceGroup, error) { + s.loadGroupsCount.Add(1) + return s.KeyspaceGroupStorage.LoadKeyspaceGroups(startID, limit) +} + +func (s *countingKeyspaceGroupStorage) LoadKeyspaceGroup(txn kv.Txn, id uint32) (*endpoint.KeyspaceGroup, error) { + s.loadGroupCount.Add(1) + return s.KeyspaceGroupStorage.LoadKeyspaceGroup(txn, id) +} + +func (s *countingKeyspaceGroupStorage) SaveKeyspaceGroup(txn kv.Txn, kg *endpoint.KeyspaceGroup) error { + s.saveGroupCount.Add(1) + return s.KeyspaceGroupStorage.SaveKeyspaceGroup(txn, kg) +} + func (s *errorKeyspaceGroupStorage) SaveKeyspaceGroup(txn kv.Txn, kg *endpoint.KeyspaceGroup) error { if s.failOnSaveID != 0 && kg.ID == s.failOnSaveID { return errSaveKeyspaceGroup @@ -50,6 +98,448 @@ func (s *errorKeyspaceGroupStorage) SaveKeyspaceGroup(txn kv.Txn, kg *endpoint.K return s.StorageEndpoint.SaveKeyspaceGroup(txn, kg) } +func TestReconcileKeyspaceGroupsDoesNotLoadHealthyGroups(t *testing.T) { + re := require.New(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + store := &countingKeyspaceGroupStorage{ + KeyspaceGroupStorage: endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil), + } + manager := NewKeyspaceGroupManager(ctx, store, nil) + defer manager.Close() + manager.nodesBalancer.Put("http://tso-1") + manager.nodesBalancer.Put("http://tso-2") + termCtx, term, _ := manager.beginKeyspaceGroupReconcileTerm(ctx) + manager.Lock() + manager.putKeyspaceGroupLocked(&endpoint.KeyspaceGroup{ + ID: constant.DefaultKeyspaceGroupID, + UserKind: endpoint.Basic.String(), + Members: []endpoint.KeyspaceGroupMember{ + {Address: "http://tso-1"}, + {Address: "http://tso-2"}, + }, + Keyspaces: make([]uint32, 1_000_000), + }) + manager.Unlock() + + var before, after runtime.MemStats + runtime.GC() + runtime.ReadMemStats(&before) + for range 10 { + re.False(manager.reconcileKeyspaceGroupsForTerm(termCtx, term)) + } + runtime.ReadMemStats(&after) + re.Zero(store.loadGroupsCount.Load()) + re.Zero(store.loadGroupCount.Load()) + re.Less((after.TotalAlloc-before.TotalAlloc)/10, uint64(4096)) +} + +func TestReconcileKeyspaceGroupsRevalidatesGroupBeforeWriting(t *testing.T) { + re := require.New(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + baseStore := endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil) + store := &countingKeyspaceGroupStorage{KeyspaceGroupStorage: baseStore} + storedGroup := &endpoint.KeyspaceGroup{ + ID: constant.DefaultKeyspaceGroupID, + UserKind: endpoint.Basic.String(), + Members: []endpoint.KeyspaceGroupMember{ + {Address: "http://tso-1", Priority: 100}, + {Address: "http://tso-2", Priority: 200}, + }, + Keyspaces: []uint32{constant.DefaultKeyspaceID}, + } + re.NoError(store.RunInTxn(ctx, func(txn kv.Txn) error { + return store.SaveKeyspaceGroup(txn, storedGroup) + })) + store.saveGroupCount.Store(0) + + manager := NewKeyspaceGroupManager(ctx, store, nil) + defer manager.Close() + manager.nodesBalancer.Put("http://tso-1") + manager.nodesBalancer.Put("http://tso-2") + termCtx, term, _ := manager.beginKeyspaceGroupReconcileTerm(ctx) + manager.Lock() + manager.putKeyspaceGroupLocked(&endpoint.KeyspaceGroup{ + ID: constant.DefaultKeyspaceGroupID, + UserKind: endpoint.Basic.String(), + Members: []endpoint.KeyspaceGroupMember{ + {Address: "http://tso-1", Priority: 100}, + }, + Keyspaces: []uint32{constant.DefaultKeyspaceID}, + }) + manager.Unlock() + + re.False(manager.reconcileKeyspaceGroupsForTerm(termCtx, term)) + re.Zero(store.saveGroupCount.Load()) + var got *endpoint.KeyspaceGroup + re.NoError(baseStore.RunInTxn(ctx, func(txn kv.Txn) error { + var err error + got, err = baseStore.LoadKeyspaceGroup(txn, constant.DefaultKeyspaceGroupID) + return err + })) + re.Equal(storedGroup.Members, got.Members) +} + +func TestKeyspaceGroupWatcherDoesNotOverwriteLocalCommit(t *testing.T) { + re := require.New(t) + manager := NewKeyspaceGroupManager(context.Background(), endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil), nil) + defer manager.Close() + + oldGroup := &endpoint.KeyspaceGroup{ + ID: 1, + UserKind: endpoint.Standard.String(), + Members: []endpoint.KeyspaceGroupMember{{Address: "http://tso-1"}}, + Keyspaces: []uint32{1}, + } + newGroup := &endpoint.KeyspaceGroup{ + ID: 1, + UserKind: endpoint.Standard.String(), + Members: []endpoint.KeyspaceGroupMember{{Address: "http://tso-1"}, {Address: "http://tso-2"}}, + Keyspaces: []uint32{1, 2}, + } + manager.Lock() + manager.groupWatcherTerm = 1 + manager.putKeyspaceGroupLocked(newGroup) + manager.Unlock() + + oldValue, err := json.Marshal(oldGroup) + re.NoError(err) + _, _, err = manager.applyKeyspaceGroupMembership(1, &mvccpb.KeyValue{Value: oldValue}, false) + re.NoError(err) + re.Equal(newGroup.Keyspaces, manager.groups[endpoint.Standard].Get(1).Keyspaces) + re.Equal(oldGroup.Members, manager.groupMemberships[1]) + + newValue, err := json.Marshal(newGroup) + re.NoError(err) + _, _, err = manager.applyKeyspaceGroupMembership(1, &mvccpb.KeyValue{Value: newValue}, false) + re.NoError(err) + re.Equal(newGroup.Keyspaces, manager.groups[endpoint.Standard].Get(1).Keyspaces) + re.Equal(newGroup.Members, manager.groupMemberships[1]) + + manager.Lock() + manager.groupWatcherTerm = 2 + manager.Unlock() + _, _, err = manager.applyKeyspaceGroupMembership(1, &mvccpb.KeyValue{Value: oldValue}, false) + re.NoError(err) + re.Equal(newGroup.Keyspaces, manager.groups[endpoint.Standard].Get(1).Keyspaces) + re.Equal(newGroup.Members, manager.groupMemberships[1]) + + re.NoError(manager.applyKeyspaceGroupMembershipDelete(2, &mvccpb.KeyValue{ + Key: []byte(keypath.KeyspaceGroupIDPath(1)), + })) + re.NotContains(manager.groupMemberships, uint32(1)) + re.Equal(newGroup.Keyspaces, manager.groups[endpoint.Standard].Get(1).Keyspaces) +} + +func TestKeyspaceGroupReconcileTermFencesPreviousLeader(t *testing.T) { + re := require.New(t) + manager := NewKeyspaceGroupManager(context.Background(), endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil), nil) + defer manager.Close() + + oldLeaderCtx, cancelOldLeader := context.WithCancel(context.Background()) + defer cancelOldLeader() + oldTermCtx, oldTerm, oldReconcileCh := manager.beginKeyspaceGroupReconcileTerm(oldLeaderCtx) + newLeaderCtx, cancelNewLeader := context.WithCancel(context.Background()) + defer cancelNewLeader() + _, newTerm, newReconcileCh := manager.beginKeyspaceGroupReconcileTerm(newLeaderCtx) + + re.Greater(newTerm, oldTerm) + select { + case <-oldTermCtx.Done(): + case <-time.After(time.Second): + t.Fatal("the previous leader term was not canceled") + } + manager.notifyKeyspaceGroupReconcile() + select { + case <-newReconcileCh: + case <-time.After(time.Second): + t.Fatal("reconcile notification was not sent to the active leader term") + } + select { + case <-oldReconcileCh: + t.Fatal("the previous leader term received a new reconcile notification") + default: + } +} + +func TestCreateKeyspaceGroupTriggersNodeAllocation(t *testing.T) { + re := require.New(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + store := endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil) + manager := NewKeyspaceGroupManager(ctx, store, nil) + defer manager.Close() + manager.nodesBalancer.Put("http://tso-1") + manager.nodesBalancer.Put("http://tso-2") + termCtx, term, reconcileCh := manager.beginKeyspaceGroupReconcileTerm(ctx) + manager.wg.Add(1) + go manager.allocNodesToAllKeyspaceGroups(termCtx, term, reconcileCh) + + re.NoError(manager.CreateKeyspaceGroups([]*endpoint.KeyspaceGroup{{ + ID: 1, + UserKind: endpoint.Standard.String(), + }})) + testutil.Eventually(re, func() bool { + group, err := manager.GetKeyspaceGroupByID(1) + return err == nil && group != nil && len(group.Members) == mcs.DefaultKeyspaceGroupReplicaCount + }, testutil.WithWaitFor(500*time.Millisecond), testutil.WithTickInterval(10*time.Millisecond)) +} + +func TestKeyspaceGroupWatcherTriggersNodeAllocation(t *testing.T) { + re := require.New(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + _, client, clean := etcdutil.NewTestEtcdCluster(t, 1, nil) + defer clean() + + store := endpoint.NewStorageEndpoint(kv.NewEtcdKVBase(client), nil) + manager := NewKeyspaceGroupManager(ctx, store, client) + defer manager.Close() + manager.nodesBalancer.Put("http://tso-1") + manager.nodesBalancer.Put("http://tso-2") + re.NoError(manager.Bootstrap(ctx)) + + re.NoError(store.RunInTxn(ctx, func(txn kv.Txn) error { + return store.SaveKeyspaceGroup(txn, &endpoint.KeyspaceGroup{ + ID: 1, + UserKind: endpoint.Standard.String(), + Keyspaces: []uint32{1}, + }) + })) + + testutil.Eventually(re, func() bool { + group, err := manager.GetKeyspaceGroupByID(1) + return err == nil && group != nil && len(group.Members) == mcs.DefaultKeyspaceGroupReplicaCount + }, testutil.WithWaitFor(time.Second), testutil.WithTickInterval(10*time.Millisecond)) +} + +func TestKeyspaceGroupWatcherRetriesInitialFullLoad(t *testing.T) { + re := require.New(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + _, client, clean := etcdutil.NewTestEtcdCluster(t, 1, nil) + defer clean() + + staleGroup := &endpoint.KeyspaceGroup{ + ID: 1, + UserKind: endpoint.Standard.String(), + Members: []endpoint.KeyspaceGroupMember{{Address: "http://tso-1"}}, + Keyspaces: []uint32{1}, + } + staleValue, err := json.Marshal(staleGroup) + re.NoError(err) + _, err = client.Put(ctx, keypath.KeyspaceGroupIDPath(staleGroup.ID), string(staleValue)) + re.NoError(err) + + groupKey := keypath.KeyspaceGroupIDPath(2) + _, err = client.Put(ctx, groupKey, "{") + re.NoError(err) + + store := endpoint.NewStorageEndpoint(kv.NewEtcdKVBase(client), nil) + manager := NewKeyspaceGroupManager(ctx, store, client) + defer manager.Close() + termCtx, term, reconcileCh := manager.beginKeyspaceGroupReconcileTerm(ctx) + loadResult := make(chan error, 1) + go func() { + loadResult <- manager.initKeyspaceGroupsWatcher(termCtx, term, reconcileCh) + }() + + // The first load sees invalid JSON and reaches its post-load callback before retrying. + select { + case <-reconcileCh: + case <-time.After(time.Second): + t.Fatal("the first keyspace group load did not finish") + } + _, err = client.Delete(ctx, keypath.KeyspaceGroupIDPath(staleGroup.ID)) + re.NoError(err) + group := &endpoint.KeyspaceGroup{ + ID: 2, + UserKind: endpoint.Standard.String(), + Members: []endpoint.KeyspaceGroupMember{{Address: "http://tso-1"}}, + Keyspaces: []uint32{1, 2}, + } + value, err := json.Marshal(group) + re.NoError(err) + _, err = client.Put(ctx, groupKey, string(value)) + re.NoError(err) + + select { + case err := <-loadResult: + re.NoError(err) + case <-time.After(3 * time.Second): + t.Fatal("the keyspace group watcher did not finish retrying its initial load") + } + manager.RLock() + loaded := manager.groups[endpoint.Standard].Get(group.ID) + stale := manager.groups[endpoint.Standard].Get(staleGroup.ID) + manager.RUnlock() + re.Equal(group, loaded) + re.Nil(stale) +} + +func TestAllocNodesToAllKeyspaceGroupsRetriesTransientFailure(t *testing.T) { + re := require.New(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + store := &flakyKeyspaceGroupStorage{ + KeyspaceGroupStorage: endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil), + loadFailureCh: make(chan struct{}, 1), + } + manager := NewKeyspaceGroupManager(ctx, store, nil) + defer manager.Close() + manager.nodesBalancer.Put("http://tso-1") + manager.nodesBalancer.Put("http://tso-2") + termCtx, term, reconcileCh := manager.beginKeyspaceGroupReconcileTerm(ctx) + re.NoError(manager.CreateKeyspaceGroups([]*endpoint.KeyspaceGroup{{ + ID: 1, + UserKind: endpoint.Standard.String(), + }})) + store.loadFailures.Store(1) + + manager.wg.Add(1) + go manager.allocNodesToAllKeyspaceGroups(termCtx, term, reconcileCh) + select { + case <-store.loadFailureCh: + case <-time.After(time.Second): + t.Fatal("reconcile did not attempt to load the under-replicated group") + } + + testutil.Eventually(re, func() bool { + group, err := manager.GetKeyspaceGroupByID(1) + return err == nil && group != nil && len(group.Members) == mcs.DefaultKeyspaceGroupReplicaCount + }, testutil.WithWaitFor(2*time.Second), testutil.WithTickInterval(10*time.Millisecond)) +} + +func TestTSONodesWatcherTriggersNodeAllocation(t *testing.T) { + re := require.New(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + _, client, clean := etcdutil.NewTestEtcdCluster(t, 1, nil) + defer clean() + + registerNode := func(address string) { + value, err := (&discovery.ServiceRegistryEntry{ServiceAddr: address}).Serialize() + re.NoError(err) + _, err = client.Put(ctx, keypath.RegistryPath(mcs.TSOServiceName, address), value) + re.NoError(err) + } + registerNode("http://tso-1") + + store := endpoint.NewStorageEndpoint(kv.NewEtcdKVBase(client), nil) + group := &endpoint.KeyspaceGroup{ + ID: 1, + UserKind: endpoint.Standard.String(), + Members: []endpoint.KeyspaceGroupMember{ + {Address: "http://tso-1"}, + }, + Keyspaces: []uint32{1}, + } + re.NoError(store.RunInTxn(ctx, func(txn kv.Txn) error { + return store.SaveKeyspaceGroup(txn, group) + })) + + manager := NewKeyspaceGroupManager(ctx, store, client) + defer manager.Close() + testutil.Eventually(re, func() bool { + return manager.GetNodesCount() == 1 + }) + termCtx, term, reconcileCh := manager.beginKeyspaceGroupReconcileTerm(ctx) + manager.Lock() + manager.putKeyspaceGroupLocked(group) + manager.Unlock() + manager.wg.Add(1) + go manager.allocNodesToAllKeyspaceGroups(termCtx, term, reconcileCh) + + registerNode("http://tso-2") + testutil.Eventually(re, func() bool { + group, err := manager.GetKeyspaceGroupByID(1) + return err == nil && group != nil && len(group.Members) == mcs.DefaultKeyspaceGroupReplicaCount + }, testutil.WithWaitFor(time.Second), testutil.WithTickInterval(10*time.Millisecond)) +} + +func BenchmarkReconcileKeyspaceGroupsMillionKeyspaces(b *testing.B) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + keyspaces := make([]uint32, 1_000_000) + for i := range keyspaces { + keyspaces[i] = uint32(i) + } + group := &endpoint.KeyspaceGroup{ + ID: constant.DefaultKeyspaceGroupID, + UserKind: endpoint.Basic.String(), + Members: []endpoint.KeyspaceGroupMember{ + {Address: "http://tso-1"}, + {Address: "http://tso-2"}, + }, + Keyspaces: keyspaces, + } + payload, err := json.Marshal(group) + if err != nil { + b.Fatal(err) + } + + store := endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil) + if err := store.RunInTxn(ctx, func(txn kv.Txn) error { + return store.SaveKeyspaceGroup(txn, group) + }); err != nil { + b.Fatal(err) + } + + b.Run("legacy-full-load", func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + groups, err := store.LoadKeyspaceGroups(constant.DefaultKeyspaceGroupID, 0) + if err != nil || len(groups) != 1 { + b.Fatalf("failed to load keyspace groups: groups=%d err=%v", len(groups), err) + } + } + b.ReportMetric(float64(len(payload)), "etcd-read-bytes/op") + }) + + b.Run("event-driven-healthy", func(b *testing.B) { + countingStore := &countingKeyspaceGroupStorage{KeyspaceGroupStorage: store} + manager := NewKeyspaceGroupManager(ctx, countingStore, nil) + defer manager.Close() + manager.nodesBalancer.Put("http://tso-1") + manager.nodesBalancer.Put("http://tso-2") + termCtx, term, _ := manager.beginKeyspaceGroupReconcileTerm(ctx) + manager.Lock() + manager.putKeyspaceGroupLocked(group) + manager.Unlock() + b.ReportAllocs() + for b.Loop() { + if manager.reconcileKeyspaceGroupsForTerm(termCtx, term) { + b.Fatal("healthy groups should not need a retry") + } + } + if countingStore.loadGroupsCount.Load() != 0 || countingStore.loadGroupCount.Load() != 0 { + b.Fatal("healthy reconciliation read keyspace groups from storage") + } + b.ReportMetric(0, "etcd-read-bytes/op") + }) + + b.Run("watch-membership-only", func(b *testing.B) { + manager := NewKeyspaceGroupManager(ctx, store, nil) + manager.groupWatcherTerm = 1 + manager.Lock() + manager.putKeyspaceGroupLocked(group) + manager.Unlock() + kv := &mvccpb.KeyValue{Value: payload} + b.ReportAllocs() + b.SetBytes(int64(len(payload))) + for b.Loop() { + if _, _, err := manager.applyKeyspaceGroupMembership(1, kv, false); err != nil { + b.Fatal(err) + } + } + }) +} + type keyspaceGroupTestSuite struct { suite.Suite ctx context.Context diff --git a/pkg/storage/kv/etcd_kv.go b/pkg/storage/kv/etcd_kv.go index 440757f9144..78d7baa52e9 100644 --- a/pkg/storage/kv/etcd_kv.go +++ b/pkg/storage/kv/etcd_kv.go @@ -239,7 +239,7 @@ func (txn *etcdTxn) Remove(key string) error { // Load loads the target value from etcd and puts a comparator into conditions. func (txn *etcdTxn) Load(key string) (string, error) { - resp, err := etcdutil.EtcdKVGet(txn.kv.client, key) + resp, err := etcdutil.EtcdKVGetWithContext(txn.ctx, txn.kv.client, key) if err != nil { return "", err } diff --git a/pkg/storage/kv/kv_test.go b/pkg/storage/kv/kv_test.go index a382aa240bf..fc5036e8384 100644 --- a/pkg/storage/kv/kv_test.go +++ b/pkg/storage/kv/kv_test.go @@ -43,6 +43,21 @@ func TestEtcd(t *testing.T) { testSaveMultiple(re, kv, 20) testLoadConflict(re, kv) testRawTxn(re, kv) + testCanceledTxn(re, kv) +} + +func testCanceledTxn(re *require.Assertions, kv Base) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + err := kv.RunInTxn(ctx, func(txn Txn) error { + _, loadErr := txn.Load("canceled-key") + re.ErrorIs(loadErr, context.Canceled) + return txn.Save("canceled-key", "value") + }) + re.ErrorIs(err, context.Canceled) + value, err := kv.Load("canceled-key") + re.NoError(err) + re.Empty(value) } func TestLevelDB(t *testing.T) { diff --git a/pkg/utils/etcdutil/etcdutil.go b/pkg/utils/etcdutil/etcdutil.go index aa9f729aa7a..aa2d9bc192c 100644 --- a/pkg/utils/etcdutil/etcdutil.go +++ b/pkg/utils/etcdutil/etcdutil.go @@ -147,7 +147,12 @@ func RemoveEtcdMember(client *clientv3.Client, id uint64) (*clientv3.MemberRemov // EtcdKVGet returns the etcd GetResponse by given key or key prefix func EtcdKVGet(c *clientv3.Client, key string, opts ...clientv3.OpOption) (*clientv3.GetResponse, error) { - ctx, cancel := context.WithTimeout(c.Ctx(), DefaultRequestTimeout) + return EtcdKVGetWithContext(c.Ctx(), c, key, opts...) +} + +// EtcdKVGetWithContext returns the etcd GetResponse using the given context. +func EtcdKVGetWithContext(ctx context.Context, c *clientv3.Client, key string, opts ...clientv3.OpOption) (*clientv3.GetResponse, error) { + ctx, cancel := context.WithTimeout(ctx, DefaultRequestTimeout) defer cancel() start := time.Now()