region/cell feature draft

go/vt/discovery/tablet_stats_cache.go

@@ -41,6 +41,9 @@ type TabletStatsCache struct {
 	// Note we keep track of all master tablets in all cells.
 	cell string
 
+	// cell to region mapping function
+	cellToRegion func(cell string) string
+
 	// mu protects the entries map. It does not protect individual
 	// entries in the map.
 	mu sync.RWMutex
@@ -65,8 +68,8 @@ type tabletStatsCacheEntry struct {
 // Note we do the registration in this code to guarantee we call
 // SetListener with sendDownEvents=true, as we need these events
 // to maintain the integrity of our cache.
-func NewTabletStatsCache(hc HealthCheck, cell string) *TabletStatsCache {
-	return newTabletStatsCache(hc, cell, true /* setListener */)
+func NewTabletStatsCache(hc HealthCheck, cell string, cellToRegion func(cell string) string) *TabletStatsCache {
+	return newTabletStatsCache(hc, cell, cellToRegion, true /* setListener */)
 }
 
 // NewTabletStatsCacheDoNotSetListener is identical to NewTabletStatsCache
@@ -76,14 +79,22 @@ func NewTabletStatsCache(hc HealthCheck, cell string) *TabletStatsCache {
 // When the caller sets its own listener on "hc", they must make sure that they
 // set the parameter  "sendDownEvents" to "true" or this cache won't properly
 // remove tablets whose tablet type changes.
-func NewTabletStatsCacheDoNotSetListener(cell string) *TabletStatsCache {
-	return newTabletStatsCache(nil, cell, false /* setListener */)
+func NewTabletStatsCacheDoNotSetListener(cell string, cellToRegion func(cell string) string) *TabletStatsCache {
+	return newTabletStatsCache(nil, cell, cellToRegion, false /* setListener */)
+}
+
+// UpdateCellsToRegions is mainly for testing purpose, the `cellsToRegions` mapping should be provided in the constructor
+func (tc *TabletStatsCache) UpdateCellsToRegions(cellsToRegions map[string]string) {
+	tc.cellToRegion = func(cell string) string {
+		return cellsToRegions[cell]
+	}
 }
 
-func newTabletStatsCache(hc HealthCheck, cell string, setListener bool) *TabletStatsCache {
+func newTabletStatsCache(hc HealthCheck, cell string, cellToRegion func(cell string) string, setListener bool) *TabletStatsCache {
 	tc := &TabletStatsCache{
-		cell:    cell,
-		entries: make(map[string]map[string]map[topodatapb.TabletType]*tabletStatsCacheEntry),
+		cell:         cell,
+		cellToRegion: cellToRegion,
+		entries:      make(map[string]map[string]map[topodatapb.TabletType]*tabletStatsCacheEntry),
 	}
 
 	if setListener {
@@ -144,8 +155,8 @@ func (tc *TabletStatsCache) getOrCreateEntry(target *querypb.Target) *tabletStat
 
 // StatsUpdate is part of the HealthCheckStatsListener interface.
 func (tc *TabletStatsCache) StatsUpdate(ts *TabletStats) {
-	if ts.Target.TabletType != topodatapb.TabletType_MASTER && ts.Tablet.Alias.Cell != tc.cell {
-		// this is for a non-master tablet in a different cell, drop it
+	if ts.Target.TabletType != topodatapb.TabletType_MASTER && ts.Tablet.Alias.Cell != tc.cell && tc.cellToRegion(ts.Tablet.Alias.Cell) != tc.cellToRegion(tc.cell) {
+		// this is for a non-master tablet in a different cell and a different region, drop it
 		return
 	}
 

go/vt/discovery/tablet_stats_cache_wait_test.go

@@ -131,7 +131,7 @@ func TestWaitForTablets(t *testing.T) {
 	createFakeConn(tablet, input)
 
 	hc := NewHealthCheck(1*time.Millisecond, 1*time.Hour)
-	tsc := NewTabletStatsCache(hc, "cell")
+	tsc := NewTabletStatsCache(hc, "cell", func(cell string) string { return cell })
 	hc.AddTablet(tablet, "")
 
 	// this should time out

go/vt/topo/cell_info.go

@@ -80,6 +80,16 @@ func (ts *Server) GetCellInfo(ctx context.Context, cell string, strongRead bool)
 	return ci, nil
 }
 
+// GetRegionByCell gets the region from a given cell
+func (ts *Server) GetRegionByCell(ctx context.Context, cell string) (string, error) {
+
+	cellInfo, err := ts.GetCellInfo(ctx, cell, false)
+	if err == nil {
+		return cellInfo.Region, nil
+	}
+	return "", err
+}
+
 // CreateCellInfo creates a new CellInfo with the provided content.
 func (ts *Server) CreateCellInfo(ctx context.Context, cell string, ci *topodatapb.CellInfo) error {
 	// Pack the content.

go/vt/topo/server.go

@@ -294,6 +294,28 @@ func (ts *Server) ConnForCell(ctx context.Context, cell string) (Conn, error) {
 	return conn, nil
 }
 
+// CellToRegionMapper function is a wrapper around topo.Server#GetRegionByCell with caching and error handling
+func (ts *Server) CellToRegionMapper() func(cell string) string {
+
+	memoize := make(map[string]string)
+	ctx := context.Background()
+
+	return func(cell string) string {
+		if ts == nil {
+			return cell
+		}
+		if region, ok := memoize[cell]; ok {
+			return region
+		}
+		if region, err := ts.GetRegionByCell(ctx, cell); err == nil && region != "" {
+			memoize[cell] = region
+			return region
+		}
+		// for backward compatibility, when region isn't available, it's the same as given cell
+		return cell
+	}
+}
+
 // Close will close all connections to underlying topo Server.
 // It will nil all member variables, so any further access will panic.
 func (ts *Server) Close() {