ClientCollectionCache.cs

//------------------------------------------------------------
// Copyright (c) Microsoft Corporation.  All rights reserved.
//------------------------------------------------------------

namespace Microsoft.Azure.Cosmos.Routing
{
    using System;
    using System.Threading;
    using System.Threading.Tasks;
    using Microsoft.Azure.Cosmos.Common;
    using Microsoft.Azure.Cosmos.Telemetry;
    using Microsoft.Azure.Cosmos.Telemetry.Collector;
    using Microsoft.Azure.Cosmos.Tracing;
    using Microsoft.Azure.Cosmos.Tracing.TraceData;
    using Microsoft.Azure.Documents;
    using Microsoft.Azure.Documents.Collections;

    /// <summary>
    /// Caches collection information.
    /// </summary>
    internal class ClientCollectionCache : CollectionCache
    {
        private const string TelemetrySourceName = "ClientCollectionCache";

        private readonly IStoreModel storeModel;
        private readonly ICosmosAuthorizationTokenProvider tokenProvider;
        private readonly IRetryPolicyFactory retryPolicy;
        private readonly ISessionContainer sessionContainer;
        private readonly TelemetryToServiceHelper telemetryToServiceHelper;

        public ClientCollectionCache(
            ISessionContainer sessionContainer,
            IStoreModel storeModel,
            ICosmosAuthorizationTokenProvider tokenProvider,
            IRetryPolicyFactory retryPolicy,
            TelemetryToServiceHelper telemetryToServiceHelper,
            bool enableAsyncCacheExceptionNoSharing = true)
            : base(enableAsyncCacheExceptionNoSharing)
        {
            this.storeModel = storeModel ?? throw new ArgumentNullException("storeModel");
            this.tokenProvider = tokenProvider;
            this.retryPolicy = retryPolicy;
            this.sessionContainer = sessionContainer;
            this.telemetryToServiceHelper = telemetryToServiceHelper;
        }

        protected override Task<ContainerProperties> GetByRidAsync(string apiVersion,
                                                    string collectionRid,
                                                    ITrace trace,
                                                    IClientSideRequestStatistics clientSideRequestStatistics,
                                                    CancellationToken cancellationToken)
        {
            cancellationToken.ThrowIfCancellationRequested();
            IDocumentClientRetryPolicy retryPolicyInstance = new ClearingSessionContainerClientRetryPolicy(
                this.sessionContainer, this.retryPolicy.GetRequestPolicy());
            return MetadataRetryHelper.ExecuteAsync(
                  (ct) => this.ReadCollectionAsync(
                      PathsHelper.GeneratePath(ResourceType.Collection, collectionRid, false),
                      retryPolicyInstance,
                      trace,
                      clientSideRequestStatistics,
                      ct),
                  retryPolicyInstance,
                  cancellationToken);
        }

        protected override Task<ContainerProperties> GetByNameAsync(string apiVersion,
                                                string resourceAddress,
                                                ITrace trace,
                                                IClientSideRequestStatistics clientSideRequestStatistics,
                                                CancellationToken cancellationToken)
        {
            cancellationToken.ThrowIfCancellationRequested();
            IDocumentClientRetryPolicy retryPolicyInstance = new ClearingSessionContainerClientRetryPolicy(
                this.sessionContainer, this.retryPolicy.GetRequestPolicy());
            return MetadataRetryHelper.ExecuteAsync(
                (ct) => this.ReadCollectionAsync(
                    resourceAddress, retryPolicyInstance, trace, clientSideRequestStatistics, ct),
                retryPolicyInstance,
                cancellationToken);
        }

        internal override Task<ContainerProperties> ResolveByNameAsync(
            string apiVersion,
            string resourceAddress,
            bool forceRefesh,
            ITrace trace,
            IClientSideRequestStatistics clientSideRequestStatistics,
            CancellationToken cancellationToken)
        {
            if (forceRefesh && this.sessionContainer != null)
            {
                return TaskHelper.InlineIfPossible(
                    async () =>
                    {
                        string oldRid = (await base.ResolveByNameAsync(
                            apiVersion,
                            resourceAddress,
                            forceRefesh: false,
                            trace,
                            clientSideRequestStatistics,
                            cancellationToken))?.ResourceId;

                        ContainerProperties propertiesAfterRefresh = await base.ResolveByNameAsync(
                            apiVersion,
                            resourceAddress,
                            forceRefesh,
                            trace,
                            clientSideRequestStatistics,
                            cancellationToken);

                        if (oldRid != null && oldRid != propertiesAfterRefresh?.ResourceId)
                        {
                            string resourceFullName = PathsHelper.GetCollectionPath(resourceAddress);
                            this.sessionContainer.ClearTokenByCollectionFullname(resourceFullName);
                        }

                        return propertiesAfterRefresh;
                    },
                    retryPolicy: null,
                    cancellationToken);
            }

            return TaskHelper.InlineIfPossible(
                () => base.ResolveByNameAsync(
                    apiVersion, resourceAddress, forceRefesh, trace, clientSideRequestStatistics, cancellationToken),
                retryPolicy: null,
                cancellationToken);
        }

        public override Task<ContainerProperties> ResolveCollectionAsync(
            DocumentServiceRequest request, CancellationToken cancellationToken, ITrace trace)
        {
            return TaskHelper.InlineIfPossible(
                () => this.ResolveCollectionWithSessionContainerCleanupAsync(
                    request,
                    () => base.ResolveCollectionAsync(request, cancellationToken, trace)),
                retryPolicy: null,
                cancellationToken);
        }

        public override Task<ContainerProperties> ResolveCollectionAsync(
            DocumentServiceRequest request, TimeSpan refreshAfter, CancellationToken cancellationToken, ITrace trace)
        {
            return TaskHelper.InlineIfPossible(
                () => this.ResolveCollectionWithSessionContainerCleanupAsync(
                    request,
                    () => base.ResolveCollectionAsync(request, refreshAfter, cancellationToken, trace)),
                retryPolicy: null,
                cancellationToken);
        }

        private async Task<ContainerProperties> ResolveCollectionWithSessionContainerCleanupAsync(
            DocumentServiceRequest request,
            Func<Task<ContainerProperties>> resolveContainerProvider)
        {
            string previouslyResolvedCollectionRid = request?.RequestContext?.ResolvedCollectionRid;

            try
            {
                ContainerProperties properties = await resolveContainerProvider();

                if (this.sessionContainer != null &&
                    previouslyResolvedCollectionRid != null &&
                    previouslyResolvedCollectionRid != properties.ResourceId)
                {
                    this.sessionContainer.ClearTokenByResourceId(previouslyResolvedCollectionRid);
                }

                return properties;
            }
            catch (DocumentClientException ex)
            {
                // When collection resolution fails with 404, set the appropriate substatus code.
                // - For document/item operations: Set substatus 1003 (OwnerResourceNotFound) to distinguish
                //   "container doesn't exist" from "item doesn't exist" (substatus 0).
                // - For container operations: Leave substatus as 0 because the container itself is the
                //   resource that doesn't exist (not an "owner" resource issue).
                if (ex.StatusCode == System.Net.HttpStatusCode.NotFound && 
                    ex.GetSubStatus() == SubStatusCodes.Unknown &&
                    request?.ResourceType != ResourceType.Collection)
                {
                    ex.Headers[WFConstants.BackendHeaders.SubStatus] = ((uint)SubStatusCodes.OwnerResourceNotFound).ToString(System.Globalization.CultureInfo.InvariantCulture);
                }

                throw;
            }
        }

        private async Task<ContainerProperties> ReadCollectionAsync(
            string collectionLink,
            IDocumentClientRetryPolicy retryPolicyInstance,
            ITrace trace,
            IClientSideRequestStatistics clientSideRequestStatistics,
            CancellationToken cancellationToken)
        {
            using (ITrace childTrace = trace.StartChild("Read Collection", TraceComponent.Transport, TraceLevel.Info))
            {
                cancellationToken.ThrowIfCancellationRequested();

                RequestNameValueCollection headers = new RequestNameValueCollection();
                using (DocumentServiceRequest request = DocumentServiceRequest.Create(
                       OperationType.Read,
                       ResourceType.Collection,
                       collectionLink,
                       AuthorizationTokenType.PrimaryMasterKey,
                       headers))
                {
                    headers.XDate = Rfc1123DateTimeCache.UtcNow();

                    request.RequestContext.ClientRequestStatistics = clientSideRequestStatistics ?? new ClientSideRequestStatisticsTraceDatum(DateTime.UtcNow, trace);
                    if (clientSideRequestStatistics == null)
                    {
                        childTrace.AddDatum(
                            "Client Side Request Stats",
                            request.RequestContext.ClientRequestStatistics);
                    }

                    string authorizationToken = await this.tokenProvider.GetUserAuthorizationTokenAsync(
                        request.ResourceAddress,
                        PathsHelper.GetResourcePath(request.ResourceType),
                        HttpConstants.HttpMethods.Get,
                        request.Headers,
                        AuthorizationTokenType.PrimaryMasterKey,
                        childTrace);

                    headers.Authorization = authorizationToken;

                    using (new ActivityScope(Guid.NewGuid()))
                    {
                        retryPolicyInstance?.OnBeforeSendRequest(request);

                        try
                        {
                            using (DocumentServiceResponse response =
                                await this.storeModel.ProcessMessageAsync(request))
                            {
                                ContainerProperties containerProperties = CosmosResource.FromStream<ContainerProperties>(response);
                                
                                this.telemetryToServiceHelper.GetCollector().CollectCacheInfo(
                                     ClientCollectionCache.TelemetrySourceName,
                                     () => new TelemetryInformation
                                     {
                                         RegionsContactedList = response.RequestStats.RegionsContacted,
                                         RequestLatency = response.RequestStats.RequestLatency,
                                         StatusCode = response.StatusCode,
                                         OperationType = request.OperationType,
                                         ResourceType = request.ResourceType,
                                         SubStatusCode = response.SubStatusCode,
                                         CollectionLink = collectionLink
                                     });

                                return containerProperties;
                            }
                        }
                        catch (DocumentClientException ex)
                        {
                            childTrace.AddDatum("Exception Message", ex.Message);
                            throw;
                        }
                    }
                }
            }
        }
    }
}