Skip to content

Commit 175c7e2

Browse files
xinlian12Annie LiangCopilotazure-sdkCopilot
authored
[SparkConnector]IncludeOperationStatusCodeHistoryInStaleProgressLogs (#48022)
* Add status code tracking to bulk operations with compressed consecutive identical (statusCode, subStatusCode) pairs into single entries with count and time range. --------- Co-authored-by: Annie Liang <anniemac@Annies-MacBook-Pro.local> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Azure SDK Bot <53356347+azure-sdk@users.noreply.github.com> Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com> Co-authored-by: kushagraThapar <14034156+kushagraThapar@users.noreply.github.com> Co-authored-by: Kushagra Thapar <kuthapar@microsoft.com>
1 parent 6bd680f commit 175c7e2

18 files changed

Lines changed: 556 additions & 152 deletions

File tree

sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
* Fixed an issue where `TransientIOErrorsRetryingIterator` would trigger extra query during retries and on close. - See [PR 47996](https://github.com/Azure/azure-sdk-for-java/pull/47996)
1111

1212
#### Other Changes
13+
* Added status code history in `BulkWriterNoProgressException` error message. - See [PR 48022](https://github.com/Azure/azure-sdk-for-java/pull/48022)
1314

1415
### 4.43.0 (2026-02-10)
1516

sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
* Fixed an issue where `TransientIOErrorsRetryingIterator` would trigger extra query during retries and on close. - See [PR 47996](https://github.com/Azure/azure-sdk-for-java/pull/47996)
1111

1212
#### Other Changes
13+
* Added status code history in `BulkWriterNoProgressException` error message. - See [PR 48022](https://github.com/Azure/azure-sdk-for-java/pull/48022)
1314

1415
### 4.43.0 (2026-02-10)
1516

sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
* Fixed an issue where `TransientIOErrorsRetryingIterator` would trigger extra query during retries and on close. - See [PR 47996](https://github.com/Azure/azure-sdk-for-java/pull/47996)
1111

1212
#### Other Changes
13+
* Added status code history in `BulkWriterNoProgressException` error message. - See [PR 48022](https://github.com/Azure/azure-sdk-for-java/pull/48022)
1314

1415
### 4.43.0 (2026-02-10)
1516

sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
* Fixed an issue where `TransientIOErrorsRetryingIterator` would trigger extra query during retries and on close. - See [PR 47996](https://github.com/Azure/azure-sdk-for-java/pull/47996)
1111

1212
#### Other Changes
13+
* Added status code history in `BulkWriterNoProgressException` error message. - See [PR 48022](https://github.com/Azure/azure-sdk-for-java/pull/48022)
1314

1415
### 4.43.0 (2026-02-10)
1516

sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/BulkWriter.scala

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ package com.azure.cosmos.spark
55
// scalastyle:off underscore.import
66
import com.azure.cosmos.implementation.apachecommons.lang.StringUtils
77
import com.azure.cosmos.implementation.batch.{BatchRequestResponseConstants, BulkExecutorDiagnosticsTracker, ItemBulkOperation}
8-
import com.azure.cosmos.implementation.{CosmosDaemonThreadFactory, UUIDs}
8+
import com.azure.cosmos.implementation.{CosmosDaemonThreadFactory, ImplementationBridgeHelpers, UUIDs}
99
import com.azure.cosmos.models._
1010
import com.azure.cosmos.spark.BulkWriter.{BulkOperationFailedException, bulkWriterInputBoundedElastic, bulkWriterRequestsBoundedElastic, bulkWriterResponsesBoundedElastic, getThreadInfo, readManyBoundedElastic}
1111
import com.azure.cosmos.spark.diagnostics.DefaultDiagnostics
@@ -888,7 +888,7 @@ private class BulkWriter
888888

889889
val message = s"All retries exhausted for '${itemOperation.getOperationType}' bulk operation - " +
890890
s"statusCode=[$effectiveStatusCode:$effectiveSubStatusCode] " +
891-
s"itemId=[${context.itemId}], partitionKeyValue=[${context.partitionKeyValue}]"
891+
s"itemId=[${context.itemId}], partitionKeyValue=[${context.partitionKeyValue}], attemptNumber=${context.attemptNumber}"
892892

893893
val exceptionToBeThrown = responseException match {
894894
case Some(e) =>
@@ -929,6 +929,11 @@ private class BulkWriter
929929
sb.append("->")
930930
val ctx = itemOperation.getContext[OperationContext]
931931
sb.append(s"${ctx.partitionKeyValue}/${ctx.itemId}/${ctx.eTag}(${ctx.attemptNumber})")
932+
itemOperation match {
933+
case op: ItemBulkOperation[_, _] =>
934+
sb.append(s", statusHistory=${op.getStatusTracker.toString}")
935+
case _ =>
936+
}
932937
})
933938

934939
// add readMany snapshot logs

sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala

Lines changed: 32 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
package com.azure.cosmos.spark
44

55
// scalastyle:off underscore.import
6-
import com.azure.cosmos.implementation.batch.{BulkExecutorDiagnosticsTracker, CosmosBulkTransactionalBatchResponse, TransactionalBulkExecutor}
6+
import com.azure.cosmos.implementation.batch.{BulkExecutorDiagnosticsTracker, CosmosBatchBulkOperation, CosmosBulkTransactionalBatchResponse, TransactionalBulkExecutor}
77
import com.azure.cosmos.implementation.{CosmosTransactionalBulkExecutionOptionsImpl, UUIDs}
88
import com.azure.cosmos.models.{CosmosBatch, CosmosBatchResponse}
99
import com.azure.cosmos.spark.BulkWriter.getThreadInfo
@@ -93,7 +93,7 @@ private class TransactionalBulkWriter
9393
private val activeBatches = new ConcurrentHashMap[PartitionKey, CosmosBatchOperation].asScala
9494
private val errorCaptureFirstException = new AtomicReference[Throwable]()
9595
private val transactionalBulkInputEmitter: Sinks.Many[TransactionalBulkItem] = Sinks.many().unicast().onBackpressureBuffer()
96-
private val transactionalBatchInputEmitter: Sinks.Many[CosmosBatch] = Sinks.many().unicast().onBackpressureBuffer()
96+
private val transactionalBatchInputEmitter: Sinks.Many[CosmosBatchBulkOperation] = Sinks.many().unicast().onBackpressureBuffer()
9797

9898
// for transactional batch, all rows/items from the dataframe should be grouped as one cosmos batch
9999
private val transactionalBatchPartitionKeyScheduled = java.util.concurrent.ConcurrentHashMap.newKeySet[PartitionKey]().asScala
@@ -268,7 +268,7 @@ private class TransactionalBulkWriter
268268
try {
269269
// all the operations in the batch will have the same partition key value
270270
// get the partition key value from the first result
271-
val partitionKeyValue = resp.getCosmosBatch.getPartitionKeyValue
271+
val partitionKeyValue = resp.getCosmosBatchBulkOperation.getPartitionKeyValue
272272
val activeBatchOperationOpt = activeBatches.remove(partitionKeyValue)
273273
val pendingBatchOperationRetriesOpt = pendingBatchRetries.remove(partitionKeyValue)
274274

@@ -291,7 +291,7 @@ private class TransactionalBulkWriter
291291
case Some(cosmosException: CosmosException) =>
292292
handleNonSuccessfulStatusCode(
293293
batchOperation.operationContext,
294-
batchOperation.cosmosBatch,
294+
batchOperation.cosmosBatchBulkOperation,
295295
None,
296296
isGettingRetried,
297297
Some(cosmosException))
@@ -307,7 +307,7 @@ private class TransactionalBulkWriter
307307
} else if (!resp.getResponse.isSuccessStatusCode) {
308308
handleNonSuccessfulStatusCode(
309309
batchOperation.operationContext,
310-
batchOperation.cosmosBatch,
310+
batchOperation.cosmosBatchBulkOperation,
311311
Some(resp.getResponse),
312312
isGettingRetried,
313313
None)
@@ -359,6 +359,7 @@ private class TransactionalBulkWriter
359359
throwIfCapturedExceptionExists()
360360

361361
val activeTasksSemaphoreTimeout = 10
362+
val cosmosBatchBulkOperation = new CosmosBatchBulkOperation(cosmosBatch)
362363
val operationContext =
363364
new OperationContext(
364365
cosmosBatch.getPartitionKeyValue,
@@ -402,7 +403,7 @@ private class TransactionalBulkWriter
402403
val cnt = totalScheduledMetrics.getAndAdd(cosmosBatch.getOperations.size())
403404
log.logTrace(s"total scheduled $cnt, Context: ${operationContext.toString} $getThreadInfo")
404405

405-
scheduleBatchInternal(CosmosBatchOperation(cosmosBatch, operationContext))
406+
scheduleBatchInternal(CosmosBatchOperation(cosmosBatchBulkOperation, operationContext))
406407
}
407408

408409
private def scheduleBatchInternal(cosmosBatchOperation: CosmosBatchOperation): Unit = {
@@ -414,17 +415,17 @@ private class TransactionalBulkWriter
414415
}
415416

416417
activeBatches.put(
417-
cosmosBatchOperation.cosmosBatch.getPartitionKeyValue,
418+
cosmosBatchOperation.cosmosBatchBulkOperation.getPartitionKeyValue,
418419
cosmosBatchOperation)
419-
transactionalBatchInputEmitter.emitNext(cosmosBatchOperation.cosmosBatch, emitFailureHandler)
420+
transactionalBatchInputEmitter.emitNext(cosmosBatchOperation.cosmosBatchBulkOperation, emitFailureHandler)
420421
}
421422

422423
//scalastyle:off method.length
423424
//scalastyle:off cyclomatic.complexity
424425
private[this] def handleNonSuccessfulStatusCode
425426
(
426427
operationContext: OperationContext,
427-
cosmosBatch: CosmosBatch,
428+
cosmosBatchBulkOperation: CosmosBatchBulkOperation,
428429
cosmosBatchResponse: Option[CosmosBatchResponse],
429430
isGettingRetried: AtomicBoolean,
430431
responseException: Option[CosmosException]
@@ -466,16 +467,16 @@ private class TransactionalBulkWriter
466467
s"Context: {${operationContext.toString}} $getThreadInfo")
467468

468469
val batchOperationRetry = CosmosBatchOperation(
469-
cosmosBatch,
470+
cosmosBatchBulkOperation,
470471
new OperationContext(
471472
operationContext.partitionKeyValueInput,
472473
operationContext.attemptNumber + 1,
473474
operationContext.sequenceNumber)
474475
)
475476

476477
this.scheduleRetry(
477-
trackPendingRetryAction = () => pendingBatchRetries.put(cosmosBatch.getPartitionKeyValue, batchOperationRetry).isEmpty,
478-
clearPendingRetryAction = () => pendingBatchRetries.remove(cosmosBatch.getPartitionKeyValue).isDefined,
478+
trackPendingRetryAction = () => pendingBatchRetries.put(cosmosBatchBulkOperation.getPartitionKeyValue, batchOperationRetry).isEmpty,
479+
clearPendingRetryAction = () => pendingBatchRetries.remove(cosmosBatchBulkOperation.getPartitionKeyValue).isDefined,
479480
batchOperationRetry,
480481
effectiveStatusCode)
481482
isGettingRetried.set(true)
@@ -519,18 +520,27 @@ private class TransactionalBulkWriter
519520
// flatten the batches
520521
activeOperationsSnapshot
521522
.values
522-
.flatMap(batchOperation =>
523-
batchOperation.cosmosBatch.getOperations.asScala.map(itemOperation => (itemOperation, batchOperation.operationContext.attemptNumber)))
523+
.flatMap(batchOperation => {
524+
val statusTracker = batchOperation.cosmosBatchBulkOperation.getStatusTracker
525+
val statusHistory = if (statusTracker != null) {
526+
Some(statusTracker.toString)
527+
} else {
528+
None
529+
}
530+
batchOperation.cosmosBatchBulkOperation.getCosmosBatch.getOperations.asScala.map(itemOperation =>
531+
(itemOperation, batchOperation.operationContext.attemptNumber, statusHistory))
532+
})
524533
.toList
525534
.take(TransactionalBulkWriter.maxItemOperationsToShowInErrorMessage)
526-
.foreach(itemOperationAttemptPair => {
535+
.foreach(itemOperationTuple => {
527536
if (sb.nonEmpty) {
528537
sb.append(", ")
529538
}
530539

531-
sb.append(itemOperationAttemptPair._1.getOperationType)
540+
sb.append(itemOperationTuple._1.getOperationType)
532541
sb.append("->")
533-
sb.append(s"${itemOperationAttemptPair._1.getId}/${itemOperationAttemptPair._1.getPartitionKeyValue}/(${itemOperationAttemptPair._2})")
542+
sb.append(s"${itemOperationTuple._1.getId}/${itemOperationTuple._1.getPartitionKeyValue}/(${itemOperationTuple._2})")
543+
itemOperationTuple._3.foreach(history => sb.append(s", statusHistory=$history"))
534544
})
535545

536546
sb.toString()
@@ -548,8 +558,8 @@ private class TransactionalBulkWriter
548558
snapshot.keys.forall(partitionKey => {
549559
if (current.contains(partitionKey)) {
550560

551-
snapshot(partitionKey).cosmosBatch.getOperations.asScala.forall(itemOperationSnapshot => {
552-
current(partitionKey).cosmosBatch.getOperations.asScala.exists(currentOperation =>
561+
snapshot(partitionKey).cosmosBatchBulkOperation.getCosmosBatch.getOperations.asScala.forall(itemOperationSnapshot => {
562+
current(partitionKey).cosmosBatchBulkOperation.getCosmosBatch.getOperations.asScala.exists(currentOperation =>
553563
itemOperationSnapshot.getOperationType == currentOperation.getOperationType
554564
&& itemOperationSnapshot.getPartitionKeyValue == currentOperation.getPartitionKeyValue
555565
&& Objects.equals(itemOperationSnapshot.getId, currentOperation.getId)
@@ -607,7 +617,7 @@ private class TransactionalBulkWriter
607617
(pendingRetriesSnapshot ++ activeOperationsSnapshot)
608618
.toList
609619
.sortBy(op => op._2.operationContext.sequenceNumber)
610-
.map(batchOperationPartitionKeyPair => batchOperationPartitionKeyPair._2.cosmosBatch)
620+
.map(batchOperationPartitionKeyPair => batchOperationPartitionKeyPair._2.cosmosBatchBulkOperation.getCosmosBatch)
611621
.flatMap(batch => batch.getOperations.asScala)
612622
)
613623
} else {
@@ -702,7 +712,7 @@ private class TransactionalBulkWriter
702712
// re-validating whether the operation is still active - if so, just re-enqueue another retry
703713
// this is harmless - because all bulkItemOperations from Spark connector are always idempotent
704714
// For FAIL_NON_SERIALIZED, will keep retry, while for other errors, use the default behavior
705-
transactionalBatchInputEmitter.emitNext(operationPartitionKeyPair._2.cosmosBatch, TransactionalBulkWriter.emitFailureHandler)
715+
transactionalBatchInputEmitter.emitNext(operationPartitionKeyPair._2.cosmosBatchBulkOperation, TransactionalBulkWriter.emitFailureHandler)
706716
log.logWarning(s"Re-enqueued a retry for pending active batch task "
707717
+ s"(${operationPartitionKeyPair._1})' "
708718
+ s"- Attempt: ${numberOfIntervalsWithIdenticalActiveOperationSnapshots.get} - "
@@ -850,7 +860,7 @@ private class TransactionalBulkWriter
850860
}
851861
}
852862

853-
private case class CosmosBatchOperation(cosmosBatch: CosmosBatch, operationContext: OperationContext)
863+
private case class CosmosBatchOperation(cosmosBatchBulkOperation: CosmosBatchBulkOperation, operationContext: OperationContext)
854864
private case class TransactionalBulkItem(partitionKey: PartitionKey, objectNode: ObjectNode)
855865
}
856866

sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
* Fixed an issue where `TransientIOErrorsRetryingIterator` would trigger extra query during retries and on close. - See [PR 47996](https://github.com/Azure/azure-sdk-for-java/pull/47996)
1111

1212
#### Other Changes
13+
* Added status code history in `BulkWriterNoProgressException` error message. - See [PR 48022](https://github.com/Azure/azure-sdk-for-java/pull/48022)
1314

1415
### 4.43.0 (2026-02-10)
1516

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosBulkAsyncTest.java

Lines changed: 66 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,17 @@
44
package com.azure.cosmos;
55

66
import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
7+
import com.azure.cosmos.implementation.batch.BulkOperationStatusTracker;
8+
import com.azure.cosmos.implementation.batch.ItemBulkOperation;
79
import com.azure.cosmos.models.CosmosBulkExecutionOptions;
10+
import com.azure.cosmos.models.CosmosBulkOperationResponse;
811
import com.azure.cosmos.models.CosmosBulkOperations;
912
import com.azure.cosmos.models.CosmosContainerProperties;
13+
import com.azure.cosmos.models.CosmosItemOperation;
1014
import com.azure.cosmos.models.PartitionKey;
1115
import com.azure.cosmos.models.PartitionKeyDefinition;
1216
import com.azure.cosmos.models.ThroughputProperties;
17+
import com.azure.cosmos.test.faultinjection.CosmosFaultInjectionHelper;
1318
import com.azure.cosmos.test.faultinjection.FaultInjectionCondition;
1419
import com.azure.cosmos.test.faultinjection.FaultInjectionConditionBuilder;
1520
import com.azure.cosmos.test.faultinjection.FaultInjectionConnectionType;
@@ -757,6 +762,54 @@ public void replaceItem_withBulk() {
757762
assertThat(processedDoc.get()).isEqualTo(totalRequest);
758763
}
759764

765+
@Test(groups = {"fast"}, timeOut = TIMEOUT)
766+
public void createItem_withBulk_tooManyRequest_recordStatusHistory() {
767+
768+
List<CosmosItemOperation> cosmosItemOperations = new ArrayList<>();
769+
String partitionKey = UUID.randomUUID().toString();
770+
TestDoc testDoc = this.populateTestDoc(partitionKey);
771+
cosmosItemOperations.add(CosmosBulkOperations.getCreateItemOperation(testDoc, new PartitionKey(partitionKey)));
772+
773+
FaultInjectionRule tooManyRequestRule = injectBatchFailure(
774+
"statusTracker-429-" + UUID.randomUUID(),
775+
FaultInjectionServerErrorType.TOO_MANY_REQUEST,
776+
2,
777+
Duration.ZERO);
778+
779+
try {
780+
CosmosFaultInjectionHelper
781+
.configureFaultInjectionRules(bulkAsyncContainer, Collections.singletonList(tooManyRequestRule))
782+
.block();
783+
784+
List<CosmosBulkOperationResponse<Object>> responses =
785+
bulkAsyncContainer
786+
.<Object>executeBulkOperations(Flux.fromIterable(cosmosItemOperations), new CosmosBulkExecutionOptions())
787+
.collectList()
788+
.block();
789+
790+
assertThat(responses).isNotNull();
791+
assertThat(responses.size()).isEqualTo(1);
792+
793+
for (CosmosBulkOperationResponse<Object> response : responses) {
794+
assertThat(response.getResponse()).isNotNull();
795+
assertThat(response.getResponse().getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code());
796+
797+
CosmosItemOperation operation = response.getOperation();
798+
assertThat(operation).isInstanceOf(ItemBulkOperation.class);
799+
800+
BulkOperationStatusTracker statusTracker =
801+
((ItemBulkOperation<?, ?>) operation).getStatusTracker();
802+
assertThat(statusTracker).isNotNull();
803+
804+
String statusHistory = statusTracker.toString();
805+
assertThat(statusHistory).contains("429/");
806+
assertThat(statusHistory).contains("count=2");
807+
}
808+
} finally {
809+
tooManyRequestRule.disable();
810+
}
811+
}
812+
760813
private void createItemsAndVerify(List<com.azure.cosmos.models.CosmosItemOperation> cosmosItemOperations) {
761814
CosmosBulkExecutionOptions cosmosBulkExecutionOptions = new CosmosBulkExecutionOptions();
762815

@@ -806,8 +859,19 @@ public Object[][] faultInjectionProvider() {
806859
};
807860
}
808861

809-
private FaultInjectionRule injectBatchFailure(String id, FaultInjectionServerErrorType serverErrorType, int hitLimit) {
862+
private FaultInjectionRule injectBatchFailure(
863+
String id,
864+
FaultInjectionServerErrorType serverErrorType,
865+
int hitLimit) {
866+
867+
return injectBatchFailure(id, serverErrorType, hitLimit, Duration.ofSeconds(1));
868+
}
810869

870+
private FaultInjectionRule injectBatchFailure(
871+
String id,
872+
FaultInjectionServerErrorType serverErrorType,
873+
int hitLimit,
874+
Duration startDelay) {
811875

812876
FaultInjectionServerErrorResultBuilder faultInjectionResultBuilder = FaultInjectionResultBuilders
813877
.getResultBuilder(serverErrorType)
@@ -833,7 +897,7 @@ private FaultInjectionRule injectBatchFailure(String id, FaultInjectionServerErr
833897
return new FaultInjectionRuleBuilder(id)
834898
.condition(condition)
835899
.result(result)
836-
.startDelay(Duration.ofSeconds(1))
900+
.startDelay(startDelay)
837901
.hitLimit(hitLimit)
838902
.build();
839903
}

0 commit comments

Comments
 (0)