// 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)
}