Add flag to select tx throttler tablet type (vitessio#12174)

* Add flag to select tx throttler tablet type Signed-off-by: Tim Vaillancourt <[email protected]> * REPLICA and/or RDONLY only Signed-off-by: Tim Vaillancourt <[email protected]> * Update flag help msg Signed-off-by: Tim Vaillancourt <[email protected]> * Lowercase types in help/doc Signed-off-by: Tim Vaillancourt <[email protected]> * Help update Signed-off-by: Tim Vaillancourt <[email protected]> * fix test Signed-off-by: Tim Vaillancourt <[email protected]> * No underscores in flag Signed-off-by: Tim Vaillancourt <[email protected]> * Fix test Signed-off-by: Tim Vaillancourt <[email protected]> * Fix merge Signed-off-by: Tim Vaillancourt <[email protected]> * PR suggestion, consolidate config logic Signed-off-by: Tim Vaillancourt <[email protected]> * Update go/vt/vttablet/tabletserver/tabletenv/config.go Co-authored-by: Andrew Mason <[email protected]> Signed-off-by: Tim Vaillancourt <[email protected]> * Use topoproto.TabletTypeListFlag to handle flag Signed-off-by: Tim Vaillancourt <[email protected]> * Fix unit test Signed-off-by: Tim Vaillancourt <[email protected]> * Update go/vt/vttablet/tabletserver/tabletenv/config.go Co-authored-by: Andrew Mason <[email protected]> Signed-off-by: Tim Vaillancourt <[email protected]> * improve test Signed-off-by: Tim Vaillancourt <[email protected]> * pr suggestions Signed-off-by: Tim Vaillancourt <[email protected]> * go fmt Signed-off-by: Tim Vaillancourt <[email protected]> --------- Signed-off-by: Tim Vaillancourt <[email protected]> Co-authored-by: Andrew Mason <[email protected]>
slackhq · May 16, 2023 · 7eb3009 · 7eb3009
1 parent 9604c05
commit 7eb3009
Show file tree

Hide file tree

Showing 7 changed files with 102 additions and 18 deletions.
diff --git a/config/tablet/default.yaml b/config/tablet/default.yaml
@@ -117,6 +117,7 @@ cacheResultFields: true                  # enable-query-plan-field-caching
 # enable-tx-throttler
 # tx-throttler-config
 # tx-throttler-healthcheck-cells
+# tx-throttler-tablet-types
 # enable_transaction_limit
 # enable_transaction_limit_dry_run
 # transaction_limit_per_user

diff --git a/doc/ReplicationLagBasedThrottlingOfTransactions.md b/doc/ReplicationLagBasedThrottlingOfTransactions.md
@@ -30,7 +30,13 @@ If this is not specified a [default](https://github.com/vitessio/vitess/tree/mai
 * *tx-throttler-healthcheck-cells*
 
 A comma separated list of datacenter cells. The throttler will only monitor
-the non-RDONLY replicas found in these cells for replication lag.
+the replicas found in these cells for replication lag.
+
+* *tx-throttler-tablet-types*
+
+A comma separated list of tablet types. The throttler will only monitor tablets
+with these types. Only `replica` and/or `rdonly` types are supported. The default
+is `replica`.
 
 # Caveats and Known Issues
 * The throttler keeps trying to explore the maximum rate possible while keeping
@@ -39,4 +45,3 @@ lag limit may occasionally be slightly violated.
 
 * Transactions are considered homogeneous. There is currently no support
 for specifying how `expensive` a transaction is.
-
diff --git a/go/flags/endtoend/vttablet.txt b/go/flags/endtoend/vttablet.txt
@@ -939,6 +939,8 @@ max_rate_approach_threshold: 0.9
 	Default priority assigned to queries that lack priority information.
   --tx-throttler-healthcheck-cells value
 	Synonym to -tx_throttler_healthcheck_cells
+  --tx-throttler-tablet-types string
+	A comma-separated list of tablet types. Only tablets of this type are monitored for replication lag by the transaction throttler. Supported types are replica and/or rdonly. (default replica)
   --tx_throttler_config string
 	The configuration of the transaction throttler as a text formatted throttlerdata.Configuration protocol buffer message (default target_replication_lag_sec: 2
 max_replication_lag_sec: 10

diff --git a/go/vt/vttablet/tabletserver/tabletenv/config.go b/go/vt/vttablet/tabletserver/tabletenv/config.go
@@ -29,7 +29,13 @@ import (
 	"vitess.io/vitess/go/streamlog"
 	"vitess.io/vitess/go/vt/dbconfigs"
 	"vitess.io/vitess/go/vt/log"
+
+	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
+	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
+	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/throttler"
+	"vitess.io/vitess/go/vt/topo/topoproto"
+	"vitess.io/vitess/go/vt/vterrors"
 )
 
 // These constants represent values for various config parameters.
@@ -81,6 +87,7 @@ var (
 	unhealthyThreshold           time.Duration
 	transitionGracePeriod        time.Duration
 	enableReplicationReporter    bool
+	txThrottlerTabletTypes       string
 )
 
 func init() {
@@ -135,6 +142,7 @@ func init() {
 	flagutil.DualFormatStringVar(&currentConfig.TxThrottlerConfig, "tx_throttler_config", defaultConfig.TxThrottlerConfig, "The configuration of the transaction throttler as a text formatted throttlerdata.Configuration protocol buffer message")
 	flagutil.DualFormatStringListVar(&currentConfig.TxThrottlerHealthCheckCells, "tx_throttler_healthcheck_cells", defaultConfig.TxThrottlerHealthCheckCells, "A comma-separated list of cells. Only tabletservers running in these cells will be monitored for replication lag by the transaction throttler.")
 	flag.IntVar(&currentConfig.TxThrottlerDefaultPriority, "tx-throttler-default-priority", defaultConfig.TxThrottlerDefaultPriority, "Default priority assigned to queries that lack priority information.")
+	flag.StringVar(&txThrottlerTabletTypes, "tx-throttler-tablet-types", "replica", "A comma-separated list of tablet types. Only tablets of this type are monitored for replication lag by the transaction throttler. Supported types are replica and/or rdonly.")
 
 	flag.BoolVar(&enableHotRowProtection, "enable_hot_row_protection", false, "If true, incoming transactions for the same row (range) will be queued and cannot consume all txpool slots.")
 	flag.BoolVar(&enableHotRowProtectionDryRun, "enable_hot_row_protection_dry_run", false, "If true, hot row protection is not enforced but logs if transactions would have been queued.")
@@ -241,6 +249,13 @@ func Init() {
 	if *txLogHandler != "" {
 		TxLogger.ServeLogs(*txLogHandler, streamlog.GetFormatter(TxLogger))
 	}
+
+	if txThrottlerTabletTypes != "" {
+		var err error
+		if currentConfig.TxThrottlerTabletTypes, err = topoproto.ParseTabletTypes(txThrottlerTabletTypes); err != nil {
+			log.Exit("Invalid -tx-throttler-tablet-types provided")
+		}
+	}
 }
 
 // TabletConfig contains all the configuration for query service
@@ -289,10 +304,11 @@ type TabletConfig struct {
 	TwoPCCoordinatorAddress string  `json:"-"`
 	TwoPCAbandonAge         Seconds `json:"-"`
 
-	EnableTxThrottler           bool     `json:"-"`
-	TxThrottlerConfig           string   `json:"-"`
-	TxThrottlerHealthCheckCells []string `json:"-"`
-	TxThrottlerDefaultPriority  int      `json:"-"`
+	EnableTxThrottler           bool                    `json:"-"`
+	TxThrottlerConfig           string                  `json:"-"`
+	TxThrottlerHealthCheckCells []string                `json:"-"`
+	TxThrottlerDefaultPriority  int                     `json:"-"`
+	TxThrottlerTabletTypes      []topodatapb.TabletType `json:"-"`
 
 	EnableLagThrottler bool `json:"-"`
 
@@ -397,6 +413,9 @@ func (c *TabletConfig) Verify() error {
 	if err := c.verifyTransactionLimitConfig(); err != nil {
 		return err
 	}
+	if err := c.verifyTxThrottlerConfig(); err != nil {
+		return err
+	}
 	if v := c.HotRowProtection.MaxQueueSize; v <= 0 {
 		return fmt.Errorf("-hot_row_protection_max_queue_size must be > 0 (specified value: %v)", v)
 	}
@@ -442,6 +461,22 @@ func (c *TabletConfig) verifyTransactionLimitConfig() error {
 	return nil
 }
 
+// verifyTxThrottlerConfig checks the TxThrottler related config for sanity.
+func (c *TabletConfig) verifyTxThrottlerConfig() error {
+	if c.TxThrottlerTabletTypes == nil || len(c.TxThrottlerTabletTypes) == 0 {
+		return vterrors.New(vtrpcpb.Code_FAILED_PRECONDITION, "--tx-throttler-tablet-types must be defined when transaction throttler is enabled")
+	}
+	for _, tabletType := range c.TxThrottlerTabletTypes {
+		switch tabletType {
+		case topodatapb.TabletType_REPLICA, topodatapb.TabletType_RDONLY:
+			continue
+		default:
+			return vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "unsupported tablet type %q", tabletType)
+		}
+	}
+	return nil
+}
+
 // Some of these values are for documentation purposes.
 // They actually get overwritten during Init.
 var defaultConfig = TabletConfig{
@@ -505,7 +540,8 @@ var defaultConfig = TabletConfig{
 	EnableTxThrottler:           false,
 	TxThrottlerConfig:           defaultTxThrottlerConfig(),
 	TxThrottlerHealthCheckCells: []string{},
-	TxThrottlerDefaultPriority:  0, // This leads to all queries being candidates to throttle
+	TxThrottlerDefaultPriority:  sqlparser.MaxPriorityValue, // This leads to all queries being candidates to throttle
+	TxThrottlerTabletTypes:      []topodatapb.TabletType{},
 
 	EnableLagThrottler: false, // Feature flag; to switch to 'true' at some stage in the future
 

diff --git a/go/vt/vttablet/tabletserver/tabletenv/config_test.go b/go/vt/vttablet/tabletserver/tabletenv/config_test.go
@@ -27,6 +27,10 @@ import (
 
 	"vitess.io/vitess/go/cache"
 	"vitess.io/vitess/go/vt/dbconfigs"
+	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
+	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
+	"vitess.io/vitess/go/vt/topo/topoproto"
+	"vitess.io/vitess/go/vt/vterrors"
 	"vitess.io/vitess/go/yaml2"
 )
 
@@ -363,3 +367,32 @@ func TestFlags(t *testing.T) {
 	want.SanitizeLogMessages = true
 	assert.Equal(t, want, currentConfig)
 }
+
+func TestVerifyTxThrottlerConfig(t *testing.T) {
+	{
+		// default config (replica)
+		assert.Nil(t, currentConfig.verifyTxThrottlerConfig())
+	}
+	{
+		// replica + rdonly (allowed)
+		currentConfig.TxThrottlerTabletTypes = &topoproto.TabletTypeListFlag{
+			topodatapb.TabletType_REPLICA,
+			topodatapb.TabletType_RDONLY,
+		}
+		assert.Nil(t, currentConfig.verifyTxThrottlerConfig())
+	}
+	{
+		// no tablet types
+		currentConfig.TxThrottlerTabletTypes = &topoproto.TabletTypeListFlag{}
+		err := currentConfig.verifyTxThrottlerConfig()
+		assert.NotNil(t, err)
+		assert.Equal(t, vtrpcpb.Code_FAILED_PRECONDITION, vterrors.Code(err))
+	}
+	{
+		// disallowed tablet type
+		currentConfig.TxThrottlerTabletTypes = &topoproto.TabletTypeListFlag{topodatapb.TabletType_DRAINED}
+		err := currentConfig.verifyTxThrottlerConfig()
+		assert.NotNil(t, err)
+		assert.Equal(t, vtrpcpb.Code_INVALID_ARGUMENT, vterrors.Code(err))
+	}
+}
diff --git a/go/vt/vttablet/tabletserver/txthrottler/tx_throttler.go b/go/vt/vttablet/tabletserver/txthrottler/tx_throttler.go
@@ -178,6 +178,7 @@ func tryCreateTxThrottler(env tabletenv.Env, topoServer *topo.Server) (*TxThrott
 	return newTxThrottler(env, &txThrottlerConfig{
 		enabled:          true,
 		topoServer:       topoServer,
+		tabletTypes:      env.Config().TxThrottlerTabletTypes,
 		throttlerConfig:  &throttlerConfig,
 		healthCheckCells: healthCheckCells,
 	})
@@ -196,10 +197,15 @@ type txThrottlerConfig struct {
 	// healthCheckCells stores the cell names in which running vttablets will be monitored for
 	// replication lag.
 	healthCheckCells []string
+
+	// tabletTypes stores the tablet types for throttling
+	tabletTypes []topodatapb.TabletType
 }
 
 // txThrottlerState holds the state of an open TxThrottler object.
 type txThrottlerState struct {
+	config *txThrottlerConfig
+
 	// throttleMu serializes calls to throttler.Throttler.Throttle(threadId).
 	// That method is required to be called in serial for each threadId.
 	throttleMu sync.Mutex
@@ -294,6 +300,7 @@ func newTxThrottlerState(config *txThrottlerConfig, keyspace, shard string,
 		return nil, err
 	}
 	result := &txThrottlerState{
+		config:    config,
 		throttler: t,
 	}
 	result.healthCheck = healthCheckFactory()
@@ -344,16 +351,14 @@ func (ts *txThrottlerState) deallocateResources() {
 	ts.throttler = nil
 }
 
-// StatsUpdate is part of the LegacyHealthCheckStatsListener interface.
+// StatsUpdate updates the health of a tablet with the given healthcheck.
 func (ts *txThrottlerState) StatsUpdate(tabletStats *discovery.LegacyTabletStats) {
-	// Ignore PRIMARY and RDONLY stats.
-	// We currently do not monitor RDONLY tablets for replication lag. RDONLY tablets are not
-	// candidates for becoming primary during failover, and it's acceptable to serve somewhat
-	// stale date from these.
-	// TODO(erez): If this becomes necessary, we can add a configuration option that would
-	// determine whether we consider RDONLY tablets here, as well.
-	if tabletStats.Target.TabletType != topodatapb.TabletType_REPLICA {
-		return
+	// Monitor tablets for replication lag if they have a tablet
+	// type specified by the --tx_throttler_tablet_types flag.
+	for _, expectedTabletType := range ts.config.tabletTypes {
+		if tabletStats.Target.TabletType == expectedTabletType {
+			ts.throttler.RecordReplicationLag(time.Now(), tabletStats)
+			return
+		}
 	}
-	ts.throttler.RecordReplicationLag(time.Now(), tabletStats)
 }
diff --git a/go/vt/vttablet/tabletserver/txthrottler/tx_throttler_test.go b/go/vt/vttablet/tabletserver/txthrottler/tx_throttler_test.go
@@ -32,6 +32,7 @@ import (
 	"vitess.io/vitess/go/vt/throttler"
 	"vitess.io/vitess/go/vt/topo"
 	"vitess.io/vitess/go/vt/topo/memorytopo"
+	"vitess.io/vitess/go/vt/topo/topoproto"
 	"vitess.io/vitess/go/vt/vttablet/tabletserver/tabletenv"
 
 	querypb "vitess.io/vitess/go/vt/proto/query"
@@ -112,8 +113,9 @@ func TestEnabledThrottler(t *testing.T) {
 	config := tabletenv.NewDefaultConfig()
 	config.EnableTxThrottler = true
 	config.TxThrottlerHealthCheckCells = []string{"cell1", "cell2"}
-	env := tabletenv.NewEnv(config, t.Name())
+	config.TxThrottlerTabletTypes = &topoproto.TabletTypeListFlag{topodatapb.TabletType_REPLICA}
 
+	env := tabletenv.NewEnv(config, t.Name())
 	throttler, err := tryCreateTxThrottler(env, ts)
 	assert.Nil(t, err)
 	throttler.InitDBConfig(&querypb.Target{