Change healthcheck to only reload specific keyspace shard tablets

go/vt/discovery/fake_healthcheck.go

@@ -192,7 +192,7 @@ func (fhc *FakeHealthCheck) Unsubscribe(c chan *TabletHealth) {
 }
 
 // GetLoadTabletsTrigger is not implemented.
-func (fhc *FakeHealthCheck) GetLoadTabletsTrigger() chan struct{} {
+func (fhc *FakeHealthCheck) GetLoadTabletsTrigger() chan topo.KeyspaceShard {
 	return nil
 }
 

go/vt/discovery/healthcheck.go

@@ -247,7 +247,7 @@ type HealthCheck interface {
 	Unsubscribe(c chan *TabletHealth)
 
 	// GetLoadTabletsTrigger returns a channel that is used to inform when to load tablets.
-	GetLoadTabletsTrigger() chan struct{}
+	GetLoadTabletsTrigger() chan topo.KeyspaceShard
 }
 
 var _ HealthCheck = (*HealthCheckImpl)(nil)
@@ -297,8 +297,8 @@ type HealthCheckImpl struct {
 	subMu sync.Mutex
 	// subscribers
 	subscribers map[chan *TabletHealth]struct{}
-	// loadTablets trigger is used to immediately load a new primary tablet when the current one has been demoted
-	loadTabletsTrigger chan struct{}
+	// loadTabletsTrigger is used to immediately load information about tablets of a specific shard.
+	loadTabletsTrigger chan topo.KeyspaceShard
 }
 
 // NewVTGateHealthCheckFilters returns healthcheck filters for vtgate.
@@ -363,7 +363,7 @@ func NewHealthCheck(ctx context.Context, retryDelay, healthCheckTimeout time.Dur
 		healthy:            make(map[KeyspaceShardTabletType][]*TabletHealth),
 		subscribers:        make(map[chan *TabletHealth]struct{}),
 		cellAliases:        make(map[string]string),
-		loadTabletsTrigger: make(chan struct{}, 1),
+		loadTabletsTrigger: make(chan topo.KeyspaceShard, 1024),
 	}
 	var topoWatchers []*TopologyWatcher
 	cells := strings.Split(cellsToWatch, ",")
@@ -535,18 +535,21 @@ func (hc *HealthCheckImpl) updateHealth(th *TabletHealth, prevTarget *query.Targ
 		}
 
 		// If the previous tablet type was primary, we need to check if the next new primary has already been assigned.
-		// If no new primary has been assigned, we will trigger a `loadTablets` call to immediately redirect traffic to the new primary.
+		// If no new primary has been assigned, we will trigger loading of tablets for this keyspace shard to immediately redirect traffic to the new primary.
 		//
 		// This is to avoid a situation where a newly primary tablet for a shard has just been started and the tableRefreshInterval has not yet passed,
 		// causing an interruption where no primary is assigned to the shard.
 		if prevTarget.TabletType == topodata.TabletType_PRIMARY {
 			if primaries := hc.healthData[oldTargetKey]; len(primaries) == 0 {
 				log.Infof("We will have no health data for the next new primary tablet after demoting the tablet: %v, so start loading tablets now", topotools.TabletIdent(th.Tablet))
-				// We want to trigger a loadTablets call, but if the channel is not empty
-				// then a trigger is already scheduled, we don't need to trigger another one.
-				// This also prevents the code from deadlocking as described in https://github.com/vitessio/vitess/issues/16994.
+				// We want to trigger a call to load tablets for this keyspace-shard,
+				// but we want this to be non-blocking to prevent the code from deadlocking as described in https://github.com/vitessio/vitess/issues/16994.
+				// If the buffer is exhausted, then we'll just receive the update when all the tablets are loaded on the ticker.
 				select {
-				case hc.loadTabletsTrigger <- struct{}{}:
+				case hc.loadTabletsTrigger <- topo.KeyspaceShard{
+					Keyspace: prevTarget.Keyspace,
+					Shard:    prevTarget.Shard,
+				}:
 				default:
 				}
 			}
@@ -662,7 +665,7 @@ func (hc *HealthCheckImpl) broadcast(th *TabletHealth) {
 }
 
 // GetLoadTabletsTrigger returns a channel that is used to inform when to load tablets.
-func (hc *HealthCheckImpl) GetLoadTabletsTrigger() chan struct{} {
+func (hc *HealthCheckImpl) GetLoadTabletsTrigger() chan topo.KeyspaceShard {
 	return hc.loadTabletsTrigger
 }
 

go/vt/discovery/healthcheck_test.go

@@ -1227,6 +1227,68 @@ func TestPrimaryInOtherCell(t *testing.T) {
 	mustMatch(t, want, a[0], "Expecting healthy primary")
 }
 
+// TestLoadTabletsTrigger tests that we send the correct information on the load tablets trigger.
+func TestLoadTabletsTrigger(t *testing.T) {
+	ctx := utils.LeakCheckContext(t)
+
+	// create a health check instance.
+	hc := NewHealthCheck(ctx, time.Hour, time.Hour, nil, "", "", nil)
+	defer hc.Close()
+
+	ks := "keyspace"
+	shard := "shard"
+	// Add a tablet to the topology.
+	tablet1 := &topodatapb.Tablet{
+		Alias: &topodatapb.TabletAlias{
+			Cell: "zone-1",
+			Uid:  100,
+		},
+		Type:     topodatapb.TabletType_REPLICA,
+		Hostname: "host1",
+		PortMap: map[string]int32{
+			"grpc": 123,
+		},
+		Keyspace: ks,
+		Shard:    shard,
+	}
+
+	// We want to run updateHealth with arguments that always
+	// make it trigger load Tablets.
+	th := &TabletHealth{
+		Tablet: tablet1,
+		Target: &querypb.Target{
+			Keyspace:   ks,
+			Shard:      shard,
+			TabletType: topodatapb.TabletType_REPLICA,
+		},
+	}
+	prevTarget := &querypb.Target{
+		Keyspace:   ks,
+		Shard:      shard,
+		TabletType: topodatapb.TabletType_PRIMARY,
+	}
+	hc.AddTablet(tablet1)
+
+	numTriggers := 10
+	for i := 0; i < numTriggers; i++ {
+		// Since the previous target was a primary, and there are no other
+		// primary tablets for the given keyspace shard, we will see the healtcheck
+		// send on the loadTablets trigger. We just want to verify the information
+		// there is correct.
+		hc.updateHealth(th, prevTarget, false, false)
+	}
+
+	ch := hc.GetLoadTabletsTrigger()
+	require.Len(t, ch, numTriggers)
+	for i := 0; i < numTriggers; i++ {
+		// Read from the channel and verify we indeed have the right values.
+		kss := <-ch
+		require.EqualValues(t, ks, kss.Keyspace)
+		require.EqualValues(t, shard, kss.Shard)
+	}
+	require.Len(t, ch, 0)
+}
+
 func TestReplicaInOtherCell(t *testing.T) {
 	ctx := utils.LeakCheckContext(t)
 

go/vt/discovery/topology_watcher.go

@@ -110,23 +110,52 @@ func (tw *TopologyWatcher) getTablets() ([]*topo.TabletInfo, error) {
 	return tw.topoServer.GetTabletsByCell(tw.ctx, tw.cell, nil)
 }
 
+func (tw *TopologyWatcher) getTabletsByShard(keyspace string, shard string) ([]*topo.TabletInfo, error) {
+	return tw.topoServer.GetTabletsByShardCell(tw.ctx, keyspace, shard, []string{tw.cell})
+}
+
 // Start starts the topology watcher.
 func (tw *TopologyWatcher) Start() {
-	tw.wg.Add(1)
+	tw.wg.Add(2)
+	// Goroutine to refresh the tablets list periodically.
 	go func(t *TopologyWatcher) {
 		defer t.wg.Done()
 		ticker := time.NewTicker(t.refreshInterval)
 		defer ticker.Stop()
 		for {
-			t.loadTablets()
+			// Since we are going to load all the tablets,
+			// we can clear out the entire list for reloading
+			// specific keyspace shards.
+			func() {
+				for {
+					select {
+					case <-tw.healthcheck.GetLoadTabletsTrigger():
+					default:
+						return
+					}
+				}
+			}()
+			t.loadTablets("", "")
 			select {
 			case <-t.ctx.Done():
 				return
-			case <-tw.healthcheck.GetLoadTabletsTrigger():
 			case <-ticker.C:
 			}
 		}
 	}(tw)
+	// Go routine to refresh tablets for a specific
+	// keyspace shard.
+	go func(t *TopologyWatcher) {
+		defer t.wg.Done()
+		for {
+			select {
+			case <-t.ctx.Done():
+				return
+			case kss := <-tw.healthcheck.GetLoadTabletsTrigger():
+				t.loadTablets(kss.Keyspace, kss.Shard)
+			}
+		}
+	}(tw)
 }
 
 // Stop stops the watcher. It does not clean up the tablets added to HealthCheck.
@@ -136,23 +165,36 @@ func (tw *TopologyWatcher) Stop() {
 	tw.wg.Wait()
 }
 
-func (tw *TopologyWatcher) loadTablets() {
+func (tw *TopologyWatcher) loadTablets(keyspace string, shard string) {
 	newTablets := make(map[string]*tabletInfo)
 	var partialResult bool
 
-	// First get the list of all tablets.
-	tabletInfos, err := tw.getTablets()
-	topologyWatcherOperations.Add(topologyWatcherOpListTablets, 1)
-	if err != nil {
-		topologyWatcherErrors.Add(topologyWatcherOpListTablets, 1)
-		// If we get a partial result error, we just log it and process the tablets that we did manage to fetch.
-		if topo.IsErrType(err, topo.PartialResult) {
-			log.Errorf("received partial result from getTablets for cell %v: %v", tw.cell, err)
-			partialResult = true
-		} else { // For all other errors, just return.
-			log.Errorf("error getting tablets for cell: %v: %v", tw.cell, err)
+	var tabletInfos []*topo.TabletInfo
+	var err error
+	if keyspace != "" && shard != "" {
+		tabletInfos, err = tw.getTabletsByShard(keyspace, shard)
+		if err != nil {
+			log.Errorf("error getting tablets for keyspace-shard: %v:%v: %v", keyspace, shard, err)
 			return
 		}
+		// Since we are only reading tablets for a keyspace shard,
+		// this is by default a partial result.
+		partialResult = true
+	} else {
+		// First get the list of all tablets.
+		tabletInfos, err = tw.getTablets()
+		topologyWatcherOperations.Add(topologyWatcherOpListTablets, 1)
+		if err != nil {
+			topologyWatcherErrors.Add(topologyWatcherOpListTablets, 1)
+			// If we get a partial result error, we just log it and process the tablets that we did manage to fetch.
+			if topo.IsErrType(err, topo.PartialResult) {
+				log.Errorf("received partial result from getTablets for cell %v: %v", tw.cell, err)
+				partialResult = true
+			} else { // For all other errors, just return.
+				log.Errorf("error getting tablets for cell: %v: %v", tw.cell, err)
+				return
+			}
+		}
 	}
 
 	// Accumulate a list of all known alias strings to use later