Skip to content

Commit

Permalink
Routing: Adds ExcludeRegions Feature to RequestOptions (#4128)
Browse files Browse the repository at this point in the history
* adds excludeRegions

* suggested changes

* removed unused usings

* fixed blank line error

* removed using

* update contracts

* fixed test

* reverted automatic changes to BaselineTests

* requested changes

* bug fix

* PPOF test fix
  • Loading branch information
NaluTripician authored Nov 9, 2023
1 parent d287f08 commit 966b481
Show file tree
Hide file tree
Showing 8 changed files with 303 additions and 32 deletions.
8 changes: 4 additions & 4 deletions Microsoft.Azure.Cosmos/src/ClientRetryPolicy.cs
Original file line number Diff line number Diff line change
Expand Up @@ -246,9 +246,9 @@ private async Task<ShouldRetryResult> ShouldRetryInternalAsync(
if (statusCode == HttpStatusCode.NotFound
&& subStatusCode == SubStatusCodes.ReadSessionNotAvailable)
{
return this.ShouldRetryOnSessionNotAvailable();
return this.ShouldRetryOnSessionNotAvailable(this.documentServiceRequest);
}

// Received 503 due to client connect timeout or Gateway
if (statusCode == HttpStatusCode.ServiceUnavailable)
{
Expand Down Expand Up @@ -330,7 +330,7 @@ private async Task<ShouldRetryResult> ShouldRetryOnEndpointFailureAsync(
return ShouldRetryResult.RetryAfter(retryDelay);
}

private ShouldRetryResult ShouldRetryOnSessionNotAvailable()
private ShouldRetryResult ShouldRetryOnSessionNotAvailable(DocumentServiceRequest request)
{
this.sessionTokenRetryCount++;

Expand All @@ -343,7 +343,7 @@ private ShouldRetryResult ShouldRetryOnSessionNotAvailable()
{
if (this.canUseMultipleWriteLocations)
{
ReadOnlyCollection<Uri> endpoints = this.isReadRequest ? this.globalEndpointManager.ReadEndpoints : this.globalEndpointManager.WriteEndpoints;
ReadOnlyCollection<Uri> endpoints = this.globalEndpointManager.GetApplicableEndpoints(request, this.isReadRequest);

if (this.sessionTokenRetryCount > endpoints.Count)
{
Expand Down
4 changes: 3 additions & 1 deletion Microsoft.Azure.Cosmos/src/Handler/RequestMessage.cs
Original file line number Diff line number Diff line change
Expand Up @@ -282,7 +282,8 @@ internal DocumentServiceRequest ToDocumentServiceRequest()
{
this.DocumentServiceRequest.RouteTo(this.PartitionKeyRangeId);
}


this.DocumentServiceRequest.RequestContext.ExcludeRegions = this.RequestOptions?.ExcludeRegions;
this.OnBeforeRequestHandler(this.DocumentServiceRequest);
return this.DocumentServiceRequest;
}
Expand All @@ -299,6 +300,7 @@ private static Headers CreateHeaders()

private void OnBeforeRequestHandler(DocumentServiceRequest serviceRequest)
{
serviceRequest.RequestContext.ExcludeRegions = this.RequestOptions?.ExcludeRegions;
this.OnBeforeSendRequestActions?.Invoke(serviceRequest);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
namespace Microsoft.Azure.Cosmos
{
using System;
using System.ComponentModel;
using System.Text;
using Microsoft.Azure.Cosmos.CosmosElements;
using Microsoft.Azure.Cosmos.Query;
Expand Down
7 changes: 7 additions & 0 deletions Microsoft.Azure.Cosmos/src/RequestOptions/RequestOptions.cs
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,13 @@ public class RequestOptions
/// </summary>
public CosmosThresholdOptions CosmosThresholdOptions { get; set; }

/// <summary>
/// List of regions to be excluded routing the request to.
/// This can be used to route a request to a specific region by excluding all other regions.
/// If all regions are excluded, then the request will be routed to the primary/hub region.
/// </summary>
public List<string> ExcludeRegions { get; set; }

/// <summary>
/// Gets or sets the boolean to use effective partition key routing in the cosmos db request.
/// </summary>
Expand Down
5 changes: 5 additions & 0 deletions Microsoft.Azure.Cosmos/src/Routing/GlobalEndpointManager.cs
Original file line number Diff line number Diff line change
Expand Up @@ -362,6 +362,11 @@ public string GetLocation(Uri endpoint)
return this.locationCache.GetLocation(endpoint);
}

public ReadOnlyCollection<Uri> GetApplicableEndpoints(DocumentServiceRequest request, bool isReadRequest)
{
return this.locationCache.GetApplicableEndpoints(request, isReadRequest);
}

public bool TryGetLocationForGatewayDiagnostics(Uri endpoint, out string regionName)
{
return this.locationCache.TryGetLocationForGatewayDiagnostics(endpoint, out regionName);
Expand Down
62 changes: 61 additions & 1 deletion Microsoft.Azure.Cosmos/src/Routing/LocationCache.cs
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ internal sealed class LocationCache
private readonly TimeSpan unavailableLocationsExpirationTime;
private readonly int connectionLimit;
private readonly ConcurrentDictionary<Uri, LocationUnavailabilityInfo> locationUnavailablityInfoByEndpoint;
private readonly RegionNameMapper regionNameMapper;

private DatabaseAccountLocationsInfo locationInfo;
private DateTime lastCacheUpdateTimestamp;
Expand All @@ -54,6 +55,7 @@ public LocationCache(
this.lastCacheUpdateTimestamp = DateTime.MinValue;
this.enableMultipleWriteLocations = false;
this.unavailableLocationsExpirationTime = TimeSpan.FromSeconds(LocationCache.DefaultUnavailableLocationsExpirationTimeInSeconds);
this.regionNameMapper = new RegionNameMapper();

#if !(NETSTANDARD15 || NETSTANDARD16)
#if NETSTANDARD20
Expand Down Expand Up @@ -281,14 +283,72 @@ public Uri ResolveServiceEndpoint(DocumentServiceRequest request)
}
else
{
ReadOnlyCollection<Uri> endpoints = request.OperationType.IsWriteOperation() ? this.WriteEndpoints : this.ReadEndpoints;
ReadOnlyCollection<Uri> endpoints = this.GetApplicableEndpoints(request, !request.OperationType.IsWriteOperation());
locationEndpointToRoute = endpoints[locationIndex % endpoints.Count];
}

request.RequestContext.RouteToLocation(locationEndpointToRoute);
return locationEndpointToRoute;
}

public ReadOnlyCollection<Uri> GetApplicableEndpoints(DocumentServiceRequest request, bool isReadRequest)
{
ReadOnlyCollection<Uri> endpoints = isReadRequest ? this.ReadEndpoints : this.WriteEndpoints;

if (request.RequestContext.ExcludeRegions == null || request.RequestContext.ExcludeRegions.Count == 0)
{
return endpoints;
}

return this.GetApplicableEndpoints(
endpoints,
isReadRequest ? this.locationInfo.AvailableReadEndpointByLocation : this.locationInfo.AvailableWriteEndpointByLocation,
this.defaultEndpoint,
request.RequestContext.ExcludeRegions);
}

/// <summary>
/// Gets applicable endpoints for a request, if there are no applicable endpoints, returns the fallback endpoint
/// </summary>
/// <param name="endpoints"></param>
/// <param name="regionNameByEndpoint"></param>
/// <param name="fallbackEndpoint"></param>
/// <param name="excludeRegions"></param>
/// <returns>a list of applicable endpoints for a request</returns>
private ReadOnlyCollection<Uri> GetApplicableEndpoints(
IReadOnlyList<Uri> endpoints,
ReadOnlyDictionary<string, Uri> regionNameByEndpoint,
Uri fallbackEndpoint,
IReadOnlyList<string> excludeRegions)
{
List<Uri> applicableEndpoints = new List<Uri>(endpoints.Count);
HashSet<Uri> excludeUris = new HashSet<Uri>();

foreach (string region in excludeRegions)
{
string normalizedRegionName = this.regionNameMapper.GetCosmosDBRegionName(region);
if (regionNameByEndpoint.ContainsKey(normalizedRegionName))
{
excludeUris.Add(regionNameByEndpoint[normalizedRegionName]);
}
}

foreach (Uri endpoint in endpoints)
{
if (!excludeUris.Contains(endpoint))
{
applicableEndpoints.Add(endpoint);
}
}

if (applicableEndpoints.Count == 0)
{
applicableEndpoints.Add(fallbackEndpoint);
}

return new ReadOnlyCollection<Uri>(applicableEndpoints);
}

public bool ShouldRefreshEndpoints(out bool canRefreshInBackground)
{
canRefreshInBackground = true;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7986,6 +7986,18 @@
"Attributes": [],
"MethodInfo": "System.Collections.Generic.IReadOnlyDictionary`2[System.String,System.Object] Properties;CanRead:True;CanWrite:True;System.Collections.Generic.IReadOnlyDictionary`2[System.String,System.Object] get_Properties();IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;Void set_Properties(System.Collections.Generic.IReadOnlyDictionary`2[System.String,System.Object]);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
},
"System.Collections.Generic.List`1[System.String] ExcludeRegions": {
"Type": "Property",
"Attributes": [],
"MethodInfo": "System.Collections.Generic.List`1[System.String] ExcludeRegions;CanRead:True;CanWrite:True;System.Collections.Generic.List`1[System.String] get_ExcludeRegions();IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;Void set_ExcludeRegions(System.Collections.Generic.List`1[System.String]);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
},
"System.Collections.Generic.List`1[System.String] get_ExcludeRegions()[System.Runtime.CompilerServices.CompilerGeneratedAttribute()]": {
"Type": "Method",
"Attributes": [
"CompilerGeneratedAttribute"
],
"MethodInfo": "System.Collections.Generic.List`1[System.String] get_ExcludeRegions();IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
},
"System.String get_IfMatchEtag()[System.Runtime.CompilerServices.CompilerGeneratedAttribute()]": {
"Type": "Method",
"Attributes": [
Expand Down Expand Up @@ -8029,6 +8041,13 @@
],
"MethodInfo": "Void set_CosmosThresholdOptions(Microsoft.Azure.Cosmos.CosmosThresholdOptions);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
},
"Void set_ExcludeRegions(System.Collections.Generic.List`1[System.String])[System.Runtime.CompilerServices.CompilerGeneratedAttribute()]": {
"Type": "Method",
"Attributes": [
"CompilerGeneratedAttribute"
],
"MethodInfo": "Void set_ExcludeRegions(System.Collections.Generic.List`1[System.String]);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
},
"Void set_IfMatchEtag(System.String)[System.Runtime.CompilerServices.CompilerGeneratedAttribute()]": {
"Type": "Method",
"Attributes": [
Expand Down
Loading

0 comments on commit 966b481

Please sign in to comment.