com.microsoft.azure.documentdb.internal.directconnectivity.StoreReader Maven / Gradle / Ivy

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

import java.net.URI;
import java.net.URISyntaxException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.CompletionService;
import java.util.concurrent.ExecutorCompletionService;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.concurrent.ThreadLocalRandom;

import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.math.NumberUtils;
import org.apache.http.HttpStatus;
import org.joda.time.DateTime;
import org.joda.time.DateTimeZone;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.microsoft.azure.documentdb.ConsistencyLevel;
import com.microsoft.azure.documentdb.DocumentClientException;
import com.microsoft.azure.documentdb.internal.DocumentServiceRequest;
import com.microsoft.azure.documentdb.internal.HttpConstants;
import com.microsoft.azure.documentdb.internal.OperationType;
import com.microsoft.azure.documentdb.internal.SessionContainer;
import com.microsoft.azure.documentdb.internal.SessionTokenHelper;
import com.microsoft.azure.documentdb.internal.VectorSessionToken;
import com.microsoft.azure.documentdb.ClientSideRequestStatistics;

class StoreReader {
    private final GlobalAddressResolver globalAddressResolver;
    private final TransportClient transportClient;
    private final SessionContainer sessionContainer;
    private final ExecutorService executorService;
    private String lastReadAddress;

    private final Logger logger = LoggerFactory.getLogger(StoreReader.class);
    
    public StoreReader(GlobalAddressResolver globalAddressResolver,
                       TransportClient transportClient,
                       SessionContainer sessionContainer,
                       ExecutorService executorService) {
        this.globalAddressResolver = globalAddressResolver;
        this.transportClient = transportClient;
        this.sessionContainer = sessionContainer;
        this.executorService = executorService;
    }

    private static int generateNextRandom(int maxValue) {
        return ThreadLocalRandom.current().nextInt(maxValue);
    }

    private AddressCache getAddressCache(DocumentServiceRequest request) {
        return this.globalAddressResolver.resolve(request);
    }

    StoreReadResult readEventual(DocumentServiceRequest request) throws DocumentClientException {
        request.getHeaders().remove(HttpConstants.HttpHeaders.SESSION_TOKEN);
        ArrayList resolvedEndpoints = this.getEndpointAddresses(request, true);

        if (resolvedEndpoints.size() == 0) {
            return null;
        }

        StoreReadResult result = readOneReplica(request, resolvedEndpoints, ConsistencyLevel.Eventual);
        request.getRequestChargeTracker().addCharge(result.getRequestCharge());
        return result;
    }

    StoreReadResult readSession(DocumentServiceRequest request) throws DocumentClientException {
        ArrayList resolvedEndpoints = this.getEndpointAddresses(request, true);

        if (resolvedEndpoints.size() == 0) {
            return null;
        }

        VectorSessionToken requestSessionToken = null;
        if (!request.isChangeFeedRequest()) {
            SessionTokenHelper.setPartitionLocalSessionToken(request, this.sessionContainer);
            requestSessionToken = request.getSessionToken();
        }

        StoreReadResult readExceptionResult = null;

        int endpointReadCount = 0;
        while (resolvedEndpoints.size() > 0) {
            StoreReadResult replicaReadResult = readOneReplica(request, resolvedEndpoints, ConsistencyLevel.Session);
            request.getRequestChargeTracker().addCharge(replicaReadResult.getRequestCharge());
            if (replicaReadResult.isGoneException()) {
                throw replicaReadResult.getException();
            } else if (requestSessionToken == null 
                    || (replicaReadResult.getSessionToken() != null && requestSessionToken.isValid(replicaReadResult.getSessionToken()))) {
                return replicaReadResult;
            } else if (replicaReadResult.getException() != null) {
                readExceptionResult = replicaReadResult;
            } else {
                logger.trace("StoreReadResult {} not selected: requestSessionToken = {}, replicaReadResult.getSessiontoken() = {} are incompatible",
                        endpointReadCount++, requestSessionToken.convertToString() , replicaReadResult.getSessionToken() != null ?
                                replicaReadResult.getSessionToken().convertToString() : "null");
            }
        }

        return readExceptionResult;
    }

    StoreReadResult readPrimary(DocumentServiceRequest request,
            boolean requiresValidLsn,
            ConsistencyLevel consistencyLevel) throws DocumentClientException {
        ReadReplicaResult readQuorumResult = this.readPrimaryImpl(request, requiresValidLsn, consistencyLevel);
        if (readQuorumResult.isRetryWithForceRefresh() && !request.isForceAddressRefresh()) {
            request.setForceAddressRefresh(true);
            readQuorumResult = this.readPrimaryImpl(request, requiresValidLsn, consistencyLevel);
        }

        if (readQuorumResult.getResponses().size() == 0) {
            throw new DocumentClientException(HttpStatus.SC_GONE, "The requested resource is no longer available at the server.");
        }

        return readQuorumResult.getResponses().get(0);
    }

    List readMultipleReplica(final DocumentServiceRequest request,
            boolean includePrimary, 
            int replicaCountToRead, 
            ConsistencyLevel consistencyLevel) throws DocumentClientException {
        List storeReadResults = new LinkedList();
        CompletionService completionService = new ExecutorCompletionService(this.executorService);
        ArrayList resolvedEndpoints = getEndpointAddresses(request, includePrimary);

        this.logger.debug("Reading from {} out of {} replicas", replicaCountToRead, resolvedEndpoints.size());
        
        // returning empty array if endpoints will never be able to satisfy the requested number of replicas
        if (resolvedEndpoints.size() < replicaCountToRead) {
            return storeReadResults;
        }

        // session token is not needed in case of strong consistency
        request.getHeaders().remove(HttpConstants.HttpHeaders.SESSION_TOKEN);
        
        for (int i = 0; i < replicaCountToRead; i++) {
            int uriIndex = StoreReader.generateNextRandom(resolvedEndpoints.size());
            this.logger.debug("Reading {} from replica {}", request.getOperationType(), resolvedEndpoints.get(uriIndex));
            Callable callable = this.getReadReplicaCallable(request, resolvedEndpoints.get(uriIndex), consistencyLevel);
            
            resolvedEndpoints.remove(uriIndex);
            completionService.submit(callable);
        }

        int received = 0;
        while (received < replicaCountToRead) {
            try {
                Future resultFuture = completionService.take();
                received++;
                StoreReadResult result = resultFuture.get();
                request.getRequestChargeTracker().addCharge(result.getRequestCharge());
                storeReadResults.add(result);
            } catch (Exception e) {
                throw new DocumentClientException(HttpStatus.SC_INTERNAL_SERVER_ERROR, e);
            }
        }

        return storeReadResults;
    }

    private ReadReplicaResult readPrimaryImpl(DocumentServiceRequest request,
            boolean requiresValidLsn,
            ConsistencyLevel consistencyLevel) throws DocumentClientException {
        URI primaryUri = ReplicatedResourceClient.resolvePrimaryUri(request, getAddressCache(request));
        request.getHeaders().remove(HttpConstants.HttpHeaders.SESSION_TOKEN);

        DateTime requestStartTime = DateTime.now(DateTimeZone.UTC);
        StoreReadResult storeReadResult = this.createStoreReadResult(request, primaryUri.toString(), consistencyLevel);
        recordReadResponse(storeReadResult, requestStartTime, request, primaryUri.toString());
        request.getRequestChargeTracker().addCharge(storeReadResult.getRequestCharge());
        if (storeReadResult.isGoneException()) {
            return new ReadReplicaResult(true, new ArrayList());
        }

        return new ReadReplicaResult(false, Arrays.asList(storeReadResult));
    }

    private StoreReadResult readOneReplica(DocumentServiceRequest request,
            ArrayList resolvedEndpoints,
            ConsistencyLevel consistencyLevel) throws DocumentClientException {
        if (request == null) {
            throw new IllegalArgumentException("request");
        }
        if (resolvedEndpoints == null || resolvedEndpoints.size() == 0) {
            throw new IllegalArgumentException("resolveEndpoints");
        }

        int uriIndex = StoreReader.generateNextRandom(resolvedEndpoints.size());
        String endpoint = resolvedEndpoints.get(uriIndex);
        resolvedEndpoints.remove(uriIndex);
        DateTime requestStartTime = DateTime.now(DateTimeZone.UTC);
        StoreReadResult storeReadResult = createStoreReadResult(request, endpoint, consistencyLevel);
        recordReadResponse(storeReadResult, requestStartTime, request, endpoint);
        return storeReadResult;
    }

    private ArrayList getEndpointAddresses(DocumentServiceRequest request, boolean includePrimary) throws DocumentClientException {
        AddressInformation[] addresses = getAddressCache(request).resolve(request);
        ArrayList resolvedEndpoints = new ArrayList();
        for (int i = 0; i < addresses.length; i++) {
            if (!addresses[i].isPrimary() || includePrimary) {
                resolvedEndpoints.add(addresses[i].getPhysicalUri());
            }
        }

        return resolvedEndpoints;
    }

    private Callable getReadReplicaCallable(final DocumentServiceRequest request,
            final String endpointAddress,
            final ConsistencyLevel consistencyLevel) {
        Callable callable = new Callable() {
            @Override
            public StoreReadResult call() throws DocumentClientException {
                DateTime requestStartTime = DateTime.now(DateTimeZone.UTC);
                StoreReadResult storeReadResult = createStoreReadResult(request, endpointAddress, consistencyLevel);
                recordReadResponse(storeReadResult, requestStartTime, request, endpointAddress);
                return storeReadResult;
            }
        };

        return callable;
    }

    public static StoreReadResult createStoreReadResult(StoreResponse storeResponse,
                                                        DocumentClientException storeException,
                                                        ConsistencyLevel consistencyLevel,
                                                        URI storePhysicalAddress) throws DocumentClientException {
        boolean useLocalLSNBasedHeaders = consistencyLevel != null && consistencyLevel != ConsistencyLevel.Strong
                && consistencyLevel != ConsistencyLevel.BoundedStaleness;

        long quorumAckedLSN = -1;
        int currentReplicaSetSize = -1;
        int currentWriteQuorum = -1;
        double requestCharge = 0;
        long lsn = -1;
        long globalLSN = -1;
        long numReadRegions = -1;
        long itemLSN = -1;

        if (storeException != null) {
            Map headers = storeException.getResponseHeaders();

            if (storePhysicalAddress == null) {
                try {
                    storePhysicalAddress = new URI(storeException.getResourceAddress());
                } catch (URISyntaxException e) {
                    throw new DocumentClientException(HttpStatus.SC_INTERNAL_SERVER_ERROR, e);
                }
            }

            if (headers == null) {
                return new StoreReadResult(storeResponse, storeException, lsn,
                        null, quorumAckedLSN, requestCharge, currentReplicaSetSize,
                        currentWriteQuorum, false, storePhysicalAddress, globalLSN, numReadRegions, itemLSN, null);
            }

            String quorumAckedLsnHeaderValue = headers.get(useLocalLSNBasedHeaders ? WFConstants.BackendHeaders.QuorumAckedLocalLSN : WFConstants.BackendHeaders.QuorumAckedLSN);
            if (!StringUtils.isEmpty(quorumAckedLsnHeaderValue)) {
                quorumAckedLSN = Long.parseLong(quorumAckedLsnHeaderValue);
            }

            String currentReplicaSetSizeHeaderValue = headers.get(WFConstants.BackendHeaders.CurrentReplicaSetSize);
            if (!StringUtils.isEmpty(currentReplicaSetSizeHeaderValue)) {
                currentReplicaSetSize = Integer.parseInt(currentReplicaSetSizeHeaderValue);
            }

            String currentWriteQuorumHeaderValue = headers.get(WFConstants.BackendHeaders.CurrentWriteQuorum);
            if (!StringUtils.isEmpty(currentReplicaSetSizeHeaderValue)) {
                currentWriteQuorum = Integer.parseInt(currentWriteQuorumHeaderValue);
            }

            String currentRequestChargeHeaderValue = headers.get(HttpConstants.HttpHeaders.REQUEST_CHARGE);
            if (!StringUtils.isEmpty(currentRequestChargeHeaderValue)) {
                requestCharge = Double.parseDouble(currentRequestChargeHeaderValue);
            }

            String currentRequestLSN = headers.get(useLocalLSNBasedHeaders ? WFConstants.BackendHeaders.LocalLSN : HttpConstants.HttpHeaders.LSN);
            if (!StringUtils.isEmpty(currentRequestLSN)) {
                lsn = Long.parseLong(currentRequestLSN);
            }

            globalLSN = NumberUtils.toLong(headers.get(WFConstants.BackendHeaders.GlobalCommittedLSN), globalLSN);
            itemLSN = NumberUtils.toLong(headers.get(WFConstants.BackendHeaders.ItemLSN), itemLSN);
            numReadRegions = NumberUtils.toLong(headers.get(WFConstants.BackendHeaders.NumberOfReadRegions), numReadRegions);

            VectorSessionToken sessionToken = null;
            String sessionTokenHeaderValue = headers.get(HttpConstants.HttpHeaders.SESSION_TOKEN);
            if (!StringUtils.isEmpty(sessionTokenHeaderValue)) {
                sessionToken = SessionTokenHelper.parse(sessionTokenHeaderValue);
            }

            return new StoreReadResult(storeResponse, storeException, lsn,
                    null, quorumAckedLSN, requestCharge, currentReplicaSetSize,
                    currentWriteQuorum, false, storePhysicalAddress, globalLSN, numReadRegions, itemLSN, sessionToken);
        }

        String quorumAckedLsnHeaderValue = storeResponse.getHeaderValue(useLocalLSNBasedHeaders ? WFConstants.BackendHeaders.QuorumAckedLocalLSN : WFConstants.BackendHeaders.QuorumAckedLSN);
        if (!StringUtils.isEmpty(quorumAckedLsnHeaderValue)) {
            quorumAckedLSN = Long.parseLong(quorumAckedLsnHeaderValue);
        }

        String currentReplicaSetSizeHeaderValue = storeResponse.getHeaderValue(WFConstants.BackendHeaders.CurrentReplicaSetSize);
        if (!StringUtils.isEmpty(currentReplicaSetSizeHeaderValue)) {
            currentReplicaSetSize = Integer.parseInt(currentReplicaSetSizeHeaderValue);
        }

        String currentWriteQuorumHeaderValue = storeResponse.getHeaderValue(WFConstants.BackendHeaders.CurrentWriteQuorum);
        if (!StringUtils.isEmpty(currentWriteQuorumHeaderValue)) {
            currentWriteQuorum = Integer.parseInt(currentWriteQuorumHeaderValue);
        }

        String currentRequestChargeHeaderValue = storeResponse.getHeaderValue(HttpConstants.HttpHeaders.REQUEST_CHARGE);
        if (!StringUtils.isEmpty(currentRequestChargeHeaderValue)) {
            requestCharge = Double.parseDouble(currentRequestChargeHeaderValue);
        }

        if (useLocalLSNBasedHeaders) {
            String currentRequestLSN = storeResponse.getHeaderValue(WFConstants.BackendHeaders.LocalLSN);
            if (!StringUtils.isEmpty(currentRequestLSN)) {
                lsn = Long.parseLong(currentRequestLSN);
            }
        } else {
            lsn = storeResponse.getLSN();
        }

        boolean isValid = true;

        globalLSN = NumberUtils.toLong(storeResponse.getHeaderValue(WFConstants.BackendHeaders.GlobalCommittedLSN), -1);
        numReadRegions = NumberUtils.toLong(storeResponse.getHeaderValue(WFConstants.BackendHeaders.NumberOfReadRegions), -1);
        VectorSessionToken sessionToken = null;
        String sessionTokenHeaderValue = storeResponse.getHeaderValue(HttpConstants.HttpHeaders.SESSION_TOKEN);
        if (!StringUtils.isEmpty(sessionTokenHeaderValue)) {
            sessionToken = SessionTokenHelper.parse(sessionTokenHeaderValue);
        }

        return new StoreReadResult(storeResponse, storeException, lsn,
                storeResponse.getPartitionKeyRangeId(), quorumAckedLSN, requestCharge, currentReplicaSetSize,
                currentWriteQuorum, isValid, storePhysicalAddress, globalLSN, numReadRegions, itemLSN, sessionToken);
    }

    private StoreReadResult createStoreReadResult(DocumentServiceRequest request,
            String endpointAddress,
            ConsistencyLevel consistencyLevel) throws DocumentClientException {
        StoreResponse storeResponse = null;
        DocumentClientException storeException = null;
        try {
            storeResponse = readFromStore(request, endpointAddress);
        } catch (DocumentClientException e) {
            logger.debug("Request:[operation:{} header:[{}], nameBased:{}, RID:{}, type:{}, path:{}] failed due to exception:{}", 
                    request.getOperationType(),
                    request.getHeaders(),
                    request.getIsNameBased(),
                    request.getResourceId(),
                    request.getResourceType(),
                    request.getPath(),
                    e.getMessage());
            storeException = e;
        }
        try {
            return createStoreReadResult(storeResponse, storeException, consistencyLevel, new URI(endpointAddress));
        } catch (URISyntaxException e) {
            throw new DocumentClientException(HttpStatus.SC_INTERNAL_SERVER_ERROR, e);
        }
    }

    private StoreResponse readFromStore(DocumentServiceRequest request, String address) throws DocumentClientException {
        String continuation = null;
        if (request.getOperationType() == OperationType.ReadFeed || request.getOperationType() == OperationType.Query
                || request.getOperationType() == OperationType.SqlQuery) {
            continuation = request.getHeaders().get(HttpConstants.HttpHeaders.CONTINUATION);

            if (continuation != null && continuation.contains(";")) {
                String parts[] = StringUtils.split(continuation, ';');
                if (parts.length < 3) {
                    throw new DocumentClientException(HttpStatus.SC_BAD_REQUEST, "Invalid header value");
                }

                continuation = parts[0];
            }

            request.setContinuation(continuation);
        }

        try {
            return this.transportClient.invokeResourceOperation(new URI(address), request);
        } catch (URISyntaxException e) {
            throw new DocumentClientException(HttpStatus.SC_INTERNAL_SERVER_ERROR, e);
        } finally {
            this.lastReadAddress = address;
        }
    }

    private void recordReadResponse(StoreReadResult storeReadResult, DateTime requestStartTime, DocumentServiceRequest request, String endpoint) throws DocumentClientException {
        URI targetURI;
        try {
            targetURI = new URI(endpoint);
        } catch (URISyntaxException e) {
            throw new DocumentClientException(HttpStatus.SC_INTERNAL_SERVER_ERROR, e);
        }
        if (storeReadResult.getStoreResponse() != null) {
            request.getClientSideRequestStatistics().getContactedReplicas().add(targetURI);
        }
        if (storeReadResult.getException() != null &&
                storeReadResult.getException().getStatusCode() == HttpStatus.SC_GONE) {
            request.getClientSideRequestStatistics().getFailedReplicas().add(targetURI);
        }
        ClientSideRequestStatistics clientSideRequestStatistics = request.getClientSideRequestStatistics();
        clientSideRequestStatistics.recordResponse(request, storeReadResult, requestStartTime);
        if (storeReadResult.getStoreResponse() != null) {
            storeReadResult.getStoreResponse().setClientSideRequestStatistics(clientSideRequestStatistics);
        }
    }

    public String getLastReadAddress() {
        return this.lastReadAddress;
    }

    public void setLastReadAddress(String lastReadAddress) {
        this.lastReadAddress = lastReadAddress;
    }
}