Support promotion rule override in vttablet

go/flags/endtoend/vtcombo.txt

@@ -260,6 +260,7 @@ Flags:
       --port int                                                         port for the server
       --pprof strings                                                    enable profiling
       --pprof-http                                                       enable pprof http endpoints
+      --promotion-rule topodatapb.PromotionRule                          Sets the override Promotion Rule for this tablet. Possible values 'none', 'neutral', 'prefer', 'prefer_not' and 'must_not'. Use 'must_not' with caution as it can lead to issues with reparenting. (default none)
       --proto_topo vttest.TopoData                                       vttest proto definition of the topology, encoded in compact text format. See vttest.proto for more information.
       --proxy_protocol                                                   Enable HAProxy PROXY protocol on MySQL listener socket
       --proxy_tablets                                                    Setting this true will make vtctld proxy the tablet status instead of redirecting to them

go/flags/endtoend/vttablet.txt

@@ -256,6 +256,7 @@ Flags:
       --port int                                                         port for the server
       --pprof strings                                                    enable profiling
       --pprof-http                                                       enable pprof http endpoints
+      --promotion-rule topodatapb.PromotionRule                          Sets the override Promotion Rule for this tablet. Possible values 'none', 'neutral', 'prefer', 'prefer_not' and 'must_not'. Use 'must_not' with caution as it can lead to issues with reparenting. (default none)
       --pt-osc-path string                                               override default pt-online-schema-change binary full path (default "/usr/bin/pt-online-schema-change")
       --publish_retry_interval duration                                  how long vttablet waits to retry publishing the tablet record (default 30s)
       --purge_logs_interval duration                                     how often try to remove old logs (default 1h0m0s)

go/vt/topo/topoproto/flag.go

@@ -17,6 +17,7 @@ limitations under the License.
 package topoproto
 
 import (
+	"fmt"
 	"strings"
 
 	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
@@ -59,3 +60,25 @@ func (ttf *TabletTypeFlag) Set(v string) error {
 
 // Type is part of the pflag.Value interface.
 func (*TabletTypeFlag) Type() string { return "topodatapb.TabletType" }
+
+// PromotionRuleFlag implements the pflag.Value interface, for parsing a command-line value into a PromotionRule.
+type PromotionRuleFlag topodatapb.PromotionRule
+
+// String is part of the pflag.Value interfaces.
+func (ttf *PromotionRuleFlag) String() string {
+	promotionRule := topodatapb.PromotionRule(*ttf)
+	return strings.ToLower(promotionRule.String())
+}
+
+// Set is part of the pflag.Value interfaces.
+func (ttf *PromotionRuleFlag) Set(v string) error {
+	t, err := ParsePromotionRule(v)
+	if t == topodatapb.PromotionRule_MUST {
+		return fmt.Errorf("promotion rule %v not supported yet", t)
+	}
+	*ttf = PromotionRuleFlag(t)
+	return err
+}
+
+// Type is part of the pflag.Value interface.
+func (*PromotionRuleFlag) Type() string { return "topodatapb.PromotionRule" }

go/vt/topo/topoproto/tablet.go

@@ -295,3 +295,12 @@ func IsServingType(tabletType topodatapb.TabletType) bool {
 		return false
 	}
 }
+
+// ParsePromotionRule parses the tablet promotion rule.
+func ParsePromotionRule(param string) (topodatapb.PromotionRule, error) {
+	value, ok := topodatapb.PromotionRule_value[strings.ToUpper(param)]
+	if !ok {
+		return topodatapb.PromotionRule_NONE, fmt.Errorf("unknown PromotionRule %v", param)
+	}
+	return topodatapb.PromotionRule(value), nil
+}

go/vt/topo/topoproto/tablet_test.go

@@ -129,3 +129,36 @@ func TestIsTabletsInList(t *testing.T) {
 		})
 	}
 }
+
+func TestParsePromotionRule(t *testing.T) {
+	testcases := []struct {
+		name         string
+		rule         string
+		expectedRule topodatapb.PromotionRule
+		errContains  string
+	}{
+		{
+			name:         "success",
+			rule:         "prefer_not",
+			expectedRule: topodatapb.PromotionRule_PREFER_NOT,
+		},
+		{
+			name:         "invalid",
+			rule:         "this_should_fail",
+			expectedRule: topodatapb.PromotionRule_NONE,
+			errContains:  "unknown PromotionRule",
+		},
+	}
+
+	for _, testcase := range testcases {
+		testcase := testcase
+		t.Run(testcase.name, func(t *testing.T) {
+			t.Parallel()
+			rule, err := ParsePromotionRule(testcase.rule)
+			if testcase.errContains != "" {
+				assert.ErrorContains(t, err, testcase.errContains)
+			}
+			assert.Equal(t, testcase.expectedRule, rule)
+		})
+	}
+}

go/vt/vtctl/reparentutil/durability.go

@@ -108,6 +108,14 @@ func PromotionRule(durability Durabler, tablet *topodatapb.Tablet) promotionrule
 	if tablet == nil || tablet.Alias == nil {
 		return promotionrule.MustNot
 	}
+	// Support per-tablet override.
+	if tablet.PromotionRule != topodatapb.PromotionRule_NONE {
+		promotionRule, err := promotionrule.ParseFromProto(tablet.PromotionRule)
+		if err != nil {
+			log.Errorf("failed to parse tablet promotion rule: %v", err)
+		}
+		return promotionRule
+	}
 	return durability.PromotionRule(tablet)
 }
 

go/vt/vtctl/reparentutil/durability_test.go

@@ -331,3 +331,27 @@ func TestDurabilityTest(t *testing.T) {
 		})
 	}
 }
+
+func TestPromotionRule(t *testing.T) {
+	cellName := "zone3"
+	durabler := &durabilityNone{}
+
+	// test --promotion-rule override
+	assert.Equal(t, promotionrule.MustNot, PromotionRule(durabler, &topodatapb.Tablet{
+		Alias: &topodatapb.TabletAlias{
+			Cell: cellName,
+			Uid:  4,
+		},
+		Type:          topodatapb.TabletType_REPLICA,
+		PromotionRule: topodatapb.PromotionRule_MUST_NOT,
+	}))
+
+	// default
+	assert.Equal(t, promotionrule.Neutral, PromotionRule(durabler, &topodatapb.Tablet{
+		Alias: &topodatapb.TabletAlias{
+			Cell: cellName,
+			Uid:  5,
+		},
+		Type: topodatapb.TabletType_REPLICA,
+	}))
+}

go/vt/vtctl/reparentutil/promotionrule/promotion_rule.go

@@ -17,7 +17,10 @@
 package promotionrule
 
 import (
+	"errors"
 	"fmt"
+
+	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
 )
 
 // CandidatePromotionRule describe the promotion preference/rule for an instance.
@@ -31,6 +34,8 @@ const (
 	MustNot   CandidatePromotionRule = "must_not"
 )
 
+var ErrUnsupportedPromotionRule = errors.New("unsupported promotion rule")
+
 var promotionRuleOrderMap = map[CandidatePromotionRule]int{
 	Must:      0,
 	Prefer:    1,
@@ -65,3 +70,21 @@ func Parse(ruleName string) (CandidatePromotionRule, error) {
 		return CandidatePromotionRule(""), fmt.Errorf("Invalid CandidatePromotionRule: %v", ruleName)
 	}
 }
+
+// ParseFromProto returns a *CandidatePromotionRule from a topodatapb.PromotionRule.
+func ParseFromProto(promotionRule topodatapb.PromotionRule) (CandidatePromotionRule, error) {
+	switch promotionRule {
+	case topodatapb.PromotionRule_NEUTRAL, topodatapb.PromotionRule_NONE:
+		return Neutral, nil
+	case topodatapb.PromotionRule_MUST:
+		return Must, fmt.Errorf("ParseFromProto: %v not supported yet", promotionRule)
+	case topodatapb.PromotionRule_PREFER:
+		return Prefer, nil
+	case topodatapb.PromotionRule_PREFER_NOT:
+		return PreferNot, nil
+	case topodatapb.PromotionRule_MUST_NOT:
+		return MustNot, nil
+	default:
+		return Neutral, ErrUnsupportedPromotionRule
+	}
+}

go/vt/vtctl/reparentutil/promotionrule/promotion_rule_test.go

@@ -0,0 +1,61 @@
+package promotionrule
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
+)
+
+func TestPromotionRuleParseFromProto(t *testing.T) {
+	testcases := []struct {
+		in          topodatapb.PromotionRule
+		out         CandidatePromotionRule
+		errContains string
+	}{
+		{
+			in:  topodatapb.PromotionRule_NONE,
+			out: Neutral,
+		},
+		{
+
+			in:  topodatapb.PromotionRule_NEUTRAL,
+			out: Neutral,
+		},
+		{
+			in:          topodatapb.PromotionRule_MUST,
+			out:         Must,
+			errContains: "ParseFromProto: MUST not supported yet",
+		},
+		{
+			in:  topodatapb.PromotionRule_PREFER,
+			out: Prefer,
+		},
+		{
+			in:  topodatapb.PromotionRule_PREFER_NOT,
+			out: PreferNot,
+		},
+		{
+			in:  topodatapb.PromotionRule_MUST_NOT,
+			out: MustNot,
+		},
+		{
+			in:          999,
+			out:         Neutral,
+			errContains: "unsupported promotion rule",
+		},
+	}
+
+	for _, testcase := range testcases {
+		testcase := testcase
+		t.Run(testcase.in.String(), func(t *testing.T) {
+			t.Parallel()
+			rule, err := ParseFromProto(testcase.in)
+			if testcase.errContains != "" {
+				assert.ErrorContains(t, err, testcase.errContains)
+			}
+			assert.Equal(t, testcase.out, rule)
+		})
+	}
+}

go/vt/vttablet/tabletmanager/tm_init.go

@@ -90,6 +90,7 @@ var (
 	initDbNameOverride string
 	skipBuildInfoTags  = "/.*/"
 	initTags           flagutil.StringMapValue
+	promotionRule      topodatapb.PromotionRule
 
 	initTimeout          = 1 * time.Minute
 	mysqlShutdownTimeout = mysqlctl.DefaultShutdownTimeout
@@ -105,6 +106,7 @@ func registerInitFlags(fs *pflag.FlagSet) {
 	fs.Var(&initTags, "init_tags", "(init parameter) comma separated list of key:value pairs used to tag the tablet")
 	fs.DurationVar(&initTimeout, "init_timeout", initTimeout, "(init parameter) timeout to use for the init phase.")
 	fs.DurationVar(&mysqlShutdownTimeout, "mysql-shutdown-timeout", mysqlShutdownTimeout, "timeout to use when MySQL is being shut down.")
+	fs.Var((*topoproto.PromotionRuleFlag)(&promotionRule), "promotion-rule", "Sets the override Promotion Rule for this tablet. Possible values 'none', 'neutral', 'prefer', 'prefer_not' and 'must_not'. Use 'must_not' with caution as it can lead to issues with reparenting.")
 }
 
 var (
@@ -273,6 +275,7 @@ func BuildTabletFromInput(alias *topodatapb.TabletAlias, port, grpcPort int32, d
 		DbNameOverride:       initDbNameOverride,
 		Tags:                 mergeTags(buildTags, initTags),
 		DefaultConnCollation: uint32(charset),
+		PromotionRule:        promotionRule,
 	}, nil
 }
 

proto/topodata.proto

@@ -102,6 +102,27 @@ enum TabletType {
   DRAINED = 8;
 }
 
+// PromotionRule represents a Durability Policy override for an instance of vttablet.
+enum PromotionRule {
+  // NONE represents an undefined promotion rule.
+  NONE = 0;
+
+  // MUST represents a must promotion rule.
+  MUST = 1;
+
+  // NEUTRAL represents a neutral promotion rule.
+  NEUTRAL = 2;
+
+  // PREFER represents a prefer promotion rule.
+  PREFER = 3;
+
+  // PREFER_NOT represents a prefer_not promotion rule.
+  PREFER_NOT = 4;
+
+  // MUST_NOT represents a must_not promotion rule.
+  MUST_NOT = 5;
+}
+
 // Tablet represents information about a running instance of vttablet.
 message Tablet {
   // alias is the unique name of the tablet.
@@ -160,6 +181,10 @@ message Tablet {
   // default_conn_collation is the default connection collation used by this tablet.
   uint32 default_conn_collation = 16;
 
+  // promotion_rule specified the desired Promotion Rule for this tablet. If set,
+  // this overrides the Promotion Rule of the Durability Policy.
+  PromotionRule promotion_rule = 17;
+
   // OBSOLETE: ip and tablet health information
   // string ip = 3;
   // map<string, string> health_map = 11;