Availability: Adds logic to avoid bad replica during cache refresh

Microsoft.Azure.Cosmos/src/Routing/GatewayAddressCache.cs

@@ -195,6 +195,11 @@ public async Task<PartitionAddressInformation> TryGetAddressesAsync(
                         singleValueInitFunc: (currentCachedValue) =>
                         {
                             staleAddressInfo = currentCachedValue;
+
+                            GatewayAddressCache.SetTransportAddressUrisToUnhealthy(
+                               currentCachedValue,
+                               request?.RequestContext?.FailedEndpoints);
+
                             return this.GetAddressesForRangeIdAsync(
                                 request,
                                 partitionKeyRangeIdentity.CollectionRid,
@@ -272,6 +277,32 @@ public async Task<PartitionAddressInformation> TryGetAddressesAsync(
             }
         }
 
+        private static void SetTransportAddressUrisToUnhealthy(
+            PartitionAddressInformation stalePartitionAddressInformation,
+            Lazy<HashSet<TransportAddressUri>> failedEndpoints)
+        {
+            if (stalePartitionAddressInformation == null ||
+                failedEndpoints == null ||
+                !failedEndpoints.IsValueCreated)
+            {
+                return;
+            }
+
+            IReadOnlyList<TransportAddressUri> perProtocolPartitionAddressInformation = stalePartitionAddressInformation.Get(Protocol.Tcp)?.ReplicaTransportAddressUris;
+            if (perProtocolPartitionAddressInformation == null)
+            {
+                return;
+            }
+
+            foreach (TransportAddressUri failed in perProtocolPartitionAddressInformation)
+            {
+                if (failedEndpoints.Value.Contains(failed))
+                {
+                    failed.SetUnhealthy();
+                }
+            }
+        }
+
         private static void LogPartitionCacheRefresh(
             IClientSideRequestStatistics clientSideRequestStatistics,
             PartitionAddressInformation old,

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosBadReplicaTests.cs

@@ -0,0 +1,191 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Tests
+{
+    using System;
+    using System.Collections.Generic;
+    using System.Globalization;
+    using System.IO;
+    using System.Linq;
+    using System.Net;
+    using System.Net.Http;
+    using System.Text;
+    using System.Threading;
+    using System.Threading.Tasks;
+    using global::Azure.Core;
+    using Microsoft.Azure.Cosmos.Fluent;
+    using Microsoft.VisualStudio.TestTools.UnitTesting;
+    using Moq;
+    using Newtonsoft.Json.Linq;
+
+    [TestClass]
+    public class CosmosBadReplicaTests
+    {
+        [TestMethod]
+        [Timeout(30000)]
+        public async Task TestGoneFromServiceScenarioAsync()
+        {
+            Mock<IHttpHandler> mockHttpHandler = new Mock<IHttpHandler>(MockBehavior.Strict);
+            Uri endpoint = MockSetupsHelper.SetupSingleRegionAccount(
+                "mockAccountInfo",
+                consistencyLevel: ConsistencyLevel.Session,
+                mockHttpHandler,
+                out string primaryRegionEndpoint);
+
+            string databaseName = "mockDbName";
+            string containerName = "mockContainerName";
+            string containerRid = "ccZ1ANCszwk=";
+            Documents.ResourceId cRid = Documents.ResourceId.Parse(containerRid);
+            MockSetupsHelper.SetupContainerProperties(
+                mockHttpHandler: mockHttpHandler,
+                regionEndpoint: primaryRegionEndpoint,
+                databaseName: databaseName,
+                containerName: containerName,
+                containerRid: containerRid);
+
+            MockSetupsHelper.SetupSinglePartitionKeyRange(
+                mockHttpHandler,
+                primaryRegionEndpoint,
+                cRid,
+                out IReadOnlyList<string> partitionKeyRanges);
+
+            List<string> replicaIds1 = new List<string>()
+            {
+                "11111111111111111",
+                "22222222222222222",
+                "33333333333333333",
+                "44444444444444444",
+            };
+
+            HttpResponseMessage replicaSet1 = MockSetupsHelper.CreateAddresses(
+                replicaIds1,
+                partitionKeyRanges.First(),
+                "eastus",
+                cRid);
+
+            // One replica changed on the refresh
+            List<string> replicaIds2 = new List<string>()
+            {
+                "11111111111111111",
+                "22222222222222222",
+                "33333333333333333",
+                "55555555555555555",
+            };
+
+            HttpResponseMessage replicaSet2 = MockSetupsHelper.CreateAddresses(
+                replicaIds2,
+                partitionKeyRanges.First(),
+                "eastus",
+                cRid);
+
+            bool delayCacheRefresh = true;
+            bool delayRefreshUnblocked = false;
+            mockHttpHandler.SetupSequence(x => x.SendAsync(
+                It.Is<HttpRequestMessage>(r => r.RequestUri.ToString().Contains("addresses")), It.IsAny<CancellationToken>()))
+                .Returns(Task.FromResult(replicaSet1))
+                .Returns(async ()=>
+                {
+                    //block cache refresh to verify bad replica is not visited during refresh
+                    while (delayCacheRefresh)
+                    {
+                        await Task.Delay(TimeSpan.FromMilliseconds(20));
+                    }
+
+                    delayRefreshUnblocked = true;
+                    return replicaSet2;
+                });
+
+            int callBack = 0;
+            List<Documents.TransportAddressUri> urisVisited = new List<Documents.TransportAddressUri>();
+            Mock<Documents.TransportClient> mockTransportClient = new Mock<Documents.TransportClient>(MockBehavior.Strict);
+            mockTransportClient.Setup(x => x.InvokeResourceOperationAsync(It.IsAny<Documents.TransportAddressUri>(), It.IsAny<Documents.DocumentServiceRequest>()))
+                .Callback<Documents.TransportAddressUri, Documents.DocumentServiceRequest>((t, _) => urisVisited.Add(t))
+            .Returns(() =>
+            {
+                callBack++;
+                if (callBack == 1)
+                {
+                    throw Documents.Rntbd.TransportExceptions.GetGoneException(
+                        new Uri("https://localhost:8081"),
+                        Guid.NewGuid(),
+                        new Documents.TransportException(Documents.TransportErrorCode.ConnectionBroken,
+                        null,
+                        Guid.NewGuid(),
+                        new Uri("https://localhost:8081"),
+                        "Mock",
+                        userPayload: true,
+                        payloadSent: false));
+                }
+
+                return Task.FromResult(new Documents.StoreResponse()
+                {
+                    Status = 200,
+                    Headers = new Documents.Collections.StoreResponseNameValueCollection()
+                    {
+                        ActivityId = Guid.NewGuid().ToString(),
+                        LSN = "12345",
+                        PartitionKeyRangeId = "0",
+                        GlobalCommittedLSN = "12345",
+                        SessionToken = "1#12345#1=12345"
+                    },
+                    ResponseBody = new MemoryStream()
+                });
+            });
+
+            CosmosClientOptions cosmosClientOptions = new CosmosClientOptions()
+            {
+                ConsistencyLevel = Cosmos.ConsistencyLevel.Session,
+                HttpClientFactory = () => new HttpClient(new HttpHandlerHelper(mockHttpHandler.Object)),
+                TransportClientHandlerFactory = (original) => mockTransportClient.Object,
+            };
+
+            using (CosmosClient customClient = new CosmosClient(
+                endpoint.ToString(),
+                Convert.ToBase64String(Encoding.UTF8.GetBytes(Guid.NewGuid().ToString())),
+                cosmosClientOptions))
+            {
+                try
+                {
+                    Container container = customClient.GetContainer(databaseName, containerName);
+
+                    for (int i = 0; i < 20; i++)
+                    {
+                        ResponseMessage response = await container.ReadItemStreamAsync(Guid.NewGuid().ToString(), new Cosmos.PartitionKey(Guid.NewGuid().ToString()));
+                        Assert.AreEqual(HttpStatusCode.OK, response.StatusCode);
+                    }
+
+                    mockTransportClient.VerifyAll();
+                    mockHttpHandler.VerifyAll();
+
+                    Documents.TransportAddressUri failedReplica = urisVisited.First();
+                    Assert.AreEqual(1, urisVisited.Count(x => x.Equals(failedReplica)));
+
+                    urisVisited.Clear();
+                    delayCacheRefresh = false;
+                    do 
+                    {
+                        await Task.Delay(TimeSpan.FromMilliseconds(100));
+                    }while (!delayRefreshUnblocked);
+
+                    for (int i = 0; i < 20; i++)
+                    {
+                        ResponseMessage response = await container.ReadItemStreamAsync(Guid.NewGuid().ToString(), new Cosmos.PartitionKey(Guid.NewGuid().ToString()));
+                        Assert.AreEqual(HttpStatusCode.OK, response.StatusCode);
+                    }
+
+                    Assert.AreEqual(4, urisVisited.ToHashSet().Count());
+
+                    // Clears all the setups. No network calls should be done on the next operation.
+                    mockHttpHandler.Reset();
+                    mockTransportClient.Reset();
+                }
+                finally
+                {
+                    mockTransportClient.Setup(x => x.Dispose());
+                }
+            }
+        }
+    }
+}

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/PartitionKeyRangeFailoverTests/MockSetupsHelper.cs

@@ -43,6 +43,66 @@ public static void SetupStrongAccountProperties(
                 .Returns<HttpRequestMessage, CancellationToken>((request, cancellationToken) => Task.FromResult(httpResponseMessage));
         }
 
+        public static Uri SetupSingleRegionAccount(
+          string accountName,
+          Cosmos.ConsistencyLevel consistencyLevel,
+          Mock<IHttpHandler> mockHttpHandler,
+          out string primaryRegionEndpoint)
+        {
+            primaryRegionEndpoint = $"https://{accountName}-eastus.documents.azure.com";
+            AccountRegion region = new AccountRegion()
+            {
+                Name = "East US",
+                Endpoint = primaryRegionEndpoint
+            };
+
+            AccountProperties accountProperties = new AccountProperties()
+            {
+                Id = accountName,
+                WriteLocationsInternal = new Collection<AccountRegion>()
+                {
+                    region
+                },
+                ReadLocationsInternal = new Collection<AccountRegion>()
+                {
+                    region
+                },
+                EnableMultipleWriteLocations = false,
+                Consistency = new AccountConsistency()
+                {
+                    DefaultConsistencyLevel = consistencyLevel
+                },
+                SystemReplicationPolicy = new ReplicationPolicy()
+                {
+                    MinReplicaSetSize = 3,
+                    MaxReplicaSetSize = 4
+                },
+                ReadPolicy = new ReadPolicy()
+                {
+                    PrimaryReadCoefficient = 1,
+                    SecondaryReadCoefficient = 1
+                },
+                ReplicationPolicy = new ReplicationPolicy()
+                {
+                    AsyncReplication = false,
+                    MinReplicaSetSize = 3,
+                    MaxReplicaSetSize = 4
+                }
+            };
+
+
+            Uri endpointUri = new Uri($"https://{accountName}.documents.azure.com");
+            mockHttpHandler.Setup(x => x.SendAsync(
+                It.Is<HttpRequestMessage>(x => x.RequestUri == endpointUri),
+                It.IsAny<CancellationToken>()))
+                .Returns<HttpRequestMessage, CancellationToken>((request, cancellationToken) => Task.FromResult(new HttpResponseMessage()
+                {
+                    StatusCode = HttpStatusCode.OK,
+                    Content = new StringContent(JsonConvert.SerializeObject(accountProperties))
+                }));
+            return endpointUri;
+        }
+
         public static HttpResponseMessage CreateStrongAccount(
             string accountName,
             IList<AccountRegion> writeRegions,
@@ -174,6 +234,83 @@ internal static void SetupPartitionKeyRanges(
               }));
         }
 
+        internal static void SetupSinglePartitionKeyRange(
+            Mock<IHttpHandler> mockHttpHandler,
+            string regionEndpoint,
+            ResourceId containerResourceId,
+            out IReadOnlyList<string> partitionKeyRangeIds)
+        {
+            List<Documents.PartitionKeyRange> partitionKeyRanges = new List<Documents.PartitionKeyRange>()
+            {
+                new Documents.PartitionKeyRange()
+                {
+                    MinInclusive = "",
+                    MaxExclusive = "FF",
+                    Id = "0",
+                    ResourceId = "ccZ1ANCszwkDAAAAAAAAUA==",
+                }
+            };
+
+            partitionKeyRangeIds = partitionKeyRanges.Select(x => x.Id).ToList();
+            string containerRidValue = containerResourceId.DocumentCollectionId.ToString();
+            JObject jObject = new JObject
+            {
+                { "_rid",  containerRidValue},
+                { "_count", partitionKeyRanges.Count },
+                { "PartitionKeyRanges", JArray.FromObject(partitionKeyRanges) }
+            };
+
+            Uri partitionKeyUri = new Uri($"{regionEndpoint}/dbs/{containerResourceId.DatabaseId}/colls/{containerRidValue}/pkranges");
+            mockHttpHandler.SetupSequence(x => x.SendAsync(It.Is<HttpRequestMessage>(x => x.RequestUri == partitionKeyUri), It.IsAny<CancellationToken>()))
+              .Returns(() => Task.FromResult(new HttpResponseMessage()
+              {
+                  StatusCode = HttpStatusCode.OK,
+                  Content = new StringContent(jObject.ToString())
+              }))
+              .Returns(() => Task.FromResult(new HttpResponseMessage()
+              {
+                  StatusCode = HttpStatusCode.NotModified,
+              }));
+        }
+
+        internal static HttpResponseMessage CreateAddresses(
+            List<string> replicaIds,
+            string partitionKeyRangeId,
+            string regionName,
+            ResourceId containerResourceId)
+        {
+            string basePhysicalUri = $"rntbd://cdb-ms-prod-{regionName}-fd4.documents.azure.com:14382/apps/9dc0394e-d25f-4c98-baa5-72f1c700bf3e/services/060067c7-a4e9-4465-a412-25cb0104cb58/partitions/2cda760c-f81f-4094-85d0-7bcfb2acc4e6/replicas/";
+
+            // Use the partition key range id at the end of each replica id to avoid conflicts when setting up multiple partition key ranges
+            List<Address> addresses = new List<Address>();
+            for (int i = 0; i < replicaIds.Count; i++)
+            {
+                string repliaId = replicaIds[i] + (i == 0 ? "p":"s") + "/";
+                addresses.Add(new Address()
+                {
+                    IsPrimary = i == 0,
+                    PartitionKeyRangeId = partitionKeyRangeId,
+                    PhysicalUri = basePhysicalUri + repliaId,
+                    Protocol = "rntbd",
+                    PartitionIndex = "7718513@164605136"
+                });
+            }
+
+            string containerRid = containerResourceId.DocumentCollectionId.ToString();
+            JObject jObject = new JObject
+            {
+                { "_rid", containerRid },
+                { "_count", addresses.Count },
+                { "Addresss", JArray.FromObject(addresses) }
+            };
+
+            return new HttpResponseMessage()
+            {
+                StatusCode = HttpStatusCode.OK,
+                Content = new StringContent(jObject.ToString())
+            };
+        }
+
         internal static void SetupAddresses(
             Mock<IHttpHandler> mockHttpHandler,
             string partitionKeyRangeId,
@@ -273,7 +410,6 @@ internal static void SetupCreateItemResponse(
                 .Returns<TransportAddressUri, DocumentServiceRequest>(
                 (uri, documentServiceRequest) =>
                 {
-                    Console.WriteLine($"Write Success: {physicalUri}");
                     Stream createdObject = documentServiceRequest.CloneableBody.Clone();
                     return Task.FromResult(new StoreResponse()
                     {