@azure/cosmos
Version:
Microsoft Azure Cosmos DB Service Node.js SDK for NOSQL API
154 lines • 7.39 kB
JavaScript
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
import { Constants } from "../common/constants.js";
import { StatusCodes } from "../common/statusCodes.js";
import { DiagnosticNodeInternal, DiagnosticNodeType, } from "../diagnostics/DiagnosticNodeInternal.js";
import { ErrorResponse } from "../index.js";
import { calculateObjectSizeInBytes } from "../utils/batch.js";
import { getCurrentTimestampInMs } from "../utils/time.js";
/**
* Maintains a batch of operations and dispatches it as a unit of work.
* Execution of the request is done by the @see {@link ExecuteCallback} and retry is done by the @see {@link RetryCallback}.
* @hidden
*/
export class Batcher {
batchOperationsList;
currentSize;
toBeDispatched;
executor;
retrier;
diagnosticLevel;
encryptionEnabled;
encryptionProcessor;
clientConfigDiagnostics;
limiter;
processedOperationCountRef;
constructor(limiter, executor, retrier, diagnosticLevel, encryptionEnabled, clientConfig, encryptionProcessor, processedOperationCountRef) {
this.limiter = limiter;
this.batchOperationsList = [];
this.executor = executor;
this.retrier = retrier;
this.diagnosticLevel = diagnosticLevel;
this.encryptionEnabled = encryptionEnabled;
this.encryptionProcessor = encryptionProcessor;
this.clientConfigDiagnostics = clientConfig;
this.currentSize = 0;
this.toBeDispatched = false;
this.processedOperationCountRef = processedOperationCountRef;
}
/**
* Attempts to add an operation to the current batch.
* Returns false if the batch is full or already dispatched.
*/
tryAdd(operation) {
if (this.toBeDispatched) {
return false;
}
if (!operation) {
throw new ErrorResponse("Operation is not defined");
}
if (!operation.operationContext) {
throw new ErrorResponse("Operation context is not defined");
}
if (this.batchOperationsList.length === Constants.MaxBulkOperationsCount) {
return false;
}
const currentOperationSize = calculateObjectSizeInBytes(operation.operationInput);
if (this.batchOperationsList.length > 0 &&
this.currentSize + currentOperationSize > Constants.DefaultMaxBulkRequestBodySizeInBytes) {
return false;
}
this.currentSize += currentOperationSize;
this.batchOperationsList.push(operation);
return true;
}
isEmpty() {
return this.batchOperationsList.length === 0;
}
/**
* Dispatches the current batch of operations.
* Handles retries for failed operations and updates the ordered response.
*/
async dispatch(partitionMetric) {
this.toBeDispatched = true;
const startTime = getCurrentTimestampInMs();
const diagnosticNode = new DiagnosticNodeInternal(this.diagnosticLevel, DiagnosticNodeType.BATCH_REQUEST, null);
try {
const response = await this.executor(this.batchOperationsList, diagnosticNode);
const hasThrottles = 1;
const noThrottle = 0;
const numThrottle = response?.results?.some((result) => "code" in result && result.code === StatusCodes.TooManyRequests)
? hasThrottles
: noThrottle;
const splitOrMerge = response?.results?.some((result) => "code" in result && result.code === StatusCodes.Gone)
? true
: false;
if (splitOrMerge) {
await this.limiter.pauseAndClear(StatusCodes.Gone, diagnosticNode);
}
partitionMetric.add(this.batchOperationsList.length, getCurrentTimestampInMs() - startTime, numThrottle);
for (let i = 0; i < response.operations.length; i++) {
const operation = response.operations[i];
const bulkOperationResult = response.results[i];
if (bulkOperationResult instanceof ErrorResponse) {
const shouldRetry = await operation.operationContext.retryPolicy.shouldRetry(bulkOperationResult, operation.operationContext.diagnosticNode);
if (shouldRetry) {
await this.retrier(operation, operation.operationContext.diagnosticNode);
continue;
}
}
try {
if (this.encryptionEnabled && bulkOperationResult.resourceBody) {
operation.operationContext.diagnosticNode.beginEncryptionDiagnostics(Constants.Encryption.DiagnosticsDecryptOperation);
const { body: decryptedBody, propertiesDecryptedCount } = await this.encryptionProcessor.decrypt(bulkOperationResult.resourceBody);
bulkOperationResult.resourceBody = decryptedBody;
operation.operationContext.diagnosticNode.endEncryptionDiagnostics(Constants.Encryption.DiagnosticsDecryptOperation, propertiesDecryptedCount);
}
}
catch (error) {
// if decryption fails after successful write operation, fail the operation with internal server error
if (operation.operationInput.operationType !== "Read") {
const decryptionError = new ErrorResponse(`Item ${operation.operationInput.operationType} operation was successful but response decryption failed: + ${error.message}`);
decryptionError.code = StatusCodes.ServiceUnavailable;
throw decryptionError;
}
}
operation.operationContext.addDiagnosticChild(diagnosticNode);
bulkOperationResult.diagnostics = operation.operationContext.diagnosticNode.toDiagnostic(this.clientConfigDiagnostics);
const bulkItemResponse = {
operationInput: operation.unencryptedOperationInput,
};
if (bulkOperationResult instanceof ErrorResponse) {
bulkItemResponse.error = bulkOperationResult;
}
else {
bulkItemResponse.response = bulkOperationResult;
}
operation.operationContext.complete(bulkItemResponse);
this.processedOperationCountRef.count++;
}
}
catch (error) {
// Mark all operations in the batch as failed
for (const operation of this.batchOperationsList) {
const response = {
operationInput: operation.operationInput,
error: Object.assign(new ErrorResponse(error.message), {
code: StatusCodes.InternalServerError,
diagnostics: operation.operationContext.diagnosticNode.toDiagnostic(this.clientConfigDiagnostics),
}),
};
operation.operationContext.fail(response);
this.processedOperationCountRef.count++;
}
}
finally {
// Clean up batch state
this.batchOperationsList = [];
}
}
getOperations() {
return this.batchOperationsList;
}
}
//# sourceMappingURL=Batcher.js.map