resource_group: implement async loading for resource groups

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

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

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 {

pkg/mcs/resourcemanager/server/keyspace_manager.go

@@ -67,11 +67,39 @@ type consumptionItem struct {
 	isTiFlash    bool
 }
 
+// reservedKind describes why a cached resource group entry is still
+// considered unconfirmed.
+type reservedKind int
+
+const (
+	// reservedPlaceholder marks an entry whose settings and state are both
+	// synthetic, e.g. the default group pre-inserted by
+	// ensureReservedDefaultGroupInCache before async loading runs. The async
+	// bulk merge may replace such an entry wholesale.
+	reservedPlaceholder reservedKind = iota
+	// reservedStateOnly marks an entry whose settings were confirmed from
+	// storage (and may have been modified and persisted since) but whose
+	// state failed to load. The async bulk merge must only adopt the scanned
+	// state into it, never replace its settings with the scan's older copy.
+	reservedStateOnly
+)
+
 type keyspaceResourceGroupManager struct {
 	syncutil.RWMutex
-	groups          map[string]*ResourceGroup
+	groups map[string]*ResourceGroup
+	// reservedGroups tracks names whose entry in groups is not yet fully
+	// confirmed by a storage load or a real write, together with how much of
+	// it is unconfirmed (see reservedKind). It shares the same lock as groups
+	// so a lazy load can atomically decide whether it's safe to replace the
+	// entry.
+	reservedGroups  map[string]reservedKind
 	groupRUTrackers map[string]*groupRUTracker
 	serviceLimiter  *serviceLimiter
+	// deleteGen is bumped under the write lock every time a group is removed
+	// from the cache. A lazy load snapshots it before its lock-free storage
+	// read and re-checks it under the lock before inserting, so a group
+	// deleted after the read is never resurrected by the now-stale result.
+	deleteGen uint64
 
 	keyspaceID uint32
 	storage    endpoint.ResourceGroupStorage
@@ -89,6 +117,7 @@ func newKeyspaceResourceGroupManager(
 	}
 	return &keyspaceResourceGroupManager{
 		groups:          make(map[string]*ResourceGroup),
+		reservedGroups:  make(map[string]reservedKind),
 		groupRUTrackers: make(map[string]*groupRUTracker),
 		keyspaceID:      keyspaceID,
 		storage:         storage,
@@ -159,13 +188,17 @@ 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
 	}
 
 	resourceGroup := FromProtoResourceGroup(group)
 	krgm.Lock()
 	krgm.groups[group.Name] = resourceGroup
+	delete(krgm.reservedGroups, group.Name)
 	krgm.Unlock()
 	krgm.syncBurstabilityWithServiceLimit(resourceGroup)
 	return nil
@@ -175,9 +208,20 @@ func (krgm *keyspaceResourceGroupManager) deleteResourceGroupFromCache(name stri
 	krgm.Lock()
 	delete(krgm.groups, name)
 	delete(krgm.groupRUTrackers, name)
+	delete(krgm.reservedGroups, name)
+	// Signal any in-flight lazy load that a delete happened, so it won't
+	// reinsert a copy read from storage before this deletion.
+	krgm.deleteGen++
 	krgm.Unlock()
 }
 
+// loadDeleteGen returns the current delete generation counter.
+func (krgm *keyspaceResourceGroupManager) loadDeleteGen() uint64 {
+	krgm.RLock()
+	defer krgm.RUnlock()
+	return krgm.deleteGen
+}
+
 func (krgm *keyspaceResourceGroupManager) setRawStatesIntoResourceGroup(name string, rawValue string) error {
 	tokens := &GroupStates{}
 	if err := json.Unmarshal([]byte(rawValue), tokens); err != nil {
@@ -196,8 +240,16 @@ func (krgm *keyspaceResourceGroupManager) setRawStatesIntoResourceGroup(name str
 func (krgm *keyspaceResourceGroupManager) initDefaultResourceGroup() {
 	krgm.RLock()
 	_, ok := krgm.groups[DefaultResourceGroupName]
+	kind, reserved := krgm.reservedGroups[DefaultResourceGroupName]
 	krgm.RUnlock()
-	if ok {
+	// A cached entry only makes initialization unnecessary if it's confirmed
+	// data, or at least has confirmed settings (reservedStateOnly), which must
+	// not be overwritten with synthetic ones. A reservedPlaceholder means
+	// nothing is persisted for the default group (e.g. a fresh store): it must
+	// still be created and persisted here, otherwise it would stay an
+	// unconfirmed placeholder forever, with its settings never stored and its
+	// state persistence permanently skipped.
+	if ok && (!reserved || kind == reservedStateOnly) {
 		return
 	}
 	defaultGroup := newDefaultResourceGroup()
@@ -218,6 +270,7 @@ func (krgm *keyspaceResourceGroupManager) ensureReservedDefaultGroupInCache() {
 	krgm.Lock()
 	if _, ok := krgm.groups[DefaultResourceGroupName]; !ok {
 		krgm.groups[DefaultResourceGroupName] = defaultGroup
+		krgm.reservedGroups[DefaultResourceGroupName] = reservedPlaceholder
 		inserted = true
 	}
 	krgm.Unlock()
@@ -245,6 +298,7 @@ func (krgm *keyspaceResourceGroupManager) restoreDefaultResourceGroupFromReserve
 	defaultGroup := newDefaultResourceGroup()
 	krgm.Lock()
 	krgm.groups[DefaultResourceGroupName] = defaultGroup
+	krgm.reservedGroups[DefaultResourceGroupName] = reservedPlaceholder
 	krgm.Unlock()
 	krgm.syncBurstabilityWithServiceLimit(defaultGroup)
 }
@@ -266,6 +320,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 +343,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 +388,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 +453,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",