com.microsoft.azure.documentdb.internal.query.ParallelDocumentQueryExecutionContext Maven / Gradle / Ivy

Go to download
/*
 * The MIT License (MIT)
 * Copyright (c) 2017 Microsoft Corporation
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */
package com.microsoft.azure.documentdb.internal.query;

import static com.microsoft.azure.documentdb.internal.query.ExceptionHelper.toRuntimeException;
import static com.microsoft.azure.documentdb.internal.query.ExceptionHelper.unwrap;

import java.util.Collection;
import java.util.List;
import java.util.NoSuchElementException;

import com.microsoft.azure.documentdb.Document;
import com.microsoft.azure.documentdb.DocumentQueryClientInternal;
import com.microsoft.azure.documentdb.FeedOptions;
import com.microsoft.azure.documentdb.PartitionKeyRange;
import com.microsoft.azure.documentdb.SqlQuerySpec;
import com.microsoft.azure.documentdb.internal.query.funcs.Func1;

class ParallelDocumentQueryExecutionContext extends ParallelDocumentQueryExecutionContextBase {

    private int currentDocumentProducerIndex;

    protected static ParallelDocumentQueryExecutionContext create(DocumentQueryClientInternal client,
            String collectionSelfLink,
            SqlQuerySpec querySpec, FeedOptions options, String resourceLink,
            PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, int initialPageSize) {

        try {
            ParallelDocumentQueryExecutionContext context = new ParallelDocumentQueryExecutionContext(client, collectionSelfLink, querySpec,
                    options, resourceLink, partitionedQueryExecutionInfo);

            // TODO: we can refactor this and move querying partition key ranges out of orderby/parallel execution context
            // similar to .NET
            Collection ranges = context
                    .getTargetPartitionKeyRanges(partitionedQueryExecutionInfo.getQueryRanges());
            context.initializationFuture = context.initializeAsync(partitionedQueryExecutionInfo, initialPageSize,
                    Document.class, ranges, options);
            return context;
        } catch (Exception e) {
            throw toRuntimeException(unwrap(e));
        }
    }

    private ParallelDocumentQueryExecutionContext(DocumentQueryClientInternal client,
            String collectionSelfLink,
            SqlQuerySpec querySpec,
            FeedOptions options, String resourceLink, PartitionedQueryExecutionInfo partitionedQueryExecutionInfo) {
        super(client, collectionSelfLink, querySpec, options, resourceLink, partitionedQueryExecutionInfo, Document.class);
    }

    private boolean tryMoveNextProducer(final DocumentProducer producer) throws Exception {
        final ParallelDocumentQueryExecutionContext that = this;
        return tryMoveNextProducer(producer, new Func1, DocumentProducer>() {

            @Override
            public DocumentProducer apply(DocumentProducer currentProducer) {
                return that.repairParallelContext(currentProducer);
            }
        });
    }

    private DocumentProducer repairParallelContext(DocumentProducer producer) {
        List replacementRanges = getReplacementRanges(producer.getTargetRange(), this.collectionSelfLink);

        super.repairContext(this.collectionSelfLink, this.currentDocumentProducerIndex, super.defaultComparator,
                replacementRanges, this.querySpec);

        if (this.shouldPrefetch) {
            for (int rangeIndex = 0; rangeIndex < replacementRanges.size(); rangeIndex++) {
                // capture these futures and wait on them
                documentProducers.get(rangeIndex + this.currentDocumentProducerIndex).tryScheduleFetch();
            }
        }
        return documentProducers.get(this.currentDocumentProducerIndex);
    }

    /*
     * (non-Javadoc)
     *
     * @see com.microsoft.azure.documentdb.internal.query.
     * ParallelDocumentQueryExecutionContextBase#nextInternal()
     */
    @Override
    public Document nextInternal() throws Exception {

        Document result = null;
        while (!this.isDone()) {
            DocumentProducer currentDocumentProducer = this.documentProducers
                    .get(this.currentDocumentProducerIndex);

            if (result != null) {
                break;
            }

            if (currentDocumentProducer.peek() != null) {
                result = currentDocumentProducer.peek();
            }

            if (!this.tryMoveNextProducer(currentDocumentProducer)) {
                ++this.currentDocumentProducerIndex;
                continue;
            }
        }

        if (this.isDone()) {
            super.onFinish();
        }
        // we don't know if there are more elements or not till we try.
        // For now PeekingParallelDocumentQueryExecutionContext buffers the result and fixes the issue
        // TODO: we should refactor the java code and replace .next() with drain(.)
        // similar to .Net
        if (this.isDone() && result == null) {
            throw new NoSuchElementException();
        }

        return result;
    }

    @Override
    protected boolean hasNextInternal() {
        return !isDone();
    }

    protected boolean isDone() {
        return this.currentDocumentProducerIndex >= super.documentProducers.size();
    }
}