func TestDiscoveryGatewayGetTablets(t *testing.T) {
keyspace := "ks"
shard := "0"
hc := discovery.NewFakeHealthCheck()
dg := createDiscoveryGateway(hc, topo.Server{}, nil, "local", 2).(*discoveryGateway)
// replica should only use local ones
hc.Reset()
dg.tsc.ResetForTesting()
hc.AddTestTablet("remote", "1.1.1.1", 1001, keyspace, shard, topodatapb.TabletType_REPLICA, true, 10, nil)
ep1 := hc.AddTestTablet("local", "2.2.2.2", 1001, keyspace, shard, topodatapb.TabletType_REPLICA, true, 10, nil).Tablet()
tsl := dg.tsc.GetHealthyTabletStats(keyspace, shard, topodatapb.TabletType_REPLICA)
if len(tsl) != 1 || !topo.TabletEquality(tsl[0].Tablet, ep1) {
t.Errorf("want %+v, got %+v", ep1, tsl)
}
// master should use the one with newer timestamp regardless of cell
hc.Reset()
dg.tsc.ResetForTesting()
hc.AddTestTablet("remote", "1.1.1.1", 1001, keyspace, shard, topodatapb.TabletType_MASTER, true, 5, nil)
ep1 = hc.AddTestTablet("remote", "2.2.2.2", 1001, keyspace, shard, topodatapb.TabletType_MASTER, true, 10, nil).Tablet()
tsl = dg.tsc.GetHealthyTabletStats(keyspace, shard, topodatapb.TabletType_MASTER)
if len(tsl) != 1 || !topo.TabletEquality(tsl[0].Tablet, ep1) {
t.Errorf("want %+v, got %+v", ep1, tsl)
}
}
// deleteTabletFromTargetProtected deletes the tablet for the given target.
// LOCK_REQUIRED hc.mu
func (hc *HealthCheckImpl) deleteTabletFromTargetProtected(target *querypb.Target, tablet *topodatapb.Tablet) {
shardMap, ok := hc.targetToTablets[target.Keyspace]
if !ok {
return
}
ttMap, ok := shardMap[target.Shard]
if !ok {
return
}
tList, ok := ttMap[target.TabletType]
if !ok {
return
}
for i, t := range tList {
if topo.TabletEquality(t, tablet) {
tList = append(tList[:i], tList[i+1:]...)
ttMap[target.TabletType] = tList
return
}
}
}
// addTabletToTargetProtected adds the tablet to the given target.
// LOCK_REQUIRED hc.mu
func (hc *HealthCheckImpl) addTabletToTargetProtected(target *querypb.Target, tablet *topodatapb.Tablet) {
shardMap, ok := hc.targetToTablets[target.Keyspace]
if !ok {
shardMap = make(map[string]map[topodatapb.TabletType][]*topodatapb.Tablet)
hc.targetToTablets[target.Keyspace] = shardMap
}
ttMap, ok := shardMap[target.Shard]
if !ok {
ttMap = make(map[topodatapb.TabletType][]*topodatapb.Tablet)
shardMap[target.Shard] = ttMap
}
tList, ok := ttMap[target.TabletType]
if !ok {
tList = make([]*topodatapb.Tablet, 0, 1)
}
for _, t := range tList {
if topo.TabletEquality(t, tablet) {
log.Warningf("tablet is already added: %+v", tablet)
return
}
}
ttMap[target.TabletType] = append(tList, tablet)
}