diff --git a/errors.toml b/errors.toml index 33acad6fbf..43413f79a0 100644 --- a/errors.toml +++ b/errors.toml @@ -921,6 +921,11 @@ error = ''' keyspace not found with name: %s ''' +["PD:resourcemanager:ErrResourceGroupsLoading"] +error = ''' +resource groups are still being loaded, please try again later +''' + ["PD:scatter:ErrEmptyRegion"] error = ''' empty region diff --git a/pkg/errs/errno.go b/pkg/errs/errno.go index a887d539f6..9ddb0b7b40 100644 --- a/pkg/errs/errno.go +++ b/pkg/errs/errno.go @@ -541,6 +541,7 @@ var ( ErrResourceGroupNotExists = errors.Normalize("the %s resource group does not exist", errors.RFCCodeText("PD:resourcemanager:ErrGroupNotExists")) ErrDeleteReservedGroup = errors.Normalize("cannot delete reserved group", errors.RFCCodeText("PD:resourcemanager:ErrDeleteReservedGroup")) ErrInvalidGroup = errors.Normalize("invalid group settings, please check %s", errors.RFCCodeText("PD:resourcemanager:ErrInvalidGroup")) + ErrResourceGroupsLoading = errors.Normalize("resource groups are still being loaded, please try again later", errors.RFCCodeText("PD:resourcemanager:ErrResourceGroupsLoading")) ) // Microservice errors diff --git a/pkg/mcs/resourcemanager/server/grpc_service.go b/pkg/mcs/resourcemanager/server/grpc_service.go index 875299c602..0945494cdc 100644 --- a/pkg/mcs/resourcemanager/server/grpc_service.go +++ b/pkg/mcs/resourcemanager/server/grpc_service.go @@ -231,18 +231,23 @@ func (s *Service) AcquireTokenBuckets(stream rmpb.ResourceManager_AcquireTokenBu zap.Uint32("keyspace-id", keyspaceID), zap.String("resource-group", resourceGroupName), ) + // Get the resource group from manager to acquire token buckets. This also + // triggers lazy loading of the group if async loading hasn't completed yet, + // so it must happen before accessKeyspaceResourceGroupManager below. + rg, err := s.manager.GetMutableResourceGroup(keyspaceID, resourceGroupName) + if rg == nil { + if err != nil && !errors.ErrorEqual(err, errs.ErrResourceGroupNotExists.FastGenByArgs(resourceGroupName)) { + return err + } + log.Warn("resource group not found", append(requestFields, zap.Error(err))...) + continue + } // Get keyspace resource group manager to apply service limit later. krgm, err := s.manager.accessKeyspaceResourceGroupManager(keyspaceID, resourceGroupName) if krgm == nil { log.Warn("keyspace resource group manager not found", append(requestFields, zap.Error(err))...) continue } - // Get the resource group from manager to acquire token buckets. - rg, err := s.manager.GetMutableResourceGroup(keyspaceID, resourceGroupName) - if rg == nil { - log.Warn("resource group not found", append(requestFields, zap.Error(err))...) - continue - } // Send the consumption to update the metrics. err = s.manager.dispatchConsumption(req) if err != nil { diff --git a/pkg/mcs/resourcemanager/server/keyspace_manager.go b/pkg/mcs/resourcemanager/server/keyspace_manager.go index d3d8a13dae..8c900c5479 100644 --- a/pkg/mcs/resourcemanager/server/keyspace_manager.go +++ b/pkg/mcs/resourcemanager/server/keyspace_manager.go @@ -69,7 +69,13 @@ type consumptionItem struct { type keyspaceResourceGroupManager struct { syncutil.RWMutex - groups map[string]*ResourceGroup + groups map[string]*ResourceGroup + // reservedGroups tracks names whose entry in groups is still just the + // synthetic placeholder inserted by ensureReservedDefaultGroupInCache or + // restoreDefaultResourceGroupFromReserved, not yet confirmed by a storage + // load or a real write. It shares the same lock as groups so a lazy load + // can atomically decide whether it's safe to replace the placeholder. + reservedGroups map[string]struct{} groupRUTrackers map[string]*groupRUTracker serviceLimiter *serviceLimiter @@ -89,6 +95,7 @@ func newKeyspaceResourceGroupManager( } return &keyspaceResourceGroupManager{ groups: make(map[string]*ResourceGroup), + reservedGroups: make(map[string]struct{}), groupRUTrackers: make(map[string]*groupRUTracker), keyspaceID: keyspaceID, storage: storage, @@ -159,6 +166,9 @@ func (krgm *keyspaceResourceGroupManager) upsertResourceGroupFromRaw(name string zap.Uint32("keyspace-id", krgm.keyspaceID), zap.String("name", name), zap.String("raw-value", rawValue), zap.Error(err)) return err } + krgm.Lock() + delete(krgm.reservedGroups, group.Name) + krgm.Unlock() krgm.syncBurstabilityWithServiceLimit(existing) return nil } @@ -166,6 +176,7 @@ func (krgm *keyspaceResourceGroupManager) upsertResourceGroupFromRaw(name string resourceGroup := FromProtoResourceGroup(group) krgm.Lock() krgm.groups[group.Name] = resourceGroup + delete(krgm.reservedGroups, group.Name) krgm.Unlock() krgm.syncBurstabilityWithServiceLimit(resourceGroup) return nil @@ -175,6 +186,7 @@ func (krgm *keyspaceResourceGroupManager) deleteResourceGroupFromCache(name stri krgm.Lock() delete(krgm.groups, name) delete(krgm.groupRUTrackers, name) + delete(krgm.reservedGroups, name) krgm.Unlock() } @@ -218,6 +230,7 @@ func (krgm *keyspaceResourceGroupManager) ensureReservedDefaultGroupInCache() { krgm.Lock() if _, ok := krgm.groups[DefaultResourceGroupName]; !ok { krgm.groups[DefaultResourceGroupName] = defaultGroup + krgm.reservedGroups[DefaultResourceGroupName] = struct{}{} inserted = true } krgm.Unlock() @@ -245,6 +258,7 @@ func (krgm *keyspaceResourceGroupManager) restoreDefaultResourceGroupFromReserve defaultGroup := newDefaultResourceGroup() krgm.Lock() krgm.groups[DefaultResourceGroupName] = defaultGroup + krgm.reservedGroups[DefaultResourceGroupName] = struct{}{} krgm.Unlock() krgm.syncBurstabilityWithServiceLimit(defaultGroup) } @@ -266,6 +280,7 @@ func (krgm *keyspaceResourceGroupManager) addResourceGroup(grouppb *rmpb.Resourc } krgm.Lock() krgm.groups[group.Name] = group + delete(krgm.reservedGroups, group.Name) krgm.Unlock() krgm.syncBurstabilityWithServiceLimit(group) return nil @@ -288,6 +303,9 @@ func (krgm *keyspaceResourceGroupManager) modifyResourceGroup(group *rmpb.Resour if err != nil { return err } + // Deliberately not clearing reservedGroups here: modifying only patches + // settings, it never establishes the group's state, so it must not make + // a state-unconfirmed entry look fully confirmed. return curGroup.persistSettings(krgm.keyspaceID, krgm.storage) } @@ -330,6 +348,17 @@ func (krgm *keyspaceResourceGroupManager) deleteResourceGroup(name string) error return nil } +// isReserved reports whether name's cached entry is still just the synthetic +// placeholder inserted by ensureReservedDefaultGroupInCache or +// restoreDefaultResourceGroupFromReserved, not yet confirmed by a storage +// load or a real write. +func (krgm *keyspaceResourceGroupManager) isReserved(name string) bool { + krgm.RLock() + defer krgm.RUnlock() + _, ok := krgm.reservedGroups[name] + return ok +} + func (krgm *keyspaceResourceGroupManager) getResourceGroup(name string, withStats bool) *ResourceGroup { krgm.RLock() defer krgm.RUnlock() @@ -384,6 +413,14 @@ func (krgm *keyspaceResourceGroupManager) persistResourceGroupRunningState() { for idx := range keys { krgm.RLock() group, ok := krgm.groups[keys[idx]] + _, reserved := krgm.reservedGroups[keys[idx]] + if ok && reserved { + // The entry is still just an unconfirmed placeholder (e.g. the + // synthetic default installed before async loading completes); + // persisting its fresh state would permanently overwrite any + // real persisted state still waiting to be loaded. + ok = false + } if ok { if err := group.persistStates(krgm.keyspaceID, krgm.storage); err != nil { log.Error("persist keyspace resource group state failed", diff --git a/pkg/mcs/resourcemanager/server/manager.go b/pkg/mcs/resourcemanager/server/manager.go index 4f811a0392..e45be2cad3 100644 --- a/pkg/mcs/resourcemanager/server/manager.go +++ b/pkg/mcs/resourcemanager/server/manager.go @@ -21,6 +21,7 @@ import ( "os" "strings" "sync" + "sync/atomic" "time" "github.com/prometheus/client_golang/prometheus/push" @@ -114,8 +115,22 @@ type Manager struct { metrics *metrics // ruCollector is used to collect the RU metering data. ruCollector *ruCollector + // async loading state management + loadingState int32 // atomic access + // syncLoadedGroups records groups that were loaded synchronously (e.g., by lazy loading) + syncLoadedGroups map[trackerKey]bool } +// LoadingState represents the current loading state of resource groups +const ( + // LoadingStateNotStarted means resource groups haven't started loading + LoadingStateNotStarted int32 = iota + // LoadingStateInProgress means resource groups are being loaded asynchronously + LoadingStateInProgress + // LoadingStateCompleted means all resource groups have been loaded + LoadingStateCompleted +) + // factoryProvider is a factory provider for the manager, which injects some specialized functions // that need to be retrieved from the `bs.Server` instance without interacting with its interface. type factoryProvider interface { @@ -149,6 +164,8 @@ func newManagerBase(controllerConfig *ControllerConfig, writeRole ResourceGroupW keyspaceIDLookup: make(map[string]uint32), metrics: newMetrics(), ruCollector: newRUCollector(), + loadingState: LoadingStateNotStarted, + syncLoadedGroups: make(map[trackerKey]bool), } } @@ -194,7 +211,10 @@ func NewMetadataOnlyManager[T metadataFactoryProvider](srv bs.Server) (*Manager, nil, ) m.srv = srv - if err := m.initMetadata(); err != nil { + if err := m.initControllerConfig(); err != nil { + return nil, err + } + if err := m.loadKeyspaceResourceGroups(); err != nil { return nil, err } return m, nil @@ -275,6 +295,13 @@ func (m *Manager) GetRUVersionPolicy() *RUVersionPolicy { return m.controllerConfig.RUVersionPolicy.Clone() } +// getOrCreateKeyspaceResourceGroupManager returns the keyspace resource group +// manager for keyspaceID, creating it if needed. When initDefault is true, it +// also ensures the default resource group is present. While async loading is +// still in progress this goes through loadResourceGroupIfNeeded, which tries +// a storage point load first so a customized default is never clobbered by a +// blindly synthesized one. Once loading has completed, any default missing +// from the cache truly doesn't exist anywhere, so it's synthesized directly. func (m *Manager) getOrCreateKeyspaceResourceGroupManager(keyspaceID uint32, initDefault bool) *keyspaceResourceGroupManager { m.Lock() krgm, ok := m.krgms[keyspaceID] @@ -283,9 +310,15 @@ func (m *Manager) getOrCreateKeyspaceResourceGroupManager(keyspaceID uint32, ini m.krgms[keyspaceID] = krgm } m.Unlock() - // Init the default resource group if needed. if initDefault { - krgm.initDefaultResourceGroup() + if atomic.LoadInt32(&m.loadingState) == LoadingStateCompleted { + // Async loading (if any) has already finished, so if the default + // group isn't cached yet it truly doesn't exist anywhere; it's + // safe to synthesize and persist it directly. + krgm.initDefaultResourceGroup() + } else if err := m.loadResourceGroupIfNeeded(keyspaceID, DefaultResourceGroupName); err != nil { + log.Debug("failed to load default resource group", zap.Uint32("keyspace-id", keyspaceID), zap.Error(err)) + } } return krgm } @@ -325,13 +358,16 @@ func (m *Manager) Init(ctx context.Context) error { m.wg.Wait() return err } + atomic.StoreInt32(&m.loadingState, LoadingStateCompleted) } else { - if err := m.initMetadata(); err != nil { - return err - } // This context is derived from the leader/primary context, it will be canceled // from the outside loop when the leader/primary step down. ctx, m.cancel = context.WithCancel(ctx) + if err := m.initMetadata(ctx); err != nil { + m.cancel() + m.wg.Wait() + return err + } } m.wg.Add(1) // Start the background metrics flusher. @@ -369,34 +405,145 @@ func (m *Manager) initControllerConfig() error { return nil } -func (m *Manager) initMetadata() error { +func (m *Manager) initMetadata(ctx context.Context) error { if err := m.initControllerConfig(); err != nil { return err } - // Load keyspace resource groups from the storage. - return m.loadKeyspaceResourceGroups() + m.Lock() + m.krgms = make(map[uint32]*keyspaceResourceGroupManager) + m.syncLoadedGroups = make(map[trackerKey]bool) + atomic.StoreInt32(&m.loadingState, LoadingStateNotStarted) + m.Unlock() + + m.initReservedInCache() + if err := m.loadServiceLimits(); err != nil { + return err + } + + m.wg.Add(1) + go m.asyncLoadResourceGroups(ctx) + return nil +} + +func (m *Manager) loadServiceLimits() error { + return m.storage.LoadServiceLimits(func(keyspaceID uint32, serviceLimit float64) { + m.getOrCreateKeyspaceResourceGroupManager(keyspaceID, false).setServiceLimitFromStorage(serviceLimit) + }) } func (m *Manager) loadKeyspaceResourceGroups() error { - // Empty the keyspace resource group manager map before the loading. + tempKrgms, err := m.loadKeyspaceResourceGroupsFromStorage() + if err != nil { + return err + } m.Lock() - m.krgms = make(map[uint32]*keyspaceResourceGroupManager) + m.krgms = tempKrgms + m.syncLoadedGroups = nil + atomic.StoreInt32(&m.loadingState, LoadingStateCompleted) m.Unlock() - // Load keyspace resource group meta info from the storage. + m.initReserved() + return m.loadServiceLimits() +} + +func (m *Manager) asyncLoadResourceGroups(ctx context.Context) { + defer logutil.LogPanic() + defer m.wg.Done() + + const retryInterval = 10 * time.Second + retry := 0 + for { + select { + case <-ctx.Done(): + log.Info("async loading resource groups cancelled") + return + default: + } + if retry > 0 { + log.Info("retrying async loading resource groups", zap.Int("retry", retry)) + timer := time.NewTimer(retryInterval) + select { + case <-ctx.Done(): + timer.Stop() + log.Info("async loading resource groups cancelled") + return + case <-timer.C: + } + } + + atomic.StoreInt32(&m.loadingState, LoadingStateInProgress) + startTime := time.Now() + tempKrgms, err := m.loadKeyspaceResourceGroupsFromStorage() + if err != nil { + log.Error("failed to load resource groups", zap.Error(err), zap.Int("retry", retry)) + atomic.StoreInt32(&m.loadingState, LoadingStateNotStarted) + retry++ + continue + } + + loaded := 0 + m.Lock() + for keyspaceID, tempKrgm := range tempKrgms { + krgm := m.krgms[keyspaceID] + if krgm == nil { + krgm = newKeyspaceResourceGroupManager(keyspaceID, m.storage, m.writeRole) + m.krgms[keyspaceID] = krgm + } + groupsToSync := make([]*ResourceGroup, 0) + tempKrgm.RLock() + krgm.Lock() + for name, group := range tempKrgm.groups { + key := trackerKey{keyspaceID: keyspaceID, groupName: name} + if !m.syncLoadedGroups[key] { + krgm.groups[name] = group + // This group is now confirmed, fully-loaded data (settings + // and state); it must no longer be treated as an + // unconfirmed placeholder by loadResourceGroupIfNeeded or + // skipped by the state persist loop. + delete(krgm.reservedGroups, name) + groupsToSync = append(groupsToSync, group) + loaded++ + } + } + krgm.Unlock() + tempKrgm.RUnlock() + for _, group := range groupsToSync { + krgm.syncBurstabilityWithServiceLimit(group) + } + } + m.syncLoadedGroups = nil + m.Unlock() + + m.initReserved() + atomic.StoreInt32(&m.loadingState, LoadingStateCompleted) + duration := time.Since(startTime) + asyncLoadGroupDuration.Observe(duration.Seconds()) + log.Info("async loading resource groups completed", zap.Int("loaded-groups", loaded), zap.Duration("duration", duration)) + return + } +} + +func (m *Manager) loadKeyspaceResourceGroupsFromStorage() (map[uint32]*keyspaceResourceGroupManager, error) { + tempKrgms := make(map[uint32]*keyspaceResourceGroupManager) + getOrCreateTempKrgm := func(keyspaceID uint32) *keyspaceResourceGroupManager { + krgm, ok := tempKrgms[keyspaceID] + if !ok { + krgm = newKeyspaceResourceGroupManager(keyspaceID, m.storage, m.writeRole) + tempKrgms[keyspaceID] = krgm + } + return krgm + } if err := m.storage.LoadResourceGroupSettings(func(keyspaceID uint32, name string, rawValue string) { - // Since the default resource group might be loaded from the storage, we don't need to initialize it here. - err := m.getOrCreateKeyspaceResourceGroupManager(keyspaceID, false).addResourceGroupFromRaw(name, rawValue) + err := getOrCreateTempKrgm(keyspaceID).addResourceGroupFromRaw(name, rawValue) if err != nil { log.Error("failed to add resource group to the keyspace resource group manager", zap.Uint32("keyspace-id", keyspaceID), zap.String("group-name", name), zap.Error(err)) } }); err != nil { - return err + return nil, err } - // Load keyspace resource group states from the storage. if err := m.storage.LoadResourceGroupStates(func(keyspaceID uint32, name string, rawValue string) { - krgm := m.getKeyspaceResourceGroupManager(keyspaceID) + krgm := tempKrgms[keyspaceID] if krgm == nil { log.Warn("failed to get the corresponding keyspace resource group manager", zap.Uint32("keyspace-id", keyspaceID), zap.String("group-name", name)) @@ -408,18 +555,119 @@ func (m *Manager) loadKeyspaceResourceGroups() error { zap.Uint32("keyspace-id", keyspaceID), zap.String("group-name", name), zap.Error(err)) } }); err != nil { + return nil, err + } + return tempKrgms, nil +} + +// loadResourceGroup loads a single resource group from storage. The returned +// stateLoaded reports whether the group's persisted state was successfully +// read; the caller must not mark such a group as sync-loaded, so that a +// concurrent or later async bulk load can still fill in its real state. +func (m *Manager) loadResourceGroup(keyspaceID uint32, name string) (group *ResourceGroup, stateLoaded bool, err error) { + rawValue, err := m.storage.LoadResourceGroupSetting(keyspaceID, name) + if err != nil { + return nil, false, err + } + if rawValue == "" { + return nil, false, errs.ErrResourceGroupNotExists.FastGenByArgs(name) + } + krgm := newKeyspaceResourceGroupManager(keyspaceID, m.storage, m.writeRole) + if err := krgm.addResourceGroupFromRaw(name, rawValue); err != nil { + return nil, false, err + } + state, err := m.storage.LoadResourceGroupState(keyspaceID, name) + if err != nil { + log.Warn("failed to load resource group state, continuing without state", + zap.Uint32("keyspace-id", keyspaceID), + zap.String("group-name", name), + zap.Error(err)) + return krgm.getMutableResourceGroup(name), false, nil + } + if state != "" { + if err := krgm.setRawStatesIntoResourceGroup(name, state); err != nil { + return nil, false, err + } + } + return krgm.getMutableResourceGroup(name), true, nil +} + +func (m *Manager) loadResourceGroupIfNeeded(keyspaceID uint32, name string) error { + if atomic.LoadInt32(&m.loadingState) == LoadingStateCompleted { + return nil + } + krgm := m.getKeyspaceResourceGroupManager(keyspaceID) + if krgm != nil { + // A cached entry only satisfies this call if it's confirmed data, not + // just a synthetic placeholder (e.g. from ensureReservedDefaultGroupInCache) + // installed before async loading had a chance to run. + if group := krgm.getMutableResourceGroup(name); group != nil && !krgm.isReserved(name) { + return nil + } + } + group, stateLoaded, err := m.loadResourceGroup(keyspaceID, name) + if err != nil { + if name == DefaultResourceGroupName && errors.ErrorEqual(err, errs.ErrResourceGroupNotExists.FastGenByArgs(name)) { + // No persisted default group settings exist yet (e.g. a brand-new + // keyspace), so it's safe to synthesize the reserved default group. + // This calls initDefaultResourceGroup directly instead of going + // through getOrCreateKeyspaceResourceGroupManager(id, true), which + // now routes back into this same function and would recurse. + m.getOrCreateKeyspaceResourceGroupManager(keyspaceID, false).initDefaultResourceGroup() + return nil + } return err } - // Initialize the reserved keyspace resource group manager and default resource groups. - m.initReserved() - // Load service limits from the storage after all resource groups are loaded. - return m.loadServiceLimits() + krgm = m.getOrCreateKeyspaceResourceGroupManager(keyspaceID, false) + inserted := false + krgm.Lock() + if _, exists := krgm.groups[name]; !exists { + krgm.groups[name] = group + inserted = true + } else if _, reserved := krgm.reservedGroups[name]; reserved { + // The existing entry is just an unconfirmed placeholder; the freshly + // loaded group is the real, confirmed data, so replacing it is safe. + krgm.groups[name] = group + inserted = true + } + // Only clear the placeholder mark once the state was actually read; a + // metadata-only group (state load failed) must stay reserved so a later + // call or the async bulk merge remains free to fill in the real state, + // instead of this partial result being treated as final forever. + if stateLoaded { + delete(krgm.reservedGroups, name) + } else { + krgm.reservedGroups[name] = struct{}{} + } + krgm.Unlock() + if inserted { + krgm.syncBurstabilityWithServiceLimit(group) + } + // Only mark the group as sync-loaded when its persisted state was actually + // read; otherwise a later async bulk load must remain free to fill in the + // real state instead of being skipped forever. + if stateLoaded { + markKey := trackerKey{keyspaceID: keyspaceID, groupName: name} + m.Lock() + if m.syncLoadedGroups != nil { + m.syncLoadedGroups[markKey] = true + } + m.Unlock() + } + syncLoadGroupCounter.Inc() + return nil } -func (m *Manager) loadServiceLimits() error { - return m.storage.LoadServiceLimits(func(keyspaceID uint32, serviceLimit float64) { - m.getOrCreateKeyspaceResourceGroupManager(keyspaceID, false).setServiceLimitFromStorage(serviceLimit) - }) +func (m *Manager) markResourceGroupSyncLoaded(keyspaceID uint32, name string) { + m.Lock() + defer m.Unlock() + if m.syncLoadedGroups != nil { + m.syncLoadedGroups[trackerKey{keyspaceID: keyspaceID, groupName: name}] = true + } +} + +func (m *Manager) isResourceGroupLoadingComplete() bool { + return atomic.LoadInt32(&m.loadingState) == LoadingStateCompleted } func cloneControllerConfig(cfg *ControllerConfig) *ControllerConfig { @@ -456,6 +704,7 @@ func (m *Manager) applyResourceGroupSettingFromRaw(keyspaceID uint32, name, rawV zap.Error(err)) return err } + m.markResourceGroupSyncLoaded(keyspaceID, name) return nil } @@ -492,6 +741,7 @@ func (m *Manager) applyResourceGroupStatesFromRaw(keyspaceID uint32, name, rawVa zap.Error(err)) return err } + m.markResourceGroupSyncLoaded(keyspaceID, name) return nil } @@ -504,6 +754,15 @@ func (m *Manager) initReserved() { } } +func (m *Manager) initReservedInCache() { + // Initialize the reserved default group in memory before async loading + // without overwriting persisted default group settings. + m.getOrCreateKeyspaceResourceGroupManager(constant.NullKeyspaceID, false).ensureReservedDefaultGroupInCache() + for _, krgm := range m.getKeyspaceResourceGroupManagers() { + krgm.ensureReservedDefaultGroupInCache() + } +} + // UpdateControllerConfigItem updates the controller config item. func (m *Manager) UpdateControllerConfigItem(key string, value any) error { if !m.writeRole.AllowsMetadataWrite() { @@ -574,7 +833,16 @@ func (m *Manager) AddResourceGroup(grouppb *rmpb.ResourceGroup) error { if krgm == nil { return errs.ErrKeyspaceNotExists.FastGenByArgs(keyspaceID) } - return krgm.addResourceGroup(grouppb) + if err := m.loadResourceGroupIfNeeded(keyspaceID, grouppb.Name); err != nil && + !errors.ErrorEqual(err, errs.ErrResourceGroupNotExists.FastGenByArgs(grouppb.Name)) { + log.Warn("failed to load resource group before add", zap.Uint32("keyspace-id", keyspaceID), zap.String("name", grouppb.Name), zap.Error(err)) + return err + } + if err := krgm.addResourceGroup(grouppb); err != nil { + return err + } + m.markResourceGroupSyncLoaded(keyspaceID, grouppb.Name) + return nil } // ModifyResourceGroup modifies an existing resource group. @@ -583,11 +851,25 @@ func (m *Manager) ModifyResourceGroup(grouppb *rmpb.ResourceGroup) error { return errMetadataWriteDisabled } keyspaceID := ExtractKeyspaceID(grouppb.GetKeyspaceId()) + if err := m.loadResourceGroupIfNeeded(keyspaceID, grouppb.Name); err != nil { + log.Debug("failed to load resource group", zap.Uint32("keyspace-id", keyspaceID), zap.String("name", grouppb.Name), zap.Error(err)) + return err + } krgm, err := m.accessKeyspaceResourceGroupManager(keyspaceID, grouppb.Name) if err != nil { return err } - return krgm.modifyResourceGroup(grouppb) + if err := krgm.modifyResourceGroup(grouppb); err != nil { + return err + } + // Modifying only patches settings, it never establishes the group's + // state. If the state still hasn't been confirmed (isReserved), marking + // it sync-loaded here would make the async bulk merge skip it forever, + // so the persisted running state would never get applied. + if !krgm.isReserved(grouppb.Name) { + m.markResourceGroupSyncLoaded(keyspaceID, grouppb.Name) + } + return nil } // DeleteResourceGroup deletes a resource group. @@ -595,16 +877,28 @@ func (m *Manager) DeleteResourceGroup(keyspaceID uint32, name string) error { if !m.writeRole.AllowsMetadataWrite() { return errMetadataWriteDisabled } + if err := m.loadResourceGroupIfNeeded(keyspaceID, name); err != nil { + log.Debug("failed to load resource group", zap.Uint32("keyspace-id", keyspaceID), zap.String("name", name), zap.Error(err)) + return err + } // "default" group can't be deleted, so there is not need to call accessKeyspaceResourceGroupManager krgm := m.getKeyspaceResourceGroupManager(keyspaceID) if krgm == nil { return errs.ErrKeyspaceNotExists.FastGenByArgs(keyspaceID) } - return krgm.deleteResourceGroup(name) + if err := krgm.deleteResourceGroup(name); err != nil { + return err + } + m.markResourceGroupSyncLoaded(keyspaceID, name) + return nil } // GetResourceGroup returns a copy of a resource group. func (m *Manager) GetResourceGroup(keyspaceID uint32, name string, withStats bool) (*ResourceGroup, error) { + if err := m.loadResourceGroupIfNeeded(keyspaceID, name); err != nil { + log.Debug("failed to load resource group", zap.Uint32("keyspace-id", keyspaceID), zap.String("name", name), zap.Error(err)) + return nil, err + } krgm, err := m.accessKeyspaceResourceGroupManager(keyspaceID, name) if err != nil { return nil, err @@ -614,6 +908,10 @@ func (m *Manager) GetResourceGroup(keyspaceID uint32, name string, withStats boo // GetMutableResourceGroup returns a mutable resource group. func (m *Manager) GetMutableResourceGroup(keyspaceID uint32, name string) (*ResourceGroup, error) { + if err := m.loadResourceGroupIfNeeded(keyspaceID, name); err != nil { + log.Debug("failed to load resource group", zap.Uint32("keyspace-id", keyspaceID), zap.String("name", name), zap.Error(err)) + return nil, err + } krgm, err := m.accessKeyspaceResourceGroupManager(keyspaceID, name) if err != nil { return nil, err @@ -622,7 +920,12 @@ func (m *Manager) GetMutableResourceGroup(keyspaceID uint32, name string) (*Reso } // GetResourceGroupList returns copies of resource group list. +// Returns error if resource groups are still being loaded asynchronously. func (m *Manager) GetResourceGroupList(keyspaceID uint32, withStats bool) ([]*ResourceGroup, error) { + if !m.isResourceGroupLoadingComplete() { + log.Debug("resource groups are still being loaded, cannot return list") + return nil, errs.ErrResourceGroupsLoading + } krgm, err := m.accessKeyspaceResourceGroupManager(keyspaceID, DefaultResourceGroupName) if err != nil { return nil, err diff --git a/pkg/mcs/resourcemanager/server/manager_async_test.go b/pkg/mcs/resourcemanager/server/manager_async_test.go new file mode 100644 index 0000000000..73db769b88 --- /dev/null +++ b/pkg/mcs/resourcemanager/server/manager_async_test.go @@ -0,0 +1,247 @@ +// Copyright 2026 TiKV Project Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package server + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/pingcap/kvproto/pkg/resource_manager" + + "github.com/tikv/pd/pkg/errs" + "github.com/tikv/pd/pkg/keyspace/constant" + "github.com/tikv/pd/pkg/storage" + "github.com/tikv/pd/pkg/utils/testutil" +) + +type blockingResourceGroupStorage struct { + storage.Storage + + once sync.Once + releaseOnce sync.Once + entered chan struct{} + release chan struct{} + + // failNextState, when true, makes the very next LoadResourceGroupState + // call fail once, then resets itself. + failNextState atomic.Bool +} + +func newBlockingResourceGroupStorage() *blockingResourceGroupStorage { + return &blockingResourceGroupStorage{ + Storage: storage.NewStorageWithMemoryBackend(), + entered: make(chan struct{}), + release: make(chan struct{}), + } +} + +func (s *blockingResourceGroupStorage) LoadResourceGroupSettings(f func(keyspaceID uint32, name, rawValue string)) error { + s.once.Do(func() { + close(s.entered) + <-s.release + }) + return s.Storage.LoadResourceGroupSettings(f) +} + +func (s *blockingResourceGroupStorage) LoadResourceGroupState(keyspaceID uint32, name string) (string, error) { + if s.failNextState.CompareAndSwap(true, false) { + return "", errors.New("injected resource group state load failure") + } + return s.Storage.LoadResourceGroupState(keyspaceID, name) +} + +func (s *blockingResourceGroupStorage) waitEntered(t *testing.T) { + t.Helper() + select { + case <-s.entered: + case <-time.After(time.Second): + t.Fatal("timed out waiting for async resource group loading") + } +} + +func (s *blockingResourceGroupStorage) unblock() { + s.releaseOnce.Do(func() { + close(s.release) + }) +} + +func newAsyncTestGroup(name string, fillRate uint64) *resource_manager.ResourceGroup { + return &resource_manager.ResourceGroup{ + Name: name, + Mode: resource_manager.GroupMode_RUMode, + Priority: middlePriority, + RUSettings: &resource_manager.GroupRequestUnitSettings{ + RU: &resource_manager.TokenBucket{ + Settings: &resource_manager.TokenLimitSettings{ + FillRate: fillRate, + BurstLimit: int64(fillRate), + }, + }, + }, + } +} + +func stopAsyncTestManager(m *Manager) { + if m.cancel != nil { + m.cancel() + } + m.wg.Wait() +} + +func TestAsyncLoadResourceGroupsLazyGet(t *testing.T) { + re := require.New(t) + store := newBlockingResourceGroupStorage() + re.NoError(store.SaveResourceGroupSetting(1, "lazy-group", newAsyncTestGroup("lazy-group", 100))) + + m := NewManager[*mockConfigProvider](&mockConfigProvider{}) + m.storage = store + re.NoError(m.Init(context.Background())) + defer stopAsyncTestManager(m) + // Unblock the async loader first (LIFO) so stopAsyncTestManager's wg.Wait() + // cannot hang if a later assertion aborts the test before the explicit + // store.unblock() call below is reached. + defer store.unblock() + + store.waitEntered(t) + + _, err := m.GetResourceGroupList(1, false) + re.ErrorIs(err, errs.ErrResourceGroupsLoading) + + group, err := m.GetResourceGroup(1, "lazy-group", false) + re.NoError(err) + re.NotNil(group) + re.Equal("lazy-group", group.Name) + re.Equal(float64(100), group.RUSettings.RU.getFillRate()) + + store.unblock() + testutil.Eventually(re, func() bool { + groups, err := m.GetResourceGroupList(1, false) + return err == nil && len(groups) == 2 + }, testutil.WithTickInterval(20*time.Millisecond)) +} + +func TestAsyncLoadResourceGroupsDoesNotRestoreDeletedLazyGroup(t *testing.T) { + re := require.New(t) + store := newBlockingResourceGroupStorage() + re.NoError(store.SaveResourceGroupSetting(1, "deleted-group", newAsyncTestGroup("deleted-group", 100))) + + m := NewManager[*mockConfigProvider](&mockConfigProvider{}) + m.storage = store + re.NoError(m.Init(context.Background())) + defer stopAsyncTestManager(m) + // Unblock the async loader first (LIFO) so stopAsyncTestManager's wg.Wait() + // cannot hang if a later assertion aborts the test before the explicit + // store.unblock() call below is reached. + defer store.unblock() + + store.waitEntered(t) + + group, err := m.GetResourceGroup(1, "deleted-group", false) + re.NoError(err) + re.NotNil(group) + re.NoError(m.DeleteResourceGroup(1, "deleted-group")) + + store.unblock() + testutil.Eventually(re, func() bool { + groups, err := m.GetResourceGroupList(1, false) + if err != nil { + return false + } + for _, group := range groups { + if group.Name == "deleted-group" { + return false + } + } + return true + }, testutil.WithTickInterval(20*time.Millisecond)) +} + +// TestAsyncLoadResourceGroupsLazyGetLegacyKeyspace guards against the point +// loaders (LoadResourceGroupSetting/LoadResourceGroupState) diverging from +// the bulk loaders on legacy, pre-keyspace resource groups: those are saved +// under constant.NullKeyspaceID, and a lazy Get during async loading must be +// able to find one the same way the bulk scan would once it completes. +func TestAsyncLoadResourceGroupsLazyGetLegacyKeyspace(t *testing.T) { + re := require.New(t) + store := newBlockingResourceGroupStorage() + re.NoError(store.SaveResourceGroupSetting(constant.NullKeyspaceID, "legacy-group", newAsyncTestGroup("legacy-group", 100))) + + m := NewManager[*mockConfigProvider](&mockConfigProvider{}) + m.storage = store + re.NoError(m.Init(context.Background())) + defer stopAsyncTestManager(m) + defer store.unblock() + + store.waitEntered(t) + + group, err := m.GetResourceGroup(constant.NullKeyspaceID, "legacy-group", false) + re.NoError(err) + re.NotNil(group) + re.Equal("legacy-group", group.Name) + re.Equal(float64(100), group.RUSettings.RU.getFillRate()) + + store.unblock() + testutil.Eventually(re, func() bool { + group, err := m.GetResourceGroup(constant.NullKeyspaceID, "legacy-group", false) + return err == nil && group != nil + }, testutil.WithTickInterval(20*time.Millisecond)) +} + +// TestAsyncLoadResourceGroupsRecoversFromStateLoadFailure guards against a +// group getting stuck marked reserved forever after a transient +// LoadResourceGroupState failure during lazy loading: once the async bulk +// load subsequently installs the fully-loaded (settings and state) +// confirmed data for the same group, the reserved marker must be cleared, +// otherwise loadResourceGroupIfNeeded and the state persist loop would keep +// treating already-recovered, correct data as an unconfirmed placeholder. +func TestAsyncLoadResourceGroupsRecoversFromStateLoadFailure(t *testing.T) { + re := require.New(t) + store := newBlockingResourceGroupStorage() + group := newAsyncTestGroup("flaky-group", 100) + re.NoError(store.SaveResourceGroupSetting(1, "flaky-group", group)) + re.NoError(store.SaveResourceGroupStates(1, "flaky-group", FromProtoResourceGroup(group).GetGroupStates())) + + m := NewManager[*mockConfigProvider](&mockConfigProvider{}) + m.storage = store + re.NoError(m.Init(context.Background())) + defer stopAsyncTestManager(m) + defer store.unblock() + + store.waitEntered(t) + + // Make the lazy load's own state read fail once, so the group is cached + // as a metadata-only, still-reserved entry. + store.failNextState.Store(true) + fetched, err := m.GetResourceGroup(1, "flaky-group", false) + re.NoError(err) + re.NotNil(fetched) + + krgm := m.getKeyspaceResourceGroupManager(1) + re.NotNil(krgm) + re.True(krgm.isReserved("flaky-group"), "group should still be reserved after a failed state load") + + // Let the async bulk load proceed; its own state read is unaffected + // (failNextState was already consumed) and should install confirmed data. + store.unblock() + testutil.Eventually(re, func() bool { + return !krgm.isReserved("flaky-group") + }, testutil.WithTickInterval(20*time.Millisecond)) +} diff --git a/pkg/mcs/resourcemanager/server/manager_test.go b/pkg/mcs/resourcemanager/server/manager_test.go index 16ecac456d..c045597240 100644 --- a/pkg/mcs/resourcemanager/server/manager_test.go +++ b/pkg/mcs/resourcemanager/server/manager_test.go @@ -408,7 +408,9 @@ func TestInitManager(t *testing.T) { m.storage = storage err = m.Init(ctx) re.NoError(err) - re.Len(m.getKeyspaceResourceGroupManagers(), 2) + testutil.Eventually(re, func() bool { + return len(m.getKeyspaceResourceGroupManagers()) == 2 + }) // Get the default resource group. rg, err := m.GetResourceGroup(1, DefaultResourceGroupName, true) re.NoError(err) diff --git a/pkg/mcs/resourcemanager/server/metrics.go b/pkg/mcs/resourcemanager/server/metrics.go index 2a7fd14482..3bdac0f9fd 100644 --- a/pkg/mcs/resourcemanager/server/metrics.go +++ b/pkg/mcs/resourcemanager/server/metrics.go @@ -220,6 +220,22 @@ var ( Help: "The duration of pushing RU metrics to Prometheus.", Buckets: prometheus.DefBuckets, }) + + syncLoadGroupCounter = prometheus.NewCounter( + prometheus.CounterOpts{ + Namespace: namespace, + Subsystem: serverSubsystem, + Name: "sync_load_groups_total", + Help: "Total number of resource groups loaded synchronously.", + }) + + asyncLoadGroupDuration = prometheus.NewHistogram( + prometheus.HistogramOpts{ + Namespace: namespace, + Subsystem: serverSubsystem, + Name: "async_load_group_duration_seconds", + Help: "Duration of asynchronous resource group loading in seconds.", + }) ) type metrics struct { @@ -274,6 +290,8 @@ func init() { prometheus.MustRegister(overrideSettings) prometheus.MustRegister(serviceLimit) prometheus.MustRegister(pushRUMetricsDuration) + prometheus.MustRegister(syncLoadGroupCounter) + prometheus.MustRegister(asyncLoadGroupDuration) } func newMetrics() *metrics { diff --git a/pkg/storage/endpoint/resource_group.go b/pkg/storage/endpoint/resource_group.go index 68a6297772..a54eb7a271 100644 --- a/pkg/storage/endpoint/resource_group.go +++ b/pkg/storage/endpoint/resource_group.go @@ -30,9 +30,11 @@ import ( // ResourceGroupStorage defines the storage operations on the resource group. type ResourceGroupStorage interface { LoadResourceGroupSettings(f func(keyspaceID uint32, name, rawValue string)) error + LoadResourceGroupSetting(keyspaceID uint32, name string) (string, error) SaveResourceGroupSetting(keyspaceID uint32, name string, msg proto.Message) error DeleteResourceGroupSetting(keyspaceID uint32, name string) error LoadResourceGroupStates(f func(keyspaceID uint32, name, rawValue string)) error + LoadResourceGroupState(keyspaceID uint32, name string) (string, error) SaveResourceGroupStates(keyspaceID uint32, name string, obj any) error DeleteResourceGroupStates(keyspaceID uint32, name string) error SaveControllerConfig(config any) error @@ -73,6 +75,11 @@ func (se *StorageEndpoint) LoadResourceGroupSettings(f func(keyspaceID uint32, n }) } +// LoadResourceGroupSetting loads a specific resource group from storage. +func (se *StorageEndpoint) LoadResourceGroupSetting(keyspaceID uint32, name string) (string, error) { + return se.Load(keypath.KeyspaceResourceGroupSettingPath(keyspaceID, name)) +} + // SaveResourceGroupStates stores a resource group to storage. func (se *StorageEndpoint) SaveResourceGroupStates(keyspaceID uint32, name string, obj any) error { return se.saveJSON(keypath.KeyspaceResourceGroupStatePath(keyspaceID, name), obj) @@ -102,6 +109,11 @@ func (se *StorageEndpoint) LoadResourceGroupStates(f func(keyspaceID uint32, nam }) } +// LoadResourceGroupState loads a specific resource group state from storage. +func (se *StorageEndpoint) LoadResourceGroupState(keyspaceID uint32, name string) (string, error) { + return se.Load(keypath.KeyspaceResourceGroupStatePath(keyspaceID, name)) +} + // SaveControllerConfig stores the resource controller config to storage. func (se *StorageEndpoint) SaveControllerConfig(config any) error { return se.saveJSON(keypath.ControllerConfigPath(), config) diff --git a/tests/integrations/mcs/resourcemanager/resource_manager_test.go b/tests/integrations/mcs/resourcemanager/resource_manager_test.go index d30a5cb372..3d95b2d7a2 100644 --- a/tests/integrations/mcs/resourcemanager/resource_manager_test.go +++ b/tests/integrations/mcs/resourcemanager/resource_manager_test.go @@ -22,6 +22,7 @@ import ( "math/rand/v2" "net/http" "reflect" + "sort" "strconv" "strings" "sync" @@ -227,6 +228,7 @@ func (suite *resourceManagerClientTestSuite) SetupSuite() { // Ensure RM service discovery has picked up the standalone endpoint before running tests. waitResourceManagerServiceURL(re, suite.client, true) } + waitAsyncLoadResourceGroups(re, suite.client) suite.initGroups = []*rmpb.ResourceGroup{ { @@ -305,6 +307,13 @@ func waitResourceManagerServiceURL(re *require.Assertions, cli pd.Client, wantNo }) } +func waitAsyncLoadResourceGroups(re *require.Assertions, cli pd.Client) { + testutil.Eventually(re, func() bool { + _, err := cli.ListResourceGroups(context.TODO()) + return err == nil + }, testutil.WithTickInterval(100*time.Millisecond)) +} + func TestSwitchModeDuringWorkload(t *testing.T) { for _, tc := range []struct { name string @@ -647,6 +656,7 @@ func (suite *resourceManagerClientTestSuite) resignAndWaitLeader(re *require.Ass newLeader := suite.cluster.GetServer(suite.cluster.WaitLeader()) re.NotNil(newLeader) waitLeaderServingClient(re, suite.client, newLeader.GetAddr()) + waitAsyncLoadResourceGroups(re, suite.client) } func (suite *resourceManagerClientTestSuite) TestWatchResourceGroup() { @@ -727,16 +737,21 @@ func (suite *resourceManagerClientTestSuite) TestWatchResourceGroup() { re.NoError(err) re.Contains(resp, "Success!") // Make sure the resource group active - meta, err = controller.GetResourceGroup(group.Name) - re.NotNil(meta) - re.NoError(err) + testutil.Eventually(re, func() bool { + meta, err = controller.GetResourceGroup(group.Name) + if err != nil || meta == nil { + return false + } + meta = controller.GetActiveResourceGroup(group.Name) + return meta != nil + }, testutil.WithTickInterval(50*time.Millisecond)) modifySettings(group, 30000) resp, err = cli.ModifyResourceGroup(suite.ctx, group) re.NoError(err) re.Contains(resp, "Success!") testutil.Eventually(re, func() bool { meta = controller.GetActiveResourceGroup(group.Name) - return meta.RUSettings.RU.Settings.FillRate == uint64(30000) + return meta != nil && meta.RUSettings.RU.Settings.FillRate == uint64(30000) }, testutil.WithTickInterval(100*time.Millisecond)) re.NoError(failpoint.Disable("github.com/tikv/pd/client/resource_group/controller/watchStreamError")) @@ -1599,13 +1614,40 @@ func (suite *resourceManagerClientTestSuite) TestBasicResourceGroupCURD() { // re-connect client as well suite.client = suite.setupPDClient(re) cli = suite.client + expectedGroups := normalizeResourceGroupsForSettingsCompare(groups) var newGroups []*rmpb.ResourceGroup testutil.Eventually(re, func() bool { var err error newGroups, err = cli.ListResourceGroups(suite.ctx) - return err == nil - }, testutil.WithWaitFor(time.Second)) - re.Equal(groups, newGroups) + return err == nil && reflect.DeepEqual(expectedGroups, normalizeResourceGroupsForSettingsCompare(newGroups)) + }) + re.Equal(expectedGroups, normalizeResourceGroupsForSettingsCompare(newGroups)) +} + +func normalizeResourceGroupsForSettingsCompare(groups []*rmpb.ResourceGroup) []*rmpb.ResourceGroup { + normalized := make([]*rmpb.ResourceGroup, 0, len(groups)) + for _, group := range groups { + cloned := typeutil.DeepClone(group, func() *rmpb.ResourceGroup { + return &rmpb.ResourceGroup{} + }) + cloned.RUStats = nil + resetTokenBucketRuntimeState(cloned.GetRUSettings().GetRU()) + rawSettings := cloned.GetRawResourceSettings() + resetTokenBucketRuntimeState(rawSettings.GetCpu()) + resetTokenBucketRuntimeState(rawSettings.GetIoRead()) + resetTokenBucketRuntimeState(rawSettings.GetIoWrite()) + normalized = append(normalized, cloned) + } + sort.Slice(normalized, func(i, j int) bool { + return normalized[i].GetName() < normalized[j].GetName() + }) + return normalized +} + +func resetTokenBucketRuntimeState(bucket *rmpb.TokenBucket) { + if bucket != nil { + bucket.Tokens = 0 + } } func (suite *resourceManagerClientTestSuite) TestResourceGroupRUConsumption() {