func TestDiscoveryGatewayGetEndPoints(t *testing.T) {
keyspace := "ks"
shard := "0"
hc := newFakeHealthCheck()
dg := createDiscoveryGateway(hc, topo.Server{}, nil, "local", time.Millisecond, 2, time.Second, time.Second, time.Second, nil).(*discoveryGateway)
// replica should only use local ones
hc.addTestEndPoint("remote", "1.1.1.1", 1001, keyspace, shard, pbt.TabletType_REPLICA, 10, nil, nil)
ep1 := hc.addTestEndPoint("local", "2.2.2.2", 1001, keyspace, shard, pbt.TabletType_REPLICA, 10, nil, nil)
eps := dg.getEndPoints(keyspace, shard, pbt.TabletType_REPLICA)
if len(eps) != 1 || !topo.EndPointEquality(eps[0], ep1) {
t.Errorf("want %+v, got %+v", ep1, eps[0])
}
// master should use the one with newer timestamp regardless of cell
hc.addTestEndPoint("remote", "1.1.1.1", 1001, keyspace, shard, pbt.TabletType_MASTER, 5, nil, nil)
ep1 = hc.addTestEndPoint("remote", "2.2.2.2", 1001, keyspace, shard, pbt.TabletType_MASTER, 10, nil, nil)
eps = dg.getEndPoints(keyspace, shard, pbt.TabletType_MASTER)
if len(eps) != 1 || !topo.EndPointEquality(eps[0], ep1) {
t.Errorf("want %+v, got %+v", ep1, eps[0])
}
}
func updateEndpoint(ctx context.Context, ts topo.Server, cell, keyspace, shard string, tabletType pb.TabletType, endpoint *pb.EndPoint) error {
return retryUpdateEndpoints(ctx, ts, cell, keyspace, shard, tabletType, true, /* create */
func(endpoints *pb.EndPoints) bool {
// Look for an existing entry to update.
for i := range endpoints.Entries {
if endpoints.Entries[i].Uid == endpoint.Uid {
if topo.EndPointEquality(endpoints.Entries[i], endpoint) {
// The entry already exists and is the same.
return false
}
// Update an existing entry.
endpoints.Entries[i] = endpoint
return true
}
}
// The entry doesn't exist, so add it.
endpoints.Entries = append(endpoints.Entries, endpoint)
return true
})
}
func (zkts *Server) updateTabletEndpoint(oldValue string, oldStat zk.Stat, addr *topodatapb.EndPoint) (newValue string, err error) {
if oldStat == nil {
// The incoming object doesn't exist - we haven't been placed in the serving
// graph yet, so don't update. Assume the next process that rebuilds the graph
// will get the updated tablet location.
return "", errSkipUpdate
}
var addrs *topodatapb.EndPoints
if oldValue != "" {
addrs = &topodatapb.EndPoints{}
if len(oldValue) > 0 {
if err := json.Unmarshal([]byte(oldValue), addrs); err != nil {
return "", fmt.Errorf("EndPoints unmarshal failed: %v %v", oldValue, err)
}
}
foundTablet := false
for i, entry := range addrs.Entries {
if entry.Uid == addr.Uid {
foundTablet = true
if !topo.EndPointEquality(entry, addr) {
addrs.Entries[i] = addr
}
break
}
}
if !foundTablet {
addrs.Entries = append(addrs.Entries, addr)
}
} else {
addrs = topo.NewEndPoints()
addrs.Entries = append(addrs.Entries, addr)
}
data, err := json.MarshalIndent(addrs, "", " ")
if err != nil {
return "", err
}
return string(data), nil
}
// deleteEndPointFromTargetProtected deletes the endpoint for the given target.
// LOCK_REQUIRED hc.mu
func (hc *HealthCheckImpl) deleteEndPointFromTargetProtected(target *querypb.Target, endPoint *topodatapb.EndPoint) {
shardMap, ok := hc.targetToEPs[target.Keyspace]
if !ok {
return
}
ttMap, ok := shardMap[target.Shard]
if !ok {
return
}
epList, ok := ttMap[target.TabletType]
if !ok {
return
}
for i, ep := range epList {
if topo.EndPointEquality(ep, endPoint) {
epList = append(epList[:i], epList[i+1:]...)
ttMap[target.TabletType] = epList
return
}
}
}
// deleteEndPointFromTargetProtected deletes the endpoint for the given target.
// LOCK_REQUIRED hc.mu
func (hc *HealthCheckImpl) deleteEndPointFromTargetProtected(target *pbq.Target, endPoint *pbt.EndPoint) {
shardMap, ok := hc.targetToEPs[target.Keyspace]
if !ok {
return
}
ttMap, ok := shardMap[target.Shard]
if !ok {
return
}
epList, ok := ttMap[target.TabletType]
if !ok {
return
}
for i, ep := range epList {
if topo.EndPointEquality(ep, endPoint) {
epList = append(epList[:i], epList[i+1:]...)
ttMap[target.TabletType] = epList
hcConnCounters.Add([]string{target.Keyspace, target.Shard, strings.ToLower(target.TabletType.String())}, -1)
return
}
}
}
// addEndPointToTargetProtected adds the endpoint to the given target.
// LOCK_REQUIRED hc.mu
func (hc *HealthCheckImpl) addEndPointToTargetProtected(target *querypb.Target, endPoint *topodatapb.EndPoint) {
shardMap, ok := hc.targetToEPs[target.Keyspace]
if !ok {
shardMap = make(map[string]map[topodatapb.TabletType][]*topodatapb.EndPoint)
hc.targetToEPs[target.Keyspace] = shardMap
}
ttMap, ok := shardMap[target.Shard]
if !ok {
ttMap = make(map[topodatapb.TabletType][]*topodatapb.EndPoint)
shardMap[target.Shard] = ttMap
}
epList, ok := ttMap[target.TabletType]
if !ok {
epList = make([]*topodatapb.EndPoint, 0, 1)
}
for _, ep := range epList {
if topo.EndPointEquality(ep, endPoint) {
log.Warningf("endpoint is already added: %+v", endPoint)
return
}
}
ttMap[target.TabletType] = append(epList, endPoint)
}
// addEndPointToTargetProtected adds the endpoint to the given target.
// LOCK_REQUIRED hc.mu
func (hc *HealthCheckImpl) addEndPointToTargetProtected(target *pbq.Target, endPoint *pbt.EndPoint) {
shardMap, ok := hc.targetToEPs[target.Keyspace]
if !ok {
shardMap = make(map[string]map[pbt.TabletType][]*pbt.EndPoint)
hc.targetToEPs[target.Keyspace] = shardMap
}
ttMap, ok := shardMap[target.Shard]
if !ok {
ttMap = make(map[pbt.TabletType][]*pbt.EndPoint)
shardMap[target.Shard] = ttMap
}
epList, ok := ttMap[target.TabletType]
if !ok {
epList = make([]*pbt.EndPoint, 0, 1)
}
for _, ep := range epList {
if topo.EndPointEquality(ep, endPoint) {
log.Warningf("endpoint is already added: %+v", endPoint)
return
}
}
ttMap[target.TabletType] = append(epList, endPoint)
hcConnCounters.Add([]string{target.Keyspace, target.Shard, strings.ToLower(target.TabletType.String())}, 1)
}