ddl: refactor bundle[2/2] [6/6]

ddl/placement/bundle.go

@@ -19,6 +19,7 @@ import (
  "errors"
  "fmt"
  "strconv"
+ "strings"
 
  "github.com/pingcap/failpoint"
  "github.com/pingcap/parser/ast"
@@ -201,3 +202,40 @@ func (b *Bundle) Clone() *Bundle {
 func (b *Bundle) IsEmpty() bool {
  return len(b.Rules) == 0 && b.Index == 0 && !b.Override
 }
+
+// ObjectID extracts the db/table/partition ID from the group ID
+func (b *Bundle) ObjectID() (int64, error) {
+ // If the rule doesn't come from TiDB, skip it.
+ if !strings.HasPrefix(b.ID, BundleIDPrefix) {
+ return 0, ErrInvalidBundleIDFormat
+ }
+ id, err := strconv.ParseInt(b.ID[len(BundleIDPrefix):], 10, 64)
+ if err != nil {
+ return 0, fmt.Errorf("%w: %s", ErrInvalidBundleID, err)
+ }
+ if id <= 0 {
+ return 0, fmt.Errorf("%w: %s doesn't include an id", ErrInvalidBundleID, b.ID)
+ }
+ return id, nil
+}
+
+func isValidLeaderRule(rule *Rule, dcLabelKey string) bool {
+ if rule.Role == Leader && rule.Count == 1 {
+ for _, con := range rule.Constraints {
+ if con.Op == In && con.Key == dcLabelKey && len(con.Values) == 1 {
+ return true
+ }
+ }
+ }
+ return false
+}
+
+// GetLeaderDC returns the leader's DC by Bundle if found.
+func (b *Bundle) GetLeaderDC(dcLabelKey string) (string, bool) {
+ for _, rule := range b.Rules {
+ if isValidLeaderRule(rule, dcLabelKey) {
+ return rule.Constraints[0].Values[0], true
+ }
+ }
+ return "", false
+}

ddl/placement/bundle_test.go

@@ -56,6 +56,290 @@ func (s *testBundleSuite) TestClone(c *C) {
  c.Assert(newBundle, DeepEquals, &Bundle{ID: GroupID(2), Rules: []*Rule{{ID: "121"}}})
 }
 
+func (s *testBundleSuite) TestObjectID(c *C) {
+ type TestCase struct {
+ name string
+ bundleID string
+ expectedID int64
+ err error
+ }
+ tests := []TestCase{
+ {"non tidb bundle", "pd", 0, ErrInvalidBundleIDFormat},
+ {"id of words", "TiDB_DDL_foo", 0, ErrInvalidBundleID},
+ {"id of words and nums", "TiDB_DDL_3x", 0, ErrInvalidBundleID},
+ {"id of floats", "TiDB_DDL_3.0", 0, ErrInvalidBundleID},
+ {"id of negatives", "TiDB_DDL_-10", 0, ErrInvalidBundleID},
+ {"id of positive integer", "TiDB_DDL_10", 10, nil},
+ }
+ for _, t := range tests {
+ comment := Commentf("%s", t.name)
+ bundle := Bundle{ID: t.bundleID}
+ id, err := bundle.ObjectID()
+ if t.err == nil {
+ c.Assert(err, IsNil, comment)
+ c.Assert(id, Equals, t.expectedID, comment)
+ } else {
+ c.Assert(errors.Is(err, t.err), IsTrue, comment)
+ }
+ }
+}
+
+func (s *testBundleSuite) TestGetLeaderDCByBundle(c *C) {
+ testcases := []struct {
+ name string
+ bundle *Bundle
+ expectedDC string
+ }{
+ {
+ name: "only leader",
+ bundle: &Bundle{
+ ID: GroupID(1),
+ Rules: []*Rule{
+ {
+ ID: "12",
+ Role: Leader,
+ Constraints: Constraints{
+ {
+ Key: "zone",
+ Op: In,
+ Values: []string{"bj"},
+ },
+ },
+ Count: 1,
+ },
+ },
+ },
+ expectedDC: "bj",
+ },
+ {
+ name: "no leader",
+ bundle: &Bundle{
+ ID: GroupID(1),
+ Rules: []*Rule{
+ {
+ ID: "12",
+ Role: Voter,
+ Constraints: Constraints{
+ {
+ Key: "zone",
+ Op: In,
+ Values: []string{"bj"},
+ },
+ },
+ Count: 3,
+ },
+ },
+ },
+ expectedDC: "",
+ },
+ {
+ name: "voter and leader",
+ bundle: &Bundle{
+ ID: GroupID(1),
+ Rules: []*Rule{
+ {
+ ID: "11",
+ Role: Leader,
+ Constraints: Constraints{
+ {
+ Key: "zone",
+ Op: In,
+ Values: []string{"sh"},
+ },
+ },
+ Count: 1,
+ },
+ {
+ ID: "12",
+ Role: Voter,
+ Constraints: Constraints{
+ {
+ Key: "zone",
+ Op: In,
+ Values: []string{"bj"},
+ },
+ },
+ Count: 3,
+ },
+ },
+ },
+ expectedDC: "sh",
+ },
+ {
+ name: "wrong label key",
+ bundle: &Bundle{
+ ID: GroupID(1),
+ Rules: []*Rule{
+ {
+ ID: "11",
+ Role: Leader,
+ Constraints: Constraints{
+ {
+ Key: "fake",
+ Op: In,
+ Values: []string{"sh"},
+ },
+ },
+ Count: 1,
+ },
+ },
+ },
+ expectedDC: "",
+ },
+ {
+ name: "wrong operator",
+ bundle: &Bundle{
+ ID: GroupID(1),
+ Rules: []*Rule{
+ {
+ ID: "11",
+ Role: Leader,
+ Constraints: Constraints{
+ {
+ Key: "zone",
+ Op: NotIn,
+ Values: []string{"sh"},
+ },
+ },
+ Count: 1,
+ },
+ },
+ },
+ expectedDC: "",
+ },
+ {
+ name: "leader have multi values",
+ bundle: &Bundle{
+ ID: GroupID(1),
+ Rules: []*Rule{
+ {
+ ID: "11",
+ Role: Leader,
+ Constraints: Constraints{
+ {
+ Key: "zone",
+ Op: In,
+ Values: []string{"sh", "bj"},
+ },
+ },
+ Count: 1,
+ },
+ },
+ },
+ expectedDC: "",
+ },
+ {
+ name: "irrelvant rules",
+ bundle: &Bundle{
+ ID: GroupID(1),
+ Rules: []*Rule{
+ {
+ ID: "15",
+ Role: Leader,
+ Constraints: Constraints{
+ {
+ Key: EngineLabelKey,
+ Op: NotIn,
+ Values: []string{EngineLabelTiFlash},
+ },
+ },
+ Count: 1,
+ },
+ {
+ ID: "14",
+ Role: Leader,
+ Constraints: Constraints{
+ {
+ Key: "disk",
+ Op: NotIn,
+ Values: []string{"ssd", "hdd"},
+ },
+ },
+ Count: 1,
+ },
+ {
+ ID: "13",
+ Role: Leader,
+ Constraints: Constraints{
+ {
+ Key: "zone",
+ Op: In,
+ Values: []string{"bj"},
+ },
+ },
+ Count: 1,
+ },
+ },
+ },
+ expectedDC: "bj",
+ },
+ {
+ name: "multi leaders 1",
+ bundle: &Bundle{
+ ID: GroupID(1),
+ Rules: []*Rule{
+ {
+ ID: "16",
+ Role: Leader,
+ Constraints: Constraints{
+ {
+ Key: "zone",
+ Op: In,
+ Values: []string{"sh"},
+ },
+ },
+ Count: 2,
+ },
+ },
+ },
+ expectedDC: "",
+ },
+ {
+ name: "multi leaders 2",
+ bundle: &Bundle{
+ ID: GroupID(1),
+ Rules: []*Rule{
+ {
+ ID: "17",
+ Role: Leader,
+ Constraints: Constraints{
+ {
+ Key: "zone",
+ Op: In,
+ Values: []string{"sh"},
+ },
+ },
+ Count: 1,
+ },
+ {
+ ID: "18",
+ Role: Leader,
+ Constraints: Constraints{
+ {
+ Key: "zone",
+ Op: In,
+ Values: []string{"bj"},
+ },
+ },
+ Count: 1,
+ },
+ },
+ },
+ expectedDC: "sh",
+ },
+ }
+ for _, testcase := range testcases {
+ comment := Commentf("%s", testcase.name)
+ result, ok := testcase.bundle.GetLeaderDC("zone")
+ if len(testcase.expectedDC) > 0 {
+ c.Assert(ok, Equals, true, comment)
+ } else {
+ c.Assert(ok, Equals, false, comment)
+ }
+ c.Assert(result, Equals, testcase.expectedDC, comment)
+ }
+}
+
 func (s *testBundleSuite) TestApplyPlacmentSpec(c *C) {
  type TestCase struct {
  name string