// withRetry gets available connections and executes the action. If there are retryable errors,
// it retries retryCount times before failing. It does not retry if the connection is in
// the middle of a transaction. While returning the error check if it maybe a result of
// a resharding event, and set the re-resolve bit and let the upper layers
// re-resolve and retry.
func (dg *discoveryGateway) withRetry(ctx context.Context, keyspace, shard string, tabletType pbt.TabletType, action func(conn tabletconn.TabletConn) error, transactionID int64, isStreaming bool) error {
var endPointLastUsed *pbt.EndPoint
var err error
inTransaction := (transactionID != 0)
invalidEndPoints := make(map[string]bool)
for i := 0; i < dg.retryCount+1; i++ {
var endPoint *pbt.EndPoint
endPoints := dg.getEndPoints(keyspace, shard, tabletType)
if len(endPoints) == 0 {
// fail fast if there is no endpoint
err = vterrors.FromError(vtrpc.ErrorCode_INTERNAL_ERROR, fmt.Errorf("no valid endpoint"))
break
}
shuffleEndPoints(endPoints)
// skip endpoints we tried before
for _, ep := range endPoints {
if _, ok := invalidEndPoints[discovery.EndPointToMapKey(ep)]; !ok {
endPoint = ep
break
}
}
if endPoint == nil {
if err == nil {
// do not override error from last attempt.
err = vterrors.FromError(vtrpc.ErrorCode_INTERNAL_ERROR, fmt.Errorf("no available connection"))
}
break
}
// execute
endPointLastUsed = endPoint
conn := dg.hc.GetConnection(endPoint)
if conn == nil {
err = vterrors.FromError(vtrpc.ErrorCode_INTERNAL_ERROR, fmt.Errorf("no connection for %+v", endPoint))
invalidEndPoints[discovery.EndPointToMapKey(endPoint)] = true
continue
}
err = action(conn)
if dg.canRetry(ctx, err, transactionID, isStreaming) {
invalidEndPoints[discovery.EndPointToMapKey(endPoint)] = true
continue
}
break
}
return WrapError(err, keyspace, shard, tabletType, endPointLastUsed, inTransaction)
}
// GetConnection returns the TabletConn of the given endpoint.
func (fhc *fakeHealthCheck) GetConnection(endPoint *pbt.EndPoint) tabletconn.TabletConn {
key := discovery.EndPointToMapKey(endPoint)
if item := fhc.items[key]; item != nil {
return item.conn
}
return nil
}
// AddEndPoint adds the endpoint, and starts health check.
func (fhc *fakeHealthCheck) AddEndPoint(cell string, endPoint *pbt.EndPoint) {
key := discovery.EndPointToMapKey(endPoint)
item := &fhcItem{
eps: &discovery.EndPointStats{
EndPoint: endPoint,
Cell: cell,
},
}
fhc.items[key] = item
}
// AddEndPoint adds the endpoint, and starts health check.
func (fhc *fakeHealthCheck) AddEndPoint(cell, name string, endPoint *topodatapb.EndPoint) {
key := discovery.EndPointToMapKey(endPoint)
item := &fhcItem{
eps: &discovery.EndPointStats{
EndPoint: endPoint,
Cell: cell,
Name: name,
},
}
fhc.items[key] = item
}
func (fhc *fakeHealthCheck) addTestEndPoint(cell, host string, port int32, keyspace, shard string, tabletType pbt.TabletType, reparentTS int64, err error, conn tabletconn.TabletConn) *pbt.EndPoint {
ep := topo.NewEndPoint(0, host)
ep.PortMap["vt"] = port
key := discovery.EndPointToMapKey(ep)
item := fhc.items[key]
if item == nil {
fhc.AddEndPoint(cell, ep)
item = fhc.items[key]
}
item.eps.Target = &pbq.Target{Keyspace: keyspace, Shard: shard, TabletType: tabletType}
item.eps.TabletExternallyReparentedTimestamp = reparentTS
item.eps.LastError = err
item.conn = conn
return ep
}
// RemoveEndPoint removes the endpoint, and stops the health check.
func (fhc *fakeHealthCheck) RemoveEndPoint(endPoint *pbt.EndPoint) {
key := discovery.EndPointToMapKey(endPoint)
delete(fhc.items, key)
}