com.microsoft.azure.documentdb.internal.routing.PartitionKeyRangeCache Maven / Gradle / Ivy

Go to download
package com.microsoft.azure.documentdb.internal.routing;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;

import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.tuple.ImmutablePair;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.microsoft.azure.documentdb.ChangeFeedOptions;
import com.microsoft.azure.documentdb.DocumentQueryClientInternal;
import com.microsoft.azure.documentdb.FeedResponse;
import com.microsoft.azure.documentdb.PartitionKeyRange;
import com.microsoft.azure.documentdb.internal.AsyncCache;
import com.microsoft.azure.documentdb.internal.PathInfo;
import com.microsoft.azure.documentdb.internal.PathsHelper;
import com.microsoft.azure.documentdb.internal.ResourceType;

/**
 * Used internally to cache collections' partition key ranges in the Azure Cosmos DB database service.
 */
public class PartitionKeyRangeCache implements RoutingMapProvider {
    private static final Logger logger = LoggerFactory.getLogger(PartitionKeyRangeCache.class);

    private final DocumentQueryClientInternal client;
    private final AsyncCache routingMapCache;

    public PartitionKeyRangeCache(DocumentQueryClientInternal client) {
        this.client = client;
        this.routingMapCache = new AsyncCache<>(client.getExecutorService());
    }

    @Override
    public Collection getOverlappingRanges(String collectionSelfLink,
            Range range,
            boolean forceRefresh) {
        if (StringUtils.isEmpty(collectionSelfLink)) {
            throw new IllegalArgumentException("collectionSelfLink cannot be null");
        }
        if (range == null) {
            throw new IllegalArgumentException("range cannot be null");
        }

        CollectionRoutingMap routingMap;
        routingMap = tryLookUp(collectionSelfLink, null);
        if (forceRefresh && routingMap != null) {
            logger.debug("Request triggers force refresh on the PartitionKeyRangeCache. collectionLink {}, " +
                    "range {}", collectionSelfLink, range.toJson());
            routingMap = tryLookUp(collectionSelfLink, routingMap);
        }
        if (routingMap == null) {
            return Collections.emptyList();
        }
        return routingMap.getOverlappingRanges(range);
    }

    private CollectionRoutingMap tryLookUp(final String collectionLink, final CollectionRoutingMap previousValue) {
        try {
            logger.debug("Try looking up routing map for collection {}, obsoleteValue: {} stack {}",
                    collectionLink, previousValue);
            Future future = this.routingMapCache
                    .get(collectionLink, previousValue, new Callable() {
                        @Override
                        public CollectionRoutingMap call() {
                            return PartitionKeyRangeCache.this.getRoutingMapForCollection(collectionLink, previousValue);
                        }
                    });
            if (future != null) {
                return future.get();
            } else {
                logger.warn("There is no routing map lookup task for collection {}", collectionLink);
            }
        } catch (InterruptedException e) {
            logger.warn("InterruptedException while trying to look up CollectionRoutingMap", e);
        } catch (ExecutionException e) {
            logger.warn("ExecutionException while trying to look up CollectionRoutingMap", e);
        }
        return null;
    }

    @Override
    public PartitionKeyRange getPartitionKeyRangeById(String collectionSelfLink,
                                                      String partitionKeyRangeId,
                                                      boolean forceRefresh) {
        if (StringUtils.isEmpty(collectionSelfLink)) {
            throw new IllegalArgumentException("collectionSelfLink cannot be null");
        }
        if (StringUtils.isEmpty(partitionKeyRangeId)) {
            throw new IllegalArgumentException("partitionKeyRangeId cannot be null");
        }

        CollectionRoutingMap routingMap = tryLookUp(collectionSelfLink, null);
        if (forceRefresh && routingMap != null) {
            routingMap = tryLookUp(collectionSelfLink, routingMap);
        }

        if (routingMap == null) {
            return null;
        }
        return routingMap.getRangeByPartitionKeyRangeId(partitionKeyRangeId);
    }

    @Override
    public PartitionKeyRange tryGetRangeByEffectivePartitionKey(String collectionSelfLink, String effectivePartitionKey) {
        Collection ranges = this.getOverlappingRanges(collectionSelfLink,
                Range.getPointRange(effectivePartitionKey), false);

        if (ranges == null) {
            return null;
        }

        return ranges.iterator().next();
    }

    private CollectionRoutingMap getRoutingMapForCollection(String collectionLink,
                                                            CollectionRoutingMap previousRoutingMap) {
        logger.trace("Getting routing map for collection {}", collectionLink);

        // Read a change feed of partition key ranges
        ChangeFeedOptions options = new ChangeFeedOptions();
        if (previousRoutingMap == null) {
            options.setStartFromBeginning(true);
        } else {
            options.setRequestContinuation(previousRoutingMap.getChangeFeedNextIfNoneMatch());
        }
        FeedResponse deltaRangesResponse = this.client.readPartitionKeyRangesChangeFeed(collectionLink, options);

        // Building a list of  pair for the CollectionRoutingMap
        List> ranges = discardGoneRanges(deltaRangesResponse.getQueryIterable());

        CollectionRoutingMap routingMap = null;
        if (previousRoutingMap == null) {
            PathInfo pathInfo = PathsHelper.parsePathSegments(collectionLink);
            if (pathInfo == null) {
                throw new IllegalArgumentException("CollectionLink is not valid");
            }
            routingMap = InMemoryCollectionRoutingMap.tryCreateCompleteRoutingMap(ranges, pathInfo.resourceIdOrFullName);
            if(routingMap != null) {
                routingMap.setChangeFeedNextIfNoneMatch(deltaRangesResponse.getResponseContinuation());
            }
        }
        else {
            logger.debug("Combining partition key range cache with {} ranges change", ranges.size());
            routingMap = previousRoutingMap.combine(ranges, deltaRangesResponse.getResponseContinuation());
        }

        if (routingMap == null) {
            // Range information either doesn't exist or is not complete.
            throw new IllegalStateException("Cannot create complete routing map");
        }

        return routingMap;
    }

    /**
     * From the partition key range change feed response, filter out the gone ranges (the parent ranges) and create a
     * list of  pair for tryCreateCompleteRoutingMap.
     * This method is at package visibility for testing purpose.
     *
     * @param ranges the partition key range change feed response
     * @return       the list of partition key range
     */
    static List> discardGoneRanges(Iterable ranges) {
        HashMap> rangesMap = new HashMap<>();
        for (PartitionKeyRange range : ranges) {
            // The QueryIterable may have a null element at the end, due to empty last query page
            // This null need to be skipped.
            if (range == null) {
                continue;
            }
            if (range.getParents() != null) {
                // A split commit may happen between the query page payload. In that case, the parent ranges of
                // the children ranges should be discarded so the routing map is valid.
                rangesMap.keySet().removeAll(range.getParents());
            }
            rangesMap.put(range.getId(), new ImmutablePair<>(range, true));
        }
        return new ArrayList<>(rangesMap.values());
    }
}