private async Task <PartitionKeyRangeBatchExecutionResult> ExecuteAsync(
PartitionKeyRangeServerBatchRequest serverRequest,
CancellationToken cancellationToken)
{
CosmosDiagnosticsContext diagnosticsContext = new CosmosDiagnosticsContextCore();
CosmosDiagnosticScope limiterScope = diagnosticsContext.CreateScope("BatchAsyncContainerExecutor.Limiter");
SemaphoreSlim limiter = this.GetOrAddLimiterForPartitionKeyRange(serverRequest.PartitionKeyRangeId);
using (await limiter.UsingWaitAsync(cancellationToken))
{
limiterScope.Dispose();
using (Stream serverRequestPayload = serverRequest.TransferBodyStream())
{
Debug.Assert(serverRequestPayload != null, "Server request payload expected to be non-null");
ResponseMessage responseMessage = await this.cosmosClientContext.ProcessResourceOperationStreamAsync(
this.cosmosContainer.LinkUri,
ResourceType.Document,
OperationType.Batch,
new RequestOptions(),
cosmosContainerCore : this.cosmosContainer,
partitionKey : null,
streamPayload : serverRequestPayload,
requestEnricher : requestMessage => BatchAsyncContainerExecutor.AddHeadersToRequestMessage(requestMessage, serverRequest.PartitionKeyRangeId),
diagnosticsContext : diagnosticsContext,
cancellationToken : cancellationToken).ConfigureAwait(false);
using (diagnosticsContext.CreateScope("BatchAsyncContainerExecutor.ToResponse"))
{
TransactionalBatchResponse serverResponse = await TransactionalBatchResponse.FromResponseMessageAsync(
responseMessage,
serverRequest,
this.cosmosClientContext.SerializerCore,
shouldPromoteOperationStatus : true,
shouldPerformDecryption : false,
cancellationToken).ConfigureAwait(false);
return(new PartitionKeyRangeBatchExecutionResult(serverRequest.PartitionKeyRangeId, serverRequest.Operations, serverResponse));
}
}
}
}
private async Task <PartitionKeyRangeBatchExecutionResult> ExecuteAsync(
PartitionKeyRangeServerBatchRequest serverRequest,
ITrace trace,
CancellationToken cancellationToken)
{
SemaphoreSlim limiter = this.GetOrAddLimiterForPartitionKeyRange(serverRequest.PartitionKeyRangeId);
using (await limiter.UsingWaitAsync(trace, cancellationToken))
{
using (Stream serverRequestPayload = serverRequest.TransferBodyStream())
{
Debug.Assert(serverRequestPayload != null, "Server request payload expected to be non-null");
ResponseMessage responseMessage = await this.cosmosClientContext.ProcessResourceOperationStreamAsync(
this.cosmosContainer.LinkUri,
ResourceType.Document,
OperationType.Batch,
new RequestOptions(),
cosmosContainerCore : this.cosmosContainer,
feedRange : null,
streamPayload : serverRequestPayload,
requestEnricher : requestMessage => BatchAsyncContainerExecutor.AddHeadersToRequestMessage(requestMessage, serverRequest),
trace : trace,
cancellationToken : cancellationToken).ConfigureAwait(false);
TransactionalBatchResponse serverResponse = await TransactionalBatchResponse.FromResponseMessageAsync(
responseMessage,
serverRequest,
this.cosmosClientContext.SerializerCore,
shouldPromoteOperationStatus : true,
trace,
cancellationToken).ConfigureAwait(false);
return(new PartitionKeyRangeBatchExecutionResult(
serverRequest.PartitionKeyRangeId,
serverRequest.Operations,
serverResponse));
}
}
}
/// <summary>
/// Makes a single batch request to the server.
/// </summary>
/// <param name="serverRequest">A server request with a set of operations on items.</param>
/// <param name="trace">The trace.</param>
/// <param name="cancellationToken"><see cref="CancellationToken"/> representing request cancellation.</param>
/// <returns>Response from the server.</returns>
private async Task <TransactionalBatchResponse> ExecuteServerRequestAsync(
SinglePartitionKeyServerBatchRequest serverRequest,
ITrace trace,
CancellationToken cancellationToken)
{
using (ITrace executeBatchTrace = trace.StartChild("Execute Batch Request", TraceComponent.Batch, Tracing.TraceLevel.Info))
{
using (Stream serverRequestPayload = serverRequest.TransferBodyStream())
{
Debug.Assert(serverRequestPayload != null, "Server request payload expected to be non-null");
ResponseMessage responseMessage = await this.clientContext.ProcessResourceOperationStreamAsync(
this.container.LinkUri,
ResourceType.Document,
OperationType.Batch,
this.batchOptions,
this.container,
serverRequest.PartitionKey.HasValue?new FeedRangePartitionKey(serverRequest.PartitionKey.Value) : null,
serverRequestPayload,
requestMessage =>
{
requestMessage.Headers.Add(HttpConstants.HttpHeaders.IsBatchRequest, bool.TrueString);
requestMessage.Headers.Add(HttpConstants.HttpHeaders.IsBatchAtomic, bool.TrueString);
requestMessage.Headers.Add(HttpConstants.HttpHeaders.IsBatchOrdered, bool.TrueString);
},
executeBatchTrace,
cancellationToken);
return(await TransactionalBatchResponse.FromResponseMessageAsync(
responseMessage,
serverRequest,
this.clientContext.SerializerCore,
shouldPromoteOperationStatus : true,
executeBatchTrace,
cancellationToken));
}
}
}
/// <summary>
/// Initializes a new instance of the <see cref="PartitionKeyRangeBatchResponse"/> class.
/// </summary>
/// <param name="originalOperationsCount">Original operations that generated the server responses.</param>
/// <param name="serverResponse">Response from the server.</param>
/// <param name="serializerCore">Serializer to deserialize response resource body streams.</param>
internal PartitionKeyRangeBatchResponse(
int originalOperationsCount,
TransactionalBatchResponse serverResponse,
CosmosSerializerCore serializerCore)
{
this.StatusCode = serverResponse.StatusCode;
this.serverResponse = serverResponse;
this.resultsByOperationIndex = new TransactionalBatchOperationResult[originalOperationsCount];
StringBuilder errorMessageBuilder = new StringBuilder();
List <ItemBatchOperation> itemBatchOperations = new List <ItemBatchOperation>();
// We expect number of results == number of operations here
for (int index = 0; index < serverResponse.Operations.Count; index++)
{
int operationIndex = serverResponse.Operations[index].OperationIndex;
if (this.resultsByOperationIndex[operationIndex] == null ||
this.resultsByOperationIndex[operationIndex].StatusCode == (HttpStatusCode)StatusCodes.TooManyRequests)
{
this.resultsByOperationIndex[operationIndex] = serverResponse[index];
}
}
itemBatchOperations.AddRange(serverResponse.Operations);
this.Headers = serverResponse.Headers;
if (!string.IsNullOrEmpty(serverResponse.ErrorMessage))
{
errorMessageBuilder.AppendFormat("{0}; ", serverResponse.ErrorMessage);
}
this.ErrorMessage = errorMessageBuilder.Length > 2 ? errorMessageBuilder.ToString(0, errorMessageBuilder.Length - 2) : null;
this.Operations = itemBatchOperations;
this.SerializerCore = serializerCore;
}
/// <summary>
/// Makes a single batch request to the server.
/// </summary>
/// <param name="serverRequest">A server request with a set of operations on items.</param>
/// <param name="cancellationToken"><see cref="CancellationToken"/> representing request cancellation.</param>
/// <returns>Response from the server.</returns>
private async Task <TransactionalBatchResponse> ExecuteServerRequestAsync(
SinglePartitionKeyServerBatchRequest serverRequest,
CancellationToken cancellationToken)
{
using (Stream serverRequestPayload = serverRequest.TransferBodyStream())
{
Debug.Assert(serverRequestPayload != null, "Server request payload expected to be non-null");
ResponseMessage responseMessage = await this.clientContext.ProcessResourceOperationStreamAsync(
this.container.LinkUri,
ResourceType.Document,
OperationType.Batch,
this.batchOptions,
this.container,
serverRequest.PartitionKey,
serverRequestPayload,
requestMessage =>
{
requestMessage.Headers.Add(HttpConstants.HttpHeaders.IsBatchRequest, bool.TrueString);
requestMessage.Headers.Add(HttpConstants.HttpHeaders.IsBatchAtomic, bool.TrueString);
requestMessage.Headers.Add(HttpConstants.HttpHeaders.IsBatchOrdered, bool.TrueString);
},
diagnosticsContext : this.diagnosticsContext,
cancellationToken);
using (this.diagnosticsContext.CreateScope("TransactionalBatchResponse"))
{
return(await TransactionalBatchResponse.FromResponseMessageAsync(
responseMessage,
serverRequest,
this.clientContext.SerializerCore,
shouldPromoteOperationStatus : true,
shouldPerformDecryption : true,
cancellationToken));
}
}
}
private static async Task <TransactionalBatchResponse> PopulateFromContentAsync(
Stream content,
ResponseMessage responseMessage,
ServerBatchRequest serverRequest,
CosmosSerializerCore serializer,
ITrace trace,
bool shouldPromoteOperationStatus)
{
List <TransactionalBatchOperationResult> results = new List <TransactionalBatchOperationResult>();
// content is ensured to be seekable in caller.
int resizerInitialCapacity = (int)content.Length;
Result res = await content.ReadRecordIOAsync(
(Func <ReadOnlyMemory <byte>, Result>)(record =>
{
Result r = TransactionalBatchOperationResult.ReadOperationResult(record, out TransactionalBatchOperationResult operationResult);
if (r != Result.Success)
{
return(r);
}
operationResult.Trace = trace;
results.Add(operationResult);
return(r);
}),
resizer : new MemorySpanResizer <byte>(resizerInitialCapacity));
if (res != Result.Success)
{
return(null);
}
HttpStatusCode responseStatusCode = responseMessage.StatusCode;
SubStatusCodes responseSubStatusCode = responseMessage.Headers.SubStatusCode;
// Promote the operation error status as the Batch response error status if we have a MultiStatus response
// to provide users with status codes they are used to.
if ((int)responseMessage.StatusCode == (int)StatusCodes.MultiStatus &&
shouldPromoteOperationStatus)
{
foreach (TransactionalBatchOperationResult result in results)
{
if ((int)result.StatusCode != (int)StatusCodes.FailedDependency && (int)result.StatusCode >= (int)StatusCodes.StartingErrorCode)
{
responseStatusCode = result.StatusCode;
responseSubStatusCode = result.SubStatusCode;
break;
}
}
}
TransactionalBatchResponse response = new TransactionalBatchResponse(
responseStatusCode,
responseSubStatusCode,
responseMessage.ErrorMessage,
responseMessage.Headers,
trace,
serverRequest.Operations,
serializer)
{
results = results
};
return(response);
}
internal static async Task <TransactionalBatchResponse> FromResponseMessageAsync(
ResponseMessage responseMessage,
ServerBatchRequest serverRequest,
CosmosSerializerCore serializer,
bool shouldPromoteOperationStatus,
ITrace trace,
CancellationToken cancellationToken)
{
using (ITrace createResponseTrace = trace.StartChild("Create Trace", TraceComponent.Batch, TraceLevel.Info))
{
using (responseMessage)
{
TransactionalBatchResponse response = null;
if (responseMessage.Content != null)
{
Stream content = responseMessage.Content;
// Shouldn't be the case practically, but handle it for safety.
if (!responseMessage.Content.CanSeek)
{
content = new MemoryStream();
await responseMessage.Content.CopyToAsync(content);
}
if (content.ReadByte() == (int)HybridRowVersion.V1)
{
content.Position = 0;
response = await TransactionalBatchResponse.PopulateFromContentAsync(
content,
responseMessage,
serverRequest,
serializer,
trace,
shouldPromoteOperationStatus);
if (response == null)
{
// Convert any payload read failures as InternalServerError
response = new TransactionalBatchResponse(
HttpStatusCode.InternalServerError,
SubStatusCodes.Unknown,
ClientResources.ServerResponseDeserializationFailure,
responseMessage.Headers,
trace,
serverRequest.Operations,
serializer);
}
}
}
if (response == null)
{
response = new TransactionalBatchResponse(
responseMessage.StatusCode,
responseMessage.Headers.SubStatusCode,
responseMessage.ErrorMessage,
responseMessage.Headers,
trace,
serverRequest.Operations,
serializer);
}
if (response.results == null || response.results.Count != serverRequest.Operations.Count)
{
if (responseMessage.IsSuccessStatusCode)
{
// Server should be guaranteeing number of results equal to operations when
// batch request is successful - so fail as InternalServerError if this is not the case.
response = new TransactionalBatchResponse(
HttpStatusCode.InternalServerError,
SubStatusCodes.Unknown,
ClientResources.InvalidServerResponse,
responseMessage.Headers,
trace,
serverRequest.Operations,
serializer);
}
// When the overall response status code is TooManyRequests, propagate the RetryAfter into the individual operations.
int retryAfterMilliseconds = 0;
if ((int)responseMessage.StatusCode == (int)StatusCodes.TooManyRequests)
{
if (!responseMessage.Headers.TryGetValue(HttpConstants.HttpHeaders.RetryAfterInMilliseconds, out string retryAfter) ||
retryAfter == null ||
!int.TryParse(retryAfter, out retryAfterMilliseconds))
{
retryAfterMilliseconds = 0;
}
}
response.CreateAndPopulateResults(serverRequest.Operations, trace, retryAfterMilliseconds);
}
return(response);
}
}
}
internal static async Task <TransactionalBatchResponse> FromResponseMessageAsync(
ResponseMessage responseMessage,
ServerBatchRequest serverRequest,
CosmosSerializerCore serializer,
bool shouldPromoteOperationStatus,
bool shouldPerformDecryption,
CancellationToken cancellationToken)
{
using (responseMessage)
{
TransactionalBatchResponse response = null;
if (responseMessage.Content != null)
{
Stream content = responseMessage.Content;
// Shouldn't be the case practically, but handle it for safety.
if (!responseMessage.Content.CanSeek)
{
content = new MemoryStream();
await responseMessage.Content.CopyToAsync(content);
}
if (content.ReadByte() == (int)HybridRowVersion.V1)
{
content.Position = 0;
response = await TransactionalBatchResponse.PopulateFromContentAsync(
content,
responseMessage,
serverRequest,
serializer,
shouldPromoteOperationStatus);
if (response == null)
{
// Convert any payload read failures as InternalServerError
response = new TransactionalBatchResponse(
HttpStatusCode.InternalServerError,
SubStatusCodes.Unknown,
ClientResources.ServerResponseDeserializationFailure,
responseMessage.Headers.RequestCharge,
responseMessage.Headers.RetryAfter,
responseMessage.Headers.ActivityId,
responseMessage.DiagnosticsContext,
serverRequest.Operations,
serializer);
}
}
}
if (response == null)
{
response = new TransactionalBatchResponse(
responseMessage.StatusCode,
responseMessage.Headers.SubStatusCode,
responseMessage.ErrorMessage,
responseMessage.Headers.RequestCharge,
responseMessage.Headers.RetryAfter,
responseMessage.Headers.ActivityId,
responseMessage.DiagnosticsContext,
serverRequest.Operations,
serializer);
}
if (response.results == null || response.results.Count != serverRequest.Operations.Count)
{
if (responseMessage.IsSuccessStatusCode)
{
// Server should be guaranteeing number of results equal to operations when
// batch request is successful - so fail as InternalServerError if this is not the case.
response = new TransactionalBatchResponse(
HttpStatusCode.InternalServerError,
SubStatusCodes.Unknown,
ClientResources.InvalidServerResponse,
responseMessage.Headers.RequestCharge,
responseMessage.Headers.RetryAfter,
responseMessage.Headers.ActivityId,
responseMessage.DiagnosticsContext,
serverRequest.Operations,
serializer);
}
// When the overall response status code is TooManyRequests, propagate the RetryAfter into the individual operations.
int retryAfterMilliseconds = 0;
if ((int)responseMessage.StatusCode == (int)StatusCodes.TooManyRequests)
{
if (!responseMessage.Headers.TryGetValue(HttpConstants.HttpHeaders.RetryAfterInMilliseconds, out string retryAfter) ||
retryAfter == null ||
!int.TryParse(retryAfter, out retryAfterMilliseconds))
{
retryAfterMilliseconds = 0;
}
}
response.CreateAndPopulateResults(serverRequest.Operations, retryAfterMilliseconds);
}
else if (shouldPerformDecryption)
{
for (int index = 0; index < serverRequest.Operations.Count; index++)
{
ContainerCore containerCore = serverRequest.Operations[index].ContainerCore;
TransactionalBatchOperationResult result = response.results[index];
result.ResourceStream = await containerCore.ClientContext.DecryptItemAsync(
result.ResourceStream,
(DatabaseCore)containerCore.Database,
responseMessage.DiagnosticsContext,
cancellationToken);
}
}
return(response);
}
}
