@azure/cosmos
Version:
Microsoft Azure Cosmos DB Service Node.js SDK for NOSQL API
308 lines • 16.2 kB
JavaScript
"use strict";
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
Object.defineProperty(exports, "__esModule", { value: true });
exports.BulkHelper = void 0;
const ClientUtils_js_1 = require("../client/ClientUtils.js");
const constants_js_1 = require("../common/constants.js");
const helper_js_1 = require("../common/helper.js");
const statusCodes_js_1 = require("../common/statusCodes.js");
const DiagnosticNodeInternal_js_1 = require("../diagnostics/DiagnosticNodeInternal.js");
const PartitionKeyInternal_js_1 = require("../documents/PartitionKeyInternal.js");
const index_js_1 = require("../index.js");
const bulkExecutionRetryPolicy_js_1 = require("../retry/bulkExecutionRetryPolicy.js");
const resourceThrottleRetryPolicy_js_1 = require("../retry/resourceThrottleRetryPolicy.js");
const batch_js_1 = require("../utils/batch.js");
const diagnostics_js_1 = require("../utils/diagnostics.js");
const hash_js_1 = require("../utils/hashing/hash.js");
const HelperPerPartition_js_1 = require("./HelperPerPartition.js");
const index_js_2 = require("./index.js");
/**
* BulkHelper for bulk operations in a container.
* It maintains one @see {@link HelperPerPartition} for each Partition Key Range, which allows independent execution of requests. Queue based limiters @see {@link LimiterQueue}
* rate limit requestsbat the helper / Partition Key Range level, this means that we can send parallel and independent requests to different Partition Key Ranges, but for the same Range, requests
* will be limited. Two callback implementations define how a particular request should be executed, and how operations should be retried. When the helper dispatches a batch
* the batch will create a request and call the execute callback (executeRequest), if conditions are met, it might call the retry callback (reBatchOperation).
* @hidden
*/
class BulkHelper {
container;
clientContext;
partitionKeyRangeCache;
helpersByPartitionKeyRangeId;
options;
partitionKeyDefinition;
partitionKeyDefinitionPromise;
isCancelled;
processedOperationCountRef = { count: 0 };
operationPromisesList = [];
congestionControlTimer;
congestionControlDelayInMs = 1000;
staleRidError;
operationsPerSleep = 100; // Number of operations to add per sleep
intervalForPartialBatchInMs = 1000; // Sleep interval before adding partial batch to dispatch queue
/**
* @internal
*/
constructor(container, clientContext, partitionKeyRangeCache, options) {
this.container = container;
this.clientContext = clientContext;
this.partitionKeyRangeCache = partitionKeyRangeCache;
this.helpersByPartitionKeyRangeId = new Map();
this.options = options;
this.executeRequest = this.executeRequest.bind(this);
this.reBatchOperation = this.reBatchOperation.bind(this);
this.refreshPartitionKeyRangeCache = this.refreshPartitionKeyRangeCache.bind(this);
this.isCancelled = false;
this.runCongestionControlTimer();
}
/**
* adds operation(s) to the helper
* @param operationInput - bulk operation or list of bulk operations
*/
async execute(operationInput) {
const addOperationPromises = [];
const minimalPause = 0; // minimal pause (0 ms) inserted periodically during processing.
try {
for (let i = 0; i < operationInput.length; i++) {
// After every 100 operations,sleep of 0 ms is added to allow the event loop to process any pending
// callbacks/tasks such as fetching partition key definition and dispatching batches from queue. This helps
// to prevent blocking and improves overall responsiveness.
if (i % this.operationsPerSleep === 0) {
await (0, helper_js_1.sleep)(minimalPause);
}
addOperationPromises.push(this.addOperation(operationInput[i], i));
}
await Promise.allSettled(addOperationPromises);
// After processing all operations via addOperation, it's possible that the current batch in each helper is not completely full.
// In such cases, addPartialBatchToQueue is called to ensure that all the operations are added to the dispatch queue.
// while loop below waits until the count of processed operations equals the number of input operations. This is necessary because
// some operations might fail and then again get added to current batch for retry.
while (this.processedOperationCountRef.count < operationInput.length) {
this.helpersByPartitionKeyRangeId.forEach((helper) => {
helper.addPartialBatchToQueue();
});
// Pause for 1000 ms to give pending operations chance to accumulate into a batch to avoid sending multiple small batches.
await (0, helper_js_1.sleep)(this.intervalForPartialBatchInMs);
}
}
finally {
if (this.congestionControlTimer) {
clearInterval(this.congestionControlTimer);
}
}
const settledResults = await Promise.allSettled(this.operationPromisesList);
if (this.isCancelled && this.staleRidError) {
throw this.staleRidError;
}
const bulkOperationResults = settledResults.map((result) => result.status === "fulfilled" ? result.value : result.reason);
// Formatting result: if an error is present, removing the stack trace details.
const formattedResults = bulkOperationResults.map((result) => {
if (result && result.error) {
const { stack, ...otherProps } = result.error;
const trimmedError = { message: result.error.message, ...otherProps };
return {
...result,
error: trimmedError,
};
}
return result;
});
return formattedResults;
}
async addOperation(operation, idx) {
if (this.isCancelled) {
return;
}
if (!operation) {
this.operationPromisesList[idx] = Promise.resolve({
operationInput: operation,
error: Object.assign(new index_js_1.ErrorResponse("Operation cannot be null or undefined."), {
code: statusCodes_js_1.StatusCodes.InternalServerError,
}),
});
return;
}
// Checks for id and partition key in input body
if (operation.operationType === "Create" ||
operation.operationType === "Upsert" ||
operation.operationType === "Replace") {
if (!operation.resourceBody.id) {
this.operationPromisesList[idx] = Promise.resolve({
operationInput: operation,
error: Object.assign(new index_js_1.ErrorResponse(`Operation resource body must have an 'id' for ${operation.operationType} operations.`), { code: statusCodes_js_1.StatusCodes.InternalServerError }),
});
this.processedOperationCountRef.count++;
return;
}
}
if (operation.partitionKey === undefined) {
this.operationPromisesList[idx] = Promise.resolve({
operationInput: operation,
error: Object.assign(new index_js_1.ErrorResponse(`PartitionKey is required for ${operation.operationType} operations.`), { code: statusCodes_js_1.StatusCodes.InternalServerError }),
});
this.processedOperationCountRef.count++;
return;
}
let operationError;
let diagnosticNode;
let unencryptedOperation;
let partitionKeyRangeId;
try {
diagnosticNode = new DiagnosticNodeInternal_js_1.DiagnosticNodeInternal(this.clientContext.diagnosticLevel, DiagnosticNodeInternal_js_1.DiagnosticNodeType.CLIENT_REQUEST_NODE, null);
// Ensure partition key definition is available.
if (!this.partitionKeyDefinition) {
if (!this.partitionKeyDefinitionPromise) {
this.partitionKeyDefinitionPromise = (async () => {
try {
const partitionKeyDefinition = await (0, ClientUtils_js_1.readPartitionKeyDefinition)(diagnosticNode, this.container);
this.partitionKeyDefinition = partitionKeyDefinition;
return partitionKeyDefinition;
}
finally {
this.partitionKeyDefinitionPromise = null;
}
})();
}
await this.partitionKeyDefinitionPromise;
}
unencryptedOperation = (0, helper_js_1.copyObject)(operation);
// If encryption is enabled, encrypt the operation input.
if (this.clientContext.enableEncryption) {
operation = (0, helper_js_1.copyObject)(operation);
await this.container.checkAndInitializeEncryption();
diagnosticNode.beginEncryptionDiagnostics(constants_js_1.Constants.Encryption.DiagnosticsEncryptOperation);
const { operation: encryptedOp, totalPropertiesEncryptedCount } = await (0, batch_js_1.encryptOperationInput)(this.container.encryptionProcessor, operation, 0);
operation = encryptedOp;
diagnosticNode.endEncryptionDiagnostics(constants_js_1.Constants.Encryption.DiagnosticsEncryptOperation, totalPropertiesEncryptedCount);
}
// Resolve the partition key range id.
partitionKeyRangeId = await this.resolvePartitionKeyRangeId(operation, diagnosticNode);
}
catch (error) {
operationError = error;
}
// Get helper & context.
const helperForPartition = this.getHelperForPKRange(partitionKeyRangeId);
const retryPolicy = this.getRetryPolicy();
const context = new index_js_2.ItemOperationContext(partitionKeyRangeId, retryPolicy, diagnosticNode);
const itemOperation = {
unencryptedOperationInput: unencryptedOperation,
operationInput: operation,
operationContext: context,
};
// Assign the promise (ensuring position matches input order)
this.operationPromisesList[idx] = context.operationPromise;
if (operationError) {
const response = {
operationInput: unencryptedOperation,
error: Object.assign(new index_js_1.ErrorResponse(operationError.message), {
code: statusCodes_js_1.StatusCodes.InternalServerError,
diagnostics: diagnosticNode?.toDiagnostic(this.clientContext.getClientConfig()),
}),
};
context.fail(response);
this.processedOperationCountRef.count++;
return;
}
// Add the operation to the helper.
return helperForPartition.add(itemOperation);
}
async resolvePartitionKeyRangeId(operation, diagnosticNode) {
const partitionKeyRanges = (await this.partitionKeyRangeCache.onCollectionRoutingMap(this.container.url, diagnosticNode)).getOrderedParitionKeyRanges();
const partitionKey = (0, PartitionKeyInternal_js_1.convertToInternalPartitionKey)(operation.partitionKey);
const hashedKey = (0, hash_js_1.hashPartitionKey)(partitionKey, this.partitionKeyDefinition);
const matchingRange = partitionKeyRanges.find((range) => (0, batch_js_1.isKeyInRange)(range.minInclusive, range.maxExclusive, hashedKey));
if (!matchingRange) {
throw new Error("No matching partition key range found for the operation.");
}
return matchingRange.id;
}
getRetryPolicy() {
const nextRetryPolicy = new resourceThrottleRetryPolicy_js_1.ResourceThrottleRetryPolicy(this.clientContext.getRetryOptions());
return new bulkExecutionRetryPolicy_js_1.BulkExecutionRetryPolicy(nextRetryPolicy);
}
async executeRequest(operations, diagnosticNode) {
if (this.isCancelled) {
throw new index_js_1.ErrorResponse("Bulk execution cancelled due to a previous error.");
}
if (!operations.length)
return;
const pkRangeId = operations[0].operationContext.pkRangeId;
const path = (0, helper_js_1.getPathFromLink)(this.container.url, constants_js_1.ResourceType.item);
const requestBody = [];
for (const itemOperation of operations) {
requestBody.push(this.prepareOperation(itemOperation.operationInput));
}
if (!this.options.containerRid) {
this.options.containerRid = this.container._rid;
}
try {
const response = await (0, diagnostics_js_1.addDiagnosticChild)(async (childNode) => this.clientContext.bulk({
body: requestBody,
partitionKeyRangeId: pkRangeId,
path: path,
resourceId: this.container.url,
options: this.options,
diagnosticNode: childNode,
}), diagnosticNode, DiagnosticNodeInternal_js_1.DiagnosticNodeType.BATCH_REQUEST);
if (!response) {
throw new index_js_1.ErrorResponse("Failed to fetch bulk response.");
}
return index_js_2.BulkResponse.fromResponseMessage(response, operations);
}
catch (error) {
if (this.clientContext.enableEncryption) {
try {
await this.container.throwIfRequestNeedsARetryPostPolicyRefresh(error);
}
catch (err) {
await this.cancelExecution(err);
return index_js_2.BulkResponse.createEmptyResponse(operations, 0, 0, {});
}
}
return index_js_2.BulkResponse.fromResponseMessage(error, operations);
}
}
prepareOperation(operationInput) {
operationInput.partitionKey = (0, PartitionKeyInternal_js_1.convertToInternalPartitionKey)(operationInput.partitionKey);
return {
...operationInput,
partitionKey: JSON.stringify(operationInput.partitionKey),
};
}
async reBatchOperation(operation, diagnosticNode) {
const partitionKeyRangeId = await this.resolvePartitionKeyRangeId(operation.operationInput, diagnosticNode);
operation.operationContext.updatePKRangeId(partitionKeyRangeId);
const helper = this.getHelperForPKRange(partitionKeyRangeId);
await helper.add(operation);
}
async cancelExecution(error) {
this.isCancelled = true;
this.staleRidError = error;
for (const helper of this.helpersByPartitionKeyRangeId.values()) {
await helper.dispose();
}
this.helpersByPartitionKeyRangeId.clear();
}
getHelperForPKRange(pkRangeId) {
if (this.helpersByPartitionKeyRangeId.has(pkRangeId)) {
return this.helpersByPartitionKeyRangeId.get(pkRangeId);
}
const newHelper = new HelperPerPartition_js_1.HelperPerPartition(this.executeRequest, this.reBatchOperation, this.refreshPartitionKeyRangeCache, this.clientContext.diagnosticLevel, this.clientContext.enableEncryption, this.clientContext.getClientConfig(), this.container.encryptionProcessor, this.processedOperationCountRef);
this.helpersByPartitionKeyRangeId.set(pkRangeId, newHelper);
return newHelper;
}
runCongestionControlTimer() {
this.congestionControlTimer = setInterval(() => {
this.helpersByPartitionKeyRangeId.forEach((helper) => {
helper.runCongestionAlgorithm();
});
}, this.congestionControlDelayInMs);
}
async refreshPartitionKeyRangeCache(diagnosticNode) {
await this.partitionKeyRangeCache.onCollectionRoutingMap(this.container.url, diagnosticNode, true);
}
}
exports.BulkHelper = BulkHelper;
//# sourceMappingURL=BulkHelper.js.map