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

@@ -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,13 +166,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,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",