resourcemanager: add metadata watcher scaffold

pkg/mcs/resourcemanager/server/keyspace_manager.go

@@ -23,6 +23,7 @@ import (
 	"github.com/gogo/protobuf/proto"
 	"go.uber.org/zap"
 
+	"github.com/pingcap/errors"
 	rmpb "github.com/pingcap/kvproto/pkg/resource_manager"
 	"github.com/pingcap/log"
 
@@ -106,6 +107,49 @@ func (krgm *keyspaceResourceGroupManager) addResourceGroupFromRaw(name string, r
 	return nil
 }
 
+func (krgm *keyspaceResourceGroupManager) upsertResourceGroupFromRaw(name string, rawValue string) error {
+	group := &rmpb.ResourceGroup{}
+	if err := proto.Unmarshal([]byte(rawValue), group); err != nil {
+		log.Error("failed to parse the keyspace resource group meta info",
+			zap.Uint32("keyspace-id", krgm.keyspaceID), zap.String("name", name), zap.String("raw-value", rawValue), zap.Error(err))
+		return err
+	}
+	if group.Name != name {
+		err := errors.Errorf("resource group key name %s does not match payload name %s", name, group.Name)
+		log.Error("resource group name mismatch in watcher payload",
+			zap.Uint32("keyspace-id", krgm.keyspaceID),
+			zap.String("name", name),
+			zap.String("payload-name", group.Name),
+			zap.String("raw-value", rawValue),
+			zap.Error(err))
+		return err
+	}
+
+	krgm.RLock()
+	existing := krgm.groups[group.Name]
+	krgm.RUnlock()
+	if existing != nil {
+		if err := existing.ApplySettings(group); err != 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
+		}
+		return nil
+	}
+
+	krgm.Lock()
+	krgm.groups[group.Name] = FromProtoResourceGroup(group)
+	krgm.Unlock()
+	return nil
+}
+
+func (krgm *keyspaceResourceGroupManager) deleteResourceGroupFromCache(name string) {
+	krgm.Lock()
+	delete(krgm.groups, name)
+	delete(krgm.groupRUTrackers, name)
+	krgm.Unlock()
+}
+
 func (krgm *keyspaceResourceGroupManager) setRawStatesIntoResourceGroup(name string, rawValue string) error {
 	tokens := &GroupStates{}
 	if err := json.Unmarshal([]byte(rawValue), tokens); err != nil {
@@ -128,22 +172,30 @@ func (krgm *keyspaceResourceGroupManager) initDefaultResourceGroup() {
 	if ok {
 		return
 	}
-	defaultGroup := &ResourceGroup{
+	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 newDefaultResourceGroup() *ResourceGroup {
+	return &ResourceGroup{
 		Name: DefaultResourceGroupName,
 		Mode: rmpb.GroupMode_RUMode,
-		RUSettings: &RequestUnitSettings{
-			RU: &GroupTokenBucket{
-				Settings: &rmpb.TokenLimitSettings{
-					FillRate:   UnlimitedRate,
-					BurstLimit: UnlimitedBurstLimit,
-				},
+		RUSettings: NewRequestUnitSettings(DefaultResourceGroupName, &rmpb.TokenBucket{
+			Settings: &rmpb.TokenLimitSettings{
+				FillRate:   UnlimitedRate,
+				BurstLimit: UnlimitedBurstLimit,
 			},
-		},
+		}),
 		Priority: middlePriority,
 	}
-	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) restoreDefaultResourceGroupFromReserved() {
+	krgm.Lock()
+	krgm.groups[DefaultResourceGroupName] = newDefaultResourceGroup()
+	krgm.Unlock()
 }
 
 func (krgm *keyspaceResourceGroupManager) addResourceGroup(grouppb *rmpb.ResourceGroup) error {
@@ -204,12 +256,13 @@ func (krgm *keyspaceResourceGroupManager) deleteResourceGroup(name string) error
 	if !krgm.writeRole.AllowsMetadataWrite() {
 		return errMetadataWriteDisabled
 	}
+	if err := krgm.storage.DeleteResourceGroupStates(krgm.keyspaceID, name); err != nil {
+		return err
+	}
 	if err := krgm.storage.DeleteResourceGroupSetting(krgm.keyspaceID, name); err != nil {
 		return err
 	}
-	krgm.Lock()
-	delete(krgm.groups, name)
-	krgm.Unlock()
+	krgm.deleteResourceGroupFromCache(name)
 	return nil
 }
 

pkg/mcs/resourcemanager/server/keyspace_manager_test.go

@@ -179,6 +179,7 @@ func TestDeleteResourceGroup(t *testing.T) {
 	}
 	err := krgm.addResourceGroup(group)
 	re.NoError(err)
+	krgm.groupRUTrackers[group.GetName()] = newGroupRUTracker()
 
 	// Verify the group exists.
 	re.NotNil(krgm.getResourceGroup(group.GetName(), false))
@@ -189,6 +190,8 @@ func TestDeleteResourceGroup(t *testing.T) {
 
 	// Verify the group was deleted.
 	re.Nil(krgm.getResourceGroup(group.GetName(), false))
+	_, ok := krgm.groupRUTrackers[group.GetName()]
+	re.False(ok)
 
 	// Try to delete the default group.
 	krgm.initDefaultResourceGroup()
@@ -199,6 +202,55 @@ func TestDeleteResourceGroup(t *testing.T) {
 	re.NotNil(krgm.getResourceGroup(DefaultResourceGroupName, false))
 }
 
+func TestDeleteResourceGroupClearsPersistedStatesInMetadataOnlyRole(t *testing.T) {
+	re := require.New(t)
+
+	memStorage := storage.NewStorageWithMemoryBackend()
+	krgm := newKeyspaceResourceGroupManager(1, memStorage, ResourceGroupWriteRolePDMetaOnly)
+	group := &rmpb.ResourceGroup{
+		Name:     "test_group",
+		Mode:     rmpb.GroupMode_RUMode,
+		Priority: 5,
+		RUSettings: &rmpb.GroupRequestUnitSettings{
+			RU: &rmpb.TokenBucket{
+				Settings: &rmpb.TokenLimitSettings{
+					FillRate:   100,
+					BurstLimit: 200,
+				},
+			},
+		},
+	}
+
+	re.NoError(krgm.addResourceGroup(group))
+
+	now := time.Now()
+	re.NoError(memStorage.SaveResourceGroupStates(1, group.Name, &GroupStates{
+		RU: &GroupTokenBucketState{
+			Tokens:     321,
+			LastUpdate: &now,
+		},
+		RUConsumption: &rmpb.Consumption{RRU: 11, WRU: 22},
+	}))
+
+	stateCount := 0
+	re.NoError(memStorage.LoadResourceGroupStates(func(keyspaceID uint32, name, _ string) {
+		if keyspaceID == 1 && name == group.Name {
+			stateCount++
+		}
+	}))
+	re.Equal(1, stateCount)
+
+	re.NoError(krgm.deleteResourceGroup(group.Name))
+
+	stateCount = 0
+	re.NoError(memStorage.LoadResourceGroupStates(func(keyspaceID uint32, name, _ string) {
+		if keyspaceID == 1 && name == group.Name {
+			stateCount++
+		}
+	}))
+	re.Zero(stateCount)
+}
+
 func TestGetResourceGroup(t *testing.T) {
 	re := require.New(t)
 
@@ -306,6 +358,112 @@ func TestAddResourceGroupFromRaw(t *testing.T) {
 	re.Error(err)
 }
 
+func TestUpsertResourceGroupFromRaw(t *testing.T) {
+	re := require.New(t)
+
+	krgm := newKeyspaceResourceGroupManager(1, storage.NewStorageWithMemoryBackend())
+	group := &rmpb.ResourceGroup{
+		Name:     "test_group",
+		Mode:     rmpb.GroupMode_RUMode,
+		Priority: 5,
+		RUSettings: &rmpb.GroupRequestUnitSettings{
+			RU: &rmpb.TokenBucket{
+				Settings: &rmpb.TokenLimitSettings{
+					FillRate:   100,
+					BurstLimit: 200,
+				},
+			},
+		},
+	}
+	re.NoError(krgm.addResourceGroup(group))
+
+	mutableGroup := krgm.getMutableResourceGroup(group.Name)
+	re.NotNil(mutableGroup)
+	mutableGroup.RUSettings.RU.Tokens = 123.45
+	mutableGroup.RUConsumption = &rmpb.Consumption{RRU: 12, WRU: 34}
+
+	updatedGroup := &rmpb.ResourceGroup{
+		Name:     group.Name,
+		Mode:     rmpb.GroupMode_RUMode,
+		Priority: 9,
+		RUSettings: &rmpb.GroupRequestUnitSettings{
+			RU: &rmpb.TokenBucket{
+				Settings: &rmpb.TokenLimitSettings{
+					FillRate:   500,
+					BurstLimit: 600,
+				},
+			},
+		},
+	}
+	rawUpdated, err := proto.Marshal(updatedGroup)
+	re.NoError(err)
+	re.NoError(krgm.upsertResourceGroupFromRaw(group.Name, string(rawUpdated)))
+
+	current := krgm.getMutableResourceGroup(group.Name)
+	re.NotNil(current)
+	re.Equal(uint32(9), current.Priority)
+	re.Equal(float64(500), current.getFillRate())
+	re.Equal(int64(600), current.getBurstLimit())
+	re.InDelta(123.45, current.RUSettings.RU.Tokens, 0.001)
+	re.Equal(float64(12), current.RUConsumption.RRU)
+	re.Equal(float64(34), current.RUConsumption.WRU)
+
+	newGroup := &rmpb.ResourceGroup{
+		Name:     "new_group",
+		Mode:     rmpb.GroupMode_RUMode,
+		Priority: 3,
+		RUSettings: &rmpb.GroupRequestUnitSettings{
+			RU: &rmpb.TokenBucket{
+				Settings: &rmpb.TokenLimitSettings{
+					FillRate:   111,
+					BurstLimit: 222,
+				},
+			},
+		},
+	}
+	rawNew, err := proto.Marshal(newGroup)
+	re.NoError(err)
+	re.NoError(krgm.upsertResourceGroupFromRaw(newGroup.GetName(), string(rawNew)))
+	re.NotNil(krgm.getMutableResourceGroup(newGroup.GetName()))
+
+	err = krgm.upsertResourceGroupFromRaw(group.GetName(), "invalid_data")
+	re.Error(err)
+
+	mismatchedGroup := &rmpb.ResourceGroup{
+		Name:     "payload_name",
+		Mode:     rmpb.GroupMode_RUMode,
+		Priority: 1,
+	}
+	rawMismatched, err := proto.Marshal(mismatchedGroup)
+	re.NoError(err)
+	err = krgm.upsertResourceGroupFromRaw("key_name", string(rawMismatched))
+	re.Error(err)
+	re.Nil(krgm.getMutableResourceGroup("key_name"))
+	re.Nil(krgm.getMutableResourceGroup("payload_name"))
+}
+
+func TestDeleteResourceGroupFromCache(t *testing.T) {
+	re := require.New(t)
+
+	krgm := newKeyspaceResourceGroupManager(1, storage.NewStorageWithMemoryBackend())
+	group := &rmpb.ResourceGroup{
+		Name: "test_group",
+		Mode: rmpb.GroupMode_RUMode,
+		RUSettings: &rmpb.GroupRequestUnitSettings{
+			RU: &rmpb.TokenBucket{
+				Settings: &rmpb.TokenLimitSettings{},
+			},
+		},
+	}
+	re.NoError(krgm.addResourceGroup(group))
+	krgm.groupRUTrackers[group.Name] = newGroupRUTracker()
+
+	krgm.deleteResourceGroupFromCache(group.Name)
+	re.Nil(krgm.getMutableResourceGroup(group.Name))
+	_, ok := krgm.groupRUTrackers[group.Name]
+	re.False(ok)
+}
+
 func TestSetRawStatesIntoResourceGroup(t *testing.T) {
 	re := require.New(t)