@azure/cosmos
Version:
Microsoft Azure Cosmos DB Service Node.js SDK for NOSQL API
791 lines • 38.9 kB
JavaScript
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
import PriorityQueue from "priorityqueuejs";
import semaphore from "semaphore";
import { StatusCodes, SubStatusCodes } from "../common/statusCodes.js";
import { ErrorResponse } from "../request/ErrorResponse.js";
import { QueryRange } from "../routing/QueryRange.js";
import { SmartRoutingMapProvider } from "../routing/smartRoutingMapProvider.js";
import { DocumentProducer } from "./documentProducer.js";
import { getInitialHeader, mergeHeaders } from "./headerUtils.js";
import { RidSkipCountFilter } from "./queryFilteringStrategy/RidSkipCountFilter.js";
import { DiagnosticNodeInternal, DiagnosticNodeType, } from "../diagnostics/DiagnosticNodeInternal.js";
import { createParallelQueryResult } from "./parallelQueryResult.js";
/** @hidden */
export var ParallelQueryExecutionContextBaseStates;
(function (ParallelQueryExecutionContextBaseStates) {
ParallelQueryExecutionContextBaseStates["started"] = "started";
ParallelQueryExecutionContextBaseStates["inProgress"] = "inProgress";
ParallelQueryExecutionContextBaseStates["ended"] = "ended";
})(ParallelQueryExecutionContextBaseStates || (ParallelQueryExecutionContextBaseStates = {}));
/** @hidden */
export class ParallelQueryExecutionContextBase {
clientContext;
collectionLink;
query;
options;
partitionedQueryExecutionInfo;
correlatedActivityId;
rangeManager;
queryProcessingStrategy;
documentProducerComparator;
err;
state;
static STATES = ParallelQueryExecutionContextBaseStates;
routingProvider;
requestContinuation;
respHeaders;
unfilledDocumentProducersQueue;
bufferedDocumentProducersQueue;
// TODO: update type of buffer from any --> generic can be used here
buffer;
partitionDataPatchMap = new Map();
patchCounter = 0;
updatedContinuationRanges = new Map();
sem;
diagnosticNodeWrapper;
/**
* Provides the ParallelQueryExecutionContextBase.
* This is the base class that ParallelQueryExecutionContext and OrderByQueryExecutionContext will derive from.
*
* When handling a parallelized query, it instantiates one instance of
* DocumentProcuder per target partition key range and aggregates the result of each.
*
* @param clientContext - The service endpoint to use to create the client.
* @param collectionLink - The Collection Link
* @param options - Represents the feed options.
* @param partitionedQueryExecutionInfo - PartitionedQueryExecutionInfo
* @hidden
*/
constructor(clientContext, collectionLink, query, options, partitionedQueryExecutionInfo, correlatedActivityId, rangeManager, queryProcessingStrategy, documentProducerComparator) {
this.clientContext = clientContext;
this.collectionLink = collectionLink;
this.query = query;
this.options = options;
this.partitionedQueryExecutionInfo = partitionedQueryExecutionInfo;
this.correlatedActivityId = correlatedActivityId;
this.rangeManager = rangeManager;
this.queryProcessingStrategy = queryProcessingStrategy;
this.documentProducerComparator = documentProducerComparator;
this.clientContext = clientContext;
this.collectionLink = collectionLink;
this.query = query;
this.options = options;
this.partitionedQueryExecutionInfo = partitionedQueryExecutionInfo;
this.correlatedActivityId = correlatedActivityId;
this.diagnosticNodeWrapper = {
consumed: false,
diagnosticNode: new DiagnosticNodeInternal(clientContext.diagnosticLevel, DiagnosticNodeType.PARALLEL_QUERY_NODE, null),
};
this.diagnosticNodeWrapper.diagnosticNode.addData({ stateful: true });
this.err = undefined;
this.state = ParallelQueryExecutionContextBase.STATES.started;
this.routingProvider = new SmartRoutingMapProvider(this.clientContext);
this.buffer = [];
this.requestContinuation = options
? options.continuationToken || options.continuation
: undefined;
// Validate continuation token usage immediately
if (this.requestContinuation && !this.options.enableQueryControl) {
throw new Error("Continuation tokens are supported when enableQueryControl is set true in FeedOptions");
}
// response headers of undergoing operation
this.respHeaders = getInitialHeader();
// Make priority queue for documentProducers
this.unfilledDocumentProducersQueue = new PriorityQueue((a, b) => this.compareDocumentProducersByRange(a, b));
this.bufferedDocumentProducersQueue = new PriorityQueue((a, b) => this.documentProducerComparator(b, a));
// Creating the documentProducers
this.sem = semaphore(1);
this.sem.take(() => this._initializeDocumentProducers());
}
/**
* Determine if there are still remaining resources to processs based on the value of the continuation
* token or the elements remaining on the current batch in the QueryIterator.
* @returns true if there is other elements to process in the ParallelQueryExecutionContextBase.
*/
hasMoreResults() {
return (!this.err &&
(this.buffer.length > 0 || this.state !== ParallelQueryExecutionContextBase.STATES.ended));
}
/**
* Fetches more results from the query execution context.
* @param diagnosticNode - Optional diagnostic node for tracing.
* @returns A promise that resolves to the fetched results.
* @hidden
*/
async fetchMore(diagnosticNode) {
await this.bufferDocumentProducers(diagnosticNode);
await this.fillBufferFromBufferQueue();
return this.drainBufferedItems();
}
/**
* Processes buffered document producers
* @returns A promise that resolves when processing is complete.
*/
async processBufferedDocumentProducers() {
while (this.hasBufferedProducers() &&
this.shouldProcessBufferedProducers(this.isUnfilledQueueEmpty())) {
const producer = this.getNextBufferedProducer();
if (!producer)
break;
await this.processDocumentProducer(producer);
}
}
/**
* Processes a single document producer using template method pattern.
* Common structure with query-specific processing delegated to subclasses.
*/
async processDocumentProducer(producer) {
const response = await this.fetchFromProducer(producer);
this._mergeWithActiveResponseHeaders(response.headers);
if (response.result) {
this.addToBuffer(response.result);
this.handlePartitionMapping(producer, response.result);
}
// Handle producer lifecycle
if (producer.peakNextItem() !== undefined) {
this.requeueProducer(producer);
}
else if (producer.hasMoreResults()) {
this.moveToUnfilledQueue(producer);
}
}
/**
* Handles partition mapping updates - implemented in base class using template method pattern.
* Child classes provide query-specific parameters through abstract methods.
*/
handlePartitionMapping(producer, result) {
const itemCount = result?.length || 0;
const continuationToken = this.getContinuationToken(producer);
const mapping = {
itemCount,
partitionKeyRange: producer.targetPartitionKeyRange,
continuationToken,
};
this.updatePartitionMapping(mapping);
}
/**
* Gets the continuation token to use - implemented by subclasses.
*/
getContinuationToken(producer) {
const hasMoreBufferedItems = producer.peakNextItem() !== undefined;
return hasMoreBufferedItems ? producer.previousContinuationToken : producer.continuationToken;
}
/**
* Updates partition mapping - creates new entry or merges with existing for ORDER BY queries.
*/
updatePartitionMapping(mapping) {
const currentPatch = this.partitionDataPatchMap.get(this.patchCounter.toString());
const isSamePartition = currentPatch?.partitionKeyRange?.id === mapping.partitionKeyRange.id;
if (isSamePartition && currentPatch) {
currentPatch.itemCount += mapping.itemCount;
currentPatch.continuationToken = mapping.continuationToken;
return;
}
// Create new partition mapping entry
this.partitionDataPatchMap.set((++this.patchCounter).toString(), mapping);
}
/**
* Checks if the unfilled queue is empty (used by ORDER BY for processing control).
*/
isUnfilledQueueEmpty() {
return this.unfilledDocumentProducersQueue.size() === 0;
}
/**
* Initializes document producers and fills the priority queue.
* Handles both continuation token and fresh query scenarios.
*/
async _initializeDocumentProducers() {
try {
const targetPartitionRanges = await this._onTargetPartitionRanges();
const documentProducers = this.requestContinuation
? await this._createDocumentProducersFromContinuation(targetPartitionRanges)
: this._createDocumentProducersFromFresh(targetPartitionRanges);
// Fill up our priority queue with documentProducers
this._enqueueDocumentProducers(documentProducers);
this.sem.leave();
}
catch (err) {
this.err = err;
this.sem.leave();
}
}
/**
* Creates document producers from continuation token scenario.
*/
async _createDocumentProducersFromContinuation(targetPartitionRanges) {
// Parse continuation token to get range mappings and check for split/merge scenarios
const parsedToken = this._parseContinuationToken(this.requestContinuation);
const continuationRanges = await this._handlePartitionRangeChanges(parsedToken);
// Use strategy to create additional query info from parsed token
const additionalQueryInfo = this.queryProcessingStrategy.createAdditionalQueryInfo(parsedToken);
const filterResult = this.rangeManager.filterPartitionRanges(targetPartitionRanges, continuationRanges, additionalQueryInfo);
// Extract ranges and tokens from the combined result
const rangeTokenPairs = filterResult.rangeTokenPairs;
// Use strategy to create filter context for continuation token processing
const filterContext = this.queryProcessingStrategy.createFilterContext(parsedToken);
return rangeTokenPairs.map((rangeTokenPair) => this._createDocumentProducerFromRangeTokenPair(rangeTokenPair, continuationRanges, filterContext));
}
/**
* Creates document producers from fresh query scenario (no continuation token).
*/
_createDocumentProducersFromFresh(targetPartitionRanges) {
return targetPartitionRanges.map((partitionTargetRange) => this._createTargetPartitionQueryExecutionContext(partitionTargetRange, undefined));
}
/**
* Creates a document producer from a range token pair (continuation token scenario).
*/
_createDocumentProducerFromRangeTokenPair(rangeTokenPair, continuationRanges, filterContext) {
const partitionTargetRange = rangeTokenPair.range;
const continuationToken = rangeTokenPair.continuationToken;
const filterCondition = rangeTokenPair.filteringCondition || undefined;
// Find EPK ranges for this partition range from processed continuation response
const matchingContinuationRange = continuationRanges.find((cr) => cr.range.id === partitionTargetRange.id);
const startEpk = matchingContinuationRange?.epkMin;
const endEpk = matchingContinuationRange?.epkMax;
// Use strategy to determine partition-specific filter context
const targetPartitionId = continuationRanges.length > 0 && continuationRanges[continuationRanges.length - 1].range
? continuationRanges[continuationRanges.length - 1].range.id
: undefined;
const partitionFilterContext = this.queryProcessingStrategy.getPartitionFilterContext(filterContext, targetPartitionId, partitionTargetRange.id);
return this._createTargetPartitionQueryExecutionContext(partitionTargetRange, continuationToken, startEpk, endEpk, !!(startEpk && endEpk), // populateEpkRangeHeaders - true if both EPK values are present
filterCondition, partitionFilterContext);
}
/**
* Enqueues document producers into the unfilled queue.
*/
_enqueueDocumentProducers(documentProducers) {
documentProducers.forEach((documentProducer) => {
try {
this.unfilledDocumentProducersQueue.enq(documentProducer);
}
catch (e) {
this.err = e;
}
});
}
/**
* Checks if there are buffered document producers ready for processing.
* Encapsulates queue size checking.
*/
hasBufferedProducers() {
return this.bufferedDocumentProducersQueue.size() > 0;
}
/**
* Gets the next buffered document producer for processing.
* Encapsulates queue dequeuing logic.
*/
getNextBufferedProducer() {
if (this.bufferedDocumentProducersQueue.size() > 0) {
return this.bufferedDocumentProducersQueue.deq();
}
return undefined;
}
/**
* Adds items to the result buffer. Handles both single items and arrays.
*/
addToBuffer(items) {
if (Array.isArray(items)) {
if (items.length > 0) {
this.buffer.push(...items);
}
}
else if (items) {
this.buffer.push(items);
}
}
/**
* Moves a producer to the unfilled queue for later processing.
*/
moveToUnfilledQueue(producer) {
this.unfilledDocumentProducersQueue.enq(producer);
}
/**
* Re-queues a producer to the buffered queue for further processing.
*/
requeueProducer(producer) {
this.bufferedDocumentProducersQueue.enq(producer);
}
/**
* Compares two document producers based on their partition key ranges and EPK values.
* Primary comparison: minInclusive values for left-to-right range traversal
* Secondary comparison: EPK ranges when minInclusive values are identical
* @param a - First document producer
* @param b - Second document producer
* @returns Comparison result for priority queue ordering
* @hidden
*/
compareDocumentProducersByRange(a, b) {
const aMinInclusive = a.targetPartitionKeyRange.minInclusive;
const bMinInclusive = b.targetPartitionKeyRange.minInclusive;
const minInclusiveComparison = bMinInclusive.localeCompare(aMinInclusive);
// If minInclusive values are the same, check minEPK ranges if they exist
if (minInclusiveComparison === 0) {
const aMinEpk = a.startEpk;
const bMinEpk = b.startEpk;
if (aMinEpk && bMinEpk) {
return bMinEpk.localeCompare(aMinEpk);
}
}
return minInclusiveComparison;
}
/**
* Detects partition splits/merges by analyzing parsed continuation token ranges and comparing with current topology
* @param parsed - The continuation token containing range mappings to analyze
* @returns Array of processed ranges with EPK info
*/
async _handlePartitionRangeChanges(parsed) {
const processedRanges = [];
// Extract range mappings from the already parsed token
const rangeMappings = parsed.rangeMappings;
if (!rangeMappings || rangeMappings.length === 0) {
return [];
}
// Check each range mapping for potential splits/merges
for (const rangeWithToken of rangeMappings) {
// Create a new QueryRange instance from the simplified range data
const range = rangeWithToken.queryRange;
const queryRange = new QueryRange(range.min, range.max, true, // isMinInclusive - assumption: always true
false);
const rangeMin = queryRange.min;
const rangeMax = queryRange.max;
// Get current overlapping ranges for this continuation token range
const overlappingRanges = await this.routingProvider.getOverlappingRanges(this.collectionLink, [queryRange], this.getDiagnosticNode());
// Detect split/merge scenario based on the number of overlapping ranges
if (overlappingRanges.length === 0) {
continue;
}
else if (overlappingRanges.length === 1) {
// Check if it's the same range (no change) or a merge scenario
const currentRange = overlappingRanges[0];
if (currentRange.minInclusive !== rangeMin || currentRange.maxExclusive !== rangeMax) {
// Merge scenario - include EPK ranges from original continuation token range
await this._handleContinuationTokenMerge(rangeWithToken, currentRange);
processedRanges.push({
range: currentRange,
continuationToken: rangeWithToken.continuationToken,
epkMin: rangeMin, // Original range min becomes EPK min
epkMax: rangeMax, // Original range max becomes EPK max
});
}
else {
// Same range - no merge, no EPK ranges needed
processedRanges.push({
range: currentRange,
continuationToken: rangeWithToken.continuationToken,
});
}
}
else {
// Split scenario - one range from continuation token now maps to multiple ranges
await this._handleContinuationTokenSplit(rangeWithToken, overlappingRanges);
// Add all overlapping ranges with the same continuation token to processed ranges
overlappingRanges.forEach((rangeValue) => {
processedRanges.push({
range: rangeValue,
continuationToken: rangeWithToken.continuationToken,
});
});
}
}
return processedRanges;
}
/**
* Parses the continuation token based on query type
* @param continuationToken - The continuation token string to parse
* @returns Parsed continuation token object (ORDER BY or Parallel query token)
* @throws ErrorResponse when continuation token is malformed or cannot be parsed
*/
_parseContinuationToken(continuationToken) {
try {
return this.queryProcessingStrategy.parseContinuationToken(continuationToken);
}
catch (e) {
throw new ErrorResponse(`Invalid continuation token format. Expected token with rangeMappings property. ` +
`Ensure the continuation token was generated by a compatible query and has not been modified.`);
}
}
/**
* Handles partition merge scenario for continuation token ranges
*/
async _handleContinuationTokenMerge(rangeWithToken, _newMergedRange) {
const rangeKey = `${rangeWithToken.queryRange.min}-${rangeWithToken.queryRange.max}`;
this.updatedContinuationRanges.set(rangeKey, {
oldRange: {
min: rangeWithToken.queryRange.min,
max: rangeWithToken.queryRange.max,
isMinInclusive: true, // Assumption: min is always inclusive
isMaxInclusive: false, // Assumption: max is always exclusive
},
newRanges: [
{
min: rangeWithToken.queryRange.min,
max: rangeWithToken.queryRange.max,
isMinInclusive: true, // Assumption: min is always inclusive
isMaxInclusive: false, // Assumption: max is always exclusive
},
],
continuationToken: rangeWithToken.continuationToken,
});
}
/**
* Handles partition split scenario for continuation token ranges
*/
async _handleContinuationTokenSplit(rangeWithToken, overlappingRanges) {
const rangeKey = `${rangeWithToken.queryRange.min}-${rangeWithToken.queryRange.max}`;
this.updatedContinuationRanges.set(rangeKey, {
oldRange: {
min: rangeWithToken.queryRange.min,
max: rangeWithToken.queryRange.max,
isMinInclusive: true, // Assumption: min is always inclusive
isMaxInclusive: false, // Assumption: max is always exclusive
},
newRanges: overlappingRanges.map((range) => ({
min: range.minInclusive,
max: range.maxExclusive,
isMinInclusive: true,
isMaxInclusive: false,
})),
continuationToken: rangeWithToken.continuationToken,
});
}
/**
* Handles partition merge scenario for continuation token ranges
*/
_mergeWithActiveResponseHeaders(headers) {
mergeHeaders(this.respHeaders, headers);
}
_getAndResetActiveResponseHeaders() {
const ret = this.respHeaders;
this.respHeaders = getInitialHeader();
return ret;
}
getDiagnosticNode() {
return this.diagnosticNodeWrapper.diagnosticNode;
}
async _onTargetPartitionRanges() {
// invokes the callback when the target partition ranges are ready
const parsedRanges = this.partitionedQueryExecutionInfo.queryRanges;
const queryRanges = parsedRanges.map((item) => QueryRange.parseFromDict(item));
return this.routingProvider.getOverlappingRanges(this.collectionLink, queryRanges, this.getDiagnosticNode());
}
/**
* Gets the replacement ranges for a partitionkeyrange that has been split
*/
async _getReplacementPartitionKeyRanges(documentProducer, diagnosticNode) {
const partitionKeyRange = documentProducer.targetPartitionKeyRange;
// Download the new routing map
this.routingProvider = new SmartRoutingMapProvider(this.clientContext);
// Get the queryRange that relates to this partitionKeyRange
const queryRange = QueryRange.parsePartitionKeyRange(partitionKeyRange);
return this.routingProvider.getOverlappingRanges(this.collectionLink, [queryRange], diagnosticNode);
}
async _enqueueReplacementDocumentProducers(error, diagnosticNode, documentProducer) {
// Get the replacement ranges
const replacementPartitionKeyRanges = await this._getReplacementPartitionKeyRanges(documentProducer, diagnosticNode);
if (replacementPartitionKeyRanges.length === 0) {
throw error;
}
if (this.requestContinuation) {
// Update composite continuation token to handle partition split
this._updateContinuationTokenOnPartitionChange(documentProducer, replacementPartitionKeyRanges);
}
if (replacementPartitionKeyRanges.length === 1) {
// Partition is gone due to Merge
// Create the replacement documentProducer with populateEpkRangeHeaders Flag set to true to set startEpk and endEpk headers
const replacementDocumentProducer = this._createTargetPartitionQueryExecutionContext(replacementPartitionKeyRanges[0], documentProducer.continuationToken, documentProducer.startEpk, documentProducer.endEpk, true);
this.unfilledDocumentProducersQueue.enq(replacementDocumentProducer);
}
else {
// Create the replacement documentProducers
const replacementDocumentProducers = [];
replacementPartitionKeyRanges.forEach((partitionKeyRange) => {
const queryRange = QueryRange.parsePartitionKeyRange(partitionKeyRange);
// Create replacment document producers with the parent's continuationToken
const replacementDocumentProducer = this._createTargetPartitionQueryExecutionContext(partitionKeyRange, documentProducer.continuationToken, queryRange.min, queryRange.max, false);
replacementDocumentProducers.push(replacementDocumentProducer);
});
// add document producers to the queue
replacementDocumentProducers.forEach((replacementDocumentProducer) => {
if (replacementDocumentProducer.hasMoreResults()) {
this.unfilledDocumentProducersQueue.enq(replacementDocumentProducer);
}
});
}
}
_updateContinuationTokenOnPartitionChange(originalDocumentProducer, replacementPartitionKeyRanges) {
const rangeWithToken = this._createQueryRangeWithContinuationToken(originalDocumentProducer);
if (replacementPartitionKeyRanges.length === 1) {
this._handleContinuationTokenMerge(rangeWithToken, replacementPartitionKeyRanges[0]);
}
else {
this._handleContinuationTokenSplit(rangeWithToken, replacementPartitionKeyRanges);
}
}
/**
* Creates a QueryRangeWithContinuationToken object from a DocumentProducer.
* Uses the DocumentProducer's target partition key range and continuation token.
* @param documentProducer - The DocumentProducer to convert
* @returns QueryRangeWithContinuationToken object for token operations
*/
_createQueryRangeWithContinuationToken(documentProducer) {
const partitionRange = documentProducer.targetPartitionKeyRange;
// Create a simplified QueryRange using the partition key range boundaries
const simplifiedQueryRange = {
min: documentProducer.startEpk || partitionRange.minInclusive,
max: documentProducer.endEpk || partitionRange.maxExclusive,
};
return {
queryRange: simplifiedQueryRange,
continuationToken: documentProducer.continuationToken,
};
}
static _needPartitionKeyRangeCacheRefresh(error) {
// TODO: any error
return (error.code === StatusCodes.Gone &&
"substatus" in error &&
error["substatus"] === SubStatusCodes.PartitionKeyRangeGone);
}
/**
* Creates target partition range Query Execution Context
*/
_createTargetPartitionQueryExecutionContext(partitionKeyTargetRange, continuationToken, startEpk, endEpk, populateEpkRangeHeaders, filterCondition, filterContext) {
let rewrittenQuery = this.partitionedQueryExecutionInfo.queryInfo?.rewrittenQuery;
let sqlQuerySpec;
const query = this.query;
if (typeof query === "string") {
sqlQuerySpec = { query };
}
else {
sqlQuerySpec = query;
}
const formatPlaceHolder = "{documentdb-formattableorderbyquery-filter}";
if (rewrittenQuery) {
sqlQuerySpec = JSON.parse(JSON.stringify(sqlQuerySpec));
rewrittenQuery = filterCondition
? rewrittenQuery.replace(formatPlaceHolder, filterCondition)
: rewrittenQuery.replace(formatPlaceHolder, "true");
sqlQuerySpec["query"] = rewrittenQuery;
}
const options = { ...this.options };
options.continuationToken = continuationToken;
let filter;
if (filterContext) {
filter = new RidSkipCountFilter(filterContext);
}
return new DocumentProducer(this.clientContext, this.collectionLink, sqlQuerySpec, partitionKeyTargetRange, options, this.correlatedActivityId, startEpk, endEpk, populateEpkRangeHeaders, filter);
}
async drainBufferedItems() {
return new Promise((resolve, reject) => {
this.sem.take(() => {
if (this.err) {
// if there is a prior error return error
this.sem.leave();
this.err.headers = this._getAndResetActiveResponseHeaders();
reject(this.err);
return;
}
// return undefined if there is no more results
if (this.buffer.length === 0) {
this.sem.leave();
const partitionDataPatchMap = this.partitionDataPatchMap;
this.partitionDataPatchMap = new Map();
this.patchCounter = 0;
// Get and reset updated continuation ranges
const updatedContinuationRanges = Object.fromEntries(this.updatedContinuationRanges);
this.updatedContinuationRanges.clear();
const result = createParallelQueryResult([], partitionDataPatchMap, updatedContinuationRanges, undefined);
return resolve({
result: this.state === ParallelQueryExecutionContextBase.STATES.ended ? undefined : result,
headers: this._getAndResetActiveResponseHeaders(),
});
}
// draing the entire buffer object and return that in result of return object
const bufferedResults = this.buffer;
this.buffer = [];
// reset the patchToRangeMapping
const partitionDataPatchMap = this.partitionDataPatchMap;
this.partitionDataPatchMap = new Map();
this.patchCounter = 0;
// Get and reset updated continuation ranges
const updatedContinuationRanges = Object.fromEntries(this.updatedContinuationRanges);
this.updatedContinuationRanges.clear();
// release the lock before returning
this.sem.leave();
const result = createParallelQueryResult(bufferedResults, partitionDataPatchMap, updatedContinuationRanges, undefined);
return resolve({
result,
headers: this._getAndResetActiveResponseHeaders(),
});
});
});
}
/**
* Buffers document producers based on the maximum degree of parallelism.
* Moves document producers from the unfilled queue to the buffered queue.
* @param diagnosticNode - The diagnostic node for logging and tracing.
* @returns A promise that resolves when buffering is complete.
*/
async bufferDocumentProducers(diagnosticNode) {
return new Promise((resolve, reject) => {
this.sem.take(async () => {
if (this.err) {
this.sem.leave();
reject(this.err);
return;
}
this.updateStates(this.err);
if (this.state === ParallelQueryExecutionContextBase.STATES.ended) {
this.sem.leave();
resolve();
return;
}
if (this.unfilledDocumentProducersQueue.size() === 0) {
this.sem.leave();
resolve();
return;
}
try {
const maxDegreeOfParallelism = this.options.maxDegreeOfParallelism === undefined ||
this.options.maxDegreeOfParallelism < 1
? this.unfilledDocumentProducersQueue.size() // number of partitions
: Math.min(this.options.maxDegreeOfParallelism, this.unfilledDocumentProducersQueue.size());
const documentProducers = [];
while (documentProducers.length < maxDegreeOfParallelism &&
this.unfilledDocumentProducersQueue.size() > 0) {
let documentProducer;
try {
documentProducer = this.unfilledDocumentProducersQueue.deq();
}
catch (e) {
this.err = e;
this.err.headers = this._getAndResetActiveResponseHeaders();
reject(this.err);
return;
}
documentProducers.push(documentProducer);
}
const bufferDocumentProducer = async (documentProducer) => {
try {
const headers = await documentProducer.bufferMore(diagnosticNode);
this._mergeWithActiveResponseHeaders(headers);
// Always track this document producer in patchToRangeMapping, even if it has no results
// This ensures we maintain a record of all partition ranges that were scanned
const nextItem = documentProducer.peakNextItem();
if (nextItem !== undefined) {
this.bufferedDocumentProducersQueue.enq(documentProducer);
}
else {
// Track document producer with no results in patchToRangeMapping
// This represents a scanned partition that yielded no results
// IMPORTANT: Only include if continuation token is NOT null/exhausted
// Document producers with no data in buffer and no continuation token are exhausted and should not be added to partitionDataPatchMap to prevent infinite loops in order by queries
if (documentProducer.continuationToken &&
documentProducer.continuationToken !== "" &&
documentProducer.continuationToken.toLowerCase() !== "null") {
const patchKey = `empty-${documentProducer.targetPartitionKeyRange.id}-${documentProducer.targetPartitionKeyRange.minInclusive}`;
this.partitionDataPatchMap.set(patchKey, {
itemCount: 0, // 0 items for empty result set
partitionKeyRange: documentProducer.targetPartitionKeyRange,
continuationToken: documentProducer.continuationToken,
});
}
if (documentProducer.hasMoreResults()) {
this.unfilledDocumentProducersQueue.enq(documentProducer);
}
}
}
catch (err) {
if (ParallelQueryExecutionContextBase._needPartitionKeyRangeCacheRefresh(err)) {
// We want the document producer enqueued
// So that later parts of the code can repair the execution context
// refresh the partition key ranges and ctreate new document producers and add it to the queue
await this._enqueueReplacementDocumentProducers(err, diagnosticNode, documentProducer);
resolve();
}
else {
this.err = err;
this.err.headers = this._getAndResetActiveResponseHeaders();
reject(err);
}
}
};
try {
await Promise.all(documentProducers.map((producer) => bufferDocumentProducer(producer)));
}
catch (err) {
this.err = err;
this.err.headers = this._getAndResetActiveResponseHeaders();
reject(err);
return;
}
resolve();
}
catch (err) {
this.err = err;
this.err.headers = this._getAndResetActiveResponseHeaders();
reject(err);
}
finally {
this.sem.leave();
}
});
});
}
/**
* Drains the buffer of filled document producers and appends their items to the main buffer.
* Uses template method pattern - delegates actual processing to subclasses.
* @returns A promise that resolves when the buffer is filled.
*/
async fillBufferFromBufferQueue() {
return new Promise((resolve, reject) => {
this.sem.take(async () => {
if (this.err) {
// if there is a prior error return error
this.sem.leave();
this.err.headers = this._getAndResetActiveResponseHeaders();
reject(this.err);
return;
}
if (this.state === ParallelQueryExecutionContextBase.STATES.ended ||
this.bufferedDocumentProducersQueue.size() === 0) {
this.sem.leave();
resolve();
return;
}
try {
await this.processBufferedDocumentProducers();
this.updateStates(this.err);
}
catch (err) {
this.err = err;
this.err.headers = this._getAndResetActiveResponseHeaders();
reject(this.err);
return;
}
finally {
// release the lock before returning
this.sem.leave();
}
resolve();
return;
});
});
}
updateStates(error) {
if (error) {
this.err = error;
this.state = ParallelQueryExecutionContextBase.STATES.ended;
return;
}
if (this.state === ParallelQueryExecutionContextBase.STATES.started) {
this.state = ParallelQueryExecutionContextBase.STATES.inProgress;
}
const hasNoActiveProducers = this.unfilledDocumentProducersQueue.size() === 0 &&
this.bufferedDocumentProducersQueue.size() === 0;
if (hasNoActiveProducers) {
this.state = ParallelQueryExecutionContextBase.STATES.ended;
}
}
}
//# sourceMappingURL=parallelQueryExecutionContextBase.js.map