Skip to content

Commit 7a8960f

Browse files
kundadebdattakundadebdatta
committed
Code changes to add more tests to validate endpoint unavailable scenario.
1 parent 0f367b3 commit 7a8960f

5 files changed

Lines changed: 280 additions & 3 deletions

File tree

‎Microsoft.Azure.Cosmos/src/ClientRetryPolicy.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -227,7 +227,7 @@ public void OnBeforeSendRequest(DocumentServiceRequest request)
227227
// Resolve the endpoint for the request and pin the resolution to the resolved endpoint
228228
// This enables marking the endpoint unavailability on endpoint failover/unreachability
229229
this.locationEndpoint = ConfigurationManager.IsThinClientEnabled(defaultValue: false)
230-
&& !LocationCache.IsMetaData(request)
230+
&& ThinClientStoreModel.IsOperationSupportedByThinClient(request)
231231
? this.globalEndpointManager.ResolveThinClientEndpoint(request)
232232
: this.globalEndpointManager.ResolveServiceEndpoint(request);
233233

‎Microsoft.Azure.Cosmos/src/Routing/LocationCache.cs‎

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -128,12 +128,38 @@ public ReadOnlyCollection<Uri> WriteEndpoints
128128
/// <summary>
129129
/// Gets the list of thin client read endpoints.
130130
/// </summary>
131-
public ReadOnlyCollection<Uri> ThinClientReadEndpoints => this.locationInfo.ThinClientReadEndpoints;
131+
public ReadOnlyCollection<Uri> ThinClientReadEndpoints
132+
{
133+
get
134+
{
135+
// Hot-path: avoid ConcurrentDictionary methods which acquire locks
136+
if (DateTime.UtcNow - this.lastCacheUpdateTimestamp > this.unavailableLocationsExpirationTime
137+
&& this.locationUnavailablityInfoByEndpoint.Any())
138+
{
139+
this.UpdateLocationCache();
140+
}
141+
142+
return this.locationInfo.ThinClientReadEndpoints;
143+
}
144+
}
132145

133146
/// <summary>
134147
/// Gets the list of thin client write endpoints.
135148
/// </summary>
136-
public ReadOnlyCollection<Uri> ThinClientWriteEndpoints => this.locationInfo.ThinClientWriteEndpoints;
149+
public ReadOnlyCollection<Uri> ThinClientWriteEndpoints
150+
{
151+
get
152+
{
153+
// Hot-path: avoid ConcurrentDictionary methods which acquire locks
154+
if (DateTime.UtcNow - this.lastCacheUpdateTimestamp > this.unavailableLocationsExpirationTime
155+
&& this.locationUnavailablityInfoByEndpoint.Any())
156+
{
157+
this.UpdateLocationCache();
158+
}
159+
160+
return this.locationInfo.ThinClientWriteEndpoints;
161+
}
162+
}
137163

138164
public ReadOnlyCollection<string> EffectivePreferredLocations => this.locationInfo.EffectivePreferredLocations;
139165

‎Microsoft.Azure.Cosmos/src/ThinClientStoreModel.cs‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,11 @@ public override async Task<DocumentServiceResponse> ProcessMessageAsync(
4747
DocumentServiceRequest request,
4848
CancellationToken cancellationToken = default)
4949
{
50+
if (!ThinClientStoreModel.IsOperationSupportedByThinClient(request))
51+
{
52+
return await base.ProcessMessageAsync(request, cancellationToken);
53+
}
54+
5055
await GatewayStoreModel.ApplySessionTokenAsync(
5156
request,
5257
base.defaultConsistencyLevel,
@@ -101,6 +106,21 @@ await this.CaptureSessionTokenAndHandleSplitAsync(
101106
return response;
102107
}
103108

109+
public static bool IsOperationSupportedByThinClient(
110+
DocumentServiceRequest request)
111+
{
112+
// Thin proxy supports the following operations for Document resources.
113+
return request.ResourceType == ResourceType.Document
114+
&& (request.OperationType == OperationType.Batch
115+
|| request.OperationType == OperationType.Patch
116+
|| request.OperationType == OperationType.Create
117+
|| request.OperationType == OperationType.Read
118+
|| request.OperationType == OperationType.Upsert
119+
|| request.OperationType == OperationType.Replace
120+
|| request.OperationType == OperationType.Delete
121+
|| request.OperationType == OperationType.Query);
122+
}
123+
104124
private async Task<AccountProperties> GetDatabaseAccountPropertiesAsync()
105125
{
106126
try

‎Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/PartitionKeyRangeFailoverTests/RegionFailoverTests.cs‎

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,15 +5,21 @@
55
namespace Microsoft.Azure.Cosmos.Tests
66
{
77
using System;
8+
using System.Collections.Concurrent;
89
using System.Collections.Generic;
10+
using System.Collections.ObjectModel;
911
using System.IO;
1012
using System.Linq;
1113
using System.Net;
1214
using System.Net.Http;
15+
using System.Reflection;
16+
using System.Security.AccessControl;
1317
using System.Text;
1418
using System.Threading;
1519
using System.Threading.Tasks;
1620
using Microsoft.Azure.Cosmos.Diagnostics;
21+
using Microsoft.Azure.Cosmos.Routing;
22+
using Microsoft.Azure.Cosmos.Serialization.HybridRow;
1723
using Microsoft.Azure.Documents;
1824
using Microsoft.VisualStudio.TestTools.UnitTesting;
1925
using Moq;
@@ -212,6 +218,7 @@ public async Task TestHttpRequestExceptionScenarioAsync()
212218
}
213219

214220
[TestMethod]
221+
[Owner("dkunda")]
215222
[DataRow(false, DisplayName = "Read Item Scenario without PPAF and PPCB.")]
216223
[DataRow(true, DisplayName = "Read Item Scenario with PPAF and PPCB.")]
217224
public async Task ReadItemAsync_WithThinClientEnabledAndServiceUnavailableReceived_ShouldRetryOnNextPreferredRegions(
@@ -358,6 +365,152 @@ public async Task ReadItemAsync_WithThinClientEnabledAndServiceUnavailableReceiv
358365
}
359366

360367
[TestMethod]
368+
[Owner("dkunda")]
369+
public async Task ReadItemAsync_WithThinClientEnabledAndHttpRequestExceptionReceived_ShouldMarkEndpointUnavailable()
370+
{
371+
try
372+
{
373+
// testhost.dll.config sets it to 2 seconds which causes it to always expire before retrying. Remove the override.
374+
System.Configuration.ConfigurationManager.AppSettings["UnavailableLocationsExpirationTimeInSeconds"] = "500";
375+
376+
Environment.SetEnvironmentVariable(ConfigurationManager.ThinClientModeEnabled, "True");
377+
string accountName = nameof(TestHttpRequestExceptionScenarioAsync);
378+
string primaryRegionNameForUri = "eastus";
379+
string secondaryRegionNameForUri = "westus";
380+
string globalEndpoint = $"https://{accountName}.documents.azure.com:443/";
381+
Uri globalEndpointUri = new Uri(globalEndpoint);
382+
string primaryRegionEndpoint = $"https://{accountName}-{primaryRegionNameForUri}.documents.azure.com";
383+
string secondaryRegionEndpoint = $"https://{accountName}-{secondaryRegionNameForUri}.documents.azure.com";
384+
string databaseName = "testDb";
385+
string containerName = "testContainer";
386+
string containerRid = "ccZ1ANCszwk=";
387+
ResourceId containerResourceId = ResourceId.Parse(containerRid);
388+
389+
List<AccountRegion> writeRegion = new List<AccountRegion>()
390+
{
391+
new AccountRegion()
392+
{
393+
Name = "East US",
394+
Endpoint = $"{primaryRegionEndpoint}:443/"
395+
}
396+
};
397+
398+
List<AccountRegion> readRegions = new List<AccountRegion>()
399+
{
400+
new AccountRegion()
401+
{
402+
Name = "East US",
403+
Endpoint = $"{primaryRegionEndpoint}:443/"
404+
},
405+
new AccountRegion()
406+
{
407+
Name = "West US",
408+
Endpoint = $"{secondaryRegionEndpoint}:443/"
409+
}
410+
};
411+
412+
// Create a mock http handler to inject proxy responses.
413+
// MockBehavior.Strict ensures that only the mocked APIs get called
414+
List<string> regionsVisited = new List<string>();
415+
Mock<IHttpHandler> mockHttpHandler = new Mock<IHttpHandler>(MockBehavior.Strict);
416+
string readResponseHexStringWith200Status = "2a000000C80000000000000000000000000000000000000035000201000000000000011c000200000000480000007b22636f6465223a2022343039222c226d657373616765223a2022416e206572726f72206f63637572726564207768696c6520726f7574696e67207468652072657175657374227d";
417+
mockHttpHandler.Setup(x => x.SendAsync(
418+
It.Is<HttpRequestMessage>(m => m.RequestUri == globalEndpointUri || m.RequestUri.ToString().Contains(primaryRegionNameForUri) || m.RequestUri.ToString().Contains(secondaryRegionNameForUri)),
419+
It.IsAny<CancellationToken>()))
420+
.Returns<HttpRequestMessage, CancellationToken>((request, cancellationToken) =>
421+
{
422+
if (request.Version == new Version(2, 0))
423+
{
424+
if (request.RequestUri.ToString().Contains("eastus"))
425+
{
426+
throw new HttpRequestException();
427+
}
428+
else if (request.RequestUri.ToString().Contains("westus"))
429+
{
430+
regionsVisited.Add(Regions.WestUS);
431+
return Task.FromResult(new HttpResponseMessage(HttpStatusCode.OK)
432+
{
433+
RequestMessage = request,
434+
Content = new StreamContent(new MemoryStream(Convert.FromHexString(readResponseHexStringWith200Status)))
435+
});
436+
}
437+
}
438+
439+
return Task.FromResult(MockSetupsHelper.CreateStrongAccount(accountName, writeRegion, readRegions, shouldEnableThinClient: true, shouldEnablePPAF: false));
440+
});
441+
442+
MockSetupsHelper.SetupContainerProperties(
443+
mockHttpHandler: mockHttpHandler,
444+
regionEndpoint: primaryRegionEndpoint,
445+
databaseName: databaseName,
446+
containerName: containerName,
447+
containerRid: containerRid);
448+
449+
CosmosClientOptions cosmosClientOptions = new CosmosClientOptions()
450+
{
451+
ConsistencyLevel = Cosmos.ConsistencyLevel.Strong,
452+
ApplicationPreferredRegions = new List<string>()
453+
{
454+
Regions.EastUS,
455+
Regions.WestUS
456+
},
457+
ConnectionMode = ConnectionMode.Gateway,
458+
HttpClientFactory = () => new HttpClient(new HttpHandlerHelper(mockHttpHandler.Object)),
459+
};
460+
461+
using (CosmosClient customClient = new CosmosClient(
462+
globalEndpoint,
463+
Convert.ToBase64String(Encoding.UTF8.GetBytes(Guid.NewGuid().ToString())),
464+
cosmosClientOptions))
465+
{
466+
Container container = customClient.GetContainer(databaseName, containerName);
467+
ToDoActivity toDoActivity = new ToDoActivity()
468+
{
469+
Id = "TestItem",
470+
Pk = "TestPk"
471+
};
472+
473+
ItemResponse<ToDoActivity> readResponse = await container.ReadItemAsync<ToDoActivity>(toDoActivity.Id, new Cosmos.PartitionKey(toDoActivity.Pk));
474+
Console.WriteLine($"{readResponse.Diagnostics}");
475+
Assert.AreEqual(HttpStatusCode.OK, readResponse.StatusCode);
476+
Assert.IsTrue(regionsVisited.Count == 1);
477+
Assert.AreEqual(Regions.WestUS, regionsVisited[0]);
478+
479+
GlobalEndpointManager endpointManager = customClient.DocumentClient.GlobalEndpointManager;
480+
481+
FieldInfo fieldInfo = endpointManager
482+
.GetType()
483+
.GetField(
484+
name: "locationCache",
485+
bindingAttr: BindingFlags.Instance | BindingFlags.NonPublic);
486+
487+
LocationCache locationCache = (LocationCache)fieldInfo
488+
.GetValue(
489+
obj: endpointManager);
490+
491+
MethodInfo method = locationCache.GetType().GetMethod("IsEndpointUnavailable", BindingFlags.NonPublic | BindingFlags.Instance);
492+
493+
if (method != null)
494+
{
495+
bool isEastUsAvailable = (bool)method.Invoke(locationCache, new object[] { endpointManager.ThinClientReadEndpoints[0], OperationType.Read });
496+
bool isWestUsAvailable = (bool)method.Invoke(locationCache, new object[] { endpointManager.ThinClientReadEndpoints[1], OperationType.Read });
497+
498+
Assert.IsTrue(isWestUsAvailable, "Since West US was never marked unavailable, this endpoint is expected to be available.");
499+
Assert.IsFalse(isEastUsAvailable, "Since East US was marked unavailable, this endpoint is expected to be unavailable.");
500+
}
501+
502+
mockHttpHandler.VerifyAll();
503+
}
504+
}
505+
finally
506+
{
507+
// Reset the environment variable to avoid impacting other tests.
508+
Environment.SetEnvironmentVariable(ConfigurationManager.ThinClientModeEnabled, null);
509+
}
510+
}
511+
512+
[TestMethod]
513+
[Owner("dkunda")]
361514
[DataRow(false, DisplayName = "When PPAF is disabled, Create Item Scenario should not retry on other regions on a single master write account.")]
362515
public async Task CreateItemAsync_WithThinClientEnabledAndServiceUnavailableReceived_ShouldNotRetryOnOtherRegions(
363516
bool enablePartitionLevelFailover)

‎Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ThinClientStoreModelTests.cs‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,84 @@ public async Task ProcessMessageAsync_Success_ShouldReturnDocumentServiceRespons
151151
Assert.AreEqual(HttpStatusCode.Created, response.StatusCode);
152152
}
153153

154+
[TestMethod]
155+
[Owner("dkunda")]
156+
public async Task ProcessMessageAsync_WithUnsupportedOperations_ShouldFallbackToGatewayModeAndReturnDocumentServiceResponse()
157+
{
158+
// Arrange
159+
// A single base64-encoded RNTBD response representing an HTTP 201 (Created)
160+
HttpResponseMessage successResponse = new HttpResponseMessage(HttpStatusCode.OK) { Content = new StringContent("Response") };
161+
HttpRequestMessage capturedRequest = null;
162+
Mock<CosmosHttpClient> mockCosmosHttpClient = new Mock<CosmosHttpClient>();
163+
mockCosmosHttpClient.Setup(client => client.SendHttpAsync(
164+
It.IsAny<Func<ValueTask<HttpRequestMessage>>>(),
165+
It.IsAny<ResourceType>(),
166+
It.IsAny<HttpTimeoutPolicy>(),
167+
It.IsAny<IClientSideRequestStatistics>(),
168+
It.IsAny<CancellationToken>(),
169+
It.IsAny<DocumentServiceRequest>()))
170+
.Callback<Func<ValueTask<HttpRequestMessage>>, ResourceType, HttpTimeoutPolicy, IClientSideRequestStatistics, CancellationToken, DocumentServiceRequest>(
171+
async (requestFactory, _, _, _, _, _) =>
172+
capturedRequest = await requestFactory())
173+
.ReturnsAsync(successResponse);
174+
175+
DocumentServiceRequest request = DocumentServiceRequest.Create(
176+
operationType: OperationType.QueryPlan,
177+
resourceType: ResourceType.Document,
178+
resourceId: "NH1uAJ6ANm0=",
179+
body: null,
180+
authorizationTokenType: AuthorizationTokenType.PrimaryMasterKey);
181+
182+
Mock<IDocumentClientInternal> docClientMulti = new Mock<IDocumentClientInternal>();
183+
docClientMulti.Setup(c => c.ServiceEndpoint).Returns(new Uri("http://localhost"));
184+
185+
AccountProperties validAccountProperties = new AccountProperties();
186+
187+
docClientMulti
188+
.Setup(c => c.GetDatabaseAccountInternalAsync(It.IsAny<Uri>(), It.IsAny<CancellationToken>()))
189+
.ReturnsAsync(validAccountProperties);
190+
191+
ConnectionPolicy policy = new ConnectionPolicy
192+
{
193+
UseMultipleWriteLocations = true
194+
};
195+
196+
GlobalEndpointManager multiEndpointMgr = new GlobalEndpointManager(docClientMulti.Object, policy);
197+
198+
ThinClientStoreModel storeModel = new ThinClientStoreModel(
199+
endpointManager: multiEndpointMgr,
200+
globalPartitionEndpointManager: GlobalPartitionEndpointManagerNoOp.Instance,
201+
sessionContainer: this.sessionContainer,
202+
defaultConsistencyLevel: (Cosmos.ConsistencyLevel)this.defaultConsistencyLevel,
203+
eventSource: new DocumentClientEventSource(),
204+
serializerSettings: null,
205+
httpClient: mockCosmosHttpClient.Object);
206+
207+
ClientCollectionCache clientCollectionCache = new Mock<ClientCollectionCache>(
208+
this.sessionContainer,
209+
storeModel,
210+
null,
211+
null,
212+
null,
213+
false).Object;
214+
215+
PartitionKeyRangeCache partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(
216+
null,
217+
storeModel,
218+
clientCollectionCache,
219+
multiEndpointMgr,
220+
false).Object;
221+
222+
storeModel.SetCaches(partitionKeyRangeCache, clientCollectionCache);
223+
224+
// Act
225+
DocumentServiceResponse response = await storeModel.ProcessMessageAsync(request);
226+
227+
// Assert
228+
Assert.IsNotNull(response);
229+
Assert.AreEqual(HttpStatusCode.OK, response.StatusCode);
230+
}
231+
154232
[TestMethod]
155233
public void Dispose_ShouldDisposeThinClientStoreClient()
156234
{

0 commit comments

Comments
 (0)