Skip to content

Commit dd621d3

Browse files
Merge branch 'master' into ananth/openai-exception-handling-asynccache
2 parents 8ab0077 + 1101959 commit dd621d3

16 files changed

Lines changed: 1412 additions & 30 deletions

Microsoft.Azure.Cosmos/src/CosmosClientOptions.cs

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -959,7 +959,12 @@ public IFaultInjector FaultInjector
959959
/// If <see cref="AllowBulkExecution"/> is set to true in CosmosClientOptions, throughput bucket can only be set at client level.
960960
/// </remarks>
961961
/// <seealso href="https://aka.ms/cosmsodb-bucketing"/>
962-
internal int? ThroughputBucket { get; set; }
962+
#if PREVIEW
963+
public
964+
#else
965+
internal
966+
#endif
967+
int? ThroughputBucket { get; set; }
963968

964969
internal IChaosInterceptorFactory ChaosInterceptorFactory { get; set; }
965970

Microsoft.Azure.Cosmos/src/Fluent/CosmosClientBuilder.cs

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -808,7 +808,12 @@ public CosmosClientBuilder WithClientTelemetryOptions(CosmosClientTelemetryOptio
808808
/// <param name="throughputBucket">The desired throughput bucket for the client.</param>
809809
/// <returns>The current <see cref="CosmosClientBuilder"/>.</returns>
810810
/// <seealso href="https://aka.ms/cosmsodb-bucketing"/>
811-
internal CosmosClientBuilder WithThroughputBucket(int throughputBucket)
811+
#if PREVIEW
812+
public
813+
#else
814+
internal
815+
#endif
816+
CosmosClientBuilder WithThroughputBucket(int throughputBucket)
812817
{
813818
this.clientOptions.ThroughputBucket = throughputBucket;
814819
return this;

Microsoft.Azure.Cosmos/src/GatewayStoreClient.cs

Lines changed: 19 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ namespace Microsoft.Azure.Cosmos
1414
using System.Text;
1515
using System.Threading;
1616
using System.Threading.Tasks;
17+
using Microsoft.Azure.Cosmos.Routing;
1718
using Microsoft.Azure.Cosmos.Tracing.TraceData;
1819
using Microsoft.Azure.Documents;
1920
using Microsoft.Azure.Documents.Collections;
@@ -22,8 +23,9 @@ namespace Microsoft.Azure.Cosmos
2223
internal class GatewayStoreClient : TransportClient
2324
{
2425
private readonly ICommunicationEventSource eventSource;
25-
private readonly CosmosHttpClient httpClient;
26-
private readonly JsonSerializerSettings SerializerSettings;
26+
protected readonly CosmosHttpClient httpClient;
27+
protected readonly JsonSerializerSettings SerializerSettings;
28+
2729
private static readonly HttpMethod httpPatchMethod = new HttpMethod(HttpConstants.HttpMethods.Patch);
2830

2931
public GatewayStoreClient(
@@ -48,6 +50,18 @@ public async Task<DocumentServiceResponse> InvokeAsync(
4850
}
4951
}
5052

53+
public virtual Task<DocumentServiceResponse> InvokeAsync(
54+
DocumentServiceRequest request,
55+
ResourceType resourceType,
56+
Uri physicalAddress,
57+
Uri endpoint,
58+
string globalDatabaseAccountName,
59+
ClientCollectionCache clientCollectionCache,
60+
CancellationToken cancellationToken)
61+
{
62+
return this.InvokeAsync(request, resourceType, physicalAddress, cancellationToken);
63+
}
64+
5165
public static bool IsFeedRequest(OperationType requestOperationType)
5266
{
5367
return requestOperationType == OperationType.Create ||
@@ -249,7 +263,7 @@ private static async Task<Stream> BufferContentIfAvailableAsync(HttpResponseMess
249263
}
250264

251265
[SuppressMessage("Microsoft.Reliability", "CA2000:DisposeObjectsBeforeLosingScope", Justification = "Disposable object returned by method")]
252-
private async ValueTask<HttpRequestMessage> PrepareRequestMessageAsync(
266+
internal async ValueTask<HttpRequestMessage> PrepareRequestMessageAsync(
253267
DocumentServiceRequest request,
254268
Uri physicalAddress)
255269
{
@@ -347,7 +361,7 @@ private async ValueTask<HttpRequestMessage> PrepareRequestMessageAsync(
347361
{
348362
requestMessage.Properties.Add(ClientSideRequestStatisticsTraceDatum.HttpRequestRegionNameProperty, regionName);
349363
}
350-
364+
351365
return requestMessage;
352366
}
353367

@@ -367,4 +381,4 @@ private Task<HttpResponseMessage> InvokeClientAsync(
367381
request);
368382
}
369383
}
370-
}
384+
}

Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs

Lines changed: 12 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -26,16 +26,16 @@ internal class GatewayStoreModel : IStoreModelExtension, IDisposable
2626
{
2727
private static readonly string sessionConsistencyAsString = ConsistencyLevel.Session.ToString();
2828

29-
private readonly GlobalEndpointManager endpointManager;
29+
internal readonly GlobalEndpointManager endpointManager;
3030
private readonly DocumentClientEventSource eventSource;
31-
private readonly ISessionContainer sessionContainer;
32-
private readonly ConsistencyLevel defaultConsistencyLevel;
31+
internal readonly ConsistencyLevel defaultConsistencyLevel;
3332

3433
private GatewayStoreClient gatewayStoreClient;
3534

3635
// Caches to resolve the PartitionKeyRange from request. For Session Token Optimization.
37-
private ClientCollectionCache clientCollectionCache;
38-
private PartitionKeyRangeCache partitionKeyRangeCache;
36+
protected PartitionKeyRangeCache partitionKeyRangeCache;
37+
protected ClientCollectionCache clientCollectionCache;
38+
protected ISessionContainer sessionContainer;
3939

4040
public GatewayStoreModel(
4141
GlobalEndpointManager endpointManager,
@@ -113,8 +113,7 @@ public virtual async Task<AccountProperties> GetDatabaseAccountAsync(Func<ValueT
113113
}
114114

115115
long longValue;
116-
IEnumerable<string> headerValues;
117-
if (responseMessage.Headers.TryGetValues(HttpConstants.HttpHeaders.MaxMediaStorageUsageInMB, out headerValues) &&
116+
if (responseMessage.Headers.TryGetValues(HttpConstants.HttpHeaders.MaxMediaStorageUsageInMB, out IEnumerable<string> headerValues) &&
118117
(headerValues.Count() != 0))
119118
{
120119
if (long.TryParse(headerValues.First(), out longValue))
@@ -163,8 +162,8 @@ public virtual async Task<AccountProperties> GetDatabaseAccountAsync(Func<ValueT
163162
return databaseAccount;
164163
}
165164

166-
public void SetCaches(PartitionKeyRangeCache partitionKeyRangeCache,
167-
ClientCollectionCache clientCollectionCache)
165+
public void SetCaches(PartitionKeyRangeCache partitionKeyRangeCache,
166+
ClientCollectionCache clientCollectionCache)
168167
{
169168
this.clientCollectionCache = clientCollectionCache;
170169
this.partitionKeyRangeCache = partitionKeyRangeCache;
@@ -175,7 +174,7 @@ public void Dispose()
175174
this.Dispose(true);
176175
}
177176

178-
private async Task CaptureSessionTokenAndHandleSplitAsync(
177+
internal async Task CaptureSessionTokenAndHandleSplitAsync(
179178
HttpStatusCode? statusCode,
180179
SubStatusCodes subStatusCode,
181180
DocumentServiceRequest request,
@@ -445,7 +444,7 @@ internal static bool IsStoredProcedureCrudOperation(
445444
operationType != Documents.OperationType.ExecuteJavaScript;
446445
}
447446

448-
private void Dispose(bool disposing)
447+
protected virtual void Dispose(bool disposing)
449448
{
450449
if (disposing)
451450
{
@@ -466,7 +465,7 @@ private void Dispose(bool disposing)
466465
}
467466
}
468467

469-
private Uri GetEntityUri(DocumentServiceRequest entity)
468+
internal Uri GetEntityUri(DocumentServiceRequest entity)
470469
{
471470
string contentLocation = entity.Headers[HttpConstants.HttpHeaders.ContentLocation];
472471

@@ -478,7 +477,7 @@ private Uri GetEntityUri(DocumentServiceRequest entity)
478477
return new Uri(this.endpointManager.ResolveServiceEndpoint(entity), PathsHelper.GeneratePath(entity.ResourceType, entity, false));
479478
}
480479

481-
private Uri GetFeedUri(DocumentServiceRequest request)
480+
internal Uri GetFeedUri(DocumentServiceRequest request)
482481
{
483482
return new Uri(this.endpointManager.ResolveServiceEndpoint(request), PathsHelper.GeneratePath(request.ResourceType, request, true));
484483
}

Microsoft.Azure.Cosmos/src/RequestOptions/RequestOptions.cs

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -124,7 +124,12 @@ public class RequestOptions
124124
/// <see cref="RequestOptions.ThroughputBucket"/> cannot be set in RequestOptions.
125125
/// </remarks>
126126
/// <seealso href="https://aka.ms/cosmsodb-bucketing"/>
127-
internal int? ThroughputBucket { get; set; }
127+
#if PREVIEW
128+
public
129+
#else
130+
internal
131+
#endif
132+
int? ThroughputBucket { get; set; }
128133

129134
/// <summary>
130135
/// Fill the CosmosRequestMessage headers with the set properties

Microsoft.Azure.Cosmos/src/Resource/Settings/AccountProperties.cs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -241,6 +241,11 @@ private IDictionary<string, object> QueryStringToDictConverter()
241241
}
242242
}
243243

244+
/// <summary>
245+
/// This contains the thinclient endpoint value.
246+
/// </summary>
247+
internal Uri ThinClientEndpoint { get; set; }
248+
244249
/// <summary>
245250
/// This contains additional values for scenarios where the SDK is not aware of new fields.
246251
/// This ensures that if resource is read and updated none of the fields will be lost in the process.
Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
//------------------------------------------------------------
2+
// Copyright (c) Microsoft Corporation. All rights reserved.
3+
//------------------------------------------------------------
4+
namespace Microsoft.Azure.Cosmos
5+
{
6+
using System;
7+
using System.Collections.Generic;
8+
using System.Text;
9+
10+
internal static class ThinClientConstants
11+
{
12+
public const string RoutedViaProxy = "x-ms-thinclient-route-via-proxy";
13+
public const string ProxyStartEpk = "x-ms-thinclient-range-min";
14+
public const string ProxyEndEpk = "x-ms-thinclient-range-max";
15+
16+
public const string ProxyOperationType = "x-ms-thinclient-proxy-operation-type";
17+
public const string ProxyResourceType = "x-ms-thinclient-proxy-resource-type";
18+
public const string EffectivePartitionKey = "x-ms-effective-partition-key";
19+
}
20+
}
Lines changed: 168 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,168 @@
1+
//------------------------------------------------------------
2+
// Copyright (c) Microsoft Corporation. All rights reserved.
3+
//------------------------------------------------------------
4+
5+
namespace Microsoft.Azure.Cosmos
6+
{
7+
using System;
8+
using System.Collections.Concurrent;
9+
using System.Diagnostics.CodeAnalysis;
10+
using System.IO;
11+
using System.Net.Http;
12+
using System.Threading;
13+
using System.Threading.Tasks;
14+
using Microsoft.Azure.Cosmos.Core.Trace;
15+
using Microsoft.Azure.Cosmos.Routing;
16+
using Microsoft.Azure.Documents;
17+
using Newtonsoft.Json;
18+
using static Microsoft.Azure.Cosmos.ThinClientTransportSerializer;
19+
20+
/// <summary>
21+
/// A TransportClient that sends requests to proxy endpoint.
22+
/// And then processes the response back into DocumentServiceResponse objects.
23+
/// </summary>
24+
internal class ThinClientStoreClient : GatewayStoreClient
25+
{
26+
private readonly ObjectPool<BufferProviderWrapper> bufferProviderWrapperPool;
27+
28+
public ThinClientStoreClient(
29+
CosmosHttpClient httpClient,
30+
ICommunicationEventSource eventSource,
31+
JsonSerializerSettings serializerSettings = null)
32+
: base(httpClient,
33+
eventSource,
34+
serializerSettings)
35+
{
36+
this.bufferProviderWrapperPool = new ObjectPool<BufferProviderWrapper>(() => new BufferProviderWrapper());
37+
}
38+
39+
public override async Task<DocumentServiceResponse> InvokeAsync(
40+
DocumentServiceRequest request,
41+
ResourceType resourceType,
42+
Uri physicalAddress,
43+
Uri thinClientEndpoint,
44+
string globalDatabaseAccountName,
45+
ClientCollectionCache clientCollectionCache,
46+
CancellationToken cancellationToken)
47+
{
48+
using (HttpResponseMessage responseMessage = await this.InvokeClientAsync(
49+
request,
50+
resourceType,
51+
physicalAddress,
52+
thinClientEndpoint,
53+
globalDatabaseAccountName,
54+
clientCollectionCache,
55+
cancellationToken))
56+
{
57+
HttpResponseMessage proxyResponse = await ThinClientTransportSerializer.ConvertProxyResponseAsync(responseMessage);
58+
return await ThinClientStoreClient.ParseResponseAsync(proxyResponse, request.SerializerSettings ?? base.SerializerSettings, request);
59+
}
60+
}
61+
62+
internal override async Task<StoreResponse> InvokeStoreAsync(Uri baseAddress, ResourceOperation resourceOperation, DocumentServiceRequest request)
63+
{
64+
Uri physicalAddress = ThinClientStoreClient.IsFeedRequest(request.OperationType) ?
65+
HttpTransportClient.GetResourceFeedUri(resourceOperation.resourceType, baseAddress, request) :
66+
HttpTransportClient.GetResourceEntryUri(resourceOperation.resourceType, baseAddress, request);
67+
68+
using (HttpResponseMessage responseMessage = await this.InvokeClientAsync(
69+
request,
70+
resourceOperation.resourceType,
71+
physicalAddress,
72+
default,
73+
default,
74+
default,
75+
default))
76+
{
77+
return await HttpTransportClient.ProcessHttpResponse(request.ResourceAddress, string.Empty, responseMessage, physicalAddress, request);
78+
}
79+
}
80+
81+
private async ValueTask<HttpRequestMessage> PrepareRequestForProxyAsync(
82+
DocumentServiceRequest request,
83+
Uri physicalAddress,
84+
Uri thinClientEndpoint,
85+
string globalDatabaseAccountName,
86+
ClientCollectionCache clientCollectionCache)
87+
{
88+
HttpRequestMessage requestMessage = base.PrepareRequestMessageAsync(request, physicalAddress).Result;
89+
requestMessage.Version = new Version(2, 0);
90+
91+
BufferProviderWrapper bufferProviderWrapper = this.bufferProviderWrapperPool.Get();
92+
try
93+
{
94+
requestMessage.Headers.TryAddWithoutValidation(
95+
ThinClientConstants.ProxyOperationType,
96+
request.OperationType.ToOperationTypeString());
97+
98+
requestMessage.Headers.TryAddWithoutValidation(
99+
ThinClientConstants.ProxyResourceType,
100+
request.ResourceType.ToResourceTypeString());
101+
102+
Stream contentStream = await ThinClientTransportSerializer.SerializeProxyRequestAsync(
103+
bufferProviderWrapper,
104+
globalDatabaseAccountName,
105+
clientCollectionCache,
106+
requestMessage);
107+
108+
if (!contentStream.CanSeek)
109+
{
110+
throw new InvalidOperationException(
111+
$"The serializer returned a non-seekable stream ({contentStream.GetType().FullName}).");
112+
}
113+
114+
requestMessage.Content = new StreamContent(contentStream);
115+
requestMessage.Content.Headers.ContentLength = contentStream.Length;
116+
requestMessage.Headers.Clear();
117+
requestMessage.RequestUri = thinClientEndpoint;
118+
requestMessage.Method = HttpMethod.Post;
119+
120+
return requestMessage;
121+
}
122+
finally
123+
{
124+
this.bufferProviderWrapperPool.Return(bufferProviderWrapper);
125+
}
126+
}
127+
128+
private Task<HttpResponseMessage> InvokeClientAsync(
129+
DocumentServiceRequest request,
130+
ResourceType resourceType,
131+
Uri physicalAddress,
132+
Uri thinClientEndpoint,
133+
string globalDatabaseAccountName,
134+
ClientCollectionCache clientCollectionCache,
135+
CancellationToken cancellationToken)
136+
{
137+
DefaultTrace.TraceInformation("In {0}, OperationType: {1}, ResourceType: {2}", nameof(ThinClientStoreClient), request.OperationType, request.ResourceType);
138+
return base.httpClient.SendHttpAsync(
139+
() => this.PrepareRequestForProxyAsync(request, physicalAddress, thinClientEndpoint, globalDatabaseAccountName, clientCollectionCache),
140+
resourceType,
141+
HttpTimeoutPolicy.GetTimeoutPolicy(request),
142+
request.RequestContext.ClientRequestStatistics,
143+
cancellationToken);
144+
}
145+
146+
internal class ObjectPool<T>
147+
{
148+
private readonly ConcurrentBag<T> Objects;
149+
private readonly Func<T> ObjectGenerator;
150+
151+
public ObjectPool(Func<T> objectGenerator)
152+
{
153+
this.ObjectGenerator = objectGenerator ?? throw new ArgumentNullException(nameof(objectGenerator));
154+
this.Objects = new ConcurrentBag<T>();
155+
}
156+
157+
public T Get()
158+
{
159+
return this.Objects.TryTake(out T item) ? item : this.ObjectGenerator();
160+
}
161+
162+
public void Return(T item)
163+
{
164+
this.Objects.Add(item);
165+
}
166+
}
167+
}
168+
}

0 commit comments

Comments
 (0)