com.azure.cosmos.implementation.query.ParallelDocumentQueryExecutionContext Maven / Gradle / Ivy

Go to download
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
package com.azure.cosmos.implementation.query;

import com.azure.cosmos.BridgeInternal;
import com.azure.cosmos.CosmosClientException;
import com.azure.cosmos.FeedOptions;
import com.azure.cosmos.FeedResponse;
import com.azure.cosmos.Resource;
import com.azure.cosmos.SqlQuerySpec;
import com.azure.cosmos.implementation.Configs;
import com.azure.cosmos.implementation.HttpConstants;
import com.azure.cosmos.implementation.IDocumentClientRetryPolicy;
import com.azure.cosmos.implementation.PartitionKeyRange;
import com.azure.cosmos.implementation.RequestChargeTracker;
import com.azure.cosmos.implementation.ResourceType;
import com.azure.cosmos.implementation.RxDocumentServiceRequest;
import com.azure.cosmos.implementation.Utils;
import com.azure.cosmos.implementation.Utils.ValueHolder;
import org.apache.commons.lang3.tuple.ImmutablePair;
import reactor.core.publisher.Flux;
import reactor.util.concurrent.Queues;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.Callable;
import java.util.function.Function;
import java.util.stream.Collectors;
/**
 * While this class is public, but it is not part of our published public APIs.
 * This is meant to be internally used only by our sdk.
 */
public class ParallelDocumentQueryExecutionContext
        extends ParallelDocumentQueryExecutionContextBase {
    private FeedOptions feedOptions;
    
    private ParallelDocumentQueryExecutionContext(
            IDocumentQueryClient client,
            List partitionKeyRanges,
            ResourceType resourceTypeEnum,
            Class resourceType,
            SqlQuerySpec query,
            FeedOptions feedOptions,
            String resourceLink,
            String rewrittenQuery,
            String collectionRid,
            boolean isContinuationExpected,
            boolean getLazyFeedResponse,
            UUID correlatedActivityId) {
        super(client, partitionKeyRanges, resourceTypeEnum, resourceType, query, feedOptions, resourceLink,
                rewrittenQuery, isContinuationExpected, getLazyFeedResponse, correlatedActivityId);
        this.feedOptions = feedOptions;
    }

    public static  Flux> createAsync(
            IDocumentQueryClient client,
            ResourceType resourceTypeEnum,
            Class resourceType,
            SqlQuerySpec query,
            FeedOptions feedOptions,
            String resourceLink,
            String collectionRid,
            PartitionedQueryExecutionInfo partitionedQueryExecutionInfo,
            List targetRanges,
            int initialPageSize,
            boolean isContinuationExpected,
            boolean getLazyFeedResponse,
            UUID correlatedActivityId) {

        ParallelDocumentQueryExecutionContext context = new ParallelDocumentQueryExecutionContext(client,
                targetRanges,
                resourceTypeEnum,
                resourceType,
                query,
                feedOptions,
                resourceLink,
                partitionedQueryExecutionInfo.getQueryInfo().getRewrittenQuery(),
                collectionRid,
                isContinuationExpected,
                getLazyFeedResponse,
                correlatedActivityId);

        try {
            context.initialize(collectionRid,
                    targetRanges,
                    initialPageSize,
                    feedOptions.requestContinuation());
            return Flux.just(context);
        } catch (CosmosClientException dce) {
            return Flux.error(dce);
        }
    }

    private void initialize(
            String collectionRid,
            List targetRanges,
            int initialPageSize,
            String continuationToken) throws CosmosClientException {
        // Generate the corresponding continuation token map.
        Map partitionKeyRangeToContinuationTokenMap = new HashMap();
        if (continuationToken == null) {
            // If the user does not give a continuation token,
            // then just start the query from the first partition.
            for (PartitionKeyRange targetRange : targetRanges) {
                partitionKeyRangeToContinuationTokenMap.put(targetRange,
                        null);
            }
        } else {
            // Figure out which partitions to resume from:

            // If a continuation token is given then we need to figure out partition key
            // range it maps to
            // in order to filter the partition key ranges.
            // For example if suppliedCompositeContinuationToken.RANGE.Min ==
            // partition3.RANGE.Min,
            // then we know that partitions 0, 1, 2 are fully drained.

            // Check to see if composite continuation token is a valid JSON.
            ValueHolder outCompositeContinuationToken = new ValueHolder();
            if (!CompositeContinuationToken.tryParse(continuationToken,
                    outCompositeContinuationToken)) {
                String message = String.format("INVALID JSON in continuation token %s for Parallel~Context",
                        continuationToken);
                throw BridgeInternal.createCosmosClientException(HttpConstants.StatusCodes.BADREQUEST,
                        message);
            }

            CompositeContinuationToken compositeContinuationToken = outCompositeContinuationToken.v;

            // Get the right hand side of the query ranges:
            List filteredPartitionKeyRanges = this.getPartitionKeyRangesForContinuation(
                    compositeContinuationToken,
                    targetRanges);

            // The first partition is the one we left off on and have a backend continuation
            // token for.
            partitionKeyRangeToContinuationTokenMap.put(filteredPartitionKeyRanges.get(0),
                    compositeContinuationToken.getToken());

            // The remaining partitions we have yet to touch / have null continuation tokens
            for (int i = 1; i < filteredPartitionKeyRanges.size(); i++) {
                partitionKeyRangeToContinuationTokenMap.put(filteredPartitionKeyRanges.get(i),
                        null);
            }
        }

        super.initialize(collectionRid,
                partitionKeyRangeToContinuationTokenMap,
                initialPageSize,
                this.querySpec);
    }

    private List getPartitionKeyRangesForContinuation(
            CompositeContinuationToken compositeContinuationToken,
            List partitionKeyRanges) throws CosmosClientException {
        // Find the partition key range we left off on
        int startIndex = this.FindTargetRangeAndExtractContinuationTokens(partitionKeyRanges,
                compositeContinuationToken.getRange());

        List rightHandSideRanges = new ArrayList();
        for (int i = startIndex; i < partitionKeyRanges.size(); i++) {
            rightHandSideRanges.add(partitionKeyRanges.get(i));
        }

        return rightHandSideRanges;
    }

    private static class EmptyPagesFilterTransformer
            implements Function.DocumentProducerFeedResponse>, Flux>> {
        private final RequestChargeTracker tracker;
        private DocumentProducer.DocumentProducerFeedResponse previousPage;
        private final FeedOptions feedOptions;
        
        public EmptyPagesFilterTransformer(RequestChargeTracker tracker, FeedOptions options) {

            if (tracker == null) {
                throw new IllegalArgumentException("Request Charge Tracker must not be null.");
            }

            this.tracker = tracker;
            this.previousPage = null;
            this.feedOptions = options;
        }

        private DocumentProducer.DocumentProducerFeedResponse plusCharge(
                DocumentProducer.DocumentProducerFeedResponse documentProducerFeedResponse,
                double charge) {
            FeedResponse page = documentProducerFeedResponse.pageResult;
            Map headers = new HashMap<>(page.getResponseHeaders());
            double pageCharge = page.getRequestCharge();
            pageCharge += charge;
            headers.put(HttpConstants.HttpHeaders.REQUEST_CHARGE,
                    String.valueOf(pageCharge));
            FeedResponse newPage = BridgeInternal.createFeedResponseWithQueryMetrics(page.getResults(),
                    headers,
                BridgeInternal.queryMetricsFromFeedResponse(page));
            documentProducerFeedResponse.pageResult = newPage;
            return documentProducerFeedResponse;
        }

        private DocumentProducer.DocumentProducerFeedResponse addCompositeContinuationToken(
                DocumentProducer.DocumentProducerFeedResponse documentProducerFeedResponse,
                String compositeContinuationToken) {
            FeedResponse page = documentProducerFeedResponse.pageResult;
            Map headers = new HashMap<>(page.getResponseHeaders());
            headers.put(HttpConstants.HttpHeaders.CONTINUATION,
                    compositeContinuationToken);
            FeedResponse newPage = BridgeInternal.createFeedResponseWithQueryMetrics(page.getResults(),
                    headers,
                BridgeInternal.queryMetricsFromFeedResponse(page));
            documentProducerFeedResponse.pageResult = newPage;
            return documentProducerFeedResponse;
        }

        private static Map headerResponse(
                double requestCharge) {
            return Utils.immutableMapOf(HttpConstants.HttpHeaders.REQUEST_CHARGE,
                    String.valueOf(requestCharge));
        }

        @Override
        public Flux> apply(Flux.DocumentProducerFeedResponse> source) {
            // Emit an empty page so the downstream observables know when there are no more
            // results.
            return source.filter(documentProducerFeedResponse -> {
                if (documentProducerFeedResponse.pageResult.getResults().isEmpty()
                        && !this.feedOptions.getAllowEmptyPages()) {
                    // filter empty pages and accumulate charge
                    tracker.addCharge(documentProducerFeedResponse.pageResult.getRequestCharge());
                    return false;
                }
                return true;
            }).map(documentProducerFeedResponse -> {
                // Add the request charge
                double charge = tracker.getAndResetCharge();
                if (charge > 0) {
                    return new ValueHolder<>(plusCharge(documentProducerFeedResponse,
                            charge));
                } else {
                    return new ValueHolder<>(documentProducerFeedResponse);
                }
            }).concatWith(Flux.just(new ValueHolder<>(null))).map(heldValue -> {
                DocumentProducer.DocumentProducerFeedResponse documentProducerFeedResponse = heldValue.v;
                // CREATE pairs from the stream to allow the observables downstream to "peek"
                // 1, 2, 3, null -> (null, 1), (1, 2), (2, 3), (3, null)
                    ImmutablePair.DocumentProducerFeedResponse, DocumentProducer.DocumentProducerFeedResponse> previousCurrent = new ImmutablePair<>(
                        this.previousPage,
                        documentProducerFeedResponse);
                    this.previousPage = documentProducerFeedResponse;
                    return previousCurrent;
            }).skip(1).map(currentNext -> {
                // remove the (null, 1)
                // Add the continuation token based on the current and next page.
                DocumentProducer.DocumentProducerFeedResponse current = currentNext.left;
                DocumentProducer.DocumentProducerFeedResponse next = currentNext.right;

                String compositeContinuationToken;
                String backendContinuationToken = current.pageResult.getContinuationToken();
                if (backendContinuationToken == null) {
                    // We just finished reading the last document from a partition
                    if (next == null) {
                        // It was the last partition and we are done
                        compositeContinuationToken = null;
                    } else {
                        // It wasn't the last partition, so we need to give the next range, but with a
                        // null continuation
                        CompositeContinuationToken compositeContinuationTokenDom = new CompositeContinuationToken(null,
                                next.sourcePartitionKeyRange.toRange());
                        compositeContinuationToken = compositeContinuationTokenDom.toJson();
                    }
                } else {
                    // We are in the middle of reading a partition,
                    // so give back this partition with a backend continuation token
                    CompositeContinuationToken compositeContinuationTokenDom = new CompositeContinuationToken(
                            backendContinuationToken,
                            current.sourcePartitionKeyRange.toRange());
                    compositeContinuationToken = compositeContinuationTokenDom.toJson();
                }

                DocumentProducer.DocumentProducerFeedResponse page;
                page = current;
                page = this.addCompositeContinuationToken(page,
                        compositeContinuationToken);

                return page;
            }).map(documentProducerFeedResponse -> {
                // Unwrap the documentProducerFeedResponse and get back the feedResponse
                return documentProducerFeedResponse.pageResult;
            }).switchIfEmpty(Flux.defer(() -> {
                // create an empty page if there is no result
                return Flux.just(BridgeInternal.createFeedResponse(Utils.immutableListOf(),
                        headerResponse(tracker.getAndResetCharge())));
            }));
        }
    }

    @Override
    public Flux> drainAsync(
            int maxPageSize) {
        List.DocumentProducerFeedResponse>> obs = this.documentProducers
                // Get the stream.
                .stream()
                // Start from the left most partition first.
                .sorted(Comparator.comparing(dp -> dp.targetRange.getMinInclusive()))
                // For each partition get it's stream of results.
                .map(DocumentProducer::produceAsync)
                // Merge results from all partitions.
                .collect(Collectors.toList());

        int fluxConcurrency = fluxSequentialMergeConcurrency(feedOptions, obs.size());
        int fluxPrefetch = fluxSequentialMergePrefetch(feedOptions, obs.size(), maxPageSize, fluxConcurrency);

        logger.debug("ParallelQuery: flux mergeSequential" +
                         " concurrency {}, prefetch {}", fluxConcurrency, fluxPrefetch);
        return Flux.mergeSequential(obs, fluxConcurrency, fluxPrefetch)
            .compose(new EmptyPagesFilterTransformer<>(new RequestChargeTracker(), this.feedOptions));
    }

    @Override
    public Flux> executeAsync() {
        return this.drainAsync(feedOptions.maxItemCount());
    }

    protected DocumentProducer createDocumentProducer(
            String collectionRid,
            PartitionKeyRange targetRange,
            String initialContinuationToken,
            int initialPageSize,
            FeedOptions feedOptions,
            SqlQuerySpec querySpecForInit,
            Map commonRequestHeaders,
            TriFunction createRequestFunc,
            Function>> executeFunc,
            Callable createRetryPolicyFunc) {
        return new DocumentProducer(client,
                collectionRid,
                feedOptions,
                createRequestFunc,
                executeFunc,
                targetRange,
                collectionRid,
                () -> client.getResetSessionTokenRetryPolicy().getRequestPolicy(),
                resourceType,
                correlatedActivityId,
                initialPageSize,
                initialContinuationToken,
                top);
    }

    private int fluxSequentialMergeConcurrency(FeedOptions options, int numberOfPartitions) {
        int parallelism = options.getMaxDegreeOfParallelism();
        if (parallelism < 0) {
            parallelism = Configs.getCPUCnt();
        } else if (parallelism == 0) {
            parallelism = 1;
        }

        return Math.min(numberOfPartitions, parallelism);
    }

    private int fluxSequentialMergePrefetch(FeedOptions options, int numberOfPartitions, int pageSize, int fluxConcurrency) {
        int maxBufferedItemCount = options.getMaxBufferedItemCount();

        if (maxBufferedItemCount <= 0) {
            maxBufferedItemCount = Math.min(Configs.getCPUCnt() * numberOfPartitions * pageSize, 100_000);
        }

        int fluxPrefetch = Math.max(maxBufferedItemCount / (Math.max(fluxConcurrency * pageSize, 1)), 1);
        return Math.min(fluxPrefetch, Queues.XS_BUFFER_SIZE);
    }
}