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

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

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;

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

import com.microsoft.azure.documentdb.PartitionKeyRange;

/**
 * Used internally to cache partition key ranges of a collection in the Azure Cosmos DB database service.
 */
public class InMemoryCollectionRoutingMap implements CollectionRoutingMap {
    private final static Logger logger = LoggerFactory.getLogger(InMemoryCollectionRoutingMap.class);

    private final Map> rangeById;
    private final Map rangeByInfo;
    private final List orderedPartitionKeyRanges;
    private final List> orderedRanges;
    private final List orderedPartitionInfo;
    private String collectionUniqueId;
    private HashSet goneRangeIds;
    private String changeFeedNextIfNoneMatch;

    private InMemoryCollectionRoutingMap(Map> rangeById,
                                         Map rangeByInfo,
                                         List orderedPartitionKeyRanges,
                                         List orderedPartitionInfo,
                                         String collectionUniqueId) {
        this.rangeById = rangeById;
        this.rangeByInfo = rangeByInfo;
        this.orderedPartitionKeyRanges = orderedPartitionKeyRanges;
        this.orderedRanges = new ArrayList>(this.orderedPartitionKeyRanges.size());
        for (PartitionKeyRange range : this.orderedPartitionKeyRanges) {
            this.orderedRanges.add(range.toRange());
        }

        this.orderedPartitionInfo = orderedPartitionInfo;
        this.collectionUniqueId = collectionUniqueId;
        this.goneRangeIds = new HashSet<>();
        for (PartitionKeyRange r : orderedPartitionKeyRanges) {
            if (r.getParents() != null) {
                this.goneRangeIds.addAll(r.getParents());
            }
        }
    }

    private static  InMemoryCollectionRoutingMap tryCreateRoutingMap(
            Iterable> ranges, String collectionUniqueId) {
        Map> rangeById = new HashMap>();
        Map rangeByInfo = new HashMap();

        List> sortedRanges = new ArrayList>();
        for (ImmutablePair pair : ranges) {
            rangeById.put(pair.left.getId(), pair);
            rangeByInfo.put(pair.right, pair.left);
            sortedRanges.add(pair);
        }

        Collections.sort(sortedRanges, new MinPartitionKeyPairComparator());

        List orderedPartitionKeyRanges = new ArrayList(sortedRanges.size());
        List orderedPartitionInfo = new ArrayList(sortedRanges.size());

        for (ImmutablePair pair : sortedRanges) {
            orderedPartitionKeyRanges.add(pair.left);
            orderedPartitionInfo.add(pair.right);
        }

        return new InMemoryCollectionRoutingMap(rangeById, rangeByInfo, orderedPartitionKeyRanges,
                orderedPartitionInfo, collectionUniqueId);
    }

    /**
     * Combine the current routing map with a set of partition key ranges into a new collection routing map
     *
     * @param ranges         the new set of partition key ranges, which should come from a change feed query
     * @return               a new collection routing map
     */
    @Override
    public CollectionRoutingMap combine(List> ranges,
                                        String changeFeedNextIfNoneMatch) {
        // Add new gone ranges from the new set of ranges
        for (ImmutablePair rangePair : ranges) {
            if (rangePair.left.getParents() != null) {
                this.goneRangeIds.addAll(rangePair.left.getParents());
            }
        }

        // Remove all gone ranges from the current rangeById range map
        for (String goneRangeId : this.goneRangeIds) {
            this.rangeById.remove(goneRangeId);
        }
        // Add the new ranges that are not in the gone range list
        for (ImmutablePair rangePair : ranges) {
            if (!this.goneRangeIds.contains(rangePair.left.getId()))  {
                this.rangeById.put(rangePair.left.getId(), rangePair);
            }
        }

        InMemoryCollectionRoutingMap routingMap = tryCreateRoutingMap(this.rangeById.values(), this.collectionUniqueId);
        logger.debug("Created complete routing map with {} key range", routingMap.getOrderedPartitionKeyRanges().size());

        if (!isCompleteSetOfRanges(this.collectionUniqueId, routingMap.getOrderedPartitionKeyRanges())) {
            return null;
        }

        routingMap.setChangeFeedNextIfNoneMatch(changeFeedNextIfNoneMatch);

        return routingMap;
    }

    public static  InMemoryCollectionRoutingMap tryCreateCompleteRoutingMap(
            Iterable> ranges,
            String collectionUniqueId) {
        if (ranges == null) {
            throw new IllegalArgumentException("The key ranges list cannot be null");
        }

        InMemoryCollectionRoutingMap routingMap = tryCreateRoutingMap(ranges, collectionUniqueId);
        logger.debug("Created complete routing map with {} key range", routingMap.getOrderedPartitionKeyRanges().size());

        if (!isCompleteSetOfRanges(collectionUniqueId, routingMap.getOrderedPartitionKeyRanges())) {
            return null;
        }

        return routingMap;
    }

    private static boolean isCompleteSetOfRanges(String collectionId, List orderedRanges) {
        boolean isComplete = false;
        if (orderedRanges.size() > 0) {
            PartitionKeyRange firstRange = orderedRanges.get(0);
            PartitionKeyRange lastRange = orderedRanges.get(orderedRanges.size() - 1);
            boolean firstRangeValid = firstRange.getMinInclusive()
                    .compareTo(PartitionKeyRange.MINIMUM_INCLUSIVE_EFFECTIVE_PARTITION_KEY) == 0;
            isComplete = firstRangeValid;
            if (!firstRangeValid) {
                logger.warn("First partition key range is invalid. collectionId {}, first range {}",
                        collectionId, firstRange.toJson());
            }
            boolean lastRangeValid = lastRange.getMaxExclusive()
                    .compareTo(PartitionKeyRange.MAXIMUM_EXCLUSIVE_EFFECTIVE_PARTITION_KEY) == 0;
            isComplete &= lastRangeValid;
            if (!lastRangeValid) {
                logger.warn("Last partition key range is invalid. collectionId {}, last range {}",
                        collectionId, lastRange.toJson());
            }

            for (int i = 1; i < orderedRanges.size(); i++) {
                PartitionKeyRange previousRange = orderedRanges.get(i - 1);
                PartitionKeyRange currentRange = orderedRanges.get(i);
                boolean gapValid = previousRange.getMaxExclusive().compareTo(currentRange.getMinInclusive()) == 0;
                isComplete &= gapValid;
                if (!gapValid) {
                    logger.warn("Partition key range gap is invalid. collectionId {}, previous range {}, current range {}",
                            collectionId, previousRange.toJson(), currentRange.toJson());
                }

                if (!isComplete) {
                    if (previousRange.getMaxExclusive().compareTo(currentRange.getMinInclusive()) > 0) {
                        throw new IllegalStateException(String.format("Ranges overlap %s and %s",
                                previousRange.toJson(), currentRange.toJson()));
                    }

                    break;
                }
            }
        }

        return isComplete;
    }

    public String getCollectionUniqueId() {
        return collectionUniqueId;
    }

    public TPartitionInfo getHeadPartition() {
        return this.orderedPartitionInfo.get(0);
    }

    public TPartitionInfo getTailPartition() {
        return this.orderedPartitionInfo.get(this.orderedPartitionInfo.size() - 1);
    }

    public List getOrderedPartitionInfo() {
        return this.orderedPartitionInfo;
    }

    @Override
    public List getOrderedPartitionKeyRanges() {
        return this.orderedPartitionKeyRanges;
    }

    @Override
    public PartitionKeyRange getRangeByEffectivePartitionKey(String effectivePartitionKeyValue) {
        if (PartitionKeyRange.MINIMUM_INCLUSIVE_EFFECTIVE_PARTITION_KEY.compareTo(effectivePartitionKeyValue) == 0) {
            return this.orderedPartitionKeyRanges.get(0);
        }

        if (PartitionKeyRange.MAXIMUM_EXCLUSIVE_EFFECTIVE_PARTITION_KEY.compareTo(effectivePartitionKeyValue) == 0) {
            return null;
        }

        int index = Collections.binarySearch(this.orderedRanges, Range.getPointRange(effectivePartitionKeyValue),
                new Range.MinComparator());

        if (index < 0) {
            index = Math.max(0, -index - 2);
        }

        return this.orderedPartitionKeyRanges.get(index);
    }

    @Override
    public PartitionKeyRange getRangeByPartitionKeyRangeId(String partitionKeyRangeId) {
        ImmutablePair pair = this.rangeById.get(partitionKeyRangeId);
        return pair == null ? null : pair.left;
    }

    public TPartitionInfo getInfoByPartitionKeyRangeId(String partitionKeyRangeId) {
        ImmutablePair pair = this.rangeById.get(partitionKeyRangeId);
        return pair == null ? null : pair.right;
    }

    public PartitionKeyRange getPartitionKeyRangeByPartitionInfo(TPartitionInfo partitionInfo) {
        return this.rangeByInfo.get(partitionInfo);
    }

    @Override
    public Collection getOverlappingRanges(Range range) {
        return this.getOverlappingRanges(Arrays.asList(range));
    }

    @Override
    public Collection getOverlappingRanges(Collection> providedPartitionKeyRanges) {
        if (providedPartitionKeyRanges == null) {
            throw new IllegalArgumentException("providedPartitionKeyRanges");
        }

        Map partitionRanges = new TreeMap();

        for (Range range : providedPartitionKeyRanges) {
            int minIndex = Collections.binarySearch(this.orderedRanges, range, new Range.MinComparator());
            if (minIndex < 0) {
                minIndex = Math.max(minIndex, -minIndex - 2);
            }

            int maxIndex = Collections.binarySearch(this.orderedRanges, range, new Range.MaxComparator());
            if (maxIndex < 0) {
                maxIndex = Math.min(this.orderedRanges.size() - 1, -maxIndex - 1);
            }

            for (int i = minIndex; i <= maxIndex; ++i) {
                if (Range.checkOverlapping(this.orderedRanges.get(i), range)) {
                    PartitionKeyRange partitionKeyRange = this.orderedPartitionKeyRanges.get(i);
                    partitionRanges.put(partitionKeyRange.getMinInclusive(), partitionKeyRange);
                }
            }
        }

        return partitionRanges.values();
    }

    @Override
    public String getChangeFeedNextIfNoneMatch() {
        return this.changeFeedNextIfNoneMatch;
    }

    @Override
    public void setChangeFeedNextIfNoneMatch(String value) {
        this.changeFeedNextIfNoneMatch = value;
    }

    private static class MinPartitionKeyPairComparator
            implements Comparator> {
        public int compare(ImmutablePair pair1,
                           ImmutablePair pair2) {
            return pair1.left.getMinInclusive().compareTo(pair2.left.getMinInclusive());
        }
    }
}