client: reload RM snapshot before watch retry

Signed-off-by: bufferflies <1045931706@qq.com>

Author: bufferflies

client/resource_group/controller/global_controller.go

@@ -357,9 +357,7 @@ func (c *ResourceGroupsController) Start(ctx context.Context) {
 				}
 			case <-watchRetryTimer.C:
 				if !c.ruConfig.isSingleGroupByKeyspace && watchMetaChannel == nil {
-					// Use WithPrevKV() to get the previous key-value pair when get Delete Event.
-					prefix := pd.GroupSettingsPathPrefixBytes(c.keyspaceID)
-					watchMetaChannel, err = c.provider.Watch(ctx, prefix, opt.WithPrefix(), opt.WithPrevKV())
+					watchMetaChannel, err = c.reloadResourceGroupMetaWatch(ctx)
 					if err != nil {
 						log.Warn("watch resource group meta failed", zap.Error(err))
 						watchRetryTimer.Reset(watchRetryInterval)
@@ -491,6 +489,57 @@ func (c *ResourceGroupsController) Start(ctx context.Context) {
 	}()
 }
 
+func (c *ResourceGroupsController) reloadResourceGroupMetaWatch(
+	ctx context.Context,
+) (chan []*meta_storagepb.Event, error) {
+	groups, revision, err := c.provider.LoadResourceGroups(ctx)
+	if err != nil {
+		return nil, err
+	}
+	c.syncResourceGroupSnapshot(groups)
+	// Start from the next revision after the freshly loaded snapshot so reconnects
+	// keep the cache in sync without replaying the snapshot itself as watch events.
+	return c.provider.Watch(
+		ctx,
+		pd.GroupSettingsPathPrefixBytes(c.keyspaceID),
+		opt.WithRev(revision+1),
+		opt.WithPrefix(),
+		opt.WithPrevKV(),
+	)
+}
+
+func (c *ResourceGroupsController) syncResourceGroupSnapshot(groups []*rmpb.ResourceGroup) {
+	groupMap := make(map[string]*rmpb.ResourceGroup, len(groups))
+	for _, group := range groups {
+		groupMap[group.GetName()] = group
+	}
+
+	c.groupsController.Range(func(key, value any) bool {
+		name := key.(string)
+		group, exists := groupMap[name]
+		if !exists {
+			c.tombstoneGroupCostController(name)
+			return true
+		}
+		gc := value.(*groupCostController)
+		if !gc.tombstone.Load() {
+			gc.modifyMeta(group)
+			return true
+		}
+		newGC, err := newGroupCostController(group, c.ruConfig, c.lowTokenNotifyChan, c.tokenBucketUpdateChan)
+		if err != nil {
+			log.Warn("[resource group controller] re-create resource group cost controller from snapshot failed",
+				zap.String("name", name), zap.Error(err))
+			return true
+		}
+		if c.groupsController.CompareAndSwap(name, gc, newGC) {
+			log.Info("[resource group controller] re-create resource group cost controller from snapshot",
+				zap.String("name", name))
+		}
+		return true
+	})
+}
+
 // Stop stops ResourceGroupController service.
 func (c *ResourceGroupsController) Stop() error {
 	if c.loopCancel == nil {

client/resource_group/controller/global_controller_test.go

@@ -381,6 +381,71 @@ func TestGetResourceGroup(t *testing.T) {
 	re.Nil(gc02)
 }
 
+func TestReloadResourceGroupMetaWatch(t *testing.T) {
+	re := require.New(t)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	mockProvider := &MockResourceGroupProvider{}
+	mockProvider.On("Get", mock.Anything, mock.Anything, mock.Anything).Return(&meta_storagepb.GetResponse{}, nil)
+	controller, err := NewResourceGroupController(ctx, 1, mockProvider, nil, constants.NullKeyspaceID)
+	re.NoError(err)
+
+	defaultGroup := &rmpb.ResourceGroup{
+		Name:       defaultResourceGroupName,
+		Mode:       rmpb.GroupMode_RUMode,
+		RUSettings: &rmpb.GroupRequestUnitSettings{RU: &rmpb.TokenBucket{Settings: &rmpb.TokenLimitSettings{FillRate: 1000000}}},
+	}
+	staleGroup := &rmpb.ResourceGroup{
+		Name:       "stale-group",
+		Mode:       rmpb.GroupMode_RUMode,
+		RUSettings: &rmpb.GroupRequestUnitSettings{RU: &rmpb.TokenBucket{Settings: &rmpb.TokenLimitSettings{FillRate: 100}}},
+	}
+	updatedDefaultGroup := &rmpb.ResourceGroup{
+		Name:       defaultResourceGroupName,
+		Mode:       rmpb.GroupMode_RUMode,
+		RUSettings: &rmpb.GroupRequestUnitSettings{RU: &rmpb.TokenBucket{Settings: &rmpb.TokenLimitSettings{FillRate: 2000}}},
+	}
+
+	mockProvider.On("GetResourceGroup", mock.Anything, defaultResourceGroupName, mock.Anything).Return(defaultGroup, nil)
+	mockProvider.On("GetResourceGroup", mock.Anything, "stale-group", mock.Anything).Return(staleGroup, nil)
+
+	defaultGC, err := controller.tryGetResourceGroupController(ctx, defaultResourceGroupName, false)
+	re.NoError(err)
+	re.NotNil(defaultGC)
+
+	staleGC, err := controller.tryGetResourceGroupController(ctx, "stale-group", false)
+	re.NoError(err)
+	re.NotNil(staleGC)
+	re.False(staleGC.tombstone.Load())
+
+	watchCh := make(chan []*meta_storagepb.Event)
+	mockProvider.On("LoadResourceGroups", mock.Anything).
+		Return([]*rmpb.ResourceGroup{updatedDefaultGroup}, int64(88), nil).
+		Once()
+	mockProvider.On("Watch", mock.Anything, pd.GroupSettingsPathPrefixBytes(constants.NullKeyspaceID), mock.Anything).
+		Run(func(args mock.Arguments) {
+			opts := args.Get(2).([]opt.MetaStorageOption)
+			metaOp := &opt.MetaStorageOp{}
+			for _, apply := range opts {
+				apply(metaOp)
+			}
+			re.Equal(int64(89), metaOp.Revision)
+			re.True(metaOp.IsOptsWithPrefix)
+			re.True(metaOp.PrevKv)
+		}).
+		Return(watchCh, nil).
+		Once()
+
+	reloadedWatchCh, err := controller.reloadResourceGroupMetaWatch(ctx)
+	re.NoError(err)
+	re.Equal(watchCh, reloadedWatchCh)
+	re.Equal(updatedDefaultGroup, defaultGC.getMeta())
+	staleGC, err = controller.tryGetResourceGroupController(ctx, "stale-group", true)
+	re.NoError(err)
+	re.True(staleGC.tombstone.Load())
+}
+
 func TestTokenBucketsRequestWithKeyspaceID(t *testing.T) {
 	re := require.New(t)
 	ctx, cancel := context.WithCancel(context.Background())