com.microsoft.azure.documentdb.internal.query.DocumentProducer 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 java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Callable;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.Semaphore;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

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

import com.microsoft.azure.documentdb.Document;
import com.microsoft.azure.documentdb.PartitionKeyRange;
import com.microsoft.azure.documentdb.internal.DocumentServiceRequest;
import com.microsoft.azure.documentdb.internal.DocumentServiceResponse;
import com.microsoft.azure.documentdb.internal.HttpConstants;
import com.microsoft.azure.documentdb.internal.query.funcs.Callback3;
import com.microsoft.azure.documentdb.internal.query.funcs.Func1;
import com.microsoft.azure.documentdb.internal.query.funcs.Func2;

public class DocumentProducer {
    private static final Logger LOGGER = LoggerFactory.getLogger(DocumentProducer.class);
    private static final double ITEM_BUFFER_THRESHOLD = 0.1;

    private final AtomicInteger fetchInvocationCount = new AtomicInteger(0);
    private final AtomicInteger moveNextInvocationCount = new AtomicInteger(0);

    private final Func1 executeRequestFunc;
    @SuppressWarnings("rawtypes")
    private final BlockingQueue fetchResultBuffer;
    private final Func2 createRequestFunc;
    private final PartitionKeyRange targetRange;
    private final Class deserializationClass;
    private final Semaphore fetchStateSemaphore = new Semaphore(1);
    protected final AtomicBoolean isFetching = new AtomicBoolean(false);

    private final AtomicInteger bufferedDocumentsCount = new AtomicInteger(0);
    private final FetchScheduler fetchScheduler;
    private final Callback3, Integer, Double> produceCompleteCallback;

    private Iterator currentIterator;
    private Document currentDocument;
    private boolean hasStarted;
    private int previousResponseItemCount;
    private Map previousResponseHeaders;
    public volatile boolean isDone;
    private String responseContinuation;
    private String previousResponseContinuation;
    private int pageSize; // TODO populate page size
    private int itemsTillNextContinuationBoundary;
    private volatile String currentBackendContinuationToken;
    private boolean isAtContinuationBoundary;

    @SuppressWarnings("rawtypes")
    public DocumentProducer(Func1 executeRequestFunc,
            Func2 createRequestFunc, PartitionKeyRange targetRange,
            Class deserializationClass, FetchScheduler fetchScheduler, int initialPageSize,
            String initialContinuationToken,
            Callback3, Integer, Double> produceCompleteCallback) {
        this.fetchResultBuffer = new LinkedBlockingQueue();
        this.executeRequestFunc = executeRequestFunc;
        this.createRequestFunc = createRequestFunc;
        this.targetRange = targetRange;
        this.deserializationClass = deserializationClass;
        this.currentDocument = null;
        this.hasStarted = false;
        this.pageSize = initialPageSize;
        this.currentBackendContinuationToken = initialContinuationToken;
        this.fetchScheduler = fetchScheduler;
        this.produceCompleteCallback = produceCompleteCallback;
    }

    boolean shouldFetchInternal() {
        return (this.itemsTillNextContinuationBoundary - 1) < this.normalizedPageSize() * ITEM_BUFFER_THRESHOLD
                && this.fetchResultBuffer.size() <= 0;
    }

    public int normalizedPageSize() {
        return this.pageSize == -1 ? 1000 : this.pageSize;
    }

    public boolean isAtContinuationBoundary() {
        return isAtContinuationBoundary;
    }

    public int getItemsTillNextContinuationBoundary() {
        return itemsTillNextContinuationBoundary;
    }

    private boolean shouldFetch() throws InterruptedException {
        if (this.fetchedAll()) {
            return false;
        }

        if (this.shouldFetchInternal()) {
            this.fetchStateSemaphore.acquire();
            try {
                return this.shouldFetchInternal() && !isFetching.get();
            } finally {
                this.fetchStateSemaphore.release();
            }
        }
        return false;
    }

    private void updateRequestContinuationToken(String continuationToken) {
        this.currentBackendContinuationToken = continuationToken;
        this.hasStarted = true;
    }

    private void completeFetch(List docs, Map headerResponse)
            throws InterruptedException {
        this.fetchStateSemaphore.acquire();
        try {
            if (docs.size() > 0) {
                this.fetchResultBuffer.add(new FetchResult(docs, headerResponse));
                bufferedDocumentsCount.addAndGet(docs.size());
            }

            if (this.fetchedAll()) {
                this.fetchResultBuffer.add(FetchResult.DoneResult);
            }

            isFetching.set(false);
        } finally {
            this.fetchStateSemaphore.release();
        }
    }

    private double parseRequestCharge(DocumentServiceResponse response) {
        String requestChargeHeader = response.getResponseHeaders().get(HttpConstants.HttpHeaders.REQUEST_CHARGE);
        if (!StringUtils.isEmpty(requestChargeHeader)) {
            return Double.valueOf(requestChargeHeader.trim());
        } else {
            return 0;
        }
    }

    private void scheduleFetch() {
        LOGGER.trace("fetchAsync invoked");
        final DocumentProducer that = this;

        Callable callable = new Callable() {

            @SuppressWarnings("rawtypes")
            @Override
            public Void call() throws Exception {
                LOGGER.trace("fetchAsync callable is getting executed");

                fetchInvocationCount.incrementAndGet();

                FetchResult exceptionFetchResult = null;
                try {
                    List items = null;
                    DocumentServiceResponse response = null;
                    double requestCharge = 0;
                    do {

                        LOGGER.trace("Sending a request with continuation token {}", that.currentBackendContinuationToken);
                        DocumentServiceRequest request = that.createRequestFunc
                                .apply(that.currentBackendContinuationToken, that.pageSize);
                        response = executeRequestFunc.apply(request);


                        requestCharge += parseRequestCharge(response);

                        that.updateRequestContinuationToken(
                                response.getResponseHeaders().get(HttpConstants.HttpHeaders.CONTINUATION));

                        items = response.getQueryResponse(that.deserializationClass);
                        that.previousResponseItemCount = items.size();
                        LOGGER.trace("Producer with range Id {} fetched {} items", that.targetRange.getId(),
                                items.size());

                    } while (!that.fetchedAll() && items.size() <= 0);

                    that.completeFetch(items, response.getResponseHeaders());
                    that.produceCompleteCallback.run(that, items.size(), requestCharge);
                } catch (Exception ex) {
                    LOGGER.debug("DocumentProducer Id: {}, Exception in FetchAsync: {}", that.targetRange.getId(),
                            ex.getMessage());

                    exceptionFetchResult = new FetchResult(ex);
                }

                if (exceptionFetchResult != null) {
                    that.updateRequestContinuationToken(that.currentBackendContinuationToken);
                    that.fetchResultBuffer.add(exceptionFetchResult);
                }

                return null;
            }
        };
        this.fetchScheduler.schedule(callable);
    }

    public boolean tryScheduleFetch() {
        if (this.fetchedAll()) {
            return false;
        }

        if (!isFetching.compareAndSet(false, true)) {
            return false;
        }
        scheduleFetch();
        return true;
    }

    /**
     * Advances to the next element in the sequence, returning the result
     * asynchronously.
     * @return true if the DocumentProducer was successfully advanced
     * to the next element; false if the DocumentProducer has passed the end of
     * the sequence.
     * @throws Exception
     */
    public boolean moveNext() throws Exception {
        moveNextInvocationCount.incrementAndGet();

        if (isDone) {
            return false;
        }

        if (shouldFetch()) {
            // this one does IO and fetched results async
            tryScheduleFetch();
        }

        if (moveNextInternal()) {
            isAtContinuationBoundary = false;
            --itemsTillNextContinuationBoundary;
            return true;
        }

        @SuppressWarnings("rawtypes")
        FetchResult fetchResult = fetchResultBuffer.take();
        switch (fetchResult.type) {
        case Done:
            isDone = true;
            itemsTillNextContinuationBoundary = 0;

            return false;
        case Exception:
            throw fetchResult.exception;
        case Result:
            updateStates(fetchResult.results, fetchResult.headerResponse);
            return true;
        default:
            throw new IllegalStateException(fetchResult.type.name());
        }
    }

    private boolean moveNextInternal() {
        if (this.currentIterator == null || !this.currentIterator.hasNext()) {
            return false;
        }

        this.currentDocument = this.currentIterator.next();
        this.bufferedDocumentsCount.decrementAndGet();
        return true;
    }

    public boolean hasStarted() {
        return this.hasStarted;
    }

    public String getId() {
        return this.targetRange.getId();
    }

    @Override
    public boolean equals(Object obj) {
        if (!(obj instanceof DocumentProducer))
            return false;

        return this.getId().compareTo(((DocumentProducer) obj).getId()) == 0;
    }

    @Override
    public int hashCode() {
        return this.getId().hashCode();
    }

    public boolean fetchedAll() {
        return this.hasStarted && StringUtils.isEmpty(this.currentBackendContinuationToken);
    }

    public Document peek() {

        if (this.isDone) {
            throw new IllegalStateException("Producer is closed");
        }

        return this.currentDocument;
    }

    public int getBufferedDocumentsCount() {
        return this.bufferedDocumentsCount.get();
    }

    public int getPreviousResponseItemCount() {
        return this.previousResponseItemCount;
    }

    public Map getPreviousResponseHeaders() {
        return this.previousResponseHeaders;
    }

    public PartitionKeyRange getTargetRange() {
        return this.targetRange;
    }

    public String getCurrentBackendContinuationToken() {
        return currentBackendContinuationToken;
    }

    public int getPageSize() {
        return pageSize;
    }

    void notifyStop() {
        LOGGER.trace("notifyStop");
        this.fetchResultBuffer.add(FetchResult.DoneResult);
        this.currentBackendContinuationToken = null;
    }

    private void updateStates(List docs, Map headerResponse) {
        this.previousResponseContinuation = this.responseContinuation;
        this.responseContinuation = headerResponse.get(HttpConstants.HttpHeaders.CONTINUATION);
        this.itemsTillNextContinuationBoundary = docs.size();
        this.isAtContinuationBoundary = true;
        this.currentIterator = docs.iterator();

        LOGGER.trace("id {} Fetched Count: {}", this.getTargetRange().getId(), docs.size());
        this.moveNextInternal();
    }

    @SuppressWarnings("unchecked")
    static class FetchResult {

        public FetchResultType type;
        public List results;
        public Exception exception;

        public Map headerResponse;

        @SuppressWarnings("rawtypes")
        public final static FetchResult DoneResult = new FetchResult();

        static {
            DoneResult.type = FetchResultType.Done;
        };

        public FetchResult(List items, Map responseHeader) {
            this.results = items;
            this.headerResponse = responseHeader;
            this.type = FetchResultType.Result;
        }

        public FetchResult(Exception exception) {
            this.exception = exception;
            this.type = FetchResultType.Exception;
        }

        private FetchResult() {}
    }

    private enum FetchResultType {
        Done, Exception, Result,
    }
}