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

Go to download
/*
 * Copyright (c) Microsoft Corporation.  All rights reserved.
 */

package com.microsoft.azure.documentdb.internal.directconnectivity;

import java.net.URI;
import java.net.URISyntaxException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.RejectedExecutionException;

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.AuthorizationTokenProvider;
import com.microsoft.azure.documentdb.internal.DatabaseAccountConfigurationProvider;
import com.microsoft.azure.documentdb.internal.DocumentServiceRequest;
import com.microsoft.azure.documentdb.internal.HttpConstants;
import com.microsoft.azure.documentdb.internal.RequestChargeTracker;
import com.microsoft.azure.documentdb.internal.SessionContainer;
import com.microsoft.azure.documentdb.internal.SessionTokenHelper;
import com.microsoft.azure.documentdb.ClientSideRequestStatistics;

/**
 * Used internally to provide local quorum-acked write and globally strong write.
 */
public class ConsistencyWriter {

    private static final Logger logger = LoggerFactory.getLogger(ConsistencyWriter.class);

    private final static int MAX_NUMBER_OF_WRITE_BARRIER_READ_RETRIES = 30;
    private final static int DELAY_BETWEEN_WRITE_BARRIER_CALLS_IN_MS = 30;

    private final static int MAX_SHORT_BARRIER_RETRIES_FOR_MULTI_REGION = 4;
    private final static int SHORT_BARRIER_RETRY_INTERVAL_IN_MS_FOR_MULTI_REGION = 10;

    private final GlobalAddressResolver globalAddressResolver;
    private final StoreReader storeReader;
    private final TransportClient transportClient;
    private final DatabaseAccountConfigurationProvider configurationProvider;
    private final AuthorizationTokenProvider authorizationTokenProvider;
    private final ExecutorService executorService;
    private final SessionContainer sessionContainer;
    private final boolean useMultipleWriteLocations;

    public ConsistencyWriter(GlobalAddressResolver globalAddressResolver,
                             SessionContainer sessionContainer,
                             TransportClient transportClient,
                             DatabaseAccountConfigurationProvider configurationProvider,
                             AuthorizationTokenProvider authorizationTokenProvider,
                             ExecutorService executorService,
                             boolean useMultipleWriteLocations) {
        this.globalAddressResolver = globalAddressResolver;
        this.storeReader = new StoreReader(globalAddressResolver,
                transportClient, null, executorService);
        this.transportClient = transportClient;
        this.configurationProvider = configurationProvider;
        this.authorizationTokenProvider = authorizationTokenProvider;
        this.executorService = executorService;
        this.sessionContainer = sessionContainer;
        this.useMultipleWriteLocations = useMultipleWriteLocations;
    }

    public StoreResponse write(DocumentServiceRequest request) throws DocumentClientException {
        String sessionToken = request.getHeaders().get(HttpConstants.HttpHeaders.SESSION_TOKEN);
        try {
            return this.writePrivate(request);
        } finally {
            request.setOriginalSessionToken(sessionToken);
        }
    }

    private void startBackgroundAddressRefresh(final DocumentServiceRequest request) {
        final DocumentServiceRequest requestFinal = request;
        try {
            executorService.submit(new Runnable() {
                @Override
                public void run() {
                    try {
                        requestFinal.setForceAddressRefresh(true);
                        ReplicatedResourceClient.resolvePrimaryUri(requestFinal, ConsistencyWriter.this.globalAddressResolver.resolve(requestFinal));
                    } catch (DocumentClientException e) {
                        logger.warn("Background refresh of the primary address failed with {}", e.getMessage(), e);
                }
                }
            });
        } catch (RejectedExecutionException e) {
            logger.warn("Background refresh of the primary address failed with {}", e.getMessage(), e);
        }
    }

    private StoreResponse writePrivate(DocumentServiceRequest request) throws DocumentClientException {
        if (request.getRequestChargeTracker() == null) {
            request.setRequestChargeTracker(new RequestChargeTracker());
        }
        if (request.getClientSideRequestStatistics() == null) {
            request.setClientSideRequestStatistics(new ClientSideRequestStatistics());
        }

        if (request.getGlobalStrongWriteResponse() == null) {
            StoreResponse response;

            AddressInformation[] addressInformations = ReplicatedResourceClient.resolveAddresses(request,
                    this.globalAddressResolver.resolve(request));
            
            URI primaryUri = ReplicatedResourceClient.resolvePrimaryUri(request, addressInformations);

            if (useMultipleWriteLocations && configurationProvider.getStoreConsistencyPolicy() == ConsistencyLevel.Session) {
                // Set session token to ensure session consistency for write requests
                // when writes can be issued to multiple locations
                SessionTokenHelper.setPartitionLocalSessionToken(request, this.sessionContainer);
            } else {
                // When writes can only go to single location, there is no reason
                // to send session token to the server.
                request.getHeaders().remove(HttpConstants.HttpHeaders.SESSION_TOKEN);
            }

            DateTime startTime = DateTime.now(DateTimeZone.UTC);
            try {
                response = this.transportClient.invokeResourceOperation(primaryUri, request);
                request.getClientSideRequestStatistics().recordResponse(request,
                        StoreReader.createStoreReadResult(response, null, configurationProvider.getStoreConsistencyPolicy(), primaryUri),
                        startTime);
            } catch (DocumentClientException e) {
                request.getClientSideRequestStatistics().recordResponse(request,
                        StoreReader.createStoreReadResult(null, e, configurationProvider.getStoreConsistencyPolicy(), primaryUri),
                        startTime);
                String header = e.getResponseHeaders() != null ?
                        e.getResponseHeaders().get(HttpConstants.HttpHeaders.WRITE_REQUEST_TRIGGER_ADDRESS_REFRESH)
                        : null;
                if (StringUtils.isNotEmpty(header) && header.equalsIgnoreCase(String.valueOf(1))) {
                    startBackgroundAddressRefresh(request);
                }
                throw e;
            }

            try {
                List contactedReplicas = new ArrayList<>();
                for (AddressInformation addressInformation : addressInformations) {
                    contactedReplicas.add(new URI(addressInformation.getPhysicalUri()));

                }
                request.getClientSideRequestStatistics().setContactedReplicas(contactedReplicas);
            }
            catch (URISyntaxException e) {
                throw new DocumentClientException(HttpStatus.SC_INTERNAL_SERVER_ERROR, e);
            }

            if (ReplicatedResourceClient.GLOBAL_STRONG_ENABLED && this.isGlobalStrong(response)) {
                long lsn = NumberUtils.toLong(response.getHeaderValue(WFConstants.BackendHeaders.LSN), -1);
                long globalCommittedLsn = NumberUtils.toLong(response.getHeaderValue(WFConstants.BackendHeaders.GlobalCommittedLSN), -1);

                if (lsn == -1 || globalCommittedLsn == -1) {
                    logger.debug("ConsistencyWriter: LSN or GlobalCommittedLSN is not set for global strong request");
                    throw new DocumentClientException(HttpStatus.SC_GONE, "ConsistencyWriter: LSN or GlobalCommittedLSN is not set for global strong request");
                }

                request.setGlobalStrongWriteResponse(response);
                request.setGlobalCommittedSelectedLSN(lsn);

                if (globalCommittedLsn < lsn) {
                    DocumentServiceRequest barrierRequest = BarrierRequestHelper.create(request, this.authorizationTokenProvider);
                    if (!this.waitForWriteBarrier(barrierRequest, lsn)) {
                        logger.debug("ConsistencyWriter: Write barrier has not been met for global strong request. SelectedGlobalCommittedLSN: " + lsn);
                        throw new DocumentClientException(HttpStatus.SC_GONE, "ConsistencyWriter: Write barrier has not been met for global strong request.");
                    }
                }
            }
            //  Set calculated client side request statistics
            response.setClientSideRequestStatistics(request.getClientSideRequestStatistics());
            return response;
        } else {
            DocumentServiceRequest barrierRequest = BarrierRequestHelper.create(request, this.authorizationTokenProvider);
            if (!this.waitForWriteBarrier(barrierRequest, request.getGlobalCommittedSelectedLSN())) {
                logger.debug("ConsistencyWriter: Write barrier has not been met for global strong request. SelectedGlobalCommittedLSN: " + request.getGlobalCommittedSelectedLSN());
                throw new DocumentClientException(HttpStatus.SC_GONE, "ConsistencyWriter: Write barrier has not been met for global strong request.");
            }
        }

        StoreResponse globalStrongWriteResponse = request.getGlobalStrongWriteResponse();
        globalStrongWriteResponse.setClientSideRequestStatistics(request.getClientSideRequestStatistics());
        return globalStrongWriteResponse;
    }

    private boolean isGlobalStrong(StoreResponse response) throws DocumentClientException {
        if (this.configurationProvider.getStoreConsistencyPolicy() == ConsistencyLevel.Strong) {
            String headerValue = response.getHeaderValue(WFConstants.BackendHeaders.NumberOfReadRegions);
            int numberOfReadRegions = Integer.parseInt(headerValue);
            if (numberOfReadRegions > 0) {
                return true;
            }
        }
        return false;
    }

    private boolean waitForWriteBarrier(DocumentServiceRequest request, long barrier) throws DocumentClientException {
        int writeBarrierRetryCount = MAX_NUMBER_OF_WRITE_BARRIER_READ_RETRIES;
        long maxGlobalCommittedLSNReceived = 0;

        while (writeBarrierRetryCount-- > 0) {
            StoreReadResult response = storeReader.readEventual(request);

            if (response != null) {
                if (response.getGlobalCommittedLSN() >= barrier) {
                    return true;
                }

                if (maxGlobalCommittedLSNReceived < response.getGlobalCommittedLSN()) {
                    maxGlobalCommittedLSNReceived = response.getGlobalCommittedLSN();
                }
            }

            try {
                if ((ConsistencyWriter.MAX_NUMBER_OF_WRITE_BARRIER_READ_RETRIES - writeBarrierRetryCount) >
                        ConsistencyWriter.MAX_SHORT_BARRIER_RETRIES_FOR_MULTI_REGION) {
                    Thread.sleep(ConsistencyWriter.DELAY_BETWEEN_WRITE_BARRIER_CALLS_IN_MS);
                } else {
                    Thread.sleep(ConsistencyWriter.SHORT_BARRIER_RETRY_INTERVAL_IN_MS_FOR_MULTI_REGION);
                }
            } catch (InterruptedException e) {
                throw new IllegalStateException("Delay thread interrupted with exception: ", e);
            }
        }

        logger.trace("ConsistencyWriter: Highest global committed lsn received for write barrier call is " + maxGlobalCommittedLSNReceived);

        return false;
    }
}