Skip to content

Commit

Permalink
rule_manager: generate default rules based on whether the witness is …
Browse files Browse the repository at this point in the history 
…enabled

ref tikv#5520

Signed-off-by: Wenbo Zhang <ethercflow@gmail.com>
  • Loading branch information
@ethercflow
ethercflow committed Mar 9, 2023
1 parent afd79ed commit ae09979
Show file tree
Hide file tree
Showing 4 changed files with 79 additions and 22 deletions.
7 changes: 7 additions & 0 deletions server/config/persist_options.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -172,6 +172,13 @@ func (o *PersistOptions) SetPlacementRulesCacheEnabled(enabled bool) {
o.SetReplicationConfig(v)
}

// SetWitnessEnabled set PlacementRuleEnabled
func (o *PersistOptions) SetWitnessEnabled(enabled bool) {
v := o.GetScheduleConfig().Clone()
v.EnableWitness = enabled
o.SetScheduleConfig(v)
}

// GetStrictlyMatchLabel returns whether check label strict.
func (o *PersistOptions) GetStrictlyMatchLabel() bool {
return o.GetReplicationConfig().StrictlyMatchLabel
Expand Down
1 change: 1 addition & 0 deletions server/schedule/config/config.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,7 @@ type Config interface {
SetSplitMergeInterval(time.Duration)
SetMaxReplicas(int)
SetPlacementRulesCacheEnabled(bool)
SetWitnessEnabled(bool)
}

// StoreConfig is the interface that wraps the StoreConfig related methods.
Expand Down
46 changes: 36 additions & 10 deletions server/schedule/placement/rule_manager.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,17 +81,43 @@ func (m *RuleManager) Initialize(maxReplica int, locationLabels []string) error
}
if len(m.ruleConfig.rules) == 0 {
// migrate from old config.
defaultRule := &Rule{
GroupID: "pd",
ID: "default",
Role: Voter,
Count: maxReplica,
LocationLabels: locationLabels,
}
if err := m.storage.SaveRule(defaultRule.StoreKey(), defaultRule); err != nil {
return err
var defaultRules []*Rule
if m.conf.IsWitnessAllowed() {
defaultRules = append(defaultRules,
[]*Rule{
{
GroupID: "pd",
ID: "default",
Role: Voter,
Count: maxReplica - 1,
LocationLabels: locationLabels,
},
{
GroupID: "pd",
ID: "witness",
Role: Voter,
Count: 1,
IsWitness: true,
LocationLabels: locationLabels,
},
}...,
)
} else {
defaultRules = append(defaultRules, &Rule{
GroupID: "pd",
ID: "default",
Role: Voter,
Count: maxReplica,
LocationLabels: locationLabels,
})
}
for _, defaultRule := range defaultRules {
if err := m.storage.SaveRule(defaultRule.StoreKey(), defaultRule); err != nil {
// TODO: Need to delete the previously successfully saved Rules?
return err
}
m.ruleConfig.setRule(defaultRule)
}
m.ruleConfig.setRule(defaultRule)
}
m.ruleConfig.adjust()
ruleList, err := buildRuleList(m.ruleConfig)
Expand Down
47 changes: 35 additions & 12 deletions server/schedule/placement/rule_manager_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,19 +28,20 @@ import (
"github.com/tikv/pd/pkg/storage/kv"
)

func newTestManager(t *testing.T) (endpoint.RuleStorage, *RuleManager) {
func newTestManager(t *testing.T, enableWitness bool) (endpoint.RuleStorage, *RuleManager) {
re := require.New(t)
store := endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil)
var err error
manager := NewRuleManager(store, nil, mockconfig.NewTestOptions())
manager.conf.SetWitnessEnabled(enableWitness)
err = manager.Initialize(3, []string{"zone", "rack", "host"})
re.NoError(err)
return store, manager
}

func TestDefault(t *testing.T) {
re := require.New(t)
_, manager := newTestManager(t)
_, manager := newTestManager(t, false)
rules := manager.GetAllRules()
re.Len(rules, 1)
re.Equal("pd", rules[0].GroupID)
Expand All @@ -52,9 +53,31 @@ func TestDefault(t *testing.T) {
re.Equal([]string{"zone", "rack", "host"}, rules[0].LocationLabels)
}

func TestDefault2(t *testing.T) {
re := require.New(t)
_, manager := newTestManager(t, true)
rules := manager.GetAllRules()
re.Len(rules, 2)
re.Equal("pd", rules[0].GroupID)
re.Equal("default", rules[0].ID)
re.Equal(0, rules[0].Index)
re.Empty(rules[0].StartKey)
re.Empty(rules[0].EndKey)
re.Equal(Voter, rules[0].Role)
re.Equal([]string{"zone", "rack", "host"}, rules[0].LocationLabels)
re.Equal("pd", rules[1].GroupID)
re.Equal("witness", rules[1].ID)
re.Equal(0, rules[1].Index)
re.Empty(rules[1].StartKey)
re.Empty(rules[1].EndKey)
re.Equal(Voter, rules[1].Role)
re.Equal(true, rules[1].IsWitness)
re.Equal([]string{"zone", "rack", "host"}, rules[1].LocationLabels)
}

func TestAdjustRule(t *testing.T) {
re := require.New(t)
_, manager := newTestManager(t)
_, manager := newTestManager(t, false)
rules := []Rule{
{GroupID: "group", ID: "id", StartKeyHex: "123abc", EndKeyHex: "123abf", Role: "voter", Count: 3},
{GroupID: "", ID: "id", StartKeyHex: "123abc", EndKeyHex: "123abf", Role: "voter", Count: 3},
Expand Down Expand Up @@ -106,7 +129,7 @@ func TestAdjustRule(t *testing.T) {

func TestLeaderCheck(t *testing.T) {
re := require.New(t)
_, manager := newTestManager(t)
_, manager := newTestManager(t, false)
re.Regexp(".*needs at least one leader or voter.*", manager.SetRule(&Rule{GroupID: "pd", ID: "default", Role: "learner", Count: 3}).Error())
re.Regexp(".*define multiple leaders by count 2.*", manager.SetRule(&Rule{GroupID: "g2", ID: "33", Role: "leader", Count: 2}).Error())
re.Regexp(".*multiple leader replicas.*", manager.Batch([]RuleOp{
Expand All @@ -123,7 +146,7 @@ func TestLeaderCheck(t *testing.T) {

func TestSaveLoad(t *testing.T) {
re := require.New(t)
store, manager := newTestManager(t)
store, manager := newTestManager(t, false)
rules := []*Rule{
{GroupID: "pd", ID: "default", Role: "voter", Count: 5},
{GroupID: "foo", ID: "baz", StartKeyHex: "", EndKeyHex: "abcd", Role: "voter", Count: 1},
Expand All @@ -144,7 +167,7 @@ func TestSaveLoad(t *testing.T) {

func TestSetAfterGet(t *testing.T) {
re := require.New(t)
store, manager := newTestManager(t)
store, manager := newTestManager(t, false)
rule := manager.GetRule("pd", "default")
rule.Count = 1
manager.SetRule(rule)
Expand All @@ -166,7 +189,7 @@ func checkRules(t *testing.T, rules []*Rule, expect [][2]string) {

func TestKeys(t *testing.T) {
re := require.New(t)
_, manager := newTestManager(t)
_, manager := newTestManager(t, false)
rules := []*Rule{
{GroupID: "1", ID: "1", Role: "voter", Count: 1, StartKeyHex: "", EndKeyHex: ""},
{GroupID: "2", ID: "2", Role: "voter", Count: 1, StartKeyHex: "11", EndKeyHex: "ff"},
Expand Down Expand Up @@ -255,7 +278,7 @@ func TestKeys(t *testing.T) {
}

func TestDeleteByIDPrefix(t *testing.T) {
_, manager := newTestManager(t)
_, manager := newTestManager(t, false)
manager.SetRules([]*Rule{
{GroupID: "g1", ID: "foo1", Role: "voter", Count: 1},
{GroupID: "g2", ID: "foo1", Role: "voter", Count: 1},
Expand All @@ -275,7 +298,7 @@ func TestDeleteByIDPrefix(t *testing.T) {

func TestRangeGap(t *testing.T) {
re := require.New(t)
_, manager := newTestManager(t)
_, manager := newTestManager(t, false)
err := manager.DeleteRule("pd", "default")
re.Error(err)

Expand All @@ -298,7 +321,7 @@ func TestRangeGap(t *testing.T) {

func TestGroupConfig(t *testing.T) {
re := require.New(t)
_, manager := newTestManager(t)
_, manager := newTestManager(t, false)
pd1 := &RuleGroup{ID: "pd"}
re.Equal(pd1, manager.GetRuleGroup("pd"))

Expand Down Expand Up @@ -334,7 +357,7 @@ func TestGroupConfig(t *testing.T) {

func TestRuleVersion(t *testing.T) {
re := require.New(t)
_, manager := newTestManager(t)
_, manager := newTestManager(t, false)
rule1 := manager.GetRule("pd", "default")
re.Equal(uint64(0), rule1.Version)
// create new rule
Expand Down Expand Up @@ -425,7 +448,7 @@ func TestCheckApplyRules(t *testing.T) {

func TestCacheManager(t *testing.T) {
re := require.New(t)
_, manager := newTestManager(t)
_, manager := newTestManager(t, false)
manager.conf.SetPlacementRulesCacheEnabled(true)
rules := addExtraRules(0)
re.NoError(manager.SetRules(rules))
Expand Down

0 comments on commit ae09979

Please sign in to comment.