tx throttler: healthcheck all cells if --tx-throttler-healthcheck-cells is undefined

go/vt/vttablet/tabletserver/tabletenv/config.go

@@ -723,9 +723,6 @@ func (c *TabletConfig) verifyTxThrottlerConfig() error {
 		return vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "failed to parse throttlerdatapb.Configuration config: %v", err)
 	}
 
-	if len(c.TxThrottlerHealthCheckCells) == 0 {
-		return vterrors.Errorf(vtrpcpb.Code_FAILED_PRECONDITION, "empty healthCheckCells given: %+v", c.TxThrottlerHealthCheckCells)
-	}
 	if v := c.TxThrottlerDefaultPriority; v > sqlparser.MaxPriorityValue || v < 0 {
 		return vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "--tx-throttler-default-priority must be > 0 and < 100 (specified value: %d)", v)
 	}

go/vt/vttablet/tabletserver/tabletenv/config_test.go

@@ -382,13 +382,6 @@ func TestVerifyTxThrottlerConfig(t *testing.T) {
 			EnableTxThrottler: true,
 			TxThrottlerConfig: &TxThrottlerConfigFlag{invalidMaxReplicationLagModuleConfig},
 		},
-		{
-			// enabled without cells defined
-			Name:              "enabled without cells",
-			ExpectedErrorCode: vtrpcpb.Code_FAILED_PRECONDITION,
-			EnableTxThrottler: true,
-			TxThrottlerConfig: &TxThrottlerConfigFlag{defaultMaxReplicationLagModuleConfig},
-		},
 		{
 			// enabled with good config (default/replica tablet type)
 			Name:                        "enabled",

go/vt/vttablet/tabletserver/txthrottler/tx_throttler.go

@@ -19,6 +19,7 @@ package txthrottler
 import (
 	"context"
 	"math/rand"
+	"reflect"
 	"strings"
 	"sync"
 	"time"
@@ -32,10 +33,12 @@ import (
 	"vitess.io/vitess/go/vt/throttler"
 	"vitess.io/vitess/go/vt/topo"
 	"vitess.io/vitess/go/vt/topo/topoproto"
+	"vitess.io/vitess/go/vt/vterrors"
 	"vitess.io/vitess/go/vt/vttablet/tabletserver/tabletenv"
 
 	querypb "vitess.io/vitess/go/vt/proto/query"
 	throttlerdatapb "vitess.io/vitess/go/vt/proto/throttlerdata"
+	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
 )
 
 // These vars store the functions used to create the topo server, healthcheck,
@@ -46,9 +49,11 @@ type topologyWatcherFactoryFunc func(topoServer *topo.Server, hc discovery.Healt
 type throttlerFactoryFunc func(name, unit string, threadCount int, maxRate int64, maxReplicationLagConfig throttler.MaxReplicationLagModuleConfig) (ThrottlerInterface, error)
 
 var (
-	healthCheckFactory     healthCheckFactoryFunc
-	topologyWatcherFactory topologyWatcherFactoryFunc
-	throttlerFactory       throttlerFactoryFunc
+	healthCheckFactory       healthCheckFactoryFunc
+	topologyWatcherFactory   topologyWatcherFactoryFunc
+	throttlerFactory         throttlerFactoryFunc
+	topoCellsRefreshInterval = time.Minute
+	ErrFoundNoCells          = vterrors.New(vtrpcpb.Code_FAILED_PRECONDITION, "found no cells")
 )
 
 func resetTxThrottlerFactories() {
@@ -101,6 +106,16 @@ type TopologyWatcherInterface interface {
 // go/vt/throttler.GlobalManager.
 const TxThrottlerName = "TransactionThrottler"
 
+// fetchKnownCells gathers a list of known cells from the topology. An ErrFoundNoCells
+// error is returned if the GetKnownCells call to topoServer returns no cells.
+func fetchKnownCells(ctx context.Context, topoServer *topo.Server) ([]string, error) {
+	cells, err := topoServer.GetKnownCells(ctx)
+	if err == nil && len(cells) == 0 {
+		err = ErrFoundNoCells
+	}
+	return cells, err
+}
+
 // txThrottler implements TxThrottle for throttling transactions based on replication lag.
 // It's a thin wrapper around the throttler found in vitess/go/vt/throttler.
 // It uses a discovery.HealthCheck to send replication-lag updates to the wrapped throttler.
@@ -175,7 +190,9 @@ type txThrottlerState struct {
 	stopHealthCheck context.CancelFunc
 
 	healthCheck      discovery.HealthCheck
+	healthCheckChan  chan *discovery.TabletHealth
 	topologyWatchers []TopologyWatcherInterface
+	cellsFromTopo    bool
 }
 
 // NewTxThrottler tries to construct a txThrottler from the
@@ -286,44 +303,104 @@ func newTxThrottlerState(topoServer *topo.Server, config *txThrottlerConfig, tar
 		t.Close()
 		return nil, err
 	}
-	result := &txThrottlerState{
+	state := &txThrottlerState{
 		config:    config,
 		throttler: t,
 	}
-	createTxThrottlerHealthCheck(topoServer, config, result, target.Cell)
 
-	result.topologyWatchers = make(
-		[]TopologyWatcherInterface, 0, len(config.healthCheckCells))
-	for _, cell := range config.healthCheckCells {
-		result.topologyWatchers = append(
-			result.topologyWatchers,
+	// get cells from topo if none defined in tabletenv config
+	if len(config.healthCheckCells) == 0 {
+		ctx, cancel := context.WithTimeout(context.Background(), topo.RemoteOperationTimeout)
+		defer cancel()
+		var err error
+		if config.healthCheckCells, err = fetchKnownCells(ctx, topoServer); err != nil {
+			log.Errorf("txThrottler: failed to open throttler: %+v", err)
+			return nil, err
+		}
+		state.cellsFromTopo = true
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	state.stopHealthCheck = cancel
+	state.initHealthCheckStream(topoServer, target)
+	hcProcessor := state.healthChecksProcessorFactory(topoServer, target)
+	go hcProcessor(ctx)
+
+	return state, nil
+}
+
+func (ts *txThrottlerState) initHealthCheckStream(topoServer *topo.Server, target *querypb.Target) {
+	ts.healthCheck = healthCheckFactory(topoServer, target.Cell, ts.config.healthCheckCells)
+	ts.healthCheckChan = ts.healthCheck.Subscribe()
+	ts.topologyWatchers = make(
+		[]TopologyWatcherInterface, 0, len(ts.config.healthCheckCells))
+	for _, cell := range ts.config.healthCheckCells {
+		ts.topologyWatchers = append(
+			ts.topologyWatchers,
 			topologyWatcherFactory(
 				topoServer,
-				result.healthCheck,
+				ts.healthCheck,
 				cell,
 				target.Keyspace,
 				target.Shard,
 				discovery.DefaultTopologyWatcherRefreshInterval,
-				discovery.DefaultTopoReadConcurrency))
+				discovery.DefaultTopoReadConcurrency),
+		)
 	}
-	return result, nil
 }
 
-func createTxThrottlerHealthCheck(topoServer *topo.Server, config *txThrottlerConfig, result *txThrottlerState, cell string) {
-	ctx, cancel := context.WithCancel(context.Background())
-	result.stopHealthCheck = cancel
-	result.healthCheck = healthCheckFactory(topoServer, cell, config.healthCheckCells)
-	ch := result.healthCheck.Subscribe()
-	go func(ctx context.Context) {
+func (ts *txThrottlerState) closeHealthCheckStream() {
+	if ts.healthCheck == nil {
+		return
+	}
+	for _, watcher := range ts.topologyWatchers {
+		watcher.Stop()
+	}
+	ts.topologyWatchers = nil
+	ts.stopHealthCheck()
+	ts.healthCheck.Close()
+}
+
+func (ts *txThrottlerState) healthChecksProcessorFactory(topoServer *topo.Server, target *querypb.Target) func(ctx context.Context) {
+	if ts.cellsFromTopo {
+		return func(ctx context.Context) {
+			cellsUpdateTicker := time.NewTicker(topoCellsRefreshInterval)
+			defer cellsUpdateTicker.Stop()
+			for {
+				select {
+				case <-ctx.Done():
+					return
+				case <-cellsUpdateTicker.C:
+					ctx, cancel := context.WithTimeout(context.Background(), topo.RemoteOperationTimeout)
+					defer cancel()
+
+					cells, err := fetchKnownCells(ctx, topoServer)
+					if err != nil {
+						log.Errorf("txThrottler: failed to fetch cells from topo: %+v", err)
+						continue
+					}
+					if !reflect.DeepEqual(cells, ts.config.healthCheckCells) {
+						log.Info("txThrottler: restarting healthcheck stream due to topology cells update")
+						ts.config.healthCheckCells = cells
+						ts.closeHealthCheckStream()
+						ts.initHealthCheckStream(topoServer, target)
+					}
+				case th := <-ts.healthCheckChan:
+					ts.StatsUpdate(th)
+				}
+			}
+		}
+	}
+	return func(ctx context.Context) {
 		for {
 			select {
 			case <-ctx.Done():
 				return
-			case th := <-ch:
-				result.StatsUpdate(th)
+			case th := <-ts.healthCheckChan:
+				ts.StatsUpdate(th)
 			}
 		}
-	}(ctx)
+	}
 }
 
 func (ts *txThrottlerState) throttle() bool {
@@ -338,15 +415,8 @@ func (ts *txThrottlerState) throttle() bool {
 }
 
 func (ts *txThrottlerState) deallocateResources() {
-	// We don't really need to nil out the fields here
-	// as deallocateResources is not expected to be called
-	// more than once, but it doesn't hurt to do so.
-	for _, watcher := range ts.topologyWatchers {
-		watcher.Stop()
-	}
-	ts.topologyWatchers = nil
-
-	ts.healthCheck.Close()
+	// Close healthcheck and topo watchers
+	ts.closeHealthCheckStream()
 	ts.healthCheck = nil
 
 	// After ts.healthCheck is closed txThrottlerState.StatsUpdate() is guaranteed not

go/vt/vttablet/tabletserver/txthrottler/tx_throttler_test.go

@@ -22,6 +22,7 @@ package txthrottler
 //go:generate mockgen -destination mock_topology_watcher_test.go -package txthrottler vitess.io/vitess/go/vt/vttablet/tabletserver/txthrottler TopologyWatcherInterface
 
 import (
+	context "context"
 	"testing"
 	"time"
 
@@ -111,7 +112,6 @@ func TestEnabledThrottler(t *testing.T) {
 
 	config := tabletenv.NewDefaultConfig()
 	config.EnableTxThrottler = true
-	config.TxThrottlerHealthCheckCells = []string{"cell1", "cell2"}
 	config.TxThrottlerTabletTypes = &topoproto.TabletTypeListFlag{topodatapb.TabletType_REPLICA}
 
 	env := tabletenv.NewEnv(config, t.Name())
@@ -150,6 +150,21 @@ func TestEnabledThrottler(t *testing.T) {
 	assert.Zero(t, throttlerImpl.throttlerRunning.Get())
 }
 
+func TestFetchKnownCells(t *testing.T) {
+	{
+		ts := memorytopo.NewServer("cell1", "cell2")
+		cells, err := fetchKnownCells(context.Background(), ts)
+		assert.Nil(t, err)
+		assert.Equal(t, []string{"cell1", "cell2"}, cells)
+	}
+	{
+		ts := memorytopo.NewServer()
+		cells, err := fetchKnownCells(context.Background(), ts)
+		assert.NotNil(t, err)
+		assert.Zero(t, cells)
+	}
+}
+
 func TestNewTxThrottler(t *testing.T) {
 	config := tabletenv.NewDefaultConfig()
 	env := tabletenv.NewEnv(config, t.Name())