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 5465e5dcf6..04243a584f 100644 --- a/pkg/errs/errno.go +++ b/pkg/errs/errno.go @@ -67,6 +67,16 @@ var ( ErrGRPCRateLimitExceeded = func(err error) error { return status.Error(codes.ResourceExhausted, err.Error()) } + // ErrResourceGroupsLoadingGRPC converts the retryable + // ErrResourceGroupsLoading into codes.Unavailable, so generic + // client-side retry logic can act on it instead of seeing an opaque + // codes.Unknown. Other errors pass through unchanged. + ErrResourceGroupsLoadingGRPC = func(err error) error { + if ErrResourceGroupsLoading.Equal(err) { + return status.Error(codes.Unavailable, err.Error()) + } + return err + } // FailedPrecondition indicates operation was rejected because the // system is not in a state required for the operation's execution. @@ -538,6 +548,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/metadataapi/config_service.go b/pkg/mcs/resourcemanager/metadataapi/config_service.go index e427ab4229..14ea310699 100644 --- a/pkg/mcs/resourcemanager/metadataapi/config_service.go +++ b/pkg/mcs/resourcemanager/metadataapi/config_service.go @@ -513,6 +513,13 @@ func (*ConfigService) respondStoreReadError(c *gin.Context, err error) { c.String(http.StatusNotFound, err.Error()) return } + // Resource groups are still being loaded asynchronously: the request can + // succeed once loading completes, so report it as retryable rather than as + // an internal error. + if errs.ErrResourceGroupsLoading.Equal(err) { + c.String(http.StatusServiceUnavailable, err.Error()) + return + } c.String(http.StatusInternalServerError, err.Error()) } @@ -525,5 +532,9 @@ func (*ConfigService) respondStoreWriteError(c *gin.Context, err error) { c.String(http.StatusNotFound, err.Error()) return } + if errs.ErrResourceGroupsLoading.Equal(err) { + c.String(http.StatusServiceUnavailable, err.Error()) + return + } c.String(http.StatusInternalServerError, err.Error()) } diff --git a/pkg/mcs/resourcemanager/metadataapi/config_service_test.go b/pkg/mcs/resourcemanager/metadataapi/config_service_test.go index 966b259a0d..d11bc47b5e 100644 --- a/pkg/mcs/resourcemanager/metadataapi/config_service_test.go +++ b/pkg/mcs/resourcemanager/metadataapi/config_service_test.go @@ -226,6 +226,26 @@ func TestConfigServiceGroupCRUDAndErrorCodes(t *testing.T) { } } +// TestConfigServiceLoadingReturns503 asserts the "resource groups are still +// being loaded" error is reported as retryable. It is a transient startup state, +// not an internal error, so callers and load balancers must be able to tell the +// difference and retry. +func TestConfigServiceLoadingReturns503(t *testing.T) { + t.Parallel() + + re := require.New(t) + store := newTestStore() + handler := newTestHTTPHandler(store) + + store.listErr = pderrors.ErrResourceGroupsLoading + resp := doJSONRequest(re, handler, http.MethodGet, "/resource-manager/api/v1/config/groups", nil) + re.Equal(http.StatusServiceUnavailable, resp.Code) + + store.listErr = errors.New("boom") + resp = doJSONRequest(re, handler, http.MethodGet, "/resource-manager/api/v1/config/groups", nil) + re.Equal(http.StatusInternalServerError, resp.Code) +} + func TestConfigServiceControllerAllOrNothing(t *testing.T) { t.Parallel() @@ -351,6 +371,7 @@ type testStore struct { serviceLimits map[uint32]float64 addErr error setServiceLimitErr error + listErr error updatedControllerConfigItems []string } @@ -406,7 +427,10 @@ func (s *testStore) GetResourceGroup(keyspaceID uint32, name string, withStats b return group.Clone(withStats), nil } -func (*testStore) GetResourceGroupList(_ uint32, _ bool) ([]*rmserver.ResourceGroup, error) { +func (s *testStore) GetResourceGroupList(_ uint32, _ bool) ([]*rmserver.ResourceGroup, error) { + if s.listErr != nil { + return nil, s.listErr + } return []*rmserver.ResourceGroup{}, nil } diff --git a/pkg/mcs/resourcemanager/server/grpc_service.go b/pkg/mcs/resourcemanager/server/grpc_service.go index 8bc9472d42..4e0340b538 100644 --- a/pkg/mcs/resourcemanager/server/grpc_service.go +++ b/pkg/mcs/resourcemanager/server/grpc_service.go @@ -110,7 +110,7 @@ func (s *Service) GetResourceGroup(_ context.Context, req *rmpb.GetResourceGroup keyspaceID := ExtractKeyspaceID(req.GetKeyspaceId()) rg, err := s.manager.GetResourceGroup(keyspaceID, req.ResourceGroupName, req.WithRuStats) if err != nil { - return nil, err + return nil, errs.ErrResourceGroupsLoadingGRPC(err) } if rg == nil { return nil, errs.ErrResourceGroupNotExists.FastGenByArgs(req.ResourceGroupName) @@ -129,7 +129,7 @@ func (s *Service) ListResourceGroups(_ context.Context, req *rmpb.ListResourceGr keyspaceID := ExtractKeyspaceID(req.GetKeyspaceId()) groups, err := s.manager.GetResourceGroupList(keyspaceID, req.WithRuStats) if err != nil { - return nil, err + return nil, errs.ErrResourceGroupsLoadingGRPC(err) } resps := &rmpb.ListResourceGroupsResponse{ Groups: make([]*rmpb.ResourceGroup, 0, len(groups)), @@ -151,7 +151,7 @@ func (s *Service) AddResourceGroup(_ context.Context, req *rmpb.PutResourceGroup } err := s.manager.AddResourceGroup(req.GetGroup()) if err != nil { - return nil, err + return nil, errs.ErrResourceGroupsLoadingGRPC(err) } return &rmpb.PutResourceGroupResponse{Body: "Success!"}, nil } @@ -166,7 +166,7 @@ func (s *Service) DeleteResourceGroup(_ context.Context, req *rmpb.DeleteResourc } err := s.manager.DeleteResourceGroup(ExtractKeyspaceID(req.GetKeyspaceId()), req.ResourceGroupName) if err != nil { - return nil, err + return nil, errs.ErrResourceGroupsLoadingGRPC(err) } return &rmpb.DeleteResourceGroupResponse{Body: "Success!"}, nil } @@ -181,7 +181,7 @@ func (s *Service) ModifyResourceGroup(_ context.Context, req *rmpb.PutResourceGr } err := s.manager.ModifyResourceGroup(req.GetGroup()) if err != nil { - return nil, err + return nil, errs.ErrResourceGroupsLoadingGRPC(err) } return &rmpb.PutResourceGroupResponse{Body: "Success!"}, nil } @@ -230,18 +230,25 @@ 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. + // Lazy loading can fail transiently (a storage read error, or retries + // exhausted while async loading is still in progress). Skip just this + // request instead of returning: the error belongs to one resource group, + // while returning would tear down the whole stream and force every + // client multiplexed on it to reconnect. + rg, err := s.manager.GetMutableResourceGroup(keyspaceID, resourceGroupName) + if rg == nil { + log.Warn("resource group is unavailable", 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..0e1773f61d 100644 --- a/pkg/mcs/resourcemanager/server/keyspace_manager.go +++ b/pkg/mcs/resourcemanager/server/keyspace_manager.go @@ -69,9 +69,26 @@ 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 a synthetic + // placeholder and not yet confirmed by a storage load or a real write. + reservedGroups map[string]struct{} 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 + // defaultGroupMu serializes every path that can create or persist the + // default group: on-demand synthesis (initDefaultResourceGroup) and a + // real Add/ModifyResourceGroup targeting "default" both hold it across + // their persist step. Without this, a synthetic write and a concurrent + // customized write race independently of krgm's RWMutex (which is only + // held for the cache mutation, not the storage I/O before it), so + // whichever one's storage/cache write lands last wins even if it + // started first - silently discarding a successful customized write. + defaultGroupMu syncutil.Mutex keyspaceID uint32 storage endpoint.ResourceGroupStorage @@ -89,6 +106,7 @@ func newKeyspaceResourceGroupManager( } return &keyspaceResourceGroupManager{ groups: make(map[string]*ResourceGroup), + reservedGroups: make(map[string]struct{}), groupRUTrackers: make(map[string]*groupRUTracker), keyspaceID: keyspaceID, storage: storage, @@ -154,28 +172,62 @@ func (krgm *keyspaceResourceGroupManager) upsertResourceGroupFromRaw(name string existing := krgm.groups[group.Name] krgm.RUnlock() if existing != nil { - if err := existing.ApplySettings(group); err != nil { + // Merge the settings patch and the burst-limit sync into a single + // rg.Lock() critical section: ApplySettings can flip a group from + // bounded to unbounded (e.g. a raw burst limit going negative) + // without touching overrideBurstLimit at all, so applying them as two + // separate critical sections would let a concurrent token request + // observe the now-unbounded settings before the sync catches up and + // bypass an active keyspace service limit in between. + serviceLimit, isSet := krgm.getServiceLimit() + existing.Lock() + patchErr := existing.patchSettingsLocked(group, false) + if patchErr == nil { + applyBurstabilitySyncLocked(existing, serviceLimit, isSet) + } + existing.Unlock() + if patchErr != nil { log.Error("failed to apply the keyspace resource group settings from raw value", - zap.Uint32("keyspace-id", krgm.keyspaceID), zap.String("name", name), zap.String("raw-value", rawValue), zap.Error(err)) - return err + zap.Uint32("keyspace-id", krgm.keyspaceID), zap.String("name", name), zap.String("raw-value", rawValue), zap.Error(patchErr)) + return patchErr } - krgm.syncBurstabilityWithServiceLimit(existing) + krgm.Lock() + delete(krgm.reservedGroups, group.Name) + krgm.Unlock() return nil } resourceGroup := FromProtoResourceGroup(group) krgm.Lock() krgm.groups[group.Name] = resourceGroup + delete(krgm.reservedGroups, group.Name) + krgm.syncBurstabilityWithServiceLimitLocked(resourceGroup) krgm.Unlock() - krgm.syncBurstabilityWithServiceLimit(resourceGroup) return nil } func (krgm *keyspaceResourceGroupManager) deleteResourceGroupFromCache(name string) { krgm.Lock() + krgm.removeResourceGroupLocked(name) + krgm.Unlock() +} + +// removeResourceGroupLocked removes every cache trace of the group. The +// caller must hold the write lock. +func (krgm *keyspaceResourceGroupManager) removeResourceGroupLocked(name string) { delete(krgm.groups, name) delete(krgm.groupRUTrackers, name) - krgm.Unlock() + 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++ +} + +// 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 { @@ -193,19 +245,6 @@ func (krgm *keyspaceResourceGroupManager) setRawStatesIntoResourceGroup(name str return nil } -func (krgm *keyspaceResourceGroupManager) initDefaultResourceGroup() { - krgm.RLock() - _, ok := krgm.groups[DefaultResourceGroupName] - krgm.RUnlock() - if ok { - return - } - defaultGroup := newDefaultResourceGroup() - if err := krgm.addResourceGroup(defaultGroup.IntoProtoResourceGroup(krgm.keyspaceID)); err != nil { - log.Warn("init default group failed", zap.Uint32("keyspace-id", krgm.keyspaceID), zap.Error(err)) - } -} - func (krgm *keyspaceResourceGroupManager) ensureReservedDefaultGroupInCache() { krgm.RLock() _, ok := krgm.groups[DefaultResourceGroupName] @@ -214,16 +253,13 @@ func (krgm *keyspaceResourceGroupManager) ensureReservedDefaultGroupInCache() { return } defaultGroup := newDefaultResourceGroup() - inserted := false krgm.Lock() if _, ok := krgm.groups[DefaultResourceGroupName]; !ok { krgm.groups[DefaultResourceGroupName] = defaultGroup - inserted = true + krgm.reservedGroups[DefaultResourceGroupName] = struct{}{} + krgm.syncBurstabilityWithServiceLimitLocked(defaultGroup) } krgm.Unlock() - if inserted { - krgm.syncBurstabilityWithServiceLimit(defaultGroup) - } } func newDefaultResourceGroup() *ResourceGroup { @@ -245,53 +281,79 @@ func (krgm *keyspaceResourceGroupManager) restoreDefaultResourceGroupFromReserve defaultGroup := newDefaultResourceGroup() krgm.Lock() krgm.groups[DefaultResourceGroupName] = defaultGroup + krgm.reservedGroups[DefaultResourceGroupName] = struct{}{} + krgm.syncBurstabilityWithServiceLimitLocked(defaultGroup) krgm.Unlock() - krgm.syncBurstabilityWithServiceLimit(defaultGroup) } -func (krgm *keyspaceResourceGroupManager) addResourceGroup(grouppb *rmpb.ResourceGroup) error { +// persistResourceGroup validates grouppb, builds the in-memory group, and +// persists its settings and states to storage. It's the storage-only phase of +// an Add: the caller is responsible for publishing the returned group into +// whichever keyspace manager is current at publish time. +func (krgm *keyspaceResourceGroupManager) persistResourceGroup(grouppb *rmpb.ResourceGroup) (*ResourceGroup, error) { if err := validateResourceGroupProto(grouppb); err != nil { - return err + return nil, err } group := FromProtoResourceGroup(grouppb) if krgm.writeRole.AllowsMetadataWrite() { if err := group.persistSettings(krgm.keyspaceID, krgm.storage); err != nil { - return err + return nil, err } } if krgm.writeRole.AllowsStateWrite() { if err := group.persistStates(krgm.keyspaceID, krgm.storage); err != nil { - return err + return nil, err } } + return group, nil +} + +func (krgm *keyspaceResourceGroupManager) addResourceGroup(grouppb *rmpb.ResourceGroup) error { + group, err := krgm.persistResourceGroup(grouppb) + if err != nil { + return err + } krgm.Lock() krgm.groups[group.Name] = group + delete(krgm.reservedGroups, group.Name) + krgm.syncBurstabilityWithServiceLimitLocked(group) krgm.Unlock() - krgm.syncBurstabilityWithServiceLimit(group) return nil } -func (krgm *keyspaceResourceGroupManager) modifyResourceGroup(group *rmpb.ResourceGroup) error { +// modifyResourceGroup patches the cached group's settings and persists them, +// returning the patched group so the caller can republish it if the cache it +// came from is no longer the live one. +func (krgm *keyspaceResourceGroupManager) modifyResourceGroup(group *rmpb.ResourceGroup) (*ResourceGroup, error) { if group == nil || group.Name == "" { - return errs.ErrInvalidGroup.FastGenByArgs("the group name") + return nil, errs.ErrInvalidGroup.FastGenByArgs("the group name") } krgm.RLock() curGroup, ok := krgm.groups[group.Name] krgm.RUnlock() if !ok { - return errs.ErrResourceGroupNotExists.FastGenByArgs(group.Name) + return nil, errs.ErrResourceGroupNotExists.FastGenByArgs(group.Name) } if !krgm.writeRole.AllowsMetadataWrite() { - return errMetadataWriteDisabled + return nil, errMetadataWriteDisabled } - err := curGroup.PatchSettings(group) - if err != nil { - return err + if err := curGroup.PatchSettings(group); err != nil { + return nil, err } - return curGroup.persistSettings(krgm.keyspaceID, krgm.storage) + // 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. + if err := curGroup.persistSettings(krgm.keyspaceID, krgm.storage); err != nil { + return nil, err + } + return curGroup, nil } -func (krgm *keyspaceResourceGroupManager) deleteResourceGroup(name string) error { +// deleteResourceGroupFromStorage validates the request and removes the +// group's settings and states from storage. It's the storage-only phase of a +// Delete: the caller is responsible for removing the group from whichever +// keyspace manager cache is current at publish time. +func (krgm *keyspaceResourceGroupManager) deleteResourceGroupFromStorage(name string) error { if name == DefaultResourceGroupName { return errs.ErrDeleteReservedGroup } @@ -326,10 +388,55 @@ func (krgm *keyspaceResourceGroupManager) deleteResourceGroup(name string) error zap.Error(err)) } } + return nil +} + +func (krgm *keyspaceResourceGroupManager) deleteResourceGroup(name string) error { + if err := krgm.deleteResourceGroupFromStorage(name); err != nil { + return err + } krgm.deleteResourceGroupFromCache(name) 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 +} + +// hasConfirmedResourceGroup reports whether name is cached as confirmed data, +// i.e. present and not a reserved placeholder. It answers both questions under +// a single read lock, since it sits on the lazy-loading fast path that every +// point and token request takes while async loading is still in progress. +func (krgm *keyspaceResourceGroupManager) hasConfirmedResourceGroup(name string) bool { + krgm.RLock() + defer krgm.RUnlock() + _, ok := krgm.confirmedResourceGroupLocked(name) + return ok +} + +// confirmedResourceGroupLocked is hasConfirmedResourceGroup for a caller that +// already holds krgm's lock (read or write), returning the group itself too +// so a caller that needs both the confirmed check and the group (e.g. the +// state persist loop below) doesn't have to re-derive the same "present and +// not a reserved placeholder" logic or take a second lock round trip. +func (krgm *keyspaceResourceGroupManager) confirmedResourceGroupLocked(name string) (*ResourceGroup, bool) { + group, ok := krgm.groups[name] + if !ok { + return nil, false + } + if _, reserved := krgm.reservedGroups[name]; reserved { + return nil, false + } + return group, true +} + func (krgm *keyspaceResourceGroupManager) getResourceGroup(name string, withStats bool) *ResourceGroup { krgm.RLock() defer krgm.RUnlock() @@ -383,7 +490,11 @@ func (krgm *keyspaceResourceGroupManager) persistResourceGroupRunningState() { krgm.RUnlock() for idx := range keys { krgm.RLock() - group, ok := krgm.groups[keys[idx]] + // A reserved placeholder (e.g. the synthetic default installed before + // async loading completes) is skipped: persisting its fresh state + // would permanently overwrite any real persisted state still waiting + // to be loaded. + group, ok := krgm.confirmedResourceGroupLocked(keys[idx]) if ok { if err := group.persistStates(krgm.keyspaceID, krgm.storage); err != nil { log.Error("persist keyspace resource group state failed", @@ -432,6 +543,12 @@ func (krgm *keyspaceResourceGroupManager) getServiceLimiter() *serviceLimiter { func (krgm *keyspaceResourceGroupManager) getServiceLimit() (float64, bool) { krgm.RLock() defer krgm.RUnlock() + return krgm.getServiceLimitLocked() +} + +// getServiceLimitLocked is getServiceLimit for a caller that already holds +// krgm's lock (read or write). +func (krgm *keyspaceResourceGroupManager) getServiceLimitLocked() (float64, bool) { if krgm.serviceLimiter == nil { return 0, false } @@ -883,14 +1000,64 @@ func (krgm *keyspaceResourceGroupManager) cleanupOverrides() { // Newly loaded groups can miss the initial service-limit replay, so apply the // same baseline burst invalidation when they enter the cache. func (krgm *keyspaceResourceGroupManager) syncBurstabilityWithServiceLimit(group *ResourceGroup) { - if group == nil || group.getBurstLimit(true) >= 0 || group.getOverrideBurstLimit() >= 0 { + krgm.RLock() + defer krgm.RUnlock() + krgm.syncBurstabilityWithServiceLimitLocked(group) +} + +// syncBurstabilityWithServiceLimitLocked is syncBurstabilityWithServiceLimit +// for a caller that already holds krgm's lock (read or write). It must never +// call back into a krgm-locking helper (e.g. getServiceLimit/syncBurstabilityWithServiceLimit +// itself): krgm's lock isn't reentrant, so doing so from the same goroutine +// would deadlock. Taking group's own lock below is safe while holding krgm's: +// group has no back-reference to krgm, so this one-directional nesting +// (krgm's lock outer, group's lock inner) can't cycle with any other lock +// ordering in this package. +// +// Callers use this to make a group's cache insertion and its burst-limit +// sync take effect atomically under a single krgm.Lock() critical section: +// without it, a concurrent token request could observe the group between the +// insert and a separately-locked sync call and read its unsynced, possibly +// unlimited/negative burst setting, letting it bypass an active keyspace +// service limit until the sync catches up. +func (krgm *keyspaceResourceGroupManager) syncBurstabilityWithServiceLimitLocked(group *ResourceGroup) { + if group == nil { + return + } + serviceLimit, isSet := krgm.getServiceLimitLocked() + // Cheap short-circuit before taking group's write lock: most keyspaces + // have no active service limit at all, and this runs on every group + // insert/update across the whole system (including the async bulk + // merge's up-to-500k-group batches), so skipping the write lock entirely + // in the common no-op case avoids needless contention against concurrent + // RequestRU calls on the same, already-live group. + if !isSet || serviceLimit <= 0 { + return + } + group.Lock() + defer group.Unlock() + applyBurstabilitySyncLocked(group, serviceLimit, isSet) +} + +// applyBurstabilitySyncLocked is the group-lock-only core shared by both +// syncBurstabilityWithServiceLimitLocked (krgm-lock callers) and +// upsertResourceGroupFromRaw's in-place update path (which merges this into +// the same rg.Lock() critical section as ApplySettings, so a settings patch +// that newly makes a group unbounded can never become visible without its +// burst override already applied). The caller must already hold group's +// lock; serviceLimit/isSet come from a plain krgm.getServiceLimit() read +// taken before acquiring it - service-limit staleness of a few instructions +// here is the same self-healing eventual consistency already accepted +// elsewhere in this file (e.g. invalidateBurstability), unlike the group's +// own settings-vs-override visibility this function exists to make atomic. +func applyBurstabilitySyncLocked(group *ResourceGroup, serviceLimit float64, isSet bool) { + if group == nil || group.getBurstLimitLocked(true) >= 0 || group.getOverrideBurstLimitLocked() >= 0 { return } - serviceLimit, isSet := krgm.getServiceLimit() if !isSet || serviceLimit <= 0 { return } - group.overrideBurstLimit(int64(serviceLimit)) + group.overrideBurstLimitLocked(int64(serviceLimit)) } // Since the burstable resource groups won't require tokens from the server anymore, diff --git a/pkg/mcs/resourcemanager/server/keyspace_manager_test.go b/pkg/mcs/resourcemanager/server/keyspace_manager_test.go index 5c8f30e6a6..e3978ed07b 100644 --- a/pkg/mcs/resourcemanager/server/keyspace_manager_test.go +++ b/pkg/mcs/resourcemanager/server/keyspace_manager_test.go @@ -86,8 +86,14 @@ func TestInitDefaultResourceGroup(t *testing.T) { _, exists := krgm.groups[DefaultResourceGroupName] re.False(exists) - // Initialize the default resource group. - krgm.initDefaultResourceGroup() + // Initialize the default resource group. initDefaultResourceGroup now + // publishes through Manager.publishResourceGroupMutation, so it needs a + // Manager with krgm registered as the live entry for keyspace 1. + m := prepareManager() + m.krgms[1] = krgm + created, err := m.initDefaultResourceGroup(1, krgm, nil) + re.NoError(err) + re.True(created) // Verify the default resource group is created. defaultGroup, exists := krgm.groups[DefaultResourceGroupName] @@ -101,6 +107,41 @@ func TestInitDefaultResourceGroup(t *testing.T) { re.Equal(int64(UnlimitedBurstLimit), defaultGroup.getBurstLimit()) } +// TestSyncBurstabilityWithServiceLimitLockedSkipsGroupLockWhenNoServiceLimit +// guards against syncBurstabilityWithServiceLimitLocked unconditionally +// taking group's write lock before checking whether the keyspace even has an +// active service limit. This runs on every group insert/update across the +// system, including the async bulk merge's up-to-500k-group batches, so +// escalating to a write lock in the common no-op case (most keyspaces have +// no service limit set) is needless contention against a concurrent +// RequestRU call already holding the same group's lock. The fix checks +// isSet/serviceLimit first (cheap: it only reads krgm's already-locked +// state) and returns before ever touching group's lock when there is +// nothing to apply. +func TestSyncBurstabilityWithServiceLimitLockedSkipsGroupLockWhenNoServiceLimit(t *testing.T) { + krgm := newKeyspaceResourceGroupManager(1, storage.NewStorageWithMemoryBackend()) + group := newDefaultResourceGroup() + + // No service limit has been configured on krgm, so + // syncBurstabilityWithServiceLimitLocked has nothing to do here. Hold + // group's lock externally to prove it's never contended. + group.Lock() + defer group.Unlock() + + done := make(chan struct{}) + go func() { + krgm.RLock() + krgm.syncBurstabilityWithServiceLimitLocked(group) + krgm.RUnlock() + close(done) + }() + select { + case <-done: + case <-time.After(200 * time.Millisecond): + t.Fatal("must not try to take group's lock when the keyspace has no active service limit") + } +} + func TestAddResourceGroup(t *testing.T) { re := require.New(t) @@ -187,7 +228,7 @@ func TestModifyResourceGroup(t *testing.T) { }, }, } - err = krgm.modifyResourceGroup(modifiedGroup) + _, err = krgm.modifyResourceGroup(modifiedGroup) re.NoError(err) // Verify the group was modified. @@ -206,7 +247,7 @@ func TestModifyResourceGroup(t *testing.T) { Name: "non_existent", Mode: rmpb.GroupMode_RUMode, } - err = krgm.modifyResourceGroup(nonExistentGroup) + _, err = krgm.modifyResourceGroup(nonExistentGroup) re.Error(err) } @@ -225,7 +266,10 @@ func TestDeleteResourceGroupBehavior(t *testing.T) { _, ok := krgm.groupRUTrackers[group.GetName()] re.False(ok) - krgm.initDefaultResourceGroup() + m := prepareManager() + m.krgms[1] = krgm + _, err := m.initDefaultResourceGroup(1, krgm, nil) + re.NoError(err) re.Error(krgm.deleteResourceGroup(DefaultResourceGroupName)) re.NotNil(krgm.getResourceGroup(DefaultResourceGroupName, false)) }) @@ -361,7 +405,10 @@ func TestGetResourceGroupList(t *testing.T) { re.Equal("group2", groups[1].Name) re.Equal("group3", groups[2].Name) - krgm.initDefaultResourceGroup() + m := prepareManager() + m.krgms[1] = krgm + _, err := m.initDefaultResourceGroup(1, krgm, nil) + re.NoError(err) groups = krgm.getResourceGroupList(false, true) re.Len(groups, 4) groups = krgm.getResourceGroupList(false, false) diff --git a/pkg/mcs/resourcemanager/server/manager.go b/pkg/mcs/resourcemanager/server/manager.go index 57d1a818ff..8b65f33e84 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,40 @@ type Manager struct { metrics *metrics // ruCollector is used to collect the RU metering data. ruCollector *ruCollector + // async loading state management + loadingState atomic.Int32 + // syncLoadedGroups records groups that were loaded synchronously (e.g., by lazy loading) + syncLoadedGroups map[trackerKey]bool + // loadEpoch is bumped (under the manager lock) every time initMetadata + // resets the loading state for a new term. The async loader captures it at + // start and re-checks it before every shared-state mutation, so a loader + // from a previous term that was blocked in a storage scan can never merge + // stale data into, or publish completion for, a newer term. + loadEpoch uint64 + // serviceLimitLocks serializes SetKeyspaceServiceLimit calls per keyspace + // ID. Unlike krgm-scoped locks (e.g. defaultGroupMu), this lives on the + // Manager itself and is never reset by Init, so it keeps serializing a + // call parked mid-persist in an old term against a competing call in a + // new term for the same keyspace - without it, the old call's eventual + // publish could overwrite a newer value the new-term call already both + // persisted and published. Entries are intentionally never removed: the + // keyspace ID space is bounded, matching krgms/keyspaceNameLookup's own + // grow-only lifetime, so there's no working-set pressure to justify the + // extra bookkeeping (and its own subtleties - see LockGroup's Unlock) + // that WithRemoveEntryOnUnlock would add. + serviceLimitLocks *syncutil.LockGroup } +// 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 { @@ -140,7 +173,7 @@ type metadataWatcherProvider interface { } func newManagerBase(controllerConfig *ControllerConfig, writeRole ResourceGroupWriteRole) *Manager { - return &Manager{ + m := &Manager{ writeRole: writeRole, controllerConfig: controllerConfig, krgms: make(map[uint32]*keyspaceResourceGroupManager), @@ -149,7 +182,22 @@ func newManagerBase(controllerConfig *ControllerConfig, writeRole ResourceGroupW keyspaceIDLookup: make(map[string]uint32), metrics: newMetrics(), ruCollector: newRUCollector(), + syncLoadedGroups: make(map[trackerKey]bool), + serviceLimitLocks: syncutil.NewLockGroup(), } + m.setLoadingState(LoadingStateNotStarted) + return m +} + +// setLoadingState publishes the loading state both to the atomic field the +// serving paths read and to the gauge operators alert on. +func (m *Manager) setLoadingState(state int32) { + m.loadingState.Store(state) + resourceGroupLoadingStateGauge.Set(float64(state)) +} + +func (m *Manager) getLoadingState() int32 { + return m.loadingState.Load() } // NewManager returns a new manager base on the given server, @@ -194,7 +242,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 @@ -237,8 +288,41 @@ func (m *Manager) SetKeyspaceServiceLimit(keyspaceID uint32, serviceLimit float6 if !m.writeRole.AllowsMetadataWrite() { return errMetadataWriteDisabled } + // Serialize against every other SetKeyspaceServiceLimit call for this + // keyspace, including ones that started in a different term: krgm + // identity resets on every Init, but this lock doesn't, so it's what + // keeps a call parked mid-persist in an old term from clobbering a + // competing new-term call's result once it resumes. See the + // serviceLimitLocks field doc for why a krgm-scoped lock (like + // defaultGroupMu) can't provide this guarantee on its own. + m.serviceLimitLocks.Lock(keyspaceID) + defer m.serviceLimitLocks.Unlock(keyspaceID) + // If the keyspace is not found, create a new keyspace resource group manager. - m.getOrCreateKeyspaceResourceGroupManager(keyspaceID, true).setServiceLimit(serviceLimit) + krgm := m.getOrCreateKeyspaceResourceGroupManager(keyspaceID, true) + failpoint.InjectCall("setServiceLimitBeforeStorage") + // Storage phase: this persists synchronously inside setServiceLimit. + krgm.setServiceLimit(serviceLimit) + // Publish phase: mirror the persisted value into whichever keyspace + // manager is current now, in case Init replaced the whole krgms map for + // a new term while the storage write above was in flight - without + // this, the write still lands in storage but only updates the detached + // krgm's in-memory limiter, leaving the live serving cache (and + // GetKeyspaceServiceLimiter) showing the stale pre-write value until + // the next full reload. This mirrors the storage-then-publish shape of + // publishResourceGroupMutation, but there's no reserved/confirmed + // distinction to protect here: a service limit is a single scalar with + // no CAS at the storage layer, so nothing short of the serviceLimitLocks + // guard above tells this call whether cur already holds a newer value. + // With that guard held for the whole call, no other SetKeyspaceServiceLimit + // for this keyspace can run concurrently, so whatever's in storage now is + // still what this call itself just wrote - safe to mirror unconditionally. + m.Lock() + cur := m.getOrCreateKeyspaceResourceGroupManagerLocked(keyspaceID) + m.Unlock() + if cur != krgm { + cur.setServiceLimitFromStorage(serviceLimit) + } return nil } @@ -248,23 +332,42 @@ func (m *Manager) SetKeyspaceRUVersion(keyspaceID uint32, ruVersion int32) error return errMetadataWriteDisabled } m.Lock() - if m.controllerConfig.RUVersionPolicy == nil { + defer m.Unlock() + // Mutate a clone, not m.controllerConfig directly, and save it before + // publishing it to the field - matching UpdateControllerConfigItem's + // pattern below. Mutating the live object in place and saving it after + // unlocking (the previous approach here) left RUVersionPolicy.Overrides, + // a plain map, both mutable by a concurrent call to this same function and + // readable by this call's own unlocked marshal - an unsynchronized + // concurrent map read/write. It also let this call's save land out of + // order relative to UpdateControllerConfigItem's, which clones, saves, + // and publishes fully inside its own lock: a save delayed past unlock + // here could persist a stale snapshot over a newer config + // UpdateControllerConfigItem had already committed. Holding the lock + // across the storage write closes both: no other config mutator can run + // until this call's save and publish are done. + controllerConfig := cloneControllerConfig(m.controllerConfig) + if controllerConfig.RUVersionPolicy == nil { // DefaultRUVersion (v1) means no RU model change. // There is currently no API to modify this global default; it is // intentionally fixed so that only per-keyspace overrides drive version bumps. - m.controllerConfig.RUVersionPolicy = &RUVersionPolicy{Default: DefaultRUVersion} + controllerConfig.RUVersionPolicy = &RUVersionPolicy{Default: DefaultRUVersion} } - if m.controllerConfig.RUVersionPolicy.Overrides == nil { - m.controllerConfig.RUVersionPolicy.Overrides = make(map[uint32]RUVersion) + if controllerConfig.RUVersionPolicy.Overrides == nil { + controllerConfig.RUVersionPolicy.Overrides = make(map[uint32]RUVersion) } - defaultVersion := m.controllerConfig.RUVersionPolicy.Default + defaultVersion := controllerConfig.RUVersionPolicy.Default if ruVersion == defaultVersion { - delete(m.controllerConfig.RUVersionPolicy.Overrides, keyspaceID) + delete(controllerConfig.RUVersionPolicy.Overrides, keyspaceID) } else { - m.controllerConfig.RUVersionPolicy.Overrides[keyspaceID] = ruVersion + controllerConfig.RUVersionPolicy.Overrides[keyspaceID] = ruVersion } - m.Unlock() - return m.storage.SaveControllerConfig(m.controllerConfig) + failpoint.InjectCall("setKeyspaceRUVersionBeforeSave") + if err := m.storage.SaveControllerConfig(controllerConfig); err != nil { + return err + } + m.controllerConfig = controllerConfig + return nil } // GetRUVersionPolicy returns a deep copy of the current RU version policy from the controller config. @@ -275,18 +378,54 @@ 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 := m.getOrCreateKeyspaceResourceGroupManagerLocked(keyspaceID) + m.Unlock() + if initDefault { + // In metadata-watcher mode, LoadingStateCompleted only means the + // initial watcher bootstrap finished, not that every subsequent PD + // write has been observed yet - same caveat loadResourceGroupIfNeeded + // documents at its own entry point. Without this guard, a request + // landing right after bootstrap but before the watcher delivers a + // just-written customized default would see nothing cached, take + // this branch, and persist the built-in default over it. Always fall + // through to loadResourceGroupIfNeeded in that mode, which does a + // storage point load first instead of trusting the cache. + if !m.enableMetadataWatcher && m.getLoadingState() == LoadingStateCompleted { + // Async loading 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. This is a best-effort + // pre-warm: any failure (stale term or a real storage error) is + // already logged inside initDefaultResourceGroup, and a later + // request for the default group will retry through + // loadResourceGroupIfNeeded, which does surface such errors. + _, _ = m.initDefaultResourceGroup(keyspaceID, krgm, func() bool { + return m.isKeyspaceManagerCurrent(keyspaceID, krgm) + }) + } 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 +} + +// getOrCreateKeyspaceResourceGroupManagerLocked is the m.krgms lookup/create +// step of getOrCreateKeyspaceResourceGroupManager for a caller that already +// holds m.Lock(). +func (m *Manager) getOrCreateKeyspaceResourceGroupManagerLocked(keyspaceID uint32) *keyspaceResourceGroupManager { krgm, ok := m.krgms[keyspaceID] if !ok { krgm = newKeyspaceResourceGroupManager(keyspaceID, m.storage, m.writeRole) m.krgms[keyspaceID] = krgm } - m.Unlock() - // Init the default resource group if needed. - if initDefault { - krgm.initDefaultResourceGroup() - } return krgm } @@ -325,13 +464,22 @@ func (m *Manager) Init(ctx context.Context) error { m.wg.Wait() return err } + // No async loader runs in this mode, so nothing will ever consume the + // sync-loaded markers. Drop the map: otherwise every watcher event keeps + // adding entries that are never removed, including for deleted groups. + m.Lock() + m.syncLoadedGroups = nil + m.Unlock() + m.setLoadingState(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. @@ -356,47 +504,343 @@ func (m *Manager) initControllerConfig() error { log.Error("resource controller config load failed", zap.Error(err), zap.String("v", v)) return err } - if err = json.Unmarshal([]byte(v), &m.controllerConfig); err != nil { + // Clone, merge, save, and publish all under the same lock, matching + // SetKeyspaceRUVersion/UpdateControllerConfigItem's pattern: this makes + // initControllerConfig's own re-save fully serialized against every + // other config mutator too, instead of only against itself via ordering + // (save-before-publish). Without this, a concurrent SetKeyspaceRUVersion + // or UpdateControllerConfigItem call could commit its change to storage + // and m.controllerConfig in between this call's unlocked read/merge and + // its own unlocked save, and this call's now-stale snapshot would + // silently overwrite it - the same class of lost update fixed for + // SetKeyspaceRUVersion. Requests are gated behind IsServing(), which + // only flips true after Init (and thus this call) returns, so holding + // the lock across the save here does not add contention with live + // request traffic in the common case. + m.Lock() + defer m.Unlock() + controllerConfig := cloneControllerConfig(m.controllerConfig) + if err = json.Unmarshal([]byte(v), &controllerConfig); err != nil { log.Warn("un-marshall controller config failed, fallback to default", zap.Error(err), zap.String("v", v)) } - - // re-save the config to make sure the config has been persisted. if m.writeRole.AllowsMetadataWrite() { - if err := m.storage.SaveControllerConfig(m.controllerConfig); err != nil { + if err := m.storage.SaveControllerConfig(controllerConfig); err != nil { return err } } + m.controllerConfig = controllerConfig 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) + m.loadEpoch++ + epoch := m.loadEpoch + m.setLoadingState(LoadingStateNotStarted) + m.Unlock() + + m.initReservedInCache() + if err := m.loadServiceLimits(); err != nil { + return err + } + + m.wg.Add(1) + go m.asyncLoadResourceGroups(ctx, epoch) + return nil +} + +// maxServiceLimitReloadAttempts bounds the point re-read retries in +// loadServiceLimits below before it falls back to the bulk-scanned value. +const maxServiceLimitReloadAttempts = 3 + +func (m *Manager) loadServiceLimits() error { + return m.storage.LoadServiceLimits(func(keyspaceID uint32, bulkScannedLimit float64) { + failpoint.InjectCall("loadServiceLimitsBeforeApply", keyspaceID) + // Serialize against SetKeyspaceServiceLimit for this keyspace, and + // re-read the value from storage instead of trusting the one the + // bulk scan above already fetched: that value was read before this + // lock was acquired, so a concurrent SetKeyspaceServiceLimit call + // (from either this term or a still-in-flight previous one) can + // persist and mirror a newer value in between, which this callback + // would otherwise silently clobber with its now-stale snapshot. A + // fresh point read taken under the same lock SetKeyspaceServiceLimit + // itself uses is guaranteed to observe whichever write actually + // landed last, since setServiceLimit's storage write always + // completes before SetKeyspaceServiceLimit releases this lock. + m.serviceLimitLocks.Lock(keyspaceID) + defer m.serviceLimitLocks.Unlock(keyspaceID) + krgm := m.getOrCreateKeyspaceResourceGroupManager(keyspaceID, false) + var ( + serviceLimit float64 + err error + ) + for attempt := 1; attempt <= maxServiceLimitReloadAttempts; attempt++ { + serviceLimit, err = m.storage.LoadServiceLimit(keyspaceID) + if err == nil { + break + } + log.Warn("failed to reload service limit, retrying", + zap.Uint32("keyspace-id", keyspaceID), zap.Int("attempt", attempt), zap.Error(err)) + } + if err != nil { + // Retries exhausted, e.g. a persistent storage failure. A + // concurrent SetKeyspaceServiceLimit call for this keyspace holds + // this same serviceLimitLocks entry across both its storage write + // and its cache mirror step, so there is no partially-applied + // state it could leave behind: by the time we hold the lock, it + // has either fully landed (cache already reflects a write that is + // strictly newer than the bulk scan's pre-lock snapshot) or has + // not started yet (nothing to race with the fallback below). If + // the cache already carries such a value, leave it untouched + // instead of clobbering it with the stale bulk-scanned one. + // + // Known gap: a service limit has no separate "has this ever been + // explicitly configured" flag - getServiceLimit's isSet also + // reads as false for a limit explicitly set to 0 (unlimited), the + // same as a never-configured one. So a concurrent + // SetKeyspaceServiceLimit(keyspaceID, 0) that lands in this same + // window is indistinguishable here from "nothing set yet", and + // the fallback below would incorrectly reapply the stale + // bulk-scanned (pre-clear) value over the just-cleared limit. + // Narrower than the gap this whole retry/fallback exists to + // close - it additionally requires the point read to keep + // failing for the full retry budget - and not closed here. + if _, alreadySet := krgm.getServiceLimit(); alreadySet { + log.Warn("giving up reloading service limit; a newer value is already cached, leaving it untouched", + zap.Uint32("keyspace-id", keyspaceID), zap.Error(err)) + return + } + // Nothing newer is cached yet: fall back to the bulk-scanned + // value instead of dropping the update entirely - strictly + // better than leaving the keyspace with no service limit cached + // at all (silently allowing burstable groups to bypass it) until + // the next Init. + log.Error("giving up reloading service limit, falling back to the bulk-scanned value", + zap.Uint32("keyspace-id", keyspaceID), zap.Error(err)) + serviceLimit = bulkScannedLimit + } + krgm.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 + m.setLoadingState(LoadingStateCompleted) + epoch := m.loadEpoch m.Unlock() - // Load keyspace resource group meta info from the storage. + // This runs to completion before the manager is exposed to any request + // (NewMetadataOnlyManager's caller has not returned yet), so there is no + // live writer to race against; eagerly confirming every loaded keyspace's + // default here is safe, unlike the same backfill from the async loader. + m.initReserved(epoch) + return m.loadServiceLimits() +} + +// storeLoadingStateIfCurrent stores the loading state only if the manager has +// not been reinitialized since the loader with the given epoch started. It +// returns false when the loader is stale and must exit without touching any +// further shared state. +func (m *Manager) storeLoadingStateIfCurrent(epoch uint64, state int32) bool { + m.Lock() + defer m.Unlock() + if m.loadEpoch != epoch { + return false + } + m.setLoadingState(state) + return true +} + +func (m *Manager) asyncLoadResourceGroups(ctx context.Context, epoch uint64) { + 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: + } + } + + if !m.storeLoadingStateIfCurrent(epoch, LoadingStateInProgress) { + log.Info("async loading resource groups aborted: manager was reinitialized") + return + } + startTime := time.Now() + tempKrgms, err := m.loadKeyspaceResourceGroupsFromStorage() + // The storage scans above can block for a long time; re-check for + // cancellation before touching any shared state, so a loader whose + // term already ended doesn't pollute a newer term's state. + select { + case <-ctx.Done(): + log.Info("async loading resource groups cancelled") + return + default: + } + if err != nil { + // Use warn level since the loader retries indefinitely until it succeeds. + // The failure counter and the loading state gauge are what make a load + // that never succeeds alertable, since it no longer fails `Init` loudly. + asyncLoadGroupFailureCounter.Inc() + log.Warn("failed to load resource groups", zap.Error(err), zap.Int("retry", retry)) + if !m.storeLoadingStateIfCurrent(epoch, LoadingStateNotStarted) { + log.Info("async loading resource groups aborted: manager was reinitialized") + return + } + retry++ + continue + } + + // Flatten the loaded groups so the merge below can run in bounded + // batches. Holding m.Lock across the whole O(total groups) merge would + // block every concurrent point/token request (they need the lock to + // resolve a keyspace manager) for the entire merge, causing a large + // latency spike right when async loading completes on a cluster with + // many resource groups. + type mergeItem struct { + keyspaceID uint32 + name string + group *ResourceGroup + } + // Size the slice up front: it holds every loaded group, which is exactly + // the scale this loader exists to handle. tempKrgms is a local map this + // goroutine alone constructed and holds - not yet reachable from + // m.krgms or any other goroutine - so reading it here needs no locking; + // only the individual *ResourceGroup values get published later, into + // whichever keyspace manager is current at merge time, not tempKrgm + // itself. + totalGroups := 0 + for _, tempKrgm := range tempKrgms { + totalGroups += len(tempKrgm.groups) + } + pending := make([]mergeItem, 0, totalGroups) + for keyspaceID, tempKrgm := range tempKrgms { + for name, group := range tempKrgm.groups { + pending = append(pending, mergeItem{keyspaceID: keyspaceID, name: name, group: group}) + } + } + + const mergeBatchSize = 1024 + loaded := 0 + aborted := false + for start := 0; start < len(pending); start += mergeBatchSize { + end := min(start+mergeBatchSize, len(pending)) + m.Lock() + if m.loadEpoch != epoch { + // The manager was reinitialized for a new term while this + // loader was scanning; its result is stale and must not be + // merged. Re-checked every batch since a term change can land + // between batches. + m.Unlock() + aborted = true + break + } + for _, it := range pending[start:end] { + key := trackerKey{keyspaceID: it.keyspaceID, groupName: it.name} + if m.syncLoadedGroups[key] { + continue + } + krgm := m.getOrCreateKeyspaceResourceGroupManagerLocked(it.keyspaceID) + krgm.Lock() + krgm.groups[it.name] = it.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, it.name) + failpoint.InjectCall("mergeBeforeBurstSync", it.keyspaceID, it.name) + // Sync burstability while still holding krgm's lock, so the + // group never becomes visible to a concurrent reader with an + // unsynced burst setting - see syncBurstabilityWithServiceLimitLocked. + krgm.syncBurstabilityWithServiceLimitLocked(it.group) + krgm.Unlock() + loaded++ + } + m.Unlock() + } + if aborted { + log.Info("async loading resource groups aborted: manager was reinitialized") + return + } + m.Lock() + if m.loadEpoch != epoch { + m.Unlock() + log.Info("async loading resource groups aborted: manager was reinitialized") + return + } + m.syncLoadedGroups = nil + m.Unlock() + + // No eager reserved-default backfill runs here. An eager pass would + // re-resolve every keyspace manager and persist a synthetic default + // concurrently with live requests: once completion is published below, + // a request for a keyspace whose default was never customized can + // synthesize and publish it (via getOrCreateKeyspaceResourceGroupManager + // or loadResourceGroupIfNeeded's confirmed-not-found path) in the same + // goroutine as its own subsequent write, with no independent writer + // racing it. An out-of-band backfill loop has no such ordering guarantee + // against those on-demand writers, so it can persist a synthetic default + // after a concurrent customized write and silently discard it. A + // keyspace nobody ever queries needs no persisted default: the persist + // loop already skips unconfirmed reserved placeholders, so leaving one + // reserved indefinitely is a normal, accepted state, not a leak. + if !m.storeLoadingStateIfCurrent(epoch, LoadingStateCompleted) { + log.Info("async loading resource groups aborted: manager was reinitialized") + return + } + 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 +852,406 @@ func (m *Manager) loadKeyspaceResourceGroups() error { zap.Uint32("keyspace-id", keyspaceID), zap.String("group-name", name), zap.Error(err)) } }); err != nil { - return err + return nil, 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() + return tempKrgms, nil } -func (m *Manager) loadServiceLimits() error { - return m.storage.LoadServiceLimits(func(keyspaceID uint32, serviceLimit float64) { - m.getOrCreateKeyspaceResourceGroupManager(keyspaceID, false).setServiceLimitFromStorage(serviceLimit) +// loadResourceGroup loads a single resource group from storage. +func (m *Manager) loadResourceGroup(keyspaceID uint32, name string) (*ResourceGroup, error) { + rawValue, err := m.storage.LoadResourceGroupSetting(keyspaceID, name) + if err != nil { + return nil, err + } + if rawValue == "" { + return nil, errs.ErrResourceGroupNotExists.FastGenByArgs(name) + } + krgm := newKeyspaceResourceGroupManager(keyspaceID, m.storage, m.writeRole) + if err := krgm.addResourceGroupFromRaw(name, rawValue); err != nil { + return nil, err + } + state, err := m.storage.LoadResourceGroupState(keyspaceID, name) + if err != nil { + log.Warn("failed to load resource group state", + zap.Uint32("keyspace-id", keyspaceID), + zap.String("group-name", name), + zap.Error(err)) + return nil, err + } + if state != "" { + if err := krgm.setRawStatesIntoResourceGroup(name, state); err != nil { + return nil, err + } + } + return krgm.getMutableResourceGroup(name), nil +} + +func (m *Manager) loadResourceGroupIfNeeded(keyspaceID uint32, name string) error { + // In metadata-watcher mode the cache is only eventually consistent with + // storage: PD writes metadata directly and the watcher applies it + // asynchronously, so LoadingStateCompleted here only means the initial + // watcher bootstrap finished, not that every subsequent write has been + // observed yet. Always fall through to a point load in that mode so a + // write that outraces its own watcher event is still visible. + if !m.enableMetadataWatcher && m.getLoadingState() == 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 krgm.hasConfirmedResourceGroup(name) { + return nil + } + } + // The lock-free storage read below can be invalidated while it runs: a + // concurrent Delete of any group in the keyspace bumps deleteGen, and a + // leadership change replaces m.krgms and m.syncLoadedGroups (bumping + // loadEpoch). Both are rare: retry the read a few times against freshly + // captured state so neither makes this request spuriously fail or publish + // into a detached manager; if it keeps racing, give up without inserting + // and let a later request or the async bulk merge reload the group. + const maxLoadAttempts = 3 + for attempt := 1; ; attempt++ { + // Capture the load epoch and the current keyspace manager atomically, + // and re-capture them on every retry: publishing into a previous + // term's detached manager while marking the new term's map would make + // the new bulk merge skip a group its cache doesn't contain. + m.Lock() + epoch := m.loadEpoch + krgm = m.getOrCreateKeyspaceResourceGroupManagerLocked(keyspaceID) + m.Unlock() + // Snapshot the delete generation before the lock-free storage read, + // so a Delete that lands after the read is detected under the insert + // lock and can't be undone by the now-stale result. + deleteGen := krgm.loadDeleteGen() + group, err := m.loadResourceGroup(keyspaceID, name) + if err != nil { + if name == DefaultResourceGroupName && errs.ErrResourceGroupNotExists.Equal(err) { + m.RLock() + stale := m.loadEpoch != epoch || m.krgms[keyspaceID] != krgm + m.RUnlock() + if stale { + if attempt >= maxLoadAttempts { + // Exhausted retries without confirming absence; report + // a retryable loading error rather than a bogus + // success the caller would mistake for a load. + return errs.ErrResourceGroupsLoading + } + continue + } + // 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. + stillCurrent := func() bool { + m.RLock() + defer m.RUnlock() + return m.loadEpoch == epoch && m.krgms[keyspaceID] == krgm + } + // initDefaultResourceGroup publishes through + // publishResourceGroupMutation, which sets the sync-loaded + // marker itself atomically with the cache effect - no + // separate marker-setting step is needed here. + _, initErr := m.initDefaultResourceGroup(keyspaceID, krgm, stillCurrent) + if initErr != nil { + if errs.ErrResourceGroupsLoading.Equal(initErr) { + // stillCurrent caught a term change that landed after the + // stale check above but before the persist started; the + // group was not created anywhere, so retry against the + // fresh term instead of reporting a bogus success. + if attempt >= maxLoadAttempts { + return errs.ErrResourceGroupsLoading + } + continue + } + // A real failure persisting the default group (e.g. a storage + // write error): propagate it rather than silently reporting + // success for a group that was never actually created. + return initErr + } + return nil + } + return err + } + inserted := false + m.Lock() + if m.loadEpoch != epoch || m.krgms[keyspaceID] != krgm { + // The manager was reinitialized for a new term while the storage + // read was in flight; retry against the new term's state. + m.Unlock() + if attempt >= maxLoadAttempts { + // Exhausted retries without publishing the group; report a + // retryable loading error rather than a bogus success that + // would surface an existing group as nonexistent. + return errs.ErrResourceGroupsLoading + } + continue + } + krgm.Lock() + if krgm.deleteGen != deleteGen { + krgm.Unlock() + m.Unlock() + if attempt >= maxLoadAttempts { + // Exhausted retries without publishing the group; report a + // retryable loading error rather than a bogus success that + // would surface an existing group as nonexistent. + return errs.ErrResourceGroupsLoading + } + continue + } + 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 + } + delete(krgm.reservedGroups, name) + if inserted { + // Sync burstability while krgm's lock is still held, so the group + // never becomes visible to a concurrent reader with an unsynced + // burst setting - see syncBurstabilityWithServiceLimitLocked. + krgm.syncBurstabilityWithServiceLimitLocked(group) + } + krgm.Unlock() + m.markResourceGroupSyncLoadedLocked(keyspaceID, name) + m.Unlock() + failpoint.Inject("lazyLoadAfterCachePublish", func() {}) + syncLoadGroupCounter.Inc() + return nil + } +} + +// markResourceGroupSyncLoaded records that the group in krgm was written or +// fully loaded synchronously. The caller passes the keyspace manager it +// actually mutated: if that manager is no longer the live one (the manager was +// reinitialized for a new term while the caller was blocked on storage I/O), +// the marker is skipped, since marking the new term's map for a group its +// cache doesn't contain would make the bulk merge skip loading it. +func (m *Manager) markResourceGroupSyncLoaded(keyspaceID uint32, krgm *keyspaceResourceGroupManager, name string) { + m.Lock() + defer m.Unlock() + if m.krgms[keyspaceID] != krgm { + return + } + m.markResourceGroupSyncLoadedLocked(keyspaceID, name) +} + +// markResourceGroupSyncLoadedLocked is markResourceGroupSyncLoaded's marker +// write for a caller that already holds m.Lock() and has already verified +// krgm is still the live entry for keyspaceID. +func (m *Manager) markResourceGroupSyncLoadedLocked(keyspaceID uint32, name string) { + if m.syncLoadedGroups != nil { + m.syncLoadedGroups[trackerKey{keyspaceID: keyspaceID, groupName: name}] = true + } +} + +// publishResourceGroupMutation applies a metadata mutation's cache effect and +// its sync-loaded marker atomically with respect to the async bulk merge, +// against whichever keyspace manager is current at publish time. The merge +// holds the manager lock across its whole merge step, so the two can only be +// fully ordered: publish first and the merge skips the marked group; merge +// first and the publish overrides its stale snapshot. +// +// krgm is the keyspace manager the caller actually persisted the mutation's +// storage phase against. If it's no longer the live manager for keyspaceID +// (Init replaced the whole krgms map for a new term while the caller was +// blocked on storage I/O) AND the new term already has confirmed data for +// this group - installed by its own bulk merge, lazy load, or a live +// Add/Modify/Delete that raced ahead of this delayed one - fn's result was +// computed from data read in the old term and is stale relative to that +// confirmed data; applying it here would silently overwrite it in the cache +// and, via the sync-loaded marker, hide the group from the new term's bulk +// merge too, so the mutation is dropped instead. If the new term hasn't +// confirmed this group yet (missing, or still a reserved placeholder), there +// is nothing newer to protect, so fn's result is applied into the new term's +// manager - the two async-load tests exercise exactly this case (a mutation +// straddling a leadership change with no competing new-term write) and +// require it to still take effect. +// +// fn runs with the keyspace manager write lock held and must not do I/O; it +// returns whether to record the sync-loaded marker and, when a group was +// (re)installed, the group to sync burstability for. +// +// Known gap: this "skip if the new term already confirmed the group" rule +// assumes the confirming write's storage effect is genuinely the latest one. +// That assumption can fail for any mutation kind, not just Delete: +// +// - Delete can be parked before its storage phase (deleteResourceGroupBeforeStorage). +// If an old-term Delete is parked there, a new-term Add for the same group +// persists and publishes (getting confirmed), and only then does the old +// Delete's storage removal finally run, it deletes that newer write in +// storage - making Delete genuinely the last writer - but this function +// skips its publish because the group is already "confirmed", leaving the +// cache showing the deleted group as present. +// - Add/Modify's own storage write always completes before either of them can +// be parked (see addResourceGroupBeforePublish/modifyResourceGroupBeforePublish, +// both placed after the storage write) - but that only rules out being +// parked *after* persisting, not the storage write itself taking real +// wall-clock time to land. A new-term bulk merge or lazy load can read and +// confirm an older snapshot while an old-term Add/Modify's storage write is +// still in flight; when that write then completes, it's the latest value in +// storage, but this function still sees "already confirmed" for the older +// snapshot and drops the publish, leaving the cache stale relative to +// storage indefinitely. +// +// In both cases nothing re-syncs the cache afterward, since it reads as +// confirmed to every other path too. The Delete case also mirrors, in the +// opposite direction, a gap the old unconditional-apply code had (a delayed +// Delete publish could instead wipe a newer confirmed Add). None of this is +// fully closed without a storage-side revision to tell which write actually +// landed last; that's the same class of gap tracked for +// initDefaultResourceGroup's stillCurrent check. No regression test exercises +// either interleaving yet. +// +// TODO(#11105): close this gap with a storage-side revision/CAS check. +func (m *Manager) publishResourceGroupMutation( + keyspaceID uint32, name string, krgm *keyspaceResourceGroupManager, + fn func(krgm *keyspaceResourceGroupManager) (mark bool, synced *ResourceGroup), +) { + m.Lock() + defer m.Unlock() + cur := m.getOrCreateKeyspaceResourceGroupManagerLocked(keyspaceID) + if cur != krgm && cur.hasConfirmedResourceGroup(name) { + log.Info("skip publishing resource group mutation: a newer confirmed write already exists", + zap.Uint32("keyspace-id", keyspaceID), zap.String("group-name", name)) + return + } + cur.Lock() + mark, synced := fn(cur) + if synced != nil { + // Sync burstability while cur's lock is still held, so the group + // never becomes visible to a concurrent reader with an unsynced + // burst setting - see syncBurstabilityWithServiceLimitLocked. + cur.syncBurstabilityWithServiceLimitLocked(synced) + } + cur.Unlock() + failpoint.InjectCall("publishMutationBeforeMark") + if mark { + m.markResourceGroupSyncLoadedLocked(keyspaceID, name) + } +} + +// initDefaultResourceGroup synthesizes and persists the built-in default +// group for keyspaceID into krgm if nothing confirmed exists yet, publishing +// it through publishResourceGroupMutation - the same path a real +// Add/ModifyResourceGroup uses - so the cache effect and the sync-loaded +// marker land atomically with respect to the async bulk merge. This used to +// publish through krgm's own lock only, with the caller setting the marker +// afterward in a separate critical section; that left a window where the +// bulk merge could run in between and replace the synthesized group - along +// with any live consumption/token update applied to it in that window - with +// its own possibly-stale scanned copy, since it had no way yet to know the +// group was confirmed. Routing through publishResourceGroupMutation closes +// that window the same way it already does for Add/Modify/Delete. +// +// created reports whether it actually performed a synthesis. The three +// (created, err) outcomes need different handling and must not be collapsed +// into a single bool by the caller: (false, nil) means confirmed data +// already existed - nothing to do, safe to treat as success; (false, +// errs.ErrResourceGroupsLoading) means stillCurrent caught a term change +// before the persist started - the group was not created anywhere, so the +// caller must retry against the fresh term rather than treat this as +// success; (false, any other non-nil error) means the persist itself failed +// (e.g. a storage write error) - the caller must propagate it rather than +// silently swallow a real failure as success. +// +// defaultGroupMu only serializes callers that share the krgm instance; it +// does nothing across a term change, since Init gives the new term an +// entirely new krgm object with its own, separate defaultGroupMu. +// stillCurrent, when non-nil, is checked immediately before the persist +// (right after defaultGroupMu is acquired) to fail fast and avoid a wasted +// storage write when the caller already knows krgm is stale - typically by +// comparing it against the manager's live entry for its keyspace ID. It is +// no longer the only guard against a term change landing while the persist's +// storage write is in flight: publishResourceGroupMutation's own +// cur != krgm && cur.hasConfirmedResourceGroup(name) check now also protects +// that window (and, if the new term hasn't confirmed a default yet, still +// applies this call's result into the new term's live krgm instead of +// silently dropping it into a detached one). Pass nil when no such +// fail-fast check is needed or available (e.g. in tests that exercise a +// krgm/Manager pair with no concurrent writer to race against). +func (m *Manager) initDefaultResourceGroup(keyspaceID uint32, krgm *keyspaceResourceGroupManager, stillCurrent func() bool) (created bool, err error) { + // A confirmed cached entry means initialization is unnecessary; a missing + // or reserved-placeholder entry means nothing is persisted for the + // default group (e.g. a fresh store), so it must still be created and + // persisted, otherwise its settings are never stored and state + // persistence stays skipped. + if krgm.hasConfirmedResourceGroup(DefaultResourceGroupName) { + return false, nil + } + // Serialize against every other synthesis or real Add/ModifyResourceGroup + // targeting "default" that shares this krgm instance: see the + // defaultGroupMu doc comment on the struct. + krgm.defaultGroupMu.Lock() + defer krgm.defaultGroupMu.Unlock() + // Re-check under defaultGroupMu: while this goroutine waited for the + // lock, a real write may have already confirmed the default group, in + // which case synthesizing here would silently clobber it. + if krgm.hasConfirmedResourceGroup(DefaultResourceGroupName) { + return false, nil + } + if stillCurrent != nil && !stillCurrent() { + return false, errs.ErrResourceGroupsLoading + } + defaultGroup := newDefaultResourceGroup() + group, err := krgm.persistResourceGroup(defaultGroup.IntoProtoResourceGroup(krgm.keyspaceID)) + if err != nil { + log.Warn("init default group failed", zap.Uint32("keyspace-id", krgm.keyspaceID), zap.Error(err)) + return false, err + } + m.publishResourceGroupMutation(keyspaceID, DefaultResourceGroupName, krgm, func(cur *keyspaceResourceGroupManager) (bool, *ResourceGroup) { + cur.groups[group.Name] = group + delete(cur.reservedGroups, group.Name) + return true, group }) + return true, nil +} + +// isKeyspaceManagerCurrent reports whether krgm is still the live keyspace +// manager for keyspaceID. Used where any term change - confirmed or not - is +// reason enough to bail out, e.g. initDefaultResourceGroup's best-effort, +// no-caller-waiting synthesis: publishResourceGroupMutation would apply it +// into the new term anyway if nothing there has confirmed a default yet, so +// aborting a harmless case here just costs a later retry, not correctness. +func (m *Manager) isKeyspaceManagerCurrent(keyspaceID uint32, krgm *keyspaceResourceGroupManager) bool { + m.RLock() + defer m.RUnlock() + return m.krgms[keyspaceID] == krgm +} + +// hasNewerConfirmedWrite reports whether the current keyspace manager for +// keyspaceID is no longer krgm and already has confirmed data for name - +// i.e. whether a mutation resolved against krgm is racing a leadership +// change that already landed a newer, competing write for the same group. +// Mirrors the guard publishResourceGroupMutation applies at publish time. +// AddResourceGroup/ModifyResourceGroup/DeleteResourceGroup check this +// immediately before their storage-write phase to fail fast with a +// retryable error in that specific case, instead of writing to storage on +// behalf of a term whose result is going to be silently dropped at publish +// time anyway. Deliberately narrower than isKeyspaceManagerCurrent: a term +// change with nothing yet confirmed for this group in the new term is +// harmless and already handled gracefully - publishResourceGroupMutation +// re-resolves the current manager and applies the result there - so it must +// not raise the same reject here that only the confirmed-collision case +// warrants. +// This only catches a collision that completes before the storage write +// starts; one that lands while the write is already in flight needs the +// storage-side revision/CAS check tracked in #11105. +func (m *Manager) hasNewerConfirmedWrite(keyspaceID uint32, name string, krgm *keyspaceResourceGroupManager) bool { + m.RLock() + defer m.RUnlock() + cur, ok := m.krgms[keyspaceID] + return ok && cur != krgm && cur.hasConfirmedResourceGroup(name) +} + +func (m *Manager) isResourceGroupLoadingComplete() bool { + return m.getLoadingState() == LoadingStateCompleted } func cloneControllerConfig(cfg *ControllerConfig) *ControllerConfig { @@ -456,6 +1288,7 @@ func (m *Manager) applyResourceGroupSettingFromRaw(keyspaceID uint32, name, rawV zap.Error(err)) return err } + m.markResourceGroupSyncLoaded(keyspaceID, krgm, name) return nil } @@ -492,15 +1325,55 @@ func (m *Manager) applyResourceGroupStatesFromRaw(keyspaceID uint32, name, rawVa zap.Error(err)) return err } + m.markResourceGroupSyncLoaded(keyspaceID, krgm, name) return nil } -func (m *Manager) initReserved() { +// initReserved backfills the default resource group for every keyspace whose +// default wasn't confirmed by loading. The caller must have just verified +// epoch via storeLoadingStateIfCurrent; re-verify it here under the manager +// lock immediately before touching krgms, since that earlier check alone +// does not cover this call once its lock is released. +// +// Only call this from a path that runs before the manager serves any +// request (construction-time full loads), not from the async loader: once +// requests are flowing, a concurrent Add/ModifyResourceGroup can confirm a +// keyspace's default between this function's existence check and its +// unconditional persist, and this backfill's synthetic write would then +// silently clobber that write in both storage and the live cache. Live +// requests already synthesize a missing default on demand (via +// getOrCreateKeyspaceResourceGroupManager or loadResourceGroupIfNeeded's +// confirmed-not-found path) sequenced with their own subsequent write, so no +// out-of-band backfill is needed once serving has started. +func (m *Manager) initReserved(epoch uint64) { + m.Lock() + if m.loadEpoch != epoch { + m.Unlock() + log.Info("skip initReserved: manager was reinitialized") + return + } + m.Unlock() // Initialize the null keyspace resource group manager if it doesn't exist. m.getOrCreateKeyspaceResourceGroupManager(constant.NullKeyspaceID, true) // Initialize the default resource group respectively for each keyspace if it doesn't exist. + // No stillCurrent check is needed: this whole function only runs before + // the manager serves any request (see the doc comment above), so there is + // no concurrent Add/ModifyResourceGroup or other initDefaultResourceGroup + // call to race against. + for _, krgm := range m.getKeyspaceResourceGroupManagers() { + // Any failure is already logged inside initDefaultResourceGroup; a + // later request for the default group retries through + // loadResourceGroupIfNeeded once serving starts. + _, _ = m.initDefaultResourceGroup(krgm.keyspaceID, krgm, nil) + } +} + +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.initDefaultResourceGroup() + krgm.ensureReservedDefaultGroupInCache() } } @@ -574,7 +1447,46 @@ 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 && + !errs.ErrResourceGroupNotExists.Equal(err) { + 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 grouppb.Name == DefaultResourceGroupName { + // Serialize against a concurrent on-demand synthesis of the same + // default group (initDefaultResourceGroup, e.g. from another + // request's getOrCreateKeyspaceResourceGroupManager/ + // loadResourceGroupIfNeeded call): without this, the synthetic + // write's storage/cache commit can land after this real write's, + // silently discarding these customized settings. + krgm.defaultGroupMu.Lock() + defer krgm.defaultGroupMu.Unlock() + } + failpoint.InjectCall("addResourceGroupBeforeStorage") + // Fail fast only if a leadership change already replaced krgm AND the new + // term already has confirmed data for this group - i.e. persisting here + // would be racing a write that's already won and would just get silently + // dropped by publishResourceGroupMutation's own matching guard below. A + // harmless term change with nothing yet confirmed in the new term is not + // rejected: publishResourceGroupMutation already applies the result there + // gracefully in that case. See hasNewerConfirmedWrite. + if m.hasNewerConfirmedWrite(keyspaceID, grouppb.Name, krgm) { + return errs.ErrResourceGroupsLoading + } + // Storage phase: validate and persist. Publishing the cache effect is done + // separately below, against krgm if it's still the live manager for + // keyspaceID by then, or dropped otherwise - see publishResourceGroupMutation. + group, err := krgm.persistResourceGroup(grouppb) + if err != nil { + return err + } + failpoint.InjectCall("addResourceGroupBeforePublish") + m.publishResourceGroupMutation(keyspaceID, grouppb.Name, krgm, func(cur *keyspaceResourceGroupManager) (bool, *ResourceGroup) { + cur.groups[group.Name] = group + delete(cur.reservedGroups, group.Name) + return true, group + }) + return nil } // ModifyResourceGroup modifies an existing resource group. @@ -583,11 +1495,58 @@ 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 grouppb.Name == DefaultResourceGroupName { + // Serialize against a concurrent on-demand synthesis of the same + // default group; see the matching guard in AddResourceGroup. + krgm.defaultGroupMu.Lock() + defer krgm.defaultGroupMu.Unlock() + } + failpoint.InjectCall("modifyResourceGroupBeforeStorage") + // Fail fast only on a newer confirmed collision; see the matching guard + // in AddResourceGroup. + if m.hasNewerConfirmedWrite(keyspaceID, grouppb.Name, krgm) { + return errs.ErrResourceGroupsLoading + } + patched, err := krgm.modifyResourceGroup(grouppb) + if err != nil { + return err + } + failpoint.InjectCall("modifyResourceGroupBeforePublish") + m.publishResourceGroupMutation(keyspaceID, grouppb.Name, krgm, func(cur *keyspaceResourceGroupManager) (bool, *ResourceGroup) { + var synced *ResourceGroup + if existing := cur.groups[grouppb.Name]; existing != patched { + // A different object sits in the current term's cache: either the + // group is missing, a reserved default placeholder (with synthetic + // token/consumption state), or a pre-modification bulk-merge + // snapshot. Install `patched` wholesale rather than only patching + // settings onto it. `patched` was loaded/confirmed before it was + // modified, so it carries both the modified settings and the + // group's confirmed running state - patching the placeholder in + // place would keep its synthetic state, which the marker below + // would then freeze as confirmed and the persist loop would write + // back over the real state. + cur.groups[grouppb.Name] = patched + synced = patched + } + // The settings and state are now confirmed data, even if the entry + // started as a reserved default placeholder (the only thing + // reservedGroups ever holds). Clear the marker and record it as + // sync-loaded: otherwise the bulk merge could revert it to a + // pre-modification snapshot, or initReserved could re-synthesize a + // fresh default over it, silently dropping the just-modified settings + // from the serving cache while storage keeps the new value. + delete(cur.reservedGroups, grouppb.Name) + return true, synced + }) + return nil } // DeleteResourceGroup deletes a resource group. @@ -595,16 +1554,46 @@ 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) + failpoint.InjectCall("deleteResourceGroupBeforeStorage") + // Fail fast only on a newer confirmed collision; see the matching guard + // in AddResourceGroup. A harmless term change alone (nothing yet + // confirmed in the new term) must still let the delete proceed, storage + // removal is independent of which term's krgm initiated it, and + // publishResourceGroupMutation republishes the result against whichever + // term is current at publish time. + if m.hasNewerConfirmedWrite(keyspaceID, name, krgm) { + return errs.ErrResourceGroupsLoading + } + // Storage phase: validate and remove from storage. Publishing the cache + // effect is done separately below, against krgm if it's still the live + // manager for keyspaceID by then, or dropped otherwise (a delete + // straddling a leadership change) - see publishResourceGroupMutation. The + // storage removal above already took effect regardless. + if err := krgm.deleteResourceGroupFromStorage(name); err != nil { + return err + } + m.publishResourceGroupMutation(keyspaceID, name, krgm, func(cur *keyspaceResourceGroupManager) (bool, *ResourceGroup) { + cur.removeResourceGroupLocked(name) + return true, nil + }) + 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 +1603,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 +1615,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..cf6cd13195 --- /dev/null +++ b/pkg/mcs/resourcemanager/server/manager_async_test.go @@ -0,0 +1,1815 @@ +// 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" + "encoding/json" + "errors" + "fmt" + "io" + "math" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + + "github.com/pingcap/failpoint" + "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 + + // statePause, when armed via armStatePause, makes the next + // LoadResourceGroupState call for the armed group name signal reached and + // then block on release, so a test can hold a lazy load right after its + // storage read but before it inserts. Re-armable, and filtered by name so + // unrelated groups' loads pass through undisturbed. + statePause atomic.Pointer[statePause] + + // pauseNextStates, when true, makes the very next bulk + // LoadResourceGroupStates call signal statesReached and then block on + // statesRelease, so a test can hold an async loader after it has captured + // the settings scan but before it merges. + pauseNextStates atomic.Bool + statesReached chan struct{} + statesRelease chan struct{} + statesReleaseOnce sync.Once +} + +type statePause struct { + name string + reached chan struct{} + release chan struct{} +} + +func newBlockingResourceGroupStorage() *blockingResourceGroupStorage { + return &blockingResourceGroupStorage{ + Storage: storage.NewStorageWithMemoryBackend(), + entered: make(chan struct{}), + release: make(chan struct{}), + statesReached: make(chan struct{}), + statesRelease: make(chan struct{}), + } +} + +// armStatePause arms a one-shot pause on the next LoadResourceGroupState call +// for the given group name and returns the pause handle. The test must wait +// on reached and eventually close release. +func (s *blockingResourceGroupStorage) armStatePause(name string) *statePause { + p := &statePause{name: name, reached: make(chan struct{}), release: make(chan struct{})} + s.statePause.Store(p) + return p +} + +func waitStatePauseReached(t *testing.T, p *statePause) { + t.Helper() + select { + case <-p.reached: + case <-time.After(time.Second): + t.Fatal("timed out waiting for the lazy load to reach its state read") + } +} + +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") + } + if p := s.statePause.Load(); p != nil && p.name == name && s.statePause.CompareAndSwap(p, nil) { + close(p.reached) + <-p.release + } + return s.Storage.LoadResourceGroupState(keyspaceID, name) +} + +func (s *blockingResourceGroupStorage) LoadResourceGroupStates(f func(keyspaceID uint32, name, rawValue string)) error { + if s.pauseNextStates.CompareAndSwap(true, false) { + close(s.statesReached) + <-s.statesRelease + } + return s.Storage.LoadResourceGroupStates(f) +} + +func (s *blockingResourceGroupStorage) waitEntered(tb testing.TB) { + tb.Helper() + select { + case <-s.entered: + case <-time.After(5 * time.Second): + tb.Fatal("timed out waiting for async resource group loading") + } +} + +func (s *blockingResourceGroupStorage) unblock() { + s.releaseOnce.Do(func() { + close(s.release) + }) +} + +func (s *blockingResourceGroupStorage) unblockStates() { + s.statesReleaseOnce.Do(func() { + close(s.statesRelease) + }) +} + +// asyncTestGroupFillRate is the fill rate used by all async-loading test +// groups; kept as a named constant so the setup and the assertions stay in +// sync. +const asyncTestGroupFillRate = 100 + +func newAsyncTestGroup(name string) *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: asyncTestGroupFillRate, + BurstLimit: asyncTestGroupFillRate, + }, + }, + }, + } +} + +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"))) + + 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(asyncTestGroupFillRate), 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"))) + + 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"))) + + 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(asyncTestGroupFillRate), 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)) +} + +// TestAsyncLoadResourceGroupsDoesNotServeStateLoadFailure guards against a +// group with confirmed settings but failed state loading being exposed with a +// fresh token bucket before the async bulk loader can recover its persisted +// state. +func TestAsyncLoadResourceGroupsDoesNotServeStateLoadFailure(t *testing.T) { + re := require.New(t) + store := newBlockingResourceGroupStorage() + group := newAsyncTestGroup("flaky-group") + 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. The group must remain + // unavailable rather than being returned with a fresh token bucket state. + store.failNextState.Store(true) + fetched, err := m.GetResourceGroup(1, "flaky-group", false) + re.Error(err) + re.Nil(fetched) + + krgm := m.getKeyspaceResourceGroupManager(1) + re.NotNil(krgm) + re.Nil(krgm.getMutableResourceGroup("flaky-group"), "failed state load must not publish the group") + + // 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.getMutableResourceGroup("flaky-group") != nil && !krgm.isReserved("flaky-group") + }, testutil.WithTickInterval(20*time.Millisecond)) +} + +// TestAsyncLoadResourceGroupsDeleteRaceDoesNotResurrect reproduces the +// lazy-load vs concurrent Delete race deterministically: a lazy load reads a +// group from storage, then a Delete removes it before the lazy load inserts. +// The stale insert must be rejected (via the delete-generation check) so the +// deleted group is not resurrected for the rest of the manager's lifetime. +func TestAsyncLoadResourceGroupsDeleteRaceDoesNotResurrect(t *testing.T) { + re := require.New(t) + store := newBlockingResourceGroupStorage() + group := newAsyncTestGroup("race-group") + re.NoError(store.SaveResourceGroupSetting(1, "race-group", group)) + re.NoError(store.SaveResourceGroupStates(1, "race-group", FromProtoResourceGroup(group).GetGroupStates())) + + m := NewManager[*mockConfigProvider](&mockConfigProvider{}) + m.storage = store + re.NoError(m.Init(context.Background())) + defer stopAsyncTestManager(m) + defer store.unblock() + + // Async bulk load is blocked, so loadingState stays in progress and lazy + // loading is active. + store.waitEntered(t) + + // Start a lazy Get that will pause inside its state read, i.e. after it has + // read the group from storage but before it inserts into the cache. + pause := store.armStatePause("race-group") + var ( + gotGroup *ResourceGroup + gotErr error + ) + getDone := make(chan struct{}) + go func() { + defer close(getDone) + gotGroup, gotErr = m.GetResourceGroup(1, "race-group", false) + }() + + waitStatePauseReached(t, pause) + + // While the lazy load is paused, delete the group. Delete does its own + // (unpaused) load-then-delete, removing it from storage and cache and + // bumping the delete generation. + re.NoError(m.DeleteResourceGroup(1, "race-group")) + + // Release the paused lazy load; its now-stale insert must be rejected. + close(pause.release) + <-getDone + // The generation mismatch makes the lazy load retry its storage read, + // which now correctly observes the group as deleted. + re.ErrorContains(gotErr, "does not exist") + re.Nil(gotGroup, "the racing lazy load must observe the group as deleted") + + krgm := m.getKeyspaceResourceGroupManager(1) + re.NotNil(krgm) + re.Nil(krgm.getMutableResourceGroup("race-group"), "deleted group must not be resurrected by the racing lazy load") + + // Finishing async loading must not bring the deleted group back either. + store.unblock() + testutil.Eventually(re, func() bool { + groups, err := m.GetResourceGroupList(1, false) + if err != nil { + return false + } + for _, g := range groups { + if g.Name == "race-group" { + return false + } + } + return true + }, testutil.WithTickInterval(20*time.Millisecond)) +} + +// TestAsyncLoadResourceGroupsStaleLoaderDoesNotPolluteNewTerm reproduces the +// stale-loader race: a loader from an old term is blocked in its storage scan +// while the leadership changes and Init runs again for a new term. When the +// old loader finally wakes up, it must not merge its stale scan into the new +// term's maps, clear the new term's syncLoadedGroups, or publish completion. +func TestAsyncLoadResourceGroupsStaleLoaderDoesNotPolluteNewTerm(t *testing.T) { + re := require.New(t) + store := newBlockingResourceGroupStorage() + group := newAsyncTestGroup("stale-group") + re.NoError(store.SaveResourceGroupSetting(1, "stale-group", group)) + + m := NewManager[*mockConfigProvider](&mockConfigProvider{}) + m.storage = store + // Term 1: the loader blocks at the start of its settings scan. + re.NoError(m.Init(context.Background())) + cancelTerm1 := m.cancel + defer stopAsyncTestManager(m) + defer store.unblock() + defer store.unblockStates() + + store.waitEntered(t) + + // Let the term-1 loader run its settings scan (capturing stale-group into + // its temp result) and then block again in the states scan, i.e. after it + // has read storage but before it merges. + store.pauseNextStates.Store(true) + store.unblock() + select { + case <-store.statesReached: + case <-time.After(time.Second): + t.Fatal("timed out waiting for the term-1 loader to reach its states scan") + } + + // Leadership changes: cancel term 1 and reinitialize for term 2. The + // term-2 loader hits neither block (both were consumed) and completes. + cancelTerm1() + re.NoError(m.Init(context.Background())) + testutil.Eventually(re, func() bool { + groups, err := m.GetResourceGroupList(1, false) + return err == nil && len(groups) == 2 + }, testutil.WithTickInterval(20*time.Millisecond)) + + // Delete the group in term 2, after loading completed. + re.NoError(m.DeleteResourceGroup(1, "stale-group")) + + // Release the stale term-1 loader. It must observe its cancelled context / + // stale epoch and exit without resurrecting the deleted group or touching + // the new term's loading state. + store.unblockStates() + time.Sleep(200 * time.Millisecond) + + krgm := m.getKeyspaceResourceGroupManager(1) + re.NotNil(krgm) + re.Nil(krgm.getMutableResourceGroup("stale-group"), "stale loader must not merge into the new term") + groups, err := m.GetResourceGroupList(1, false) + re.NoError(err) + for _, g := range groups { + re.NotEqual("stale-group", g.Name) + } +} + +// TestAsyncLoadResourceGroupsLazyPublishAndMarkAreAtomic reproduces the race +// where a lazy load publishes a cache entry before recording it in +// syncLoadedGroups. A bulk merge entering that gap can overwrite mutable state +// updated by token-bucket handling with its older scan result. +func TestAsyncLoadResourceGroupsLazyPublishAndMarkAreAtomic(t *testing.T) { + re := require.New(t) + store := newBlockingResourceGroupStorage() + group := newAsyncTestGroup("atomic-group") + re.NoError(store.SaveResourceGroupSetting(1, "atomic-group", group)) + re.NoError(store.SaveResourceGroupStates(1, "atomic-group", FromProtoResourceGroup(group).GetGroupStates())) + + m := NewManager[*mockConfigProvider](&mockConfigProvider{}) + m.storage = store + re.NoError(m.Init(context.Background())) + defer stopAsyncTestManager(m) + defer store.unblock() + defer store.unblockStates() + + store.waitEntered(t) + + // Let the bulk loader capture fill rate 100, then hold it before merge. + store.pauseNextStates.Store(true) + store.unblock() + select { + case <-store.statesReached: + case <-time.After(time.Second): + t.Fatal("timed out waiting for the bulk loader to reach its states scan") + } + + re.NoError(failpoint.Enable("github.com/tikv/pd/pkg/mcs/resourcemanager/server/lazyLoadAfterCachePublish", `pause`)) + defer func() { + re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/mcs/resourcemanager/server/lazyLoadAfterCachePublish")) + }() + + getDone := make(chan struct{}) + go func() { + defer close(getDone) + _, err := m.GetResourceGroup(1, "atomic-group", false) + re.NoError(err) + }() + + var krgm *keyspaceResourceGroupManager + testutil.Eventually(re, func() bool { + krgm = m.getKeyspaceResourceGroupManager(1) + if krgm == nil { + return false + } + return krgm.getMutableResourceGroup("atomic-group") != nil + }, testutil.WithTickInterval(20*time.Millisecond)) + + krgm.getMutableResourceGroup("atomic-group").UpdateRUConsumption(&resource_manager.Consumption{RRU: 10}) + + // With the lazy load paused at the cache-publish hook, the bulk merge must + // not be able to overwrite the updated cache entry. + store.unblockStates() + testutil.Eventually(re, func() bool { + _, err := m.GetResourceGroupList(1, false) + return err == nil + }, testutil.WithTickInterval(20*time.Millisecond)) + + re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/mcs/resourcemanager/server/lazyLoadAfterCachePublish")) + <-getDone + + got, err := m.GetResourceGroup(1, "atomic-group", false) + re.NoError(err) + re.NotNil(got) + re.Equal(float64(10), krgm.getMutableResourceGroup("atomic-group").GetGroupStates().RUConsumption.RRU, + "bulk merge must not overwrite the published lazy-loaded group") +} + +// TestAsyncLoadResourceGroupsFreshStoreDefaultPersisted guards against the +// fresh-store dead end: initReservedInCache pre-inserts a synthetic default +// placeholder, and on a store with nothing persisted, the confirmed-not-found +// fallback used to bail out on its cache-exists check, leaving the default +// group an unconfirmed placeholder forever — settings never persisted and its +// state persistence permanently skipped. +func TestAsyncLoadResourceGroupsFreshStoreDefaultPersisted(t *testing.T) { + re := require.New(t) + // A completely fresh store: nothing persisted at all. + store := newBlockingResourceGroupStorage() + + m := NewManager[*mockConfigProvider](&mockConfigProvider{}) + m.storage = store + re.NoError(m.Init(context.Background())) + defer stopAsyncTestManager(m) + defer store.unblock() + + store.waitEntered(t) + + // Fetch the default group while async loading is still in progress: the + // point load confirms nothing is persisted, so the placeholder must be + // promoted to a real, persisted default group. + group, err := m.GetResourceGroup(constant.NullKeyspaceID, DefaultResourceGroupName, false) + re.NoError(err) + re.NotNil(group) + krgm := m.getKeyspaceResourceGroupManager(constant.NullKeyspaceID) + re.NotNil(krgm) + re.False(krgm.isReserved(DefaultResourceGroupName), "the default group must be confirmed after synthesis") + raw, err := store.LoadResourceGroupSetting(constant.NullKeyspaceID, DefaultResourceGroupName) + re.NoError(err) + re.NotEmpty(raw, "the synthesized default group settings must be persisted") + + // Loading completion must keep it confirmed. + store.unblock() + testutil.Eventually(re, func() bool { + _, err := m.GetResourceGroupList(constant.NullKeyspaceID, false) + return err == nil + }, testutil.WithTickInterval(20*time.Millisecond)) + re.False(krgm.isReserved(DefaultResourceGroupName)) +} + +// TestAsyncLoadResourceGroupsUnrelatedDeleteDoesNotFailLazyLoad guards against +// the delete-generation check being too coarse: deleting group B while group A +// is being lazily loaded must not make A's request spuriously report the group +// as missing — the lazy load retries its storage read and succeeds. +func TestAsyncLoadResourceGroupsUnrelatedDeleteDoesNotFailLazyLoad(t *testing.T) { + re := require.New(t) + store := newBlockingResourceGroupStorage() + re.NoError(store.SaveResourceGroupSetting(1, "group-a", newAsyncTestGroup("group-a"))) + re.NoError(store.SaveResourceGroupSetting(1, "group-b", newAsyncTestGroup("group-b"))) + + m := NewManager[*mockConfigProvider](&mockConfigProvider{}) + m.storage = store + re.NoError(m.Init(context.Background())) + defer stopAsyncTestManager(m) + defer store.unblock() + + store.waitEntered(t) + + // Start a lazy Get of group-a and pause it inside its state read, i.e. + // after it has read the group from storage but before it inserts. + pause := store.armStatePause("group-a") + var ( + gotGroup *ResourceGroup + gotErr error + ) + getDone := make(chan struct{}) + go func() { + defer close(getDone) + gotGroup, gotErr = m.GetResourceGroup(1, "group-a", false) + }() + waitStatePauseReached(t, pause) + + // Delete the unrelated group-b while group-a's lazy load is paused; this + // bumps the keyspace's delete generation. + re.NoError(m.DeleteResourceGroup(1, "group-b")) + + // Release group-a's lazy load: the generation mismatch must make it retry + // and succeed, not report group-a as missing. + close(pause.release) + <-getDone + re.NoError(gotErr) + re.NotNil(gotGroup, "an unrelated delete must not fail the lazy load") + re.Equal("group-a", gotGroup.Name) + + // After loading completes, group-a is present and group-b stays deleted. + store.unblock() + testutil.Eventually(re, func() bool { + groups, err := m.GetResourceGroupList(1, false) + if err != nil { + return false + } + foundA := false + for _, g := range groups { + if g.Name == "group-b" { + return false + } + if g.Name == "group-a" { + foundA = true + } + } + return foundA + }, testutil.WithTickInterval(20*time.Millisecond)) +} + +// TestAsyncLoadResourceGroupsStaleLazyLoadRetriesNewTerm reproduces the +// cross-term lazy load: the load captures the keyspace manager, then blocks in +// its storage read while the leadership changes and Init replaces m.krgms and +// syncLoadedGroups. On resume it must not publish into the detached old +// manager while marking the group in the new term's map (which would make the +// new bulk merge skip a group its cache doesn't contain); instead it retries +// against the freshly captured state and publishes into the new term. +func TestAsyncLoadResourceGroupsStaleLazyLoadRetriesNewTerm(t *testing.T) { + re := require.New(t) + store := newBlockingResourceGroupStorage() + re.NoError(store.SaveResourceGroupSetting(1, "cross-term", newAsyncTestGroup("cross-term"))) + + m := NewManager[*mockConfigProvider](&mockConfigProvider{}) + m.storage = store + re.NoError(m.Init(context.Background())) + cancelTerm1 := m.cancel + defer stopAsyncTestManager(m) + defer store.unblock() + + store.waitEntered(t) + + // The term-1 lazy load pauses inside its state read, holding the old + // term's keyspace manager. + pause := store.armStatePause("cross-term") + var ( + gotGroup *ResourceGroup + gotErr error + ) + getDone := make(chan struct{}) + go func() { + defer close(getDone) + gotGroup, gotErr = m.GetResourceGroup(1, "cross-term", false) + }() + waitStatePauseReached(t, pause) + + // Leadership changes: reinitialize the manager for term 2 while the + // term-1 lazy load is still blocked. Term 2's bulk loader parks on the + // same settings-scan block until store.unblock(). + cancelTerm1() + re.NoError(m.Init(context.Background())) + + // Release the stale lazy load: it must detect the term change, retry, and + // publish into the new term's manager, so the request still succeeds. + close(pause.release) + <-getDone + re.NoError(gotErr) + re.NotNil(gotGroup, "the cross-term lazy load must retry and succeed against the new term") + re.Equal("cross-term", gotGroup.Name) + + // Finish loading; the group must remain present after the term-2 bulk + // merge (it was correctly marked in the same term it was published in). + store.unblock() + testutil.Eventually(re, func() bool { + g, err := m.GetResourceGroup(1, "cross-term", false) + return err == nil && g != nil + }, testutil.WithTickInterval(20*time.Millisecond)) +} + +// TestAsyncLoadResourceGroupsCrossTermDeletePublishesToNewTerm reproduces the +// cross-term delete race: a Delete resolves its keyspace manager, then stalls +// before its storage phase while the leadership changes and the new term's +// bulk loader snapshots storage (still containing the group). When the Delete +// resumes, it removes the group from storage but its cache effect and +// sync-loaded marker must land in the *current* term — otherwise the new +// merge would reinstall its pre-deletion snapshot and the API would report +// success while the group stays in the live cache. +func TestAsyncLoadResourceGroupsCrossTermDeletePublishesToNewTerm(t *testing.T) { + re := require.New(t) + store := newBlockingResourceGroupStorage() + re.NoError(store.SaveResourceGroupSetting(1, "ct-del", newAsyncTestGroup("ct-del"))) + + m := NewManager[*mockConfigProvider](&mockConfigProvider{}) + m.storage = store + re.NoError(m.Init(context.Background())) + cancelTerm1 := m.cancel + defer stopAsyncTestManager(m) + defer store.unblock() + defer store.unblockStates() + + // Let term 1 load fully so the Delete starts against a settled term. + store.waitEntered(t) + store.unblock() + testutil.Eventually(re, func() bool { + groups, err := m.GetResourceGroupList(1, false) + return err == nil && len(groups) == 2 + }, testutil.WithTickInterval(20*time.Millisecond)) + + // Park the Delete between resolving its keyspace manager and its storage + // phase. + reached := make(chan struct{}) + release := make(chan struct{}) + re.NoError(failpoint.EnableCall("github.com/tikv/pd/pkg/mcs/resourcemanager/server/deleteResourceGroupBeforeStorage", func() { + close(reached) + <-release + })) + defer func() { + re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/mcs/resourcemanager/server/deleteResourceGroupBeforeStorage")) + }() + var delErr error + delDone := make(chan struct{}) + go func() { + defer close(delDone) + delErr = m.DeleteResourceGroup(1, "ct-del") + }() + select { + case <-reached: + case <-time.After(time.Second): + t.Fatal("timed out waiting for the delete to reach its storage phase") + } + + // Leadership changes: term 2's loader snapshots storage (the group is + // still there) and parks before merging. + cancelTerm1() + store.pauseNextStates.Store(true) + re.NoError(m.Init(context.Background())) + select { + case <-store.statesReached: + case <-time.After(time.Second): + t.Fatal("timed out waiting for the term-2 loader to snapshot storage") + } + + // Resume the Delete: storage removal proceeds, and the cache effect and + // marker must be published into term 2, not the detached term-1 manager. + close(release) + <-delDone + re.NoError(delErr) + + // Let term 2 merge its pre-deletion snapshot; the marker must make it + // skip the deleted group. + store.unblockStates() + testutil.Eventually(re, func() bool { + groups, err := m.GetResourceGroupList(1, false) + if err != nil { + return false + } + for _, g := range groups { + if g.Name == "ct-del" { + return false + } + } + return true + }, testutil.WithTickInterval(20*time.Millisecond)) + g, err := m.GetResourceGroup(1, "ct-del", false) + re.NoError(err) + re.Nil(g, "the deleted group must not be resurrected by the new term's merge") +} + +// TestResourceGroupMutationAbortsOnNewerConfirmedWriteAcrossTermChange +// guards the "parked-before-storage" side of the #11105 interleaving for all +// three mutation kinds: an old-term Add/Modify/Delete parked before its +// storage phase must abort, not touch storage, once it wakes up to find the +// new term already has confirmed data for the same group (here, from the +// new term's own reload of unchanged storage, since the mutation never +// reached its own storage phase). Proceeding would race a write that has +// already won and would just get silently dropped by +// publishResourceGroupMutation's confirmed-write guard - for Delete that +// would leave storage empty but the cache still serving the group forever, +// the exact #11105 shape; for Add/Modify it would overwrite storage with a +// stale value nobody asked for. hasNewerConfirmedWrite must catch this and +// abort before the storage write in every case, while still leaving the +// sibling "nothing confirmed yet" case +// (TestAsyncLoadResourceGroupsCrossTermDeletePublishesToNewTerm) to succeed +// normally. +func TestResourceGroupMutationAbortsOnNewerConfirmedWriteAcrossTermChange(t *testing.T) { + cases := []struct { + name string + group string + failpoint string + mutate func(m *Manager, group string) error + }{ + { + name: "delete", + group: "ct-del-abort", + failpoint: "deleteResourceGroupBeforeStorage", + mutate: func(m *Manager, group string) error { + return m.DeleteResourceGroup(1, group) + }, + }, + { + name: "add", + group: "ct-add-abort", + failpoint: "addResourceGroupBeforeStorage", + mutate: func(m *Manager, group string) error { + updated := newAsyncTestGroup(group) + updated.KeyspaceId = &resource_manager.KeyspaceIDValue{Keyspace: &resource_manager.KeyspaceIDValue_Value{Value: 1}} + updated.RUSettings.RU.Settings.FillRate = asyncTestGroupFillRate * 2 + return m.AddResourceGroup(updated) + }, + }, + { + name: "modify", + group: "ct-modify-abort", + failpoint: "modifyResourceGroupBeforeStorage", + mutate: func(m *Manager, group string) error { + modified := newAsyncTestGroup(group) + modified.KeyspaceId = &resource_manager.KeyspaceIDValue{Keyspace: &resource_manager.KeyspaceIDValue_Value{Value: 1}} + modified.RUSettings.RU.Settings.FillRate = asyncTestGroupFillRate * 2 + return m.ModifyResourceGroup(modified) + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + re := require.New(t) + store := newBlockingResourceGroupStorage() + re.NoError(store.SaveResourceGroupSetting(1, tc.group, newAsyncTestGroup(tc.group))) + + m := NewManager[*mockConfigProvider](&mockConfigProvider{}) + m.storage = store + re.NoError(m.Init(context.Background())) + cancelTerm1 := m.cancel + defer stopAsyncTestManager(m) + defer store.unblock() + + // Let term 1 load fully so the mutation starts against a settled term. + store.waitEntered(t) + store.unblock() + testutil.Eventually(re, func() bool { + groups, err := m.GetResourceGroupList(1, false) + return err == nil && len(groups) == 2 + }, testutil.WithTickInterval(20*time.Millisecond)) + + // Park the mutation between resolving its keyspace manager and its + // storage phase. + reached := make(chan struct{}) + release := make(chan struct{}) + fpPath := "github.com/tikv/pd/pkg/mcs/resourcemanager/server/" + tc.failpoint + re.NoError(failpoint.EnableCall(fpPath, func() { + close(reached) + <-release + })) + defer func() { + re.NoError(failpoint.Disable(fpPath)) + }() + var mutErr error + mutDone := make(chan struct{}) + go func() { + defer close(mutDone) + mutErr = tc.mutate(m, tc.group) + }() + select { + case <-reached: + case <-time.After(time.Second): + t.Fatal("timed out waiting for the mutation to reach its storage phase") + } + + // Leadership changes and term 2 fully completes loading. Storage is + // unchanged (the mutation never reached its storage phase), so term + // 2's own reload confirms the group in the new term's cache. + cancelTerm1() + re.NoError(m.Init(context.Background())) + testutil.Eventually(re, func() bool { + groups, err := m.GetResourceGroupList(1, false) + return err == nil && len(groups) == 2 + }, testutil.WithTickInterval(20*time.Millisecond)) + + // Resume the parked mutation: it must see term 2 already has + // confirmed data for this group and abort, instead of touching + // storage on behalf of the detached term-1 manager. + close(release) + <-mutDone + re.ErrorIs(mutErr, errs.ErrResourceGroupsLoading) + + g, err := m.GetResourceGroup(1, tc.group, false) + re.NoError(err) + re.NotNil(g, "the aborted mutation must not have removed the group from storage") + re.Equal(float64(asyncTestGroupFillRate), g.RUSettings.RU.getFillRate(), + "the aborted mutation must not have overwritten storage with a stale value") + }) + } +} + +// TestAsyncLoadResourceGroupsExhaustedRetriesReturnLoadingError guards the +// exhausted-retry path of the lazy load: when every attempt loses the +// delete-generation race, the load must fail with a retryable loading error +// instead of reporting success without publishing the group, which callers +// would misread as the group not existing. +func TestAsyncLoadResourceGroupsExhaustedRetriesReturnLoadingError(t *testing.T) { + re := require.New(t) + store := newBlockingResourceGroupStorage() + re.NoError(store.SaveResourceGroupSetting(1, "keep-a", newAsyncTestGroup("keep-a"))) + for _, name := range []string{"del-b1", "del-b2", "del-b3"} { + re.NoError(store.SaveResourceGroupSetting(1, name, newAsyncTestGroup(name))) + } + + m := NewManager[*mockConfigProvider](&mockConfigProvider{}) + m.storage = store + re.NoError(m.Init(context.Background())) + defer stopAsyncTestManager(m) + defer store.unblock() + + // Keep the bulk loader parked so lazy loading stays active. + store.waitEntered(t) + + // Park keep-a's lazy load inside each of its three read attempts, and + // delete an unrelated group while it's parked so every attempt observes a + // delete-generation change. + pause := store.armStatePause("keep-a") + var ( + gotGroup *ResourceGroup + gotErr error + ) + getDone := make(chan struct{}) + go func() { + defer close(getDone) + gotGroup, gotErr = m.GetResourceGroup(1, "keep-a", false) + }() + for _, victim := range []string{"del-b1", "del-b2", "del-b3"} { + waitStatePauseReached(t, pause) + re.NoError(m.DeleteResourceGroup(1, victim)) + next := store.armStatePause("keep-a") + close(pause.release) + pause = next + } + <-getDone + // The last arm is left unconsumed; drop it so later loads pass through. + store.statePause.Store(nil) + + re.ErrorIs(gotErr, errs.ErrResourceGroupsLoading, + "exhausted retries must surface a retryable loading error, not a bogus success") + re.Nil(gotGroup) + + // The group still exists; once loading completes it must be served again. + store.unblock() + testutil.Eventually(re, func() bool { + g, err := m.GetResourceGroup(1, "keep-a", false) + return err == nil && g != nil + }, testutil.WithTickInterval(20*time.Millisecond)) +} + +// TestAsyncLoadResourceGroupsCrossTermModifyDefaultStaysConfirmed guards the +// Modify-of-default publish path across a leadership change. A Modify patches +// and persists the default group in term 1, then stalls before publishing. +// Term 2 reinitializes the manager, giving it a fresh reserved default +// placeholder. When the Modify resumes, publishing must leave the group +// confirmed (not reserved) in the live term with the modified settings - +// otherwise it stays a reserved placeholder that the bulk merge or +// initReserved can revert to a pre-modification/synthetic default while +// storage keeps the new value. +func TestAsyncLoadResourceGroupsCrossTermModifyDefaultStaysConfirmed(t *testing.T) { + re := require.New(t) + store := newBlockingResourceGroupStorage() + // Seed a persisted default with a recognizable running state so the test + // can tell a confirmed republish (state preserved) from a synthetic + // placeholder (state reset). + seed := newAsyncTestGroup(DefaultResourceGroupName) + re.NoError(store.SaveResourceGroupSetting(constant.NullKeyspaceID, DefaultResourceGroupName, seed)) + seedStates := FromProtoResourceGroup(seed).GetGroupStates() + seedStates.RUConsumption.RRU = 777 + re.NoError(store.SaveResourceGroupStates(constant.NullKeyspaceID, DefaultResourceGroupName, seedStates)) + + m := NewManager[*mockConfigProvider](&mockConfigProvider{}) + m.storage = store + re.NoError(m.Init(context.Background())) + cancelTerm1 := m.cancel + defer stopAsyncTestManager(m) + defer store.unblock() + defer store.unblockStates() + + // Let term 1 load fully so the default group is confirmed and persisted. + store.waitEntered(t) + store.unblock() + testutil.Eventually(re, func() bool { + _, err := m.GetResourceGroupList(constant.NullKeyspaceID, false) + return err == nil + }, testutil.WithTickInterval(20*time.Millisecond)) + + // Park the Modify after it patched and persisted, before it publishes. + reached := make(chan struct{}) + release := make(chan struct{}) + re.NoError(failpoint.EnableCall("github.com/tikv/pd/pkg/mcs/resourcemanager/server/modifyResourceGroupBeforePublish", func() { + close(reached) + <-release + })) + defer func() { + re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/mcs/resourcemanager/server/modifyResourceGroupBeforePublish")) + }() + modified := newAsyncTestGroup(DefaultResourceGroupName) + modified.RUSettings.RU.Settings.FillRate = 4242 + var modErr error + modDone := make(chan struct{}) + go func() { + defer close(modDone) + modErr = m.ModifyResourceGroup(modified) + }() + select { + case <-reached: + case <-time.After(time.Second): + t.Fatal("timed out waiting for the modify to reach its publish phase") + } + + // Leadership changes: term 2 gets a fresh reserved default placeholder, + // with its loader parked before merging. + cancelTerm1() + store.pauseNextStates.Store(true) + re.NoError(m.Init(context.Background())) + select { + case <-store.statesReached: + case <-time.After(time.Second): + t.Fatal("timed out waiting for the term-2 loader to snapshot storage") + } + + // Resume the Modify's publish into term 2. + close(release) + <-modDone + re.NoError(modErr) + + // The default in term 2 must be confirmed (not a reserved placeholder), + // so neither the merge nor initReserved reverts the modified settings. + krgm := m.getKeyspaceResourceGroupManager(constant.NullKeyspaceID) + re.NotNil(krgm) + re.False(krgm.isReserved(DefaultResourceGroupName), + "a modified default must be published as confirmed data, not left reserved") + + store.unblockStates() + testutil.Eventually(re, func() bool { + _, err := m.GetResourceGroupList(constant.NullKeyspaceID, false) + return err == nil + }, testutil.WithTickInterval(20*time.Millisecond)) + g, err := m.GetResourceGroup(constant.NullKeyspaceID, DefaultResourceGroupName, false) + re.NoError(err) + re.NotNil(g) + re.Equal(float64(4242), g.RUSettings.RU.getFillRate(), + "the modified default settings must survive into the new term") + // The confirmed running state must survive too, not revert to the fresh + // synthetic placeholder state. + re.Equal(float64(777), krgm.getMutableResourceGroup(DefaultResourceGroupName).GetGroupStates().RUConsumption.RRU, + "the confirmed running state must be preserved, not reset to a synthetic placeholder") +} + +// TestInitDefaultResourceGroupMarksAtomicallyWithPublish guards against a +// window where a synthesized default group became visible in the cache +// before its sync-loaded marker was set. initDefaultResourceGroup used to +// publish through krgm's own lock directly, with the caller (loadResourceGroupIfNeeded) +// setting the sync-loaded marker afterward in a separate, later m.Lock() +// critical section. A concurrent bulk-merge batch - which only skips an item +// already marked - could run in that window and overwrite the synthesized +// group, including any live consumption/token update applied to it in the +// meantime, with its own possibly-stale scanned copy. initDefaultResourceGroup +// now publishes through publishResourceGroupMutation, which holds m.Lock() +// across both the cache-visibility change and the marker set, so nothing +// that also needs m.Lock() - including a merge batch - can ever observe one +// without the other. +func TestInitDefaultResourceGroupMarksAtomicallyWithPublish(t *testing.T) { + re := require.New(t) + m := prepareManager() + const keyspaceID = 1 + krgm := newKeyspaceResourceGroupManager(keyspaceID, m.storage, m.writeRole) + m.krgms[keyspaceID] = krgm + + reached := make(chan struct{}) + release := make(chan struct{}) + re.NoError(failpoint.EnableCall("github.com/tikv/pd/pkg/mcs/resourcemanager/server/publishMutationBeforeMark", func() { + close(reached) + <-release + })) + defer func() { + re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/mcs/resourcemanager/server/publishMutationBeforeMark")) + }() + + var initErr error + initDone := make(chan struct{}) + go func() { + defer close(initDone) + _, initErr = m.initDefaultResourceGroup(keyspaceID, krgm, nil) + }() + + select { + case <-reached: + case <-time.After(time.Second): + t.Fatal("timed out waiting for initDefaultResourceGroup to reach its publish") + } + + // While parked between the cache-visibility change and the marker set, + // nothing that needs m.Lock() (e.g. a concurrent merge batch checking + // whether to skip this group) should be able to proceed. + checkDone := make(chan bool) + go func() { + m.RLock() + _, marked := m.syncLoadedGroups[trackerKey{keyspaceID: keyspaceID, groupName: DefaultResourceGroupName}] + m.RUnlock() + checkDone <- marked + }() + + select { + case <-checkDone: + t.Fatal("a concurrent m.Lock()-holding operation must not proceed while the publish is parked between visibility and marking") + case <-time.After(100 * time.Millisecond): + } + + close(release) + <-initDone + re.NoError(initErr) + marked := <-checkDone + re.True(marked, "the marker must already be set by the time any concurrent m.Lock() holder can observe the published group") +} + +// TestAsyncLoadResourceGroupsCrossTermSetServiceLimitPublishesToNewTerm +// reproduces a leadership change straddling SetKeyspaceServiceLimit: it +// resolves a keyspace manager, then stalls before its storage phase while the +// leadership change replaces m.krgms and the new term's synchronous +// loadServiceLimits reads the still-unwritten old value. When the call +// resumes, its storage write must still land, and its cache effect must be +// mirrored into the *current* term's keyspace manager, not just the detached +// term-1 one - otherwise the new term keeps serving the stale (default zero) +// limit indefinitely, with nothing to ever re-sync it. +func TestAsyncLoadResourceGroupsCrossTermSetServiceLimitPublishesToNewTerm(t *testing.T) { + re := require.New(t) + store := newBlockingResourceGroupStorage() + + m := NewManager[*mockConfigProvider](&mockConfigProvider{}) + m.storage = store + re.NoError(m.Init(context.Background())) + cancelTerm1 := m.cancel + defer stopAsyncTestManager(m) + defer store.unblock() + + // Let term 1 load fully so the call starts against a settled term. + store.waitEntered(t) + store.unblock() + testutil.Eventually(re, func() bool { + _, err := m.GetResourceGroupList(constant.NullKeyspaceID, false) + return err == nil + }, testutil.WithTickInterval(20*time.Millisecond)) + + // Park SetKeyspaceServiceLimit between resolving its keyspace manager and + // its storage phase. + reached := make(chan struct{}) + release := make(chan struct{}) + re.NoError(failpoint.EnableCall("github.com/tikv/pd/pkg/mcs/resourcemanager/server/setServiceLimitBeforeStorage", func() { + close(reached) + <-release + })) + defer func() { + re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/mcs/resourcemanager/server/setServiceLimitBeforeStorage")) + }() + var setErr error + setDone := make(chan struct{}) + go func() { + defer close(setDone) + setErr = m.SetKeyspaceServiceLimit(constant.NullKeyspaceID, 4242) + }() + select { + case <-reached: + case <-time.After(time.Second): + t.Fatal("timed out waiting for SetKeyspaceServiceLimit to reach its storage phase") + } + + // Leadership changes while the write is parked: term 2's synchronous + // loadServiceLimits runs and completes here, reading storage before the + // parked write has persisted anything. + cancelTerm1() + re.NoError(m.Init(context.Background())) + + // Resume the write: it persists into storage now, then must publish into + // whichever keyspace manager is current (term 2), not the detached term-1 + // one. + close(release) + <-setDone + re.NoError(setErr) + + limiter := m.GetKeyspaceServiceLimiter(constant.NullKeyspaceID) + re.NotNil(limiter) + re.Equal(float64(4242), limiter.ServiceLimit, + "the service limit set during the leadership change must be visible in the new term, not stuck on the detached old one") + + raw, err := store.LoadServiceLimit(constant.NullKeyspaceID) + re.NoError(err) + re.Equal(float64(4242), raw, "the service limit must be persisted regardless of the leadership change") +} + +// TestAsyncLoadResourceGroupsCrossTermSetServiceLimitSerializesAgainstCompetingCall +// guards against the unconditional publish-phase mirror in +// SetKeyspaceServiceLimit clobbering a competing, fully-completed call for +// the same keyspace. Without serviceLimitLocks, an old-term call parked +// mid-persist could resume after a new-term call for the same keyspace +// already persisted and published its own value, and mirror its own older +// value back in - leaving the live cache stale even though storage (and +// every other observer) already moved on. serviceLimitLocks makes the +// new-term call wait for the old one to fully finish, including its own +// mirror step, before it can even start, so this interleaving can no longer +// happen. +func TestAsyncLoadResourceGroupsCrossTermSetServiceLimitSerializesAgainstCompetingCall(t *testing.T) { + re := require.New(t) + store := newBlockingResourceGroupStorage() + + m := NewManager[*mockConfigProvider](&mockConfigProvider{}) + m.storage = store + re.NoError(m.Init(context.Background())) + cancelTerm1 := m.cancel + defer stopAsyncTestManager(m) + defer store.unblock() + + store.waitEntered(t) + store.unblock() + testutil.Eventually(re, func() bool { + _, err := m.GetResourceGroupList(constant.NullKeyspaceID, false) + return err == nil + }, testutil.WithTickInterval(20*time.Millisecond)) + + // Park the old-term call between resolving its keyspace manager and its + // storage phase - it already holds serviceLimitLocks for this keyspace + // by this point. The new-term call goes through this same, still-armed + // failpoint too once serviceLimitLocks lets it proceed; parkOnce keeps + // only the first (old-term) hit actually parking, so the second one + // passes straight through instead of trying to close(reached) again. + reached := make(chan struct{}) + release := make(chan struct{}) + var parkOnce sync.Once + re.NoError(failpoint.EnableCall("github.com/tikv/pd/pkg/mcs/resourcemanager/server/setServiceLimitBeforeStorage", func() { + parkOnce.Do(func() { + close(reached) + <-release + }) + })) + defer func() { + re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/mcs/resourcemanager/server/setServiceLimitBeforeStorage")) + }() + var oldErr error + oldDone := make(chan struct{}) + go func() { + defer close(oldDone) + oldErr = m.SetKeyspaceServiceLimit(constant.NullKeyspaceID, 100) + }() + select { + case <-reached: + case <-time.After(time.Second): + t.Fatal("timed out waiting for the old-term call to reach its storage phase") + } + + // Leadership changes while the old-term call is parked. + cancelTerm1() + re.NoError(m.Init(context.Background())) + + // A new-term call for the same keyspace must not be able to run while the + // old-term call still holds serviceLimitLocks for it: it should block + // before even reaching the (still-armed) failpoint above, since + // serviceLimitLocks is acquired first. + var newErr error + newDone := make(chan struct{}) + go func() { + defer close(newDone) + newErr = m.SetKeyspaceServiceLimit(constant.NullKeyspaceID, 200) + }() + select { + case <-newDone: + t.Fatal("the new-term call must be blocked by serviceLimitLocks while the old-term call is still parked") + case <-time.After(100 * time.Millisecond): + } + + // Resume the old-term call: it persists and publishes 100, then releases + // serviceLimitLocks, letting the new-term call finally run and persist + // and publish 200. + close(release) + <-oldDone + re.NoError(oldErr) + <-newDone + re.NoError(newErr) + + limiter := m.GetKeyspaceServiceLimiter(constant.NullKeyspaceID) + re.NotNil(limiter) + re.Equal(float64(200), limiter.ServiceLimit, + "the new-term call must run - and win - only after the old-term call fully finished, not race ahead of or get clobbered by it") +} + +// TestLoadServiceLimitsDoesNotClobberConcurrentSet guards against a race +// between loadServiceLimits' bulk replay (run on every Init/leadership +// change) and a concurrent SetKeyspaceServiceLimit call for the same +// keyspace. loadServiceLimits used to apply the value its bulk storage scan +// had already read before doing any locking; if a concurrent +// SetKeyspaceServiceLimit call persisted and mirrored a newer value while the +// replay's callback for that keyspace was still in flight, the replay would +// silently overwrite the cache with its own now-stale snapshot, leaving the +// cache stuck behind storage until the next full reload. serviceLimitLocks +// alone (the fix for the sibling cross-term race above) does not close this: +// the stale value here was captured before any lock was ever taken, so +// merely serializing the two callers does not stop the replay from applying +// data that was already out of date the moment it read it. The fix instead +// re-reads the keyspace's service limit from storage under the same lock, +// discarding the bulk scan's value entirely. +func TestLoadServiceLimitsDoesNotClobberConcurrentSet(t *testing.T) { + re := require.New(t) + m := prepareManager() + const keyspaceID = 1 + re.NoError(m.storage.SaveServiceLimit(keyspaceID, 100)) + + reached := make(chan struct{}) + release := make(chan struct{}) + re.NoError(failpoint.EnableCall("github.com/tikv/pd/pkg/mcs/resourcemanager/server/loadServiceLimitsBeforeApply", func(gotKeyspaceID uint32) { + if gotKeyspaceID != keyspaceID { + return + } + close(reached) + <-release + })) + defer func() { + re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/mcs/resourcemanager/server/loadServiceLimitsBeforeApply")) + }() + + loadDone := make(chan struct{}) + go func() { + defer close(loadDone) + re.NoError(m.loadServiceLimits()) + }() + select { + case <-reached: + case <-time.After(time.Second): + t.Fatal("timed out waiting for loadServiceLimits to reach its replay callback") + } + + // While the replay is parked - holding only the (100) value its scan + // already captured - a concurrent SetKeyspaceServiceLimit call persists + // and mirrors a newer value into the same, live keyspace manager. + re.NoError(m.SetKeyspaceServiceLimit(keyspaceID, 200)) + + close(release) + <-loadDone + + limiter := m.GetKeyspaceServiceLimiter(keyspaceID) + re.NotNil(limiter) + re.Equal(float64(200), limiter.ServiceLimit, + "the replay must not overwrite a concurrently-set newer value with its own stale snapshot") +} + +// failingServiceLimitLoadStorage makes LoadServiceLimit (the point re-read +// loadServiceLimits performs under serviceLimitLocks) fail a controlled +// number of times for one keyspace, to exercise its retry-then-fallback +// behavior. LoadServiceLimits (the bulk scan) is left untouched. +type failingServiceLimitLoadStorage struct { + storage.Storage + keyspaceID uint32 + failuresLeft atomic.Int32 + calls atomic.Int32 +} + +func (s *failingServiceLimitLoadStorage) LoadServiceLimit(keyspaceID uint32) (float64, error) { + if keyspaceID == s.keyspaceID { + s.calls.Add(1) + if s.failuresLeft.Add(-1) >= 0 { + return 0, errors.New("injected service limit load failure") + } + } + return s.Storage.LoadServiceLimit(keyspaceID) +} + +// TestLoadServiceLimitsRetriesPointReadBeforeFallingBack guards against +// loadServiceLimits treating a single point-read failure as fatal. The point +// re-read added to close the race in TestLoadServiceLimitsDoesNotClobberConcurrentSet +// above can itself fail transiently (e.g. a storage blip), so it must retry a +// few times and apply the value once a retry succeeds, instead of giving up +// and falling back to the (potentially stale) bulk-scanned value on the very +// first failure. +func TestLoadServiceLimitsRetriesPointReadBeforeFallingBack(t *testing.T) { + re := require.New(t) + const keyspaceID = 1 + base := storage.NewStorageWithMemoryBackend() + re.NoError(base.SaveServiceLimit(keyspaceID, 100)) + store := &failingServiceLimitLoadStorage{Storage: base, keyspaceID: keyspaceID} + // Fail once, then succeed on the second attempt - strictly fewer calls + // than the retry budget, to prove it stops retrying once a read succeeds + // instead of always spending the full budget. + store.failuresLeft.Store(1) + + m := NewManager[*mockConfigProvider](&mockConfigProvider{}) + m.storage = store + + re.NoError(m.loadServiceLimits()) + + limiter := m.GetKeyspaceServiceLimiter(keyspaceID) + re.NotNil(limiter) + re.Equal(float64(100), limiter.ServiceLimit, + "a point read that succeeds within the retry budget must be applied") + re.EqualValues(2, store.calls.Load(), + "must retry the point read instead of falling back after only one failure, and stop once it succeeds") +} + +// TestLoadServiceLimitsFallsBackToBulkValueOnPersistentFailure guards against +// loadServiceLimits silently dropping a keyspace's service limit entirely +// when its point re-read keeps failing (e.g. a persistent storage issue). +// Without a fallback, the keyspace would run with no service limit cached at +// all - letting burstable groups bypass the configured cap indefinitely, +// since nothing else retries this load until the next Init - which is worse +// than falling back to the bulk-scanned value and re-admitting its narrow, +// already-covered staleness window. +func TestLoadServiceLimitsFallsBackToBulkValueOnPersistentFailure(t *testing.T) { + re := require.New(t) + const keyspaceID = 1 + base := storage.NewStorageWithMemoryBackend() + re.NoError(base.SaveServiceLimit(keyspaceID, 100)) + store := &failingServiceLimitLoadStorage{Storage: base, keyspaceID: keyspaceID} + // Always fail the point read, however many times it's retried. + store.failuresLeft.Store(math.MaxInt32) + + m := NewManager[*mockConfigProvider](&mockConfigProvider{}) + m.storage = store + + re.NoError(m.loadServiceLimits()) + + limiter := m.GetKeyspaceServiceLimiter(keyspaceID) + re.NotNil(limiter) + re.Equal(float64(100), limiter.ServiceLimit, + "must fall back to the bulk-scanned value rather than leaving the service limit uncached") + re.EqualValues(maxServiceLimitReloadAttempts, store.calls.Load(), + "must exhaust the retry budget before falling back") +} + +// TestLoadServiceLimitsDoesNotClobberConcurrentSetOnPersistentPointReadFailure +// guards against the bulk-scanned fallback above overwriting a newer value +// that a concurrent SetKeyspaceServiceLimit call already fully committed - to +// storage and the live cache - before loadServiceLimits' callback for this +// keyspace ever acquired serviceLimitLocks. SetKeyspaceServiceLimit holds +// that same per-keyspace lock across both its storage write and its cache +// mirror, so there is no partially-applied state the callback can observe: +// by the time it holds the lock, a concurrent call has either fully landed +// (cache already reflects a value strictly newer than the bulk scan's +// pre-lock snapshot) or has not started at all. The fix checks the cache +// before falling back and leaves an already-newer value untouched. +func TestLoadServiceLimitsDoesNotClobberConcurrentSetOnPersistentPointReadFailure(t *testing.T) { + re := require.New(t) + const keyspaceID = 1 + base := storage.NewStorageWithMemoryBackend() + re.NoError(base.SaveServiceLimit(keyspaceID, 100)) + store := &failingServiceLimitLoadStorage{Storage: base, keyspaceID: keyspaceID} + // Always fail the point read, however many times it's retried. + store.failuresLeft.Store(math.MaxInt32) + + m := NewManager[*mockConfigProvider](&mockConfigProvider{}) + m.storage = store + + reached := make(chan struct{}) + release := make(chan struct{}) + re.NoError(failpoint.EnableCall("github.com/tikv/pd/pkg/mcs/resourcemanager/server/loadServiceLimitsBeforeApply", func(gotKeyspaceID uint32) { + if gotKeyspaceID != keyspaceID { + return + } + close(reached) + <-release + })) + defer func() { + re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/mcs/resourcemanager/server/loadServiceLimitsBeforeApply")) + }() + + loadDone := make(chan struct{}) + go func() { + defer close(loadDone) + re.NoError(m.loadServiceLimits()) + }() + select { + case <-reached: + case <-time.After(time.Second): + t.Fatal("timed out waiting for loadServiceLimits to reach its replay callback") + } + + // While the replay is parked before ever acquiring serviceLimitLocks, a + // concurrent SetKeyspaceServiceLimit call fully completes - both its + // storage write and its cache mirror - for the same keyspace. + re.NoError(m.SetKeyspaceServiceLimit(keyspaceID, 200)) + + close(release) + <-loadDone + + limiter := m.GetKeyspaceServiceLimiter(keyspaceID) + re.NotNil(limiter) + re.Equal(float64(200), limiter.ServiceLimit, + "a concurrently-completed newer write must not be overwritten by the stale bulk-scanned value even when the point read keeps failing") +} + +// TestAsyncLoadResourceGroupsMergeGroupNeverVisibleUnsynced guards against a +// window in the bulk merge where a newly-loaded group became visible in +// krgm.groups (readable by any concurrent GetResourceGroup/token request) +// before its burst limit was synced against the keyspace's active service +// limit. A group with no explicit burst limit reads as unbounded until that +// sync runs, so a request landing in the gap could be granted unlimited +// burst and bypass the service limit until the merge caught up. The fix +// moved the sync inside the same krgm.Lock() critical section as the insert, +// so a concurrent reader - which also needs krgm's lock - can no longer +// observe the group until both have happened. +func TestAsyncLoadResourceGroupsMergeGroupNeverVisibleUnsynced(t *testing.T) { + re := require.New(t) + store := storage.NewStorageWithMemoryBackend() + const keyspaceID = 1 + const groupName = "burstable-rg" + unbounded := &resource_manager.ResourceGroup{ + Name: groupName, + Mode: resource_manager.GroupMode_RUMode, + Priority: middlePriority, + RUSettings: &resource_manager.GroupRequestUnitSettings{ + RU: &resource_manager.TokenBucket{ + Settings: &resource_manager.TokenLimitSettings{ + FillRate: UnlimitedRate, + BurstLimit: UnlimitedBurstLimit, + }, + }, + }, + } + re.NoError(store.SaveResourceGroupSetting(keyspaceID, groupName, unbounded)) + re.NoError(store.SaveServiceLimit(keyspaceID, 50)) + + m := NewManager[*mockConfigProvider](&mockConfigProvider{}) + m.storage = store + + reached := make(chan struct{}) + release := make(chan struct{}) + var releaseOnce sync.Once + unblock := func() { releaseOnce.Do(func() { close(release) }) } + re.NoError(failpoint.EnableCall("github.com/tikv/pd/pkg/mcs/resourcemanager/server/mergeBeforeBurstSync", func(gotKeyspaceID uint32, gotName string) { + if gotKeyspaceID != keyspaceID || gotName != groupName { + return + } + close(reached) + <-release + })) + defer func() { + re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/mcs/resourcemanager/server/mergeBeforeBurstSync")) + }() + + re.NoError(m.Init(context.Background())) + defer stopAsyncTestManager(m) + // Unblock the merge first (LIFO), so stopAsyncTestManager's wg.Wait() + // cannot hang if a later assertion aborts the test before the explicit + // unblock() call below is reached. + defer unblock() + + // loadServiceLimits (run synchronously inside Init, before the async + // merge starts) already created krgm for keyspaceID via the service + // limit saved above - capture it now, before parking, so the reader + // below can call krgm.getResourceGroup directly instead of going through + // m.getKeyspaceResourceGroupManager. The latter needs m.RLock(), which + // the merge holds for its *entire* batch regardless of this fix, so + // routing the read through it would block on the wrong lock and the test + // would pass even without the fix; reading via the captured krgm + // isolates the one lock (krgm's own) this test is actually about. + krgm := m.getKeyspaceResourceGroupManager(keyspaceID) + re.NotNil(krgm) + + select { + case <-reached: + case <-time.After(time.Second): + t.Fatal("timed out waiting for the merge to reach its burst-sync point") + } + + readDone := make(chan *ResourceGroup) + go func() { + readDone <- krgm.getResourceGroup(groupName, false) + }() + + select { + case <-readDone: + t.Fatal("a concurrent reader must not observe the group while the merge is parked between insert and burst sync - krgm's lock should still be held") + case <-time.After(100 * time.Millisecond): + } + + unblock() + group := <-readDone + re.NotNil(group) + re.GreaterOrEqual(group.getOverrideBurstLimit(), int64(0), + "the group must never become visible without its burst override already synced") +} + +// BenchmarkAsyncLoadMergeReaderStall measures the worst-case time a concurrent +// reader is blocked while the async bulk merge installs a large number of +// resource groups. The probe uses GetControllerConfig, whose only cost is the +// manager read lock - the same lock the merge takes - so it isolates how long +// the merge stalls readers. It guards against the merge holding that lock +// across the whole O(total groups) work, which would stall every point and +// token request until loading completes on a cluster with many groups. +// +// Run with: +// +// go test -run '^$' -bench BenchmarkAsyncLoadMergeReaderStall ./pkg/mcs/resourcemanager/server/ +func BenchmarkAsyncLoadMergeReaderStall(b *testing.B) { + const groupCount = 500000 + store := newBlockingResourceGroupStorage() + for i := range groupCount { + name := fmt.Sprintf("bench-group-%06d", i) + if err := store.SaveResourceGroupSetting(1, name, newAsyncTestGroup(name)); err != nil { + b.Fatal(err) + } + } + + m := NewManager[*mockConfigProvider](&mockConfigProvider{}) + m.storage = store + if err := m.Init(context.Background()); err != nil { + b.Fatal(err) + } + defer stopAsyncTestManager(m) + defer store.unblock() + + store.waitEntered(b) + + // Concurrent readers take only the manager read lock and record the + // longest single acquisition seen while the merge runs. + stop := make(chan struct{}) + var maxStall atomic.Int64 + var wg sync.WaitGroup + for range 8 { + wg.Add(1) + go func() { + defer wg.Done() + for { + select { + case <-stop: + return + default: + } + start := time.Now() + _ = m.GetControllerConfig() + if d := time.Since(start).Nanoseconds(); d > maxStall.Load() { + maxStall.Store(d) + } + } + }() + } + + b.ResetTimer() + store.unblock() + deadline := time.Now().Add(30 * time.Second) + for !m.isResourceGroupLoadingComplete() { + if time.Now().After(deadline) { + b.Fatal("timed out waiting for async loading to complete") + } + time.Sleep(time.Millisecond) + } + b.StopTimer() + + close(stop) + wg.Wait() + b.ReportMetric(float64(maxStall.Load())/1e6, "max-reader-stall-ms") +} + +// fakeTokenBucketsStream feeds a fixed set of requests to AcquireTokenBuckets +// and records what it sends back. grpc.ServerStream stays nil since +// AcquireTokenBuckets only ever calls Send/Recv on it. +type fakeTokenBucketsStream struct { + grpc.ServerStream + + requests []*resource_manager.TokenBucketsRequest + recvCnt int + sent []*resource_manager.TokenBucketsResponse +} + +// Context is needed by the metrics stream wrapper, which resolves the peer IP +// for its labels. +func (*fakeTokenBucketsStream) Context() context.Context { return context.Background() } + +func (s *fakeTokenBucketsStream) Recv() (*resource_manager.TokenBucketsRequest, error) { + if s.recvCnt >= len(s.requests) { + return nil, io.EOF + } + req := s.requests[s.recvCnt] + s.recvCnt++ + return req, nil +} + +func (s *fakeTokenBucketsStream) Send(resp *resource_manager.TokenBucketsResponse) error { + s.sent = append(s.sent, resp) + return nil +} + +func newRUTokenBucketRequest(keyspaceID uint32, name string, ru float64) *resource_manager.TokenBucketRequest { + return &resource_manager.TokenBucketRequest{ + ResourceGroupName: name, + KeyspaceId: &resource_manager.KeyspaceIDValue{Keyspace: &resource_manager.KeyspaceIDValue_Value{Value: keyspaceID}}, + Request: &resource_manager.TokenBucketRequest_RuItems{ + RuItems: &resource_manager.TokenBucketRequest_RequestRU{ + RequestRU: []*resource_manager.RequestUnitItem{ + {Type: resource_manager.RequestUnitType_RU, Value: ru}, + }, + }, + }, + ConsumptionSinceLastRequest: &resource_manager.Consumption{}, + } +} + +// TestAcquireTokenBucketsSurvivesLazyLoadFailure guards against a transient +// lazy-load failure tearing down the whole token bucket stream: the error +// belongs to a single resource group, so the other groups multiplexed on the +// same stream must still be served instead of every client being forced to +// reconnect. +func TestAcquireTokenBucketsSurvivesLazyLoadFailure(t *testing.T) { + re := require.New(t) + store := newBlockingResourceGroupStorage() + for _, name := range []string{"bad-group", "good-group"} { + group := newAsyncTestGroup(name) + re.NoError(store.SaveResourceGroupSetting(1, name, group)) + re.NoError(store.SaveResourceGroupStates(1, name, FromProtoResourceGroup(group).GetGroupStates())) + } + + m := NewManager[*mockConfigProvider](&mockConfigProvider{}) + m.storage = store + m.srv = &testBasicServer{} + re.NoError(m.Init(context.Background())) + defer stopAsyncTestManager(m) + defer store.unblock() + + // Keep the bulk loader parked so lazy loading stays active. + store.waitEntered(t) + + // Make bad-group's lazy load fail on its state read. good-group is + // requested in the same batch and must still get its tokens. + store.failNextState.Store(true) + stream := &fakeTokenBucketsStream{ + requests: []*resource_manager.TokenBucketsRequest{{ + TargetRequestPeriodMs: 1000, + ClientUniqueId: 1, + Requests: []*resource_manager.TokenBucketRequest{ + newRUTokenBucketRequest(1, "bad-group", 10), + newRUTokenBucketRequest(1, "good-group", 10), + }, + }}, + } + svc := &Service{ctx: context.Background(), manager: m} + + re.NoError(svc.AcquireTokenBuckets(stream), + "a single group's lazy-load failure must not fail the stream") + re.Len(stream.sent, 1) + re.Len(stream.sent[0].Responses, 1, "only the loadable group should be answered") + re.Equal("good-group", stream.sent[0].Responses[0].ResourceGroupName) +} + +// TestSetKeyspaceRUVersionSerializesAgainstUpdateControllerConfigItem guards +// against SetKeyspaceRUVersion mutating RUVersionPolicy.Overrides and saving +// it to storage outside its lock. That used to capture a reference to the +// live (not cloned) controller config, unlock, and only then call +// storage.SaveControllerConfig - leaving Overrides, a plain map, both +// mutable by a concurrent SetKeyspaceRUVersion call and readable by this +// call's own unlocked marshal (an unsynchronized concurrent map read/write), +// and letting this call's save land after UpdateControllerConfigItem's own +// clone-save-publish (which runs fully inside its lock), silently +// overwriting UpdateControllerConfigItem's newer write. The fix clones +// before mutating and holds the lock across the whole +// mutate-save-publish sequence, so a concurrent config mutator can't even +// start until this call is done - proven here by asserting +// UpdateControllerConfigItem blocks while SetKeyspaceRUVersion is parked +// right before its save, and that neither call's change is lost once both +// complete. +func TestSetKeyspaceRUVersionSerializesAgainstUpdateControllerConfigItem(t *testing.T) { + re := require.New(t) + m := prepareManager() + + reached := make(chan struct{}) + release := make(chan struct{}) + re.NoError(failpoint.EnableCall("github.com/tikv/pd/pkg/mcs/resourcemanager/server/setKeyspaceRUVersionBeforeSave", func() { + close(reached) + <-release + })) + defer func() { + re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/mcs/resourcemanager/server/setKeyspaceRUVersionBeforeSave")) + }() + + setDone := make(chan error, 1) + go func() { + setDone <- m.SetKeyspaceRUVersion(1, 2) + }() + select { + case <-reached: + case <-time.After(time.Second): + t.Fatal("timed out waiting for SetKeyspaceRUVersion to reach its pre-save failpoint") + } + + // While SetKeyspaceRUVersion is parked holding the lock, a concurrent + // UpdateControllerConfigItem call must block rather than run - proving + // the two are mutually exclusive across the whole save, not just the + // in-memory mutation step. + updateDone := make(chan error, 1) + go func() { + updateDone <- m.UpdateControllerConfigItem("request-unit.read-base-cost", 1.5) + }() + select { + case err := <-updateDone: + t.Fatalf("UpdateControllerConfigItem completed (err=%v) while SetKeyspaceRUVersion still held the lock", err) + case <-time.After(100 * time.Millisecond): + } + + close(release) + re.NoError(<-setDone) + re.NoError(<-updateDone) + + re.InDelta(1.5, m.controllerConfig.RequestUnit.ReadBaseCost, 0.00001, + "UpdateControllerConfigItem's change must survive") + re.Equal(RUVersion(2), m.controllerConfig.RUVersionPolicy.Overrides[1], + "SetKeyspaceRUVersion's change must not be lost to a later save landing out of order") + + raw, err := m.storage.LoadControllerConfig() + re.NoError(err) + persisted := &ControllerConfig{} + re.NoError(json.Unmarshal([]byte(raw), persisted)) + re.InDelta(1.5, persisted.RequestUnit.ReadBaseCost, 0.00001, + "UpdateControllerConfigItem's change must be persisted") + re.Equal(RUVersion(2), persisted.RUVersionPolicy.Overrides[1], + "SetKeyspaceRUVersion's change must be persisted, not clobbered by an out-of-order save") +} diff --git a/pkg/mcs/resourcemanager/server/manager_test.go b/pkg/mcs/resourcemanager/server/manager_test.go index 8d8e5f0e7c..c8385c0334 100644 --- a/pkg/mcs/resourcemanager/server/manager_test.go +++ b/pkg/mcs/resourcemanager/server/manager_test.go @@ -406,7 +406,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) @@ -799,7 +801,8 @@ func TestKeyspaceResourceGroupManagerWriteRoleGates(t *testing.T) { RU: &rmpb.TokenBucket{Settings: &rmpb.TokenLimitSettings{FillRate: 200}}, }, } - re.ErrorIs(tokenOnlyKRGM.modifyResourceGroup(modifiedGroup), errMetadataWriteDisabled) + _, err := tokenOnlyKRGM.modifyResourceGroup(modifiedGroup) + re.ErrorIs(err, errMetadataWriteDisabled) re.Equal(float64(100), tokenOnlyKRGM.getResourceGroup(group.GetName(), false).getFillRate()) re.ErrorIs(tokenOnlyKRGM.deleteResourceGroup(group.GetName()), errMetadataWriteDisabled) re.NotNil(tokenOnlyKRGM.getResourceGroup(group.GetName(), false)) diff --git a/pkg/mcs/resourcemanager/server/metadata_watcher.go b/pkg/mcs/resourcemanager/server/metadata_watcher.go index 0f8a97d697..3a2410600e 100644 --- a/pkg/mcs/resourcemanager/server/metadata_watcher.go +++ b/pkg/mcs/resourcemanager/server/metadata_watcher.go @@ -179,8 +179,14 @@ func (m *Manager) initializeMetadataWatcher(ctx context.Context) error { if err := watcher.WaitLoad(); err != nil { return err } - // Ensure reserved default groups exist even if settings were missing in storage. - m.initReserved() + // This runs synchronously before Init() returns and before + // LoadingStateCompleted is published, so no request can be in flight yet + // to race against; ensure reserved default groups exist even if settings + // were missing in storage. + m.RLock() + epoch := m.loadEpoch + m.RUnlock() + m.initReserved(epoch) return nil } diff --git a/pkg/mcs/resourcemanager/server/metadata_watcher_test.go b/pkg/mcs/resourcemanager/server/metadata_watcher_test.go index fcbcf7dbc9..bb3d8ad105 100644 --- a/pkg/mcs/resourcemanager/server/metadata_watcher_test.go +++ b/pkg/mcs/resourcemanager/server/metadata_watcher_test.go @@ -31,6 +31,7 @@ import ( "github.com/tikv/pd/pkg/keyspace/constant" "github.com/tikv/pd/pkg/storage" "github.com/tikv/pd/pkg/utils/keypath" + "github.com/tikv/pd/pkg/utils/syncutil" ) type countingServiceLimitLoadStorage struct { @@ -45,9 +46,10 @@ func (s *countingServiceLimitLoadStorage) LoadServiceLimits(f func(keyspaceID ui func newMetadataWatcherTestManager(store storage.Storage) *Manager { return &Manager{ - storage: store, - krgms: make(map[uint32]*keyspaceResourceGroupManager), - controllerConfig: &ControllerConfig{}, + storage: store, + krgms: make(map[uint32]*keyspaceResourceGroupManager), + controllerConfig: &ControllerConfig{}, + serviceLimitLocks: syncutil.NewLockGroup(), } } @@ -388,3 +390,89 @@ func TestInitializeMetadataWatcher(t *testing.T) { re.InDelta(123.5, m.GetKeyspaceServiceLimiter(10).ServiceLimit, 0.00001) }) } + +// TestMetadataWatcherModeReleasesSyncLoadedGroups guards against the +// sync-loaded markers accumulating for the process lifetime in metadata watcher +// mode: no async loader runs there, so nothing ever consumes them, and every +// watch event would otherwise add an entry that is never removed. +func TestMetadataWatcherModeReleasesSyncLoadedGroups(t *testing.T) { + re := require.New(t) + + m := newManagerBase(&ControllerConfig{}, ResourceGroupWriteRoleLegacyAll) + m.storage = storage.NewStorageWithMemoryBackend() + m.srv = &testBasicServer{} + m.enableMetadataWatcher = true + re.NotNil(m.syncLoadedGroups) + + originalFactory := newMetadataLoopWatcher + defer func() { newMetadataLoopWatcher = originalFactory }() + newMetadataLoopWatcher = func( + _ context.Context, + _ *sync.WaitGroup, + _ *clientv3.Client, + _, _ string, + _ func([]*clientv3.Event) error, + _, _ func(*mvccpb.KeyValue) error, + _ func([]*clientv3.Event) error, + _ bool, + ) metadataLoopWatcher { + return &fakeMetadataLoopWatcher{waitLoadFn: func() error { return nil }} + } + + re.NoError(m.Init(context.Background())) + defer m.close() + + re.Equal(LoadingStateCompleted, m.getLoadingState()) + m.RLock() + syncLoadedGroups := m.syncLoadedGroups + m.RUnlock() + re.Nil(syncLoadedGroups, "watcher mode must not retain sync-loaded markers") + + // A watch event after initialization must stay a no-op for the markers. + group := newMetadataWatcherResourceGroup("watched", middlePriority, 100, 100) + rawValue, err := proto.Marshal(group) + re.NoError(err) + re.NoError(m.applyResourceGroupSettingFromRaw(10, "watched", string(rawValue))) + m.RLock() + syncLoadedGroups = m.syncLoadedGroups + m.RUnlock() + re.Nil(syncLoadedGroups) + re.NotNil(m.getKeyspaceResourceGroupManager(10).getResourceGroup("watched", false)) +} + +// TestGetOrCreateKeyspaceResourceGroupManagerWatcherModeDoesNotClobberPendingDefault +// guards against trusting LoadingStateCompleted as "cache is authoritative" in +// metadata-watcher mode. There it only means the initial bootstrap finished, +// not that every write already in storage has had its watch event delivered +// to this replica's cache yet. If getOrCreateKeyspaceResourceGroupManager +// synthesized directly on that signal, a request landing after a customized +// default was persisted but before its watch event arrived would find +// nothing cached and persist the built-in default over it. +func TestGetOrCreateKeyspaceResourceGroupManagerWatcherModeDoesNotClobberPendingDefault(t *testing.T) { + re := require.New(t) + + store := storage.NewStorageWithMemoryBackend() + // A customized default is already persisted - e.g. by PD - but its watch + // event has not been delivered to this replica's cache yet. + customized := newMetadataWatcherResourceGroup(DefaultResourceGroupName, middlePriority, 555, 555) + re.NoError(store.SaveResourceGroupSetting(1, DefaultResourceGroupName, customized)) + + m := newMetadataWatcherTestManager(store) + m.enableMetadataWatcher = true + m.setLoadingState(LoadingStateCompleted) + + krgm := m.getOrCreateKeyspaceResourceGroupManager(1, true) + re.NotNil(krgm) + + group := krgm.getResourceGroup(DefaultResourceGroupName, false) + re.NotNil(group) + re.Equal(float64(555), group.getFillRate(), + "the customized default already in storage must be picked up by a point load, not overwritten by a synthesized built-in default") + + raw, err := store.LoadResourceGroupSetting(1, DefaultResourceGroupName) + re.NoError(err) + loaded := &rmpb.ResourceGroup{} + re.NoError(proto.Unmarshal([]byte(raw), loaded)) + re.Equal(uint64(555), loaded.RUSettings.RU.Settings.FillRate, + "the customized default persisted before the watcher delivered its event must survive in storage") +} diff --git a/pkg/mcs/resourcemanager/server/metrics.go b/pkg/mcs/resourcemanager/server/metrics.go index b4891fa21f..2756c0b159 100644 --- a/pkg/mcs/resourcemanager/server/metrics.go +++ b/pkg/mcs/resourcemanager/server/metrics.go @@ -216,6 +216,47 @@ 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.", + // The whole point of the async loading is that a cluster with many + // resource groups can take far longer than the default buckets' 10s + // ceiling, so use a coarse but wide exponential range instead: it + // spans sub-second up to ~27min, which keeps the slow loads this + // metric exists to observe out of the +Inf bucket. + Buckets: prometheus.ExponentialBuckets(0.1, 4, 8), + }) + + asyncLoadGroupFailureCounter = prometheus.NewCounter( + prometheus.CounterOpts{ + Namespace: namespace, + Subsystem: serverSubsystem, + Name: "async_load_group_failures_total", + Help: "Total number of failed attempts to load resource groups asynchronously.", + }) + + // resourceGroupLoadingStateGauge exposes the loading state so a load that + // keeps failing is alertable: the loader retries indefinitely, so without + // this the only symptom is a stuck state plus a periodic warning log. + resourceGroupLoadingStateGauge = prometheus.NewGauge( + prometheus.GaugeOpts{ + Namespace: namespace, + Subsystem: serverSubsystem, + Name: "resource_group_loading_state", + Help: "Current resource group loading state: 0 - not started, 1 - in progress, 2 - completed.", + }) ) type metrics struct { @@ -269,6 +310,10 @@ func init() { prometheus.MustRegister(overrideSettings) prometheus.MustRegister(serviceLimit) prometheus.MustRegister(pushRUMetricsDuration) + prometheus.MustRegister(syncLoadGroupCounter) + prometheus.MustRegister(asyncLoadGroupDuration) + prometheus.MustRegister(asyncLoadGroupFailureCounter) + prometheus.MustRegister(resourceGroupLoadingStateGauge) } func newMetrics() *metrics { diff --git a/pkg/mcs/resourcemanager/server/resource_group.go b/pkg/mcs/resourcemanager/server/resource_group.go index 5563a70ce4..4ac5aed7c0 100644 --- a/pkg/mcs/resourcemanager/server/resource_group.go +++ b/pkg/mcs/resourcemanager/server/resource_group.go @@ -178,6 +178,10 @@ func (rg *ResourceGroup) getBurstLimitLocked(ignoreOverride ...bool) int64 { func (rg *ResourceGroup) getOverrideBurstLimit() int64 { rg.RLock() defer rg.RUnlock() + return rg.getOverrideBurstLimitLocked() +} + +func (rg *ResourceGroup) getOverrideBurstLimitLocked() int64 { return rg.RUSettings.RU.overrideBurstLimit } @@ -379,7 +383,13 @@ func (rg *ResourceGroup) SetStatesIntoResourceGroup(states *GroupStates) { switch rg.Mode { case rmpb.GroupMode_RUMode: if state := states.RU; state != nil { + // The group may already be serving requests: setState writes the + // token bucket fields directly, so guard it with the group lock + // the same way RequestRU does. UpdateRUConsumption below locks + // internally. + rg.Lock() rg.RUSettings.RU.setState(state) + rg.Unlock() log.Debug("update group token bucket state", zap.String("name", rg.Name), zap.Any("state", state)) } if states.RUConsumption != nil { diff --git a/pkg/mcs/resourcemanager/server/ru_version_policy_test.go b/pkg/mcs/resourcemanager/server/ru_version_policy_test.go index 7c711e7c6f..cca2d6b968 100644 --- a/pkg/mcs/resourcemanager/server/ru_version_policy_test.go +++ b/pkg/mcs/resourcemanager/server/ru_version_policy_test.go @@ -17,6 +17,7 @@ package server import ( "context" "encoding/json" + "sync" "testing" "github.com/stretchr/testify/require" @@ -279,3 +280,36 @@ func TestManagerSetKeyspaceRUVersionResetToDefault(t *testing.T) { _, exists := policy.Overrides[1] re.False(exists) } + +// TestManagerSetKeyspaceRUVersionConcurrentWithLeadershipChange guards +// against a data race between SetKeyspaceRUVersion and a leadership change. +// SetKeyspaceRUVersion used to mutate m.controllerConfig under the lock, then +// re-read the m.controllerConfig field itself after unlocking to persist it - +// racing against initControllerConfig, which can reassign that same field +// wholesale (also under the lock) when Init runs again for a new term. Run +// under -race: any regression back to the unlocked re-read fails this test +// without needing a specific interleaving to be forced. +func TestManagerSetKeyspaceRUVersionConcurrentWithLeadershipChange(t *testing.T) { + re := require.New(t) + m := prepareManager() + re.NoError(m.Init(context.Background())) + + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + for i := range 100 { + _ = m.SetKeyspaceRUVersion(1, int32(i%2)+1) + } + }() + go func() { + defer wg.Done() + for range 20 { + m.cancel() + m.wg.Wait() + _ = m.Init(context.Background()) + } + }() + wg.Wait() + stopAsyncTestManager(m) +} 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/server/resource_group_proxy_service.go b/server/resource_group_proxy_service.go index d44be436f4..977186aa96 100644 --- a/server/resource_group_proxy_service.go +++ b/server/resource_group_proxy_service.go @@ -25,6 +25,7 @@ import ( "github.com/pingcap/kvproto/pkg/resource_manager" "github.com/pingcap/log" + "github.com/tikv/pd/pkg/errs" rm_server "github.com/tikv/pd/pkg/mcs/resourcemanager/server" "github.com/tikv/pd/pkg/mcs/utils/constant" "github.com/tikv/pd/pkg/utils/grpcutil" @@ -138,7 +139,7 @@ func (s *resourceGroupProxyServer) AddResourceGroup(ctx context.Context, req *re return nil, status.Error(codes.Internal, "resource group metadata manager is not initialized") } if err := s.metadataManager.AddResourceGroup(req.GetGroup()); err != nil { - return nil, err + return nil, errs.ErrResourceGroupsLoadingGRPC(err) } return &resource_manager.PutResourceGroupResponse{Body: "Success!"}, nil } @@ -166,7 +167,7 @@ func (s *resourceGroupProxyServer) ModifyResourceGroup(ctx context.Context, req return nil, status.Error(codes.Internal, "resource group metadata manager is not initialized") } if err := s.metadataManager.ModifyResourceGroup(req.GetGroup()); err != nil { - return nil, err + return nil, errs.ErrResourceGroupsLoadingGRPC(err) } return &resource_manager.PutResourceGroupResponse{Body: "Success!"}, nil } @@ -197,7 +198,7 @@ func (s *resourceGroupProxyServer) DeleteResourceGroup(ctx context.Context, req rm_server.ExtractKeyspaceID(req.GetKeyspaceId()), req.GetResourceGroupName(), ); err != nil { - return nil, err + return nil, errs.ErrResourceGroupsLoadingGRPC(err) } return &resource_manager.DeleteResourceGroupResponse{Body: "Success!"}, nil } diff --git a/tests/integrations/mcs/resourcemanager/resource_manager_test.go b/tests/integrations/mcs/resourcemanager/resource_manager_test.go index 84c3e1453c..d686bbc071 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")) @@ -1607,13 +1622,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() {