From 8e7391876161dfa086710ccf131c91f60cfb49a2 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Sat, 6 Dec 2025 03:38:59 +0000 Subject: [PATCH 01/16] Cosmos Spark transactional batch support --- .../cosmos/spark/ItemsWriterBuilder.scala | 95 +- .../com/azure/cosmos/spark/CosmosConfig.scala | 11 + .../azure/cosmos/spark/CosmosWriterBase.scala | 27 +- .../spark/TransactionalBulkWriter.scala | 1506 +++++++++++++++++ .../azure/cosmos/spark/BulkWriterITest.scala | 2 +- .../spark/TransactionalBatchITest.scala | 637 +++++++ .../batch/TransactionalBulkExecutor.java | 1162 +++++++++++++ 7 files changed, 3429 insertions(+), 11 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala create mode 100644 sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/TransactionalBulkExecutor.java diff --git a/sdk/cosmos/azure-cosmos-spark_3-5/src/main/scala/com/azure/cosmos/spark/ItemsWriterBuilder.scala b/sdk/cosmos/azure-cosmos-spark_3-5/src/main/scala/com/azure/cosmos/spark/ItemsWriterBuilder.scala index 091df01d5ece..1ace9707297f 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5/src/main/scala/com/azure/cosmos/spark/ItemsWriterBuilder.scala +++ b/sdk/cosmos/azure-cosmos-spark_3-5/src/main/scala/com/azure/cosmos/spark/ItemsWriterBuilder.scala @@ -2,11 +2,14 @@ // Licensed under the MIT License. package com.azure.cosmos.spark +import com.azure.cosmos.{CosmosAsyncClient, ReadConsistencyStrategy, SparkBridgeInternal} import com.azure.cosmos.spark.diagnostics.LoggerHelper import org.apache.spark.broadcast.Broadcast +import org.apache.spark.sql.connector.distributions.{Distribution, Distributions} +import org.apache.spark.sql.connector.expressions.{Expression, Expressions, NullOrdering, SortDirection, SortOrder} import org.apache.spark.sql.connector.metric.CustomMetric import org.apache.spark.sql.connector.write.streaming.StreamingWrite -import org.apache.spark.sql.connector.write.{BatchWrite, Write, WriteBuilder} +import org.apache.spark.sql.connector.write.{BatchWrite, RequiresDistributionAndOrdering, Write, WriteBuilder} import org.apache.spark.sql.types.StructType import org.apache.spark.sql.util.CaseInsensitiveStringMap @@ -46,7 +49,7 @@ private class ItemsWriterBuilder diagnosticsConfig, sparkEnvironmentInfo) - private class CosmosWrite extends Write { + private class CosmosWrite extends Write with RequiresDistributionAndOrdering { private[this] val supportedCosmosMetrics: Array[CustomMetric] = { Array( @@ -56,6 +59,15 @@ private class ItemsWriterBuilder ) } + private[this] val writeConfig = CosmosWriteConfig.parseWriteConfig( + userConfig.asCaseSensitiveMap().asScala.toMap, + inputSchema + ) + + private[this] val containerConfig = CosmosContainerConfig.parseCosmosContainerConfig( + userConfig.asCaseSensitiveMap().asScala.toMap + ) + override def toBatch(): BatchWrite = new ItemsBatchWriter( userConfig.asCaseSensitiveMap().asScala.toMap, @@ -73,5 +85,84 @@ private class ItemsWriterBuilder sparkEnvironmentInfo) override def supportedCustomMetrics(): Array[CustomMetric] = supportedCosmosMetrics + + override def requiredDistribution(): Distribution = { + if (writeConfig.bulkEnabled && writeConfig.bulkEnableTransactions) { + // For transactional writes, partition by all partition key columns + val partitionKeyPaths = getPartitionKeyColumnNames() + if (partitionKeyPaths.nonEmpty) { + // Use public Expressions.column() factory - returns NamedReference + val clustering = partitionKeyPaths.map(path => Expressions.column(path): Expression).toArray + Distributions.clustered(clustering) + } else { + Distributions.unspecified() + } + } else { + Distributions.unspecified() + } + } + + override def requiredOrdering(): Array[SortOrder] = { + if (writeConfig.bulkEnabled && writeConfig.bulkEnableTransactions) { + // For transactional writes, order by all partition key columns (ascending) + val partitionKeyPaths = getPartitionKeyColumnNames() + if (partitionKeyPaths.nonEmpty) { + partitionKeyPaths.map { path => + // Use public Expressions.sort() factory for creating SortOrder + Expressions.sort( + Expressions.column(path), + SortDirection.ASCENDING, + NullOrdering.NULLS_FIRST + ) + }.toArray + } else { + Array.empty[SortOrder] + } + } else { + Array.empty[SortOrder] + } + } + + private def getPartitionKeyColumnNames(): Seq[String] = { + try { + // Need to create a temporary container client to get partition key definition + val clientCacheItem = CosmosClientCache( + CosmosClientConfiguration( + userConfig.asCaseSensitiveMap().asScala.toMap, + ReadConsistencyStrategy.EVENTUAL, + sparkEnvironmentInfo + ), + Some(cosmosClientStateHandles.value.cosmosClientMetadataCaches), + "ItemsWriterBuilder-PKLookup" + ) + + val container = ThroughputControlHelper.getContainer( + userConfig.asCaseSensitiveMap().asScala.toMap, + containerConfig, + clientCacheItem, + None + ) + + val containerProperties = SparkBridgeInternal.getContainerPropertiesFromCollectionCache(container) + val partitionKeyDefinition = containerProperties.getPartitionKeyDefinition + + // Release the client + clientCacheItem.close() + + if (partitionKeyDefinition != null && partitionKeyDefinition.getPaths != null) { + val paths = partitionKeyDefinition.getPaths.asScala + paths.map(path => { + // Remove leading '/' from partition key path (e.g., "/pk" -> "pk") + if (path.startsWith("/")) path.substring(1) else path + }).toSeq + } else { + Seq.empty[String] + } + } catch { + case ex: Exception => + log.logWarning(s"Failed to get partition key definition for transactional writes: ${ex.getMessage}") + Seq.empty[String] + } + } } } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala index b2137fd09dcd..3c931fd85892 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala @@ -113,6 +113,7 @@ private[spark] object CosmosConfigNames { val ClientTelemetryEnabled = "spark.cosmos.clientTelemetry.enabled" // keep this to avoid breaking changes val ClientTelemetryEndpoint = "spark.cosmos.clientTelemetry.endpoint" // keep this to avoid breaking changes val WriteBulkEnabled = "spark.cosmos.write.bulk.enabled" + val WriteBulkEnableTransactions = "spark.cosmos.write.bulk.enableTransactions" val WriteBulkMaxPendingOperations = "spark.cosmos.write.bulk.maxPendingOperations" val WriteBulkMaxBatchSize = "spark.cosmos.write.bulk.maxBatchSize" val WriteBulkMinTargetBatchSize = "spark.cosmos.write.bulk.minTargetBatchSize" @@ -242,6 +243,7 @@ private[spark] object CosmosConfigNames { ClientTelemetryEnabled, ClientTelemetryEndpoint, WriteBulkEnabled, + WriteBulkEnableTransactions, WriteBulkMaxPendingOperations, WriteBulkMaxConcurrentPartitions, WriteBulkPayloadSizeInBytes, @@ -1462,6 +1464,7 @@ private case class CosmosPatchConfigs(columnConfigsMap: TrieMap[String, CosmosPa private case class CosmosWriteConfig(itemWriteStrategy: ItemWriteStrategy, maxRetryCount: Int, bulkEnabled: Boolean, + bulkEnableTransactions: Boolean = false, bulkMaxPendingOperations: Option[Int] = None, pointMaxConcurrency: Option[Int] = None, maxConcurrentCosmosPartitions: Option[Int] = None, @@ -1486,6 +1489,12 @@ private object CosmosWriteConfig { parseFromStringFunction = bulkEnabledAsString => bulkEnabledAsString.toBoolean, helpMessage = "Cosmos DB Item Write bulk enabled") + private val bulkEnableTransactions = CosmosConfigEntry[Boolean](key = CosmosConfigNames.WriteBulkEnableTransactions, + defaultValue = Option.apply(false), + mandatory = false, + parseFromStringFunction = enableTransactionsAsString => enableTransactionsAsString.toBoolean, + helpMessage = "Cosmos DB Item Write enable transactional batch - requires bulk write to be enabled and Spark 3.5+") + private val microBatchPayloadSizeInBytes = CosmosConfigEntry[Int](key = CosmosConfigNames.WriteBulkPayloadSizeInBytes, defaultValue = Option.apply(BatchRequestResponseConstants.DEFAULT_MAX_DIRECT_MODE_BATCH_REQUEST_BODY_SIZE_IN_BYTES), mandatory = false, @@ -1758,6 +1767,7 @@ private object CosmosWriteConfig { val itemWriteStrategyOpt = CosmosConfigEntry.parse(cfg, itemWriteStrategy) val maxRetryCountOpt = CosmosConfigEntry.parse(cfg, maxRetryCount) val bulkEnabledOpt = CosmosConfigEntry.parse(cfg, bulkEnabled) + val bulkEnableTransactionsOpt = CosmosConfigEntry.parse(cfg, bulkEnableTransactions) var patchConfigsOpt = Option.empty[CosmosPatchConfigs] val throughputControlConfigOpt = CosmosThroughputControlConfig.parseThroughputControlConfig(cfg) val microBatchPayloadSizeInBytesOpt = CosmosConfigEntry.parse(cfg, microBatchPayloadSizeInBytes) @@ -1788,6 +1798,7 @@ private object CosmosWriteConfig { itemWriteStrategyOpt.get, maxRetryCountOpt.get, bulkEnabled = bulkEnabledOpt.get, + bulkEnableTransactions = bulkEnableTransactionsOpt.getOrElse(false), bulkMaxPendingOperations = CosmosConfigEntry.parse(cfg, bulkMaxPendingOperations), pointMaxConcurrency = CosmosConfigEntry.parse(cfg, pointWriteConcurrency), maxConcurrentCosmosPartitions = CosmosConfigEntry.parse(cfg, bulkMaxConcurrentPartitions), diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosWriterBase.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosWriterBase.scala index 726a6fe461ac..0edfbf2170c9 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosWriterBase.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosWriterBase.scala @@ -62,14 +62,25 @@ private abstract class CosmosWriterBase( private val writer: AtomicReference[AsyncItemWriter] = new AtomicReference( if (cosmosWriteConfig.bulkEnabled) { - new BulkWriter( - container, - cosmosTargetContainerConfig, - partitionKeyDefinition, - cosmosWriteConfig, - diagnosticsConfig, - getOutputMetricsPublisher(), - commitAttempt.getAndIncrement()) + if (cosmosWriteConfig.bulkEnableTransactions) { + new TransactionalBulkWriter( + container, + cosmosTargetContainerConfig, + partitionKeyDefinition, + cosmosWriteConfig, + diagnosticsConfig, + getOutputMetricsPublisher(), + commitAttempt.getAndIncrement()) + } else { + new BulkWriter( + container, + cosmosTargetContainerConfig, + partitionKeyDefinition, + cosmosWriteConfig, + diagnosticsConfig, + getOutputMetricsPublisher(), + commitAttempt.getAndIncrement()) + } } else { new PointWriter( container, diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala new file mode 100644 index 000000000000..7d5007c78920 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala @@ -0,0 +1,1506 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.spark + +// scalastyle:off underscore.import +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils +import com.azure.cosmos.implementation.batch.{BatchRequestResponseConstants, BulkExecutorDiagnosticsTracker, ItemBulkOperation, TransactionalBulkExecutor} +import com.azure.cosmos.implementation.{CosmosDaemonThreadFactory, CosmosBulkExecutionOptionsImpl, ImplementationBridgeHelpers, UUIDs} +import com.azure.cosmos.models._ +import com.azure.cosmos.spark.TransactionalBulkWriter.{BulkOperationFailedException, bulkWriterInputBoundedElastic, bulkWriterRequestsBoundedElastic, bulkWriterResponsesBoundedElastic, getThreadInfo, readManyBoundedElastic} +import com.azure.cosmos.spark.diagnostics.DefaultDiagnostics +import com.azure.cosmos.{BridgeInternal, CosmosAsyncContainer, CosmosDiagnosticsContext, CosmosEndToEndOperationLatencyPolicyConfigBuilder, CosmosException} +import reactor.core.Scannable +import reactor.core.publisher.{Flux, Mono} +import reactor.core.scheduler.Scheduler + +import java.util +import java.util.Objects +import java.util.concurrent.{ScheduledFuture, ScheduledThreadPoolExecutor} +import scala.collection.concurrent.TrieMap +import scala.collection.convert.ImplicitConversions.`collection AsScalaIterable` +import scala.collection.mutable +import scala.concurrent.duration.Duration +// scalastyle:on underscore.import +import com.azure.cosmos.implementation.ImplementationBridgeHelpers +import com.azure.cosmos.implementation.guava25.base.Preconditions +import com.azure.cosmos.implementation.spark.{OperationContextAndListenerTuple, OperationListener} +import com.azure.cosmos.models.PartitionKey +import com.azure.cosmos.spark.TransactionalBulkWriter.{DefaultMaxPendingOperationPerCore, emitFailureHandler} +import com.azure.cosmos.spark.diagnostics.{DiagnosticsContext, DiagnosticsLoader, LoggerHelper, SparkTaskContext} +import com.fasterxml.jackson.databind.node.ObjectNode +import org.apache.spark.TaskContext +import reactor.core.Disposable +import reactor.core.publisher.Sinks +import reactor.core.publisher.Sinks.{EmitFailureHandler, EmitResult} +import reactor.core.scala.publisher.SMono.PimpJFlux +import reactor.core.scala.publisher.{SFlux, SMono} +import reactor.core.scheduler.Schedulers + +import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger, AtomicLong, AtomicReference} +import java.util.concurrent.locks.ReentrantLock +import java.util.concurrent.{Semaphore, TimeUnit} +// scalastyle:off underscore.import +import scala.collection.JavaConverters._ +// scalastyle:on underscore.import + +//scalastyle:off null +//scalastyle:off multiple.string.literals +//scalastyle:off file.size.limit +private class TransactionalBulkWriter +( + container: CosmosAsyncContainer, + containerConfig: CosmosContainerConfig, + partitionKeyDefinition: PartitionKeyDefinition, + writeConfig: CosmosWriteConfig, + diagnosticsConfig: DiagnosticsConfig, + outputMetricsPublisher: OutputMetricsPublisherTrait, + commitAttempt: Long = 1 +) extends AsyncItemWriter { + + private val log = LoggerHelper.getLogger(diagnosticsConfig, this.getClass) + + private val verboseLoggingAfterReEnqueueingRetriesEnabled = new AtomicBoolean(false) + + private val cpuCount = SparkUtils.getNumberOfHostCPUCores + // each bulk writer allows up to maxPendingOperations being buffered + // there is one bulk writer per spark task/partition + // and default config will create one executor per core on the executor host + // so multiplying by cpuCount in the default config is too aggressive + private val maxPendingOperations = writeConfig.bulkMaxPendingOperations + .getOrElse(DefaultMaxPendingOperationPerCore) + private val maxConcurrentPartitions = writeConfig.maxConcurrentCosmosPartitions match { + // using the provided maximum of concurrent partitions per Spark partition on the input data + // multiplied by 2 to leave space for partition splits during ingestion + case Some(configuredMaxConcurrentPartitions) => 2 * configuredMaxConcurrentPartitions + // using the total number of physical partitions + // multiplied by 2 to leave space for partition splits during ingestion + case None => 2 * ContainerFeedRangesCache.getFeedRanges(container, containerConfig.feedRangeRefreshIntervalInSecondsOpt).block().size + } + log.logInfo( + s"BulkWriter instantiated (Host CPU count: $cpuCount, maxPendingOperations: $maxPendingOperations, " + + s"maxConcurrentPartitions: $maxConcurrentPartitions ...") + + // Artificial operation used to signale to the bufferUntil operator that + // the buffer should be flushed. A timer-based scheduler will publish this + // dummy operation for every batchIntervalInMs ms. This operation + // is filtered out and will never be flushed to the backend + private val readManyFlushOperationSingleton = ReadManyOperation( + new CosmosItemIdentity( + new PartitionKey("ReadManyOperation.FlushSingleton"), + "ReadManyOperation.FlushSingleton" + ), + null, + null + ) + + private val closed = new AtomicBoolean(false) + private val lock = new ReentrantLock + private val pendingTasksCompleted = lock.newCondition + private val pendingRetries = new AtomicLong(0) + private val pendingBulkWriteRetries = java.util.concurrent.ConcurrentHashMap.newKeySet[CosmosItemOperation]().asScala + private val pendingReadManyRetries = java.util.concurrent.ConcurrentHashMap.newKeySet[ReadManyOperation]().asScala + private val activeTasks = new AtomicInteger(0) + private val errorCaptureFirstException = new AtomicReference[Throwable]() + private val bulkInputEmitter: Sinks.Many[CosmosItemOperation] = Sinks.many().unicast().onBackpressureBuffer() + + private val activeBulkWriteOperations =java.util.concurrent.ConcurrentHashMap.newKeySet[CosmosItemOperation]().asScala + private val activeReadManyOperations = java.util.concurrent.ConcurrentHashMap.newKeySet[ReadManyOperation]().asScala + private val semaphore = new Semaphore(maxPendingOperations) + + private val totalScheduledMetrics = new AtomicLong(0) + private val totalSuccessfulIngestionMetrics = new AtomicLong(0) + + private val maxOperationTimeout = java.time.Duration.ofSeconds(CosmosConstants.batchOperationEndToEndTimeoutInSeconds) + private val endToEndTimeoutPolicy = new CosmosEndToEndOperationLatencyPolicyConfigBuilder(maxOperationTimeout) + .enable(true) + .build + private val cosmosBulkExecutionOptions = new CosmosBulkExecutionOptions(TransactionalBulkWriter.bulkProcessingThresholds) + + private val cosmosBulkExecutionOptionsImpl = ImplementationBridgeHelpers.CosmosBulkExecutionOptionsHelper + .getCosmosBulkExecutionOptionsAccessor + .getImpl(cosmosBulkExecutionOptions) + private val monotonicOperationCounter = new AtomicLong(0) + + cosmosBulkExecutionOptionsImpl.setSchedulerOverride(bulkWriterRequestsBoundedElastic) + cosmosBulkExecutionOptionsImpl.setMaxConcurrentCosmosPartitions(maxConcurrentPartitions) + cosmosBulkExecutionOptionsImpl.setCosmosEndToEndLatencyPolicyConfig(endToEndTimeoutPolicy) + + private class ForwardingMetricTracker(val verboseLoggingEnabled: AtomicBoolean) extends BulkExecutorDiagnosticsTracker { + override def trackDiagnostics(ctx: CosmosDiagnosticsContext): Unit = { + val ctxOption = Option.apply(ctx) + outputMetricsPublisher.trackWriteOperation(0, ctxOption) + if (ctxOption.isDefined && verboseLoggingEnabled.get) { + TransactionalBulkWriter.log.logWarning(s"Track bulk operation after re-enqueued retry: ${ctxOption.get.toJson}") + } + } + + override def verboseLoggingAfterReEnqueueingRetriesEnabled(): Boolean = { + verboseLoggingEnabled.get() + } + } + + cosmosBulkExecutionOptionsImpl.setDiagnosticsTracker( + new ForwardingMetricTracker(verboseLoggingAfterReEnqueueingRetriesEnabled) + ) + + ThroughputControlHelper.populateThroughputControlGroupName(cosmosBulkExecutionOptions, writeConfig.throughputControlConfig) + + writeConfig.maxMicroBatchPayloadSizeInBytes match { + case Some(customMaxMicroBatchPayloadSizeInBytes) => + cosmosBulkExecutionOptionsImpl + .setMaxMicroBatchPayloadSizeInBytes(customMaxMicroBatchPayloadSizeInBytes) + case None => + } + + writeConfig.initialMicroBatchSize match { + case Some(customInitialMicroBatchSize) => + cosmosBulkExecutionOptions.setInitialMicroBatchSize(Math.max(1, customInitialMicroBatchSize)) + case None => + } + + writeConfig.maxMicroBatchSize match { + case Some(customMaxMicroBatchSize) => + cosmosBulkExecutionOptions.setMaxMicroBatchSize( + Math.max( + 1, + Math.min(customMaxMicroBatchSize, BatchRequestResponseConstants.MAX_OPERATIONS_IN_DIRECT_MODE_BATCH_REQUEST) + ) + ) + case None => + } + + private val operationContext = initializeOperationContext() + private val cosmosPatchHelperOpt = writeConfig.itemWriteStrategy match { + case ItemWriteStrategy.ItemPatch | ItemWriteStrategy.ItemPatchIfExists | ItemWriteStrategy.ItemBulkUpdate => + Some(new CosmosPatchHelper(diagnosticsConfig, writeConfig.patchConfigs.get)) + case _ => None + } + private val readManyInputEmitterOpt: Option[Sinks.Many[ReadManyOperation]] = { + writeConfig.itemWriteStrategy match { + case ItemWriteStrategy.ItemBulkUpdate => Some(Sinks.many().unicast().onBackpressureBuffer()) + case _ => None + } + } + + private val batchIntervalInMs = cosmosBulkExecutionOptionsImpl + .getMaxMicroBatchInterval + .toMillis + + private[this] val flushExecutorHolder: Option[(ScheduledThreadPoolExecutor, ScheduledFuture[_])] = { + writeConfig.itemWriteStrategy match { + case ItemWriteStrategy.ItemBulkUpdate => + val executor = new ScheduledThreadPoolExecutor( + 1, + new CosmosDaemonThreadFactory( + "BulkWriterReadManyFlush" + UUIDs.nonBlockingRandomUUID() + )) + executor.setExecuteExistingDelayedTasksAfterShutdownPolicy(false) + executor.setRemoveOnCancelPolicy(true) + + val future:ScheduledFuture[_] = executor.scheduleWithFixedDelay( + () => this.onFlushReadMany(), + batchIntervalInMs, + batchIntervalInMs, + TimeUnit.MILLISECONDS) + + Some(executor, future) + case _ => None + } + } + + private def initializeOperationContext(): SparkTaskContext = { + val taskContext = TaskContext.get + + val diagnosticsContext: DiagnosticsContext = DiagnosticsContext(UUIDs.nonBlockingRandomUUID(), "BulkWriter") + + if (taskContext != null) { + val taskDiagnosticsContext = SparkTaskContext(diagnosticsContext.correlationActivityId, + taskContext.stageId(), + taskContext.partitionId(), + taskContext.taskAttemptId(), + "") + + val listener: OperationListener = + DiagnosticsLoader.getDiagnosticsProvider(diagnosticsConfig).getLogger(this.getClass) + + val operationContextAndListenerTuple = new OperationContextAndListenerTuple(taskDiagnosticsContext, listener) + cosmosBulkExecutionOptionsImpl + .setOperationContextAndListenerTuple(operationContextAndListenerTuple) + + taskDiagnosticsContext + } else{ + SparkTaskContext(diagnosticsContext.correlationActivityId, + -1, + -1, + -1, + "") + } + } + + private def onFlushReadMany(): Unit = { + if (this.readManyInputEmitterOpt.isEmpty) { + throw new IllegalStateException("Callback onFlushReadMany should only be scheduled for bulk update.") + } + try { + this.readManyInputEmitterOpt.get.tryEmitNext(readManyFlushOperationSingleton) match { + case EmitResult.OK => log.logInfo("onFlushReadMany Successfully emitted flush") + case faultEmitResult => + log.logError(s"Callback invocation 'onFlush' failed with result: $faultEmitResult.") + } + } + catch { + case t: Throwable => + log.logError("Callback invocation 'onFlush' failed.", t) + } + } + + private val readManySubscriptionDisposableOpt: Option[Disposable] = { + writeConfig.itemWriteStrategy match { + case ItemWriteStrategy.ItemBulkUpdate => Some(createReadManySubscriptionDisposable()) + case _ => None + } + } + + private def createReadManySubscriptionDisposable(): Disposable = { + log.logTrace(s"readManySubscriptionDisposable, Context: ${operationContext.toString} $getThreadInfo") + + // We start from using the bulk batch size and interval and concurrency + // If in the future, there is a need to separate the configuration, can re-consider + val bulkBatchSize = writeConfig.maxMicroBatchSize match { + case Some(customMaxMicroBatchSize) => Math.min( + BatchRequestResponseConstants.MAX_OPERATIONS_IN_DIRECT_MODE_BATCH_REQUEST, + Math.max(1, customMaxMicroBatchSize)) + case None => BatchRequestResponseConstants.MAX_OPERATIONS_IN_DIRECT_MODE_BATCH_REQUEST + } + + val batchConcurrency = cosmosBulkExecutionOptionsImpl.getMaxMicroBatchConcurrency + + val firstRecordTimeStamp = new AtomicLong(-1) + val currentMicroBatchSize = new AtomicLong(0) + + readManyInputEmitterOpt + .get + .asFlux() + .publishOn(readManyBoundedElastic) + .timestamp + .bufferUntil(timestampReadManyOperationTuple => { + val timestamp = timestampReadManyOperationTuple.getT1 + val readManyOperation = timestampReadManyOperationTuple.getT2 + + if (readManyOperation eq readManyFlushOperationSingleton) { + log.logTrace(s"FlushSingletonReceived, Context: ${operationContext.toString}") + val currentMicroBatchSizeSnapshot = currentMicroBatchSize.get + if (currentMicroBatchSizeSnapshot > 0) { + firstRecordTimeStamp.set(-1) + currentMicroBatchSize.set(0) + log.logTrace(s"FlushSingletonReceived - flushing batch, Context: ${operationContext.toString}") + true + } else { + // avoid counting flush operations for the micro batch size calculation + log.logTrace(s"FlushSingletonReceived - empty buffer, nothing to flush, Context: ${operationContext.toString}") + false + } + } else { + + firstRecordTimeStamp.compareAndSet(-1, timestamp) + val age = timestamp - firstRecordTimeStamp.get + val batchSize = currentMicroBatchSize.incrementAndGet + + if (batchSize >= bulkBatchSize || age >= batchIntervalInMs) { + log.logTrace(s"BatchIntervalExpired - flushing batch, Context: ${operationContext.toString}") + firstRecordTimeStamp.set(-1) + currentMicroBatchSize.set(0) + true + } else { + false + } + } + }) + .subscribeOn(readManyBoundedElastic) + .asScala + .flatMap(timestampReadManyOperationTuples => { + val readManyOperations = timestampReadManyOperationTuples + .filter(candidate => !candidate.getT2.equals(readManyFlushOperationSingleton)) + .map(tuple => tuple.getT2) + + if (readManyOperations.isEmpty) { + Mono.empty() + } else { + val cosmosIdentitySet = readManyOperations.map(option => option.cosmosItemIdentity).toSet + + // for each batch, use readMany to read items from cosmosdb + val requestOptions = new CosmosReadManyRequestOptions() + val requestOptionsImpl = ImplementationBridgeHelpers + .CosmosReadManyRequestOptionsHelper + .getCosmosReadManyRequestOptionsAccessor.getImpl(requestOptions) + ThroughputControlHelper.populateThroughputControlGroupName( + requestOptionsImpl, + writeConfig.throughputControlConfig) + ImplementationBridgeHelpers + .CosmosAsyncContainerHelper + .getCosmosAsyncContainerAccessor + .readMany(container, cosmosIdentitySet.toList.asJava, requestOptions, classOf[ObjectNode]) + .switchIfEmpty( + // For Java SDK, empty pages will not be returned (this can happen when all the items does not exists yet) + // create a fake empty page response + Mono.just( + ImplementationBridgeHelpers + .FeedResponseHelper + .getFeedResponseAccessor + .createFeedResponse(new util.ArrayList[ObjectNode](), null, null))) + .doOnNext(feedResponse => { + // Tracking the bytes read as part of client-side patch (readMany + replace) as bytes written as well + // to have a way to indicate the additional work happening here + outputMetricsPublisher.trackWriteOperation( + 0, + Option.apply(feedResponse.getCosmosDiagnostics) match { + case Some(diagnostics) => Option.apply(diagnostics.getDiagnosticsContext) + case None => None + }) + val resultMap = new TrieMap[CosmosItemIdentity, ObjectNode]() + for (itemNode: ObjectNode <- feedResponse.getResults.asScala) { + resultMap += ( + new CosmosItemIdentity( + PartitionKeyHelper.getPartitionKeyPath(itemNode, partitionKeyDefinition), + itemNode.get(CosmosConstants.Properties.Id).asText()) -> itemNode) + } + + // It is possible that multiple cosmosPatchBulkUpdateOperations were targeting for the same item + // Currently, we are still creating one bulk item operation for each cosmosPatchBulkUpdateOperations + // for easier exception and semaphore handling + // However a consequences of it could be, the generated bulk item operation will fail due to conflicts or pre-condition failure + // If this turns out to be a problem, we can do more optimization here: merge multiple cosmosPatchBulkUpdateOperations into one bulkItemOperation + // But even the above approach can only work within the same batch but not for the whole spark partition processing. + for (readManyOperation <- readManyOperations) { + val cosmosPatchBulkUpdateOperations = + cosmosPatchHelperOpt + .get + .createCosmosPatchBulkUpdateOperations(readManyOperation.objectNode) + + val rootNode = + cosmosPatchHelperOpt + .get + .patchBulkUpdateItem(resultMap.get(readManyOperation.cosmosItemIdentity), cosmosPatchBulkUpdateOperations) + + // create bulk item operation + val etagOpt = Option.apply(rootNode.get(CosmosConstants.Properties.ETag)) + val bulkItemOperation = etagOpt match { + case Some(etag) => + CosmosBulkOperations.getReplaceItemOperation( + readManyOperation.cosmosItemIdentity.getId, + rootNode, + readManyOperation.cosmosItemIdentity.getPartitionKey, + new CosmosBulkItemRequestOptions().setIfMatchETag(etag.asText()), + new OperationContext( + readManyOperation.operationContext.itemId, + readManyOperation.operationContext.partitionKeyValue, + Some(etag.asText()), + readManyOperation.operationContext.attemptNumber, + monotonicOperationCounter.incrementAndGet(), + Some(readManyOperation.objectNode) + )) + case None => CosmosBulkOperations.getCreateItemOperation( + rootNode, + readManyOperation.cosmosItemIdentity.getPartitionKey, + new OperationContext( + readManyOperation.operationContext.itemId, + readManyOperation.operationContext.partitionKeyValue, + eTagInput = None, + readManyOperation.operationContext.attemptNumber, + monotonicOperationCounter.incrementAndGet(), + Some(readManyOperation.objectNode) + )) + } + + this.emitBulkInput(bulkItemOperation) + } + }) + .onErrorResume(throwable => { + for (readManyOperation <- readManyOperations) { + handleReadManyExceptions(throwable, readManyOperation) + } + + Mono.empty() + }) + .doFinally(_ => { + for (readManyOperation <- readManyOperations) { + val activeReadManyOperationFound = activeReadManyOperations.remove(readManyOperation) + // for ItemBulkUpdate strategy, each active task includes two stages: ReadMany + BulkWrite + // so we are not going to make task complete here + if (!activeReadManyOperationFound) { + // can't find the read-many operation in list of active operations! + logInfoOrWarning(s"Cannot find active read-many for '" + + s"${readManyOperation.cosmosItemIdentity.getPartitionKey}/" + + s"${readManyOperation.cosmosItemIdentity.getId}'. This can happen when " + + s"retries get re-enqueued.") + + if (pendingReadManyRetries.remove(readManyOperation)) { + pendingRetries.decrementAndGet() + } + } + } + }) + .`then`(Mono.empty()) + } + }, batchConcurrency) + .subscribe() + } + + private def handleReadManyExceptions(throwable: Throwable, ReadManyOperation: ReadManyOperation): Unit = { + throwable match { + case e: CosmosException => + outputMetricsPublisher.trackWriteOperation( + 0, + Option.apply(e.getDiagnostics) match { + case Some(diagnostics) => Option.apply(diagnostics.getDiagnosticsContext) + case None => None + }) + val requestOperationContext = ReadManyOperation.operationContext + if (shouldRetry(e.getStatusCode, e.getSubStatusCode, requestOperationContext)) { + log.logInfo(s"for itemId=[${requestOperationContext.itemId}], partitionKeyValue=[${requestOperationContext.partitionKeyValue}], " + + s"encountered status code '${e.getStatusCode}:${e.getSubStatusCode}' in read many, will retry! " + + s"attemptNumber=${requestOperationContext.attemptNumber}, exceptionMessage=${e.getMessage}, " + + s"Context: {${operationContext.toString}} $getThreadInfo") + + // the task will be re-queued at the beginning of the flow, so mark it complete here + markTaskCompletion() + + this.scheduleRetry( + trackPendingRetryAction = () => pendingReadManyRetries.add(ReadManyOperation), + clearPendingRetryAction = () => pendingReadManyRetries.remove(ReadManyOperation), + ReadManyOperation.cosmosItemIdentity.getPartitionKey, + ReadManyOperation.objectNode, + ReadManyOperation.operationContext, + e.getStatusCode) + + } else { + // Non-retryable exception or has exceeded the max retry count + val requestOperationContext = ReadManyOperation.operationContext + log.logError(s"for itemId=[${requestOperationContext.itemId}], partitionKeyValue=[${requestOperationContext.partitionKeyValue}], " + + s"encountered status code '${e.getStatusCode}:${e.getSubStatusCode}', all retries exhausted! " + + s"attemptNumber=${requestOperationContext.attemptNumber}, exceptionMessage=${e.getMessage}, " + + s"Context: {${operationContext.toString} $getThreadInfo") + + val message = s"All retries exhausted for readMany - " + + s"statusCode=[${e.getStatusCode}:${e.getSubStatusCode}] " + + s"itemId=[${requestOperationContext.itemId}], partitionKeyValue=[${requestOperationContext.partitionKeyValue}]" + + val exceptionToBeThrown = new BulkOperationFailedException(e.getStatusCode, e.getSubStatusCode, message, e) + captureIfFirstFailure(exceptionToBeThrown) + cancelWork() + markTaskCompletion() + } + case _ => // handle non cosmos exceptions + log.logError(s"Unexpected failure code path in Bulk ingestion readMany stage, " + + s"Context: ${operationContext.toString} $getThreadInfo", throwable) + captureIfFirstFailure(throwable) + cancelWork() + markTaskCompletion() + } + } + + private def scheduleRetry( + trackPendingRetryAction: () => Boolean, + clearPendingRetryAction: () => Boolean, + partitionKey: PartitionKey, + objectNode: ObjectNode, + operationContext: OperationContext, + statusCode: Int): Unit = { + if (trackPendingRetryAction()) { + this.pendingRetries.incrementAndGet() + } + // this is to ensure the submission will happen on a different thread in background + // and doesn't block the active thread + val deferredRetryMono = SMono.defer(() => { + scheduleWriteInternal( + partitionKey, + objectNode, + new OperationContext( + operationContext.itemId, + operationContext.partitionKeyValue, + operationContext.eTag, + operationContext.attemptNumber + 1, + operationContext.sequenceNumber)) + if (clearPendingRetryAction()) { + this.pendingRetries.decrementAndGet() + } + SMono.empty + }) + + if (Exceptions.isTimeout(statusCode)) { + deferredRetryMono + .delaySubscription( + Duration( + TransactionalBulkWriter.minDelayOn408RequestTimeoutInMs + + scala.util.Random.nextInt( + TransactionalBulkWriter.maxDelayOn408RequestTimeoutInMs - TransactionalBulkWriter.minDelayOn408RequestTimeoutInMs), + TimeUnit.MILLISECONDS), + Schedulers.boundedElastic()) + .subscribeOn(Schedulers.boundedElastic()) + .subscribe() + + } else { + deferredRetryMono + .subscribeOn(Schedulers.boundedElastic()) + .subscribe() + } + } + + private val subscriptionDisposable: Disposable = { + log.logTrace(s"subscriptionDisposable, Context: ${operationContext.toString} $getThreadInfo") + + val inputFlux = bulkInputEmitter + .asFlux() + .onBackpressureBuffer() + .publishOn(bulkWriterInputBoundedElastic) + .doOnError(t => { + log.logError(s"Input publishing flux failed, Context: ${operationContext.toString} $getThreadInfo", t) + }) + + // Use TransactionalBulkExecutor which internally uses SinglePartitionKeyServerBatchRequest + // with setAtomicBatch(true) and setShouldContinueOnError(false) for true transactional semantics + val cosmosBulkExecutionOptionsImpl = ImplementationBridgeHelpers.CosmosBulkExecutionOptionsHelper + .getCosmosBulkExecutionOptionsAccessor + .getImpl(cosmosBulkExecutionOptions) + + val transactionalExecutor = new TransactionalBulkExecutor[Object]( + container, + inputFlux, + cosmosBulkExecutionOptionsImpl) + + val bulkOperationResponseFlux: SFlux[CosmosBulkOperationResponse[Object]] = + transactionalExecutor + .execute() + .onBackpressureBuffer() + .publishOn(bulkWriterResponsesBoundedElastic) + .doOnError(t => { + log.logError(s"Transactional bulk execution flux failed, Context: ${operationContext.toString} $getThreadInfo", t) + }) + .asScala + + bulkOperationResponseFlux.subscribe( + resp => { + val isGettingRetried = new AtomicBoolean(false) + val shouldSkipTaskCompletion = new AtomicBoolean(false) + try { + val itemOperation = resp.getOperation + val itemOperationFound = activeBulkWriteOperations.remove(itemOperation) + val pendingRetriesFound = pendingBulkWriteRetries.remove(itemOperation) + + if (pendingRetriesFound) { + pendingRetries.decrementAndGet() + } + + if (!itemOperationFound) { + // can't find the item operation in list of active operations! + logInfoOrWarning(s"Cannot find active operation for '${itemOperation.getOperationType} " + + s"${itemOperation.getPartitionKeyValue}/${itemOperation.getId}'. This can happen when " + + s"retries get re-enqueued.") + shouldSkipTaskCompletion.set(true) + } + if (pendingRetriesFound || itemOperationFound) { + val context = itemOperation.getContext[OperationContext] + val itemResponse = resp.getResponse + + if (resp.getException != null) { + Option(resp.getException) match { + case Some(cosmosException: CosmosException) => + handleNonSuccessfulStatusCode( + context, itemOperation, itemResponse, isGettingRetried, Some(cosmosException)) + case _ => + log.logWarning( + s"unexpected failure: itemId=[${context.itemId}], partitionKeyValue=[" + + s"${context.partitionKeyValue}], encountered , attemptNumber=${context.attemptNumber}, " + + s"exceptionMessage=${resp.getException.getMessage}, " + + s"Context: ${operationContext.toString} $getThreadInfo", resp.getException) + captureIfFirstFailure(resp.getException) + cancelWork() + } + } else if (Option(itemResponse).isEmpty || !itemResponse.isSuccessStatusCode) { + handleNonSuccessfulStatusCode(context, itemOperation, itemResponse, isGettingRetried, None) + } else { + // no error case + outputMetricsPublisher.trackWriteOperation(1, None) + totalSuccessfulIngestionMetrics.getAndIncrement() + } + } + } + finally { + if (!isGettingRetried.get) { + semaphore.release() + } + } + + if (!shouldSkipTaskCompletion.get) { + markTaskCompletion() + } + }, + errorConsumer = Option.apply( + ex => { + log.logError(s"Unexpected failure code path in Bulk ingestion, " + + s"Context: ${operationContext.toString} $getThreadInfo", ex) + // if there is any failure this closes the bulk. + // at this point bulk api doesn't allow any retrying + // we don't know the list of failed item-operations + // they only way to retry to keep a dictionary of pending operations outside + // so we know which operations failed and which ones can be retried. + // this is currently a kill scenario. + captureIfFirstFailure(ex) + cancelWork() + markTaskCompletion() + } + ) + ) + } + + override def scheduleWrite(partitionKeyValue: PartitionKey, objectNode: ObjectNode): Unit = { + Preconditions.checkState(!closed.get()) + throwIfCapturedExceptionExists() + + val activeTasksSemaphoreTimeout = 10 + val operationContext = new OperationContext( + getId(objectNode), + partitionKeyValue, + getETag(objectNode), + 1, + monotonicOperationCounter.incrementAndGet()) + val numberOfIntervalsWithIdenticalActiveOperationSnapshots = new AtomicLong(0) + // Don't clone the activeOperations for the first iteration + // to reduce perf impact before the Semaphore has been acquired + // this means if the semaphore can't be acquired within 10 minutes + // the first attempt will always assume it wasn't stale - so effectively we + // allow staleness for ten additional minutes - which is perfectly fine + var activeBulkWriteOperationsSnapshot = mutable.Set.empty[CosmosItemOperation] + var pendingBulkWriteRetriesSnapshot = mutable.Set.empty[CosmosItemOperation] + var activeReadManyOperationsSnapshot = mutable.Set.empty[ReadManyOperation] + var pendingReadManyRetriesSnapshot = mutable.Set.empty[ReadManyOperation] + + log.logTrace( + s"Before TryAcquire ${totalScheduledMetrics.get}, Context: ${operationContext.toString} $getThreadInfo") + while (!semaphore.tryAcquire(activeTasksSemaphoreTimeout, TimeUnit.MINUTES)) { + log.logDebug(s"Not able to acquire semaphore, Context: ${operationContext.toString} $getThreadInfo") + if (subscriptionDisposable.isDisposed || + (readManySubscriptionDisposableOpt.isDefined && readManySubscriptionDisposableOpt.get.isDisposed)) { + captureIfFirstFailure( + new IllegalStateException("Can't accept any new work - BulkWriter has been disposed already")) + } + + throwIfProgressStaled( + "Semaphore acquisition", + activeBulkWriteOperationsSnapshot, + pendingBulkWriteRetriesSnapshot, + activeReadManyOperationsSnapshot, + pendingReadManyRetriesSnapshot, + numberOfIntervalsWithIdenticalActiveOperationSnapshots, + allowRetryOnNewBulkWriterInstance = false) + + activeBulkWriteOperationsSnapshot = activeBulkWriteOperations.clone() + pendingBulkWriteRetriesSnapshot = pendingBulkWriteRetries.clone() + activeReadManyOperationsSnapshot = activeReadManyOperations.clone() + pendingReadManyRetriesSnapshot = pendingReadManyRetries.clone() + } + + val cnt = totalScheduledMetrics.getAndIncrement() + log.logTrace(s"total scheduled $cnt, Context: ${operationContext.toString} $getThreadInfo") + + scheduleWriteInternal(partitionKeyValue, objectNode, operationContext) + } + + private def scheduleWriteInternal(partitionKeyValue: PartitionKey, + objectNode: ObjectNode, + operationContext: OperationContext): Unit = { + activeTasks.incrementAndGet() + if (operationContext.attemptNumber > 1) { + logInfoOrWarning(s"bulk scheduleWrite attemptCnt: ${operationContext.attemptNumber}, " + + s"Context: ${operationContext.toString} $getThreadInfo") + } + + // The handling will make sure that during retry: + // For itemBulkUpdate -> the retry will go through readMany stage -> bulkWrite stage. + // For other strategies -> the retry will only go through bulk write stage + writeConfig.itemWriteStrategy match { + case ItemWriteStrategy.ItemBulkUpdate => scheduleReadManyInternal(partitionKeyValue, objectNode, operationContext) + case _ => scheduleBulkWriteInternal(partitionKeyValue, objectNode, operationContext) + } + } + + private def scheduleReadManyInternal(partitionKeyValue: PartitionKey, + objectNode: ObjectNode, + operationContext: OperationContext): Unit = { + + // For FAIL_NON_SERIALIZED, will keep retry, while for other errors, use the default behavior + val readManyOperation = ReadManyOperation(new CosmosItemIdentity(partitionKeyValue, operationContext.itemId), objectNode, operationContext) + activeReadManyOperations.add(readManyOperation) + readManyInputEmitterOpt.get.emitNext(readManyOperation, emitFailureHandler) + } + + private def scheduleBulkWriteInternal(partitionKeyValue: PartitionKey, + objectNode: ObjectNode, + operationContext: OperationContext): Unit = { + + val bulkItemOperation = writeConfig.itemWriteStrategy match { + case ItemWriteStrategy.ItemOverwrite => + CosmosBulkOperations.getUpsertItemOperation(objectNode, partitionKeyValue, operationContext) + case ItemWriteStrategy.ItemOverwriteIfNotModified => + operationContext.eTag match { + case Some(eTag) => + CosmosBulkOperations.getReplaceItemOperation( + operationContext.itemId, + objectNode, + partitionKeyValue, + new CosmosBulkItemRequestOptions().setIfMatchETag(eTag), + operationContext) + case _ => CosmosBulkOperations.getCreateItemOperation(objectNode, partitionKeyValue, operationContext) + } + case ItemWriteStrategy.ItemAppend => + CosmosBulkOperations.getCreateItemOperation(objectNode, partitionKeyValue, operationContext) + case ItemWriteStrategy.ItemDelete => + CosmosBulkOperations.getDeleteItemOperation(operationContext.itemId, partitionKeyValue, operationContext) + case ItemWriteStrategy.ItemDeleteIfNotModified => + CosmosBulkOperations.getDeleteItemOperation( + operationContext.itemId, + partitionKeyValue, + operationContext.eTag match { + case Some(eTag) => new CosmosBulkItemRequestOptions().setIfMatchETag(eTag) + case _ => new CosmosBulkItemRequestOptions() + }, + operationContext) + case ItemWriteStrategy.ItemPatch | ItemWriteStrategy.ItemPatchIfExists => getPatchItemOperation( + operationContext.itemId, + partitionKeyValue, + partitionKeyDefinition, + objectNode, + operationContext) + case _ => + throw new RuntimeException(s"${writeConfig.itemWriteStrategy} not supported") + } + + this.emitBulkInput(bulkItemOperation) + } + + private[this] def emitBulkInput(bulkItemOperation: CosmosItemOperation): Unit = { + activeBulkWriteOperations.add(bulkItemOperation) + + // Simply emit the operation - TransactionalBulkExecutor will handle batching by partition key + bulkInputEmitter.emitNext(bulkItemOperation, emitFailureHandler) + } + + private[this] def getPatchItemOperation(itemId: String, + partitionKey: PartitionKey, + partitionKeyDefinition: PartitionKeyDefinition, + objectNode: ObjectNode, + context: OperationContext): CosmosItemOperation = { + + assert(writeConfig.patchConfigs.isDefined) + assert(cosmosPatchHelperOpt.isDefined) + val patchConfigs = writeConfig.patchConfigs.get + val cosmosPatchHelper = cosmosPatchHelperOpt.get + + val cosmosPatchOperations = cosmosPatchHelper.createCosmosPatchOperations(itemId, partitionKeyDefinition, objectNode) + + val requestOptions = new CosmosBulkPatchItemRequestOptions() + if (patchConfigs.filter.isDefined && !StringUtils.isEmpty(patchConfigs.filter.get)) { + requestOptions.setFilterPredicate(patchConfigs.filter.get) + } + + CosmosBulkOperations.getPatchItemOperation(itemId, partitionKey, cosmosPatchOperations, requestOptions, context) + } + + //scalastyle:off method.length + //scalastyle:off cyclomatic.complexity + private[this] def handleNonSuccessfulStatusCode + ( + context: OperationContext, + itemOperation: CosmosItemOperation, + itemResponse: CosmosBulkItemResponse, + isGettingRetried: AtomicBoolean, + responseException: Option[CosmosException] + ) : Unit = { + + val exceptionMessage = responseException match { + case Some(e) => e.getMessage + case None => "" + } + + val effectiveStatusCode = Option(itemResponse) match { + case Some(r) => r.getStatusCode + case None => responseException match { + case Some(e) => e.getStatusCode + case None => CosmosConstants.StatusCodes.Timeout + } + } + + val effectiveSubStatusCode = Option(itemResponse) match { + case Some(r) => r.getSubStatusCode + case None => responseException match { + case Some(e) => e.getSubStatusCode + case None => 0 + } + } + + log.logDebug(s"encountered item operation response with status code " + + s"$effectiveStatusCode:$effectiveSubStatusCode, " + + s"Context: ${operationContext.toString} $getThreadInfo") + if (shouldIgnore(effectiveStatusCode, effectiveSubStatusCode)) { + log.logDebug(s"for itemId=[${context.itemId}], partitionKeyValue=[${context.partitionKeyValue}], " + + s"ignored encountered status code '$effectiveStatusCode:$effectiveSubStatusCode', " + + s"Context: ${operationContext.toString}") + totalSuccessfulIngestionMetrics.getAndIncrement() + // work done + } else if (shouldRetry(effectiveStatusCode, effectiveSubStatusCode, context)) { + // requeue + log.logWarning(s"for itemId=[${context.itemId}], partitionKeyValue=[${context.partitionKeyValue}], " + + s"encountered status code '$effectiveStatusCode:$effectiveSubStatusCode', will retry! " + + s"attemptNumber=${context.attemptNumber}, exceptionMessage=$exceptionMessage, " + + s"Context: {${operationContext.toString}} $getThreadInfo") + + // If the write strategy is patchBulkUpdate, the OperationContext.sourceItem will not be the original objectNode, + // It is computed through read item from cosmosdb, and then patch the item locally. + // During retry, it is important to use the original objectNode (for example for preCondition failure, it requires to go through the readMany step again) + val sourceItem = itemOperation match { + case _: ItemBulkOperation[ObjectNode, OperationContext] => + context.sourceItem match { + case Some(bulkOperationSourceItem) => bulkOperationSourceItem + case None => itemOperation.getItem.asInstanceOf[ObjectNode] + } + case _ => itemOperation.getItem.asInstanceOf[ObjectNode] + } + + this.scheduleRetry( + trackPendingRetryAction = () => pendingBulkWriteRetries.add(itemOperation), + clearPendingRetryAction = () => pendingBulkWriteRetries.remove(itemOperation), + itemOperation.getPartitionKeyValue, + sourceItem, + context, + effectiveStatusCode) + isGettingRetried.set(true) + } else { + log.logError(s"for itemId=[${context.itemId}], partitionKeyValue=[${context.partitionKeyValue}], " + + s"encountered status code '$effectiveStatusCode:$effectiveSubStatusCode', all retries exhausted! " + + s"attemptNumber=${context.attemptNumber}, exceptionMessage=$exceptionMessage, " + + s"Context: {${operationContext.toString} $getThreadInfo") + + val message = s"All retries exhausted for '${itemOperation.getOperationType}' bulk operation - " + + s"statusCode=[$effectiveStatusCode:$effectiveSubStatusCode] " + + s"itemId=[${context.itemId}], partitionKeyValue=[${context.partitionKeyValue}]" + + val exceptionToBeThrown = responseException match { + case Some(e) => + new BulkOperationFailedException(effectiveStatusCode, effectiveSubStatusCode, message, e) + case None => + new BulkOperationFailedException(effectiveStatusCode, effectiveSubStatusCode, message, null) + } + + captureIfFirstFailure(exceptionToBeThrown) + cancelWork() + } + } + //scalastyle:on method.length + //scalastyle:on cyclomatic.complexity + + private[this] def throwIfCapturedExceptionExists(): Unit = { + val errorSnapshot = errorCaptureFirstException.get() + if (errorSnapshot != null) { + log.logError(s"throw captured error ${errorSnapshot.getMessage}, " + + s"Context: ${operationContext.toString} $getThreadInfo") + throw errorSnapshot + } + } + + private[this] def getActiveOperationsLog( + activeOperationsSnapshot: mutable.Set[CosmosItemOperation], + activeReadManyOperationsSnapshot: mutable.Set[ReadManyOperation]): String = { + val sb = new StringBuilder() + + activeOperationsSnapshot + .take(TransactionalBulkWriter.maxItemOperationsToShowInErrorMessage) + .foreach(itemOperation => { + if (sb.nonEmpty) { + sb.append(", ") + } + + sb.append(itemOperation.getOperationType) + sb.append("->") + val ctx = itemOperation.getContext[OperationContext] + sb.append(s"${ctx.partitionKeyValue}/${ctx.itemId}/${ctx.eTag}(${ctx.attemptNumber})") + }) + + // add readMany snapshot logs + activeReadManyOperationsSnapshot + .take(TransactionalBulkWriter.maxItemOperationsToShowInErrorMessage - activeOperationsSnapshot.size) + .foreach(readManyOperation => { + if (sb.nonEmpty) { + sb.append(", ") + } + + sb.append("ReadMany") + sb.append("->") + val ctx = readManyOperation.operationContext + sb.append(s"${ctx.partitionKeyValue}/${ctx.itemId}/${ctx.eTag}(${ctx.attemptNumber})") + }) + + sb.toString() + } + + private[this] def sameBulkWriteOperations + ( + snapshot: mutable.Set[CosmosItemOperation], + current: mutable.Set[CosmosItemOperation] + ): Boolean = { + + if (snapshot.size != current.size) { + false + } else { + snapshot.forall(snapshotOperation => { + current.exists( + currentOperation => snapshotOperation.getOperationType == currentOperation.getOperationType + && snapshotOperation.getPartitionKeyValue == currentOperation.getPartitionKeyValue + && Objects.equals(snapshotOperation.getId, currentOperation.getId) + && Objects.equals(snapshotOperation.getItem[ObjectNode], currentOperation.getItem[ObjectNode]) + ) + }) + } + } + + private[this] def sameReadManyOperations + ( + snapshot: mutable.Set[ReadManyOperation], + current: mutable.Set[ReadManyOperation] + ): Boolean = { + + if (snapshot.size != current.size) { + false + } else { + snapshot.forall(snapshotOperation => { + current.exists( + currentOperation => snapshotOperation.cosmosItemIdentity == currentOperation.cosmosItemIdentity + && Objects.equals(snapshotOperation.objectNode, currentOperation.objectNode) + ) + }) + } + } + + private[this] def throwIfProgressStaled + ( + operationName: String, + activeOperationsSnapshot: mutable.Set[CosmosItemOperation], + pendingRetriesSnapshot: mutable.Set[CosmosItemOperation], + activeReadManyOperationsSnapshot: mutable.Set[ReadManyOperation], + pendingReadManyOperationsSnapshot: mutable.Set[ReadManyOperation], + numberOfIntervalsWithIdenticalActiveOperationSnapshots: AtomicLong, + allowRetryOnNewBulkWriterInstance: Boolean + ): Unit = { + + val operationsLog = getActiveOperationsLog(activeOperationsSnapshot, activeReadManyOperationsSnapshot) + + if (sameBulkWriteOperations(pendingRetriesSnapshot ++ activeOperationsSnapshot , activeBulkWriteOperations ++ pendingBulkWriteRetries) + && sameReadManyOperations(pendingReadManyOperationsSnapshot ++ activeReadManyOperationsSnapshot , activeReadManyOperations ++ pendingReadManyRetries)) { + + numberOfIntervalsWithIdenticalActiveOperationSnapshots.incrementAndGet() + log.logWarning( + s"$operationName has been waiting $numberOfIntervalsWithIdenticalActiveOperationSnapshots " + + s"times for identical set of operations: $operationsLog " + + s"Context: ${operationContext.toString} $getThreadInfo" + ) + } else { + numberOfIntervalsWithIdenticalActiveOperationSnapshots.set(0) + logInfoOrWarning( + s"$operationName is waiting for active bulkWrite operations: $operationsLog " + + s"Context: ${operationContext.toString} $getThreadInfo" + ) + } + + val secondsWithoutProgress = numberOfIntervalsWithIdenticalActiveOperationSnapshots.get * + writeConfig.flushCloseIntervalInSeconds + val maxNoProgressIntervalInSeconds = if (commitAttempt == 1 && allowRetryOnNewBulkWriterInstance + && this.activeReadManyOperations.isEmpty && this.pendingReadManyRetries.isEmpty) { + writeConfig.maxInitialNoProgressIntervalInSeconds + } else { + writeConfig.maxRetryNoProgressIntervalInSeconds + } + val maxAllowedIntervalWithoutAnyProgressExceeded = secondsWithoutProgress >= maxNoProgressIntervalInSeconds + + if (maxAllowedIntervalWithoutAnyProgressExceeded) { + + val exception = if (activeReadManyOperationsSnapshot.isEmpty) { + val retriableRemainingOperations = if (allowRetryOnNewBulkWriterInstance) { + Some( + (pendingRetriesSnapshot ++ activeOperationsSnapshot) + .toList + .sortBy(op => op.getContext[OperationContext].sequenceNumber) + ) + } else { + None + } + + new BulkWriterNoProgressException( + s"Stale bulk ingestion identified in $operationName - the following active operations have not been " + + s"completed (first ${TransactionalBulkWriter.maxItemOperationsToShowInErrorMessage} shown) or progressed after " + + s"$maxNoProgressIntervalInSeconds seconds: $operationsLog", + commitAttempt, + retriableRemainingOperations) + } else { + new BulkWriterNoProgressException( + s"Stale bulk ingestion as well as readMany operations identified in $operationName - the following active operations have not been " + + s"completed (first ${TransactionalBulkWriter.maxItemOperationsToShowInErrorMessage} shown) or progressed after " + + s"${maxNoProgressIntervalInSeconds} : $operationsLog", + commitAttempt, + None) + } + + captureIfFirstFailure(exception) + + cancelWork() + } + + throwIfCapturedExceptionExists() + } + + // the caller has to ensure that after invoking this method scheduleWrite doesn't get invoked + // scalastyle:off method.length + // scalastyle:off cyclomatic.complexity + override def flushAndClose(): Unit = { + this.synchronized { + try { + if (!closed.get()) { + log.logInfo(s"flushAndClose invoked $getThreadInfo") + log.logInfo(s"completed so far ${totalSuccessfulIngestionMetrics.get()}, " + + s"pending bulkWrite asks ${activeBulkWriteOperations.size}, pending readMany tasks ${activeReadManyOperations.size} $getThreadInfo") + + // error handling, if there is any error and the subscription is cancelled + // the remaining tasks will not be processed hence we never reach activeTasks = 0 + // once we do error handling we should think how to cover the scenario. + lock.lock() + try { + val numberOfIntervalsWithIdenticalActiveOperationSnapshots = new AtomicLong(0) + var activeTasksSnapshot = activeTasks.get() + var pendingRetriesSnapshot = pendingRetries.get() + while ((pendingRetriesSnapshot > 0 || activeTasksSnapshot > 0) + && errorCaptureFirstException.get == null) { + + logInfoOrWarning( + s"Waiting for pending activeTasks $activeTasksSnapshot and/or pendingRetries " + + s"$pendingRetriesSnapshot, Context: ${operationContext.toString} $getThreadInfo") + val activeOperationsSnapshot = activeBulkWriteOperations.clone() + val activeReadManyOperationsSnapshot = activeReadManyOperations.clone() + val pendingOperationsSnapshot = pendingBulkWriteRetries.clone() + val pendingReadManyOperationsSnapshot = pendingReadManyRetries.clone() + val awaitCompleted = pendingTasksCompleted.await(writeConfig.flushCloseIntervalInSeconds, TimeUnit.SECONDS) + if (!awaitCompleted) { + throwIfProgressStaled( + "FlushAndClose", + activeOperationsSnapshot, + pendingOperationsSnapshot, + activeReadManyOperationsSnapshot, + pendingReadManyOperationsSnapshot, + numberOfIntervalsWithIdenticalActiveOperationSnapshots, + allowRetryOnNewBulkWriterInstance = true + ) + + if (numberOfIntervalsWithIdenticalActiveOperationSnapshots.get > 0L) { + + + val buffered = Scannable.from(bulkInputEmitter).scan(Scannable.Attr.BUFFERED) + + if (verboseLoggingAfterReEnqueueingRetriesEnabled.compareAndSet(false, true)) { + log.logWarning(s"Starting to re-enqueue retries. Enabling verbose logs. " + + s"Number of intervals with identical pending operations: " + + s"$numberOfIntervalsWithIdenticalActiveOperationSnapshots Active Bulk Operations: " + + s"$activeOperationsSnapshot, Active Read-Many Operations: $activeReadManyOperationsSnapshot, " + + s"PendingRetries: $pendingRetriesSnapshot, Buffered tasks: $buffered " + + s"Attempt: ${numberOfIntervalsWithIdenticalActiveOperationSnapshots.get} - " + + s"Context: ${operationContext.toString} $getThreadInfo") + } else if ((numberOfIntervalsWithIdenticalActiveOperationSnapshots.get % 3) == 0) { + log.logWarning(s"Reattempting to re-enqueue retries. Enabling verbose logs. " + + s"Number of intervals with identical pending operations: " + + s"$numberOfIntervalsWithIdenticalActiveOperationSnapshots Active Bulk Operations: " + + s"$activeOperationsSnapshot, Active Read-Many Operations: $activeReadManyOperationsSnapshot, " + + s"PendingRetries: $pendingRetriesSnapshot, Buffered tasks: $buffered " + + s"Attempt: ${numberOfIntervalsWithIdenticalActiveOperationSnapshots.get} - " + + s"Context: ${operationContext.toString} $getThreadInfo") + } + + activeOperationsSnapshot.foreach(operation => { + if (activeBulkWriteOperations.contains(operation)) { + // re-validating whether the operation is still active - if so, just re-enqueue another retry + // this is harmless - because all bulkItemOperations from Spark connector are always idempotent + + // For FAIL_NON_SERIALIZED, will keep retry, while for other errors, use the default behavior + // TransactionalBulkExecutor will handle batching by partition key + bulkInputEmitter.emitNext(operation, emitFailureHandler) + log.logWarning(s"Re-enqueued a retry for pending active write task '${operation.getOperationType} " + + s"(${operation.getPartitionKeyValue}/${operation.getId})' " + + s"- Attempt: ${numberOfIntervalsWithIdenticalActiveOperationSnapshots.get} - " + + s"Context: ${operationContext.toString} $getThreadInfo") + } + }) + + activeReadManyOperationsSnapshot.foreach(operation => { + if (activeReadManyOperations.contains(operation)) { + // re-validating whether the operation is still active - if so, just re-enqueue another retry + // this is harmless - because all bulkItemOperations from Spark connector are always idempotent + + // For FAIL_NON_SERIALIZED, will keep retry, while for other errors, use the default behavior + readManyInputEmitterOpt.get.emitNext(operation, emitFailureHandler) + log.logWarning(s"Re-enqueued a retry for pending active read-many task '" + + s"(${operation.cosmosItemIdentity.getPartitionKey}/${operation.cosmosItemIdentity.getId})' " + + s"- Attempt: ${numberOfIntervalsWithIdenticalActiveOperationSnapshots.get} - " + + s"Context: ${operationContext.toString} $getThreadInfo") + } + }) + } + } + + activeTasksSnapshot = activeTasks.get() + pendingRetriesSnapshot = pendingRetries.get() + val semaphoreAvailablePermitsSnapshot = semaphore.availablePermits() + + if (awaitCompleted) { + logInfoOrWarning(s"Waiting completed for pending activeTasks $activeTasksSnapshot, pendingRetries " + + s"$pendingRetriesSnapshot Context: ${operationContext.toString} $getThreadInfo") + } else { + logInfoOrWarning(s"Waiting interrupted for pending activeTasks $activeTasksSnapshot , pendingRetries " + + s"$pendingRetriesSnapshot - available permits $semaphoreAvailablePermitsSnapshot, " + + s"Context: ${operationContext.toString} $getThreadInfo") + } + } + + logInfoOrWarning(s"Waiting completed for pending activeTasks $activeTasksSnapshot, pendingRetries " + + s"$pendingRetriesSnapshot Context: ${operationContext.toString} $getThreadInfo") + } finally { + lock.unlock() + } + + logInfoOrWarning(s"invoking bulkInputEmitter.onComplete(), Context: ${operationContext.toString} $getThreadInfo") + semaphore.release(Math.max(0, activeTasks.get())) + bulkInputEmitter.emitComplete(TransactionalBulkWriter.emitFailureHandlerForComplete) + + // complete readManyInputEmitter + if (readManyInputEmitterOpt.isDefined) { + readManyInputEmitterOpt.get.emitComplete(TransactionalBulkWriter.emitFailureHandlerForComplete) + } + + throwIfCapturedExceptionExists() + + assume(activeTasks.get() <= 0) + + assume(activeBulkWriteOperations.isEmpty) + assume(activeReadManyOperations.isEmpty) + assume(semaphore.availablePermits() >= maxPendingOperations) + + if (totalScheduledMetrics.get() != totalSuccessfulIngestionMetrics.get) { + log.logWarning(s"flushAndClose completed with no error but inconsistent total success and " + + s"scheduled metrics. This indicates that successful completion was only possible after re-enqueueing " + + s"retries. totalSuccessfulIngestionMetrics=${totalSuccessfulIngestionMetrics.get()}, " + + s"totalScheduled=$totalScheduledMetrics, Context: ${operationContext.toString} $getThreadInfo") + } else { + logInfoOrWarning(s"flushAndClose completed with no error. " + + s"totalSuccessfulIngestionMetrics=${totalSuccessfulIngestionMetrics.get()}, " + + s"totalScheduled=$totalScheduledMetrics, Context: ${operationContext.toString} $getThreadInfo") + } + } + } finally { + subscriptionDisposable.dispose() + readManySubscriptionDisposableOpt match { + case Some(readManySubscriptionDisposable) => + + + readManySubscriptionDisposable.dispose() + case _ => + } + + flushExecutorHolder match { + case Some(executorAndFutureTuple) => + val executor: ScheduledThreadPoolExecutor = executorAndFutureTuple._1 + val future: ScheduledFuture[_] = executorAndFutureTuple._2 + + try { + future.cancel(true) + log.logDebug(s"Cancelled all future scheduled tasks $getThreadInfo, Context: ${operationContext.toString}") + } catch { + case e: Exception => + log.logWarning(s"Failed to cancel scheduled tasks $getThreadInfo, Context: ${operationContext.toString}", e) + } + + try { + log.logDebug(s"Shutting down the executor service, Context: ${operationContext.toString}") + executor.shutdownNow + log.logDebug(s"Successfully shut down the executor service, Context: ${operationContext.toString}") + } catch { + case e: Exception => + log.logWarning(s"Failed to shut down the executor service, Context: ${operationContext.toString}", e) + } + case _ => + } + closed.set(true) + } + } + } + // scalastyle:on method.length + // scalastyle:on cyclomatic.complexity + + private def logInfoOrWarning(msg: => String): Unit = { + if (this.verboseLoggingAfterReEnqueueingRetriesEnabled.get()) { + log.logWarning(msg) + } else { + log.logInfo(msg) + } + } + + private def markTaskCompletion(): Unit = { + lock.lock() + try { + val activeTasksLeftSnapshot = activeTasks.decrementAndGet() + val exceptionSnapshot = errorCaptureFirstException.get() + log.logTrace(s"markTaskCompletion, Active tasks left: $activeTasksLeftSnapshot, " + + s"error: $exceptionSnapshot, Context: ${operationContext.toString} $getThreadInfo") + if (activeTasksLeftSnapshot == 0 || exceptionSnapshot != null) { + pendingTasksCompleted.signal() + } + } finally { + lock.unlock() + } + } + + private def captureIfFirstFailure(throwable: Throwable): Unit = { + log.logError(s"capture failure, Context: {${operationContext.toString}} $getThreadInfo", throwable) + lock.lock() + try { + errorCaptureFirstException.compareAndSet(null, throwable) + pendingTasksCompleted.signal() + } finally { + lock.unlock() + } + } + + private def cancelWork(): Unit = { + logInfoOrWarning(s"cancelling remaining unprocessed tasks ${activeTasks.get} " + + s"[bulkWrite tasks ${activeBulkWriteOperations.size}, readMany tasks ${activeReadManyOperations.size} ]" + + s"Context: ${operationContext.toString}") + subscriptionDisposable.dispose() + if (readManySubscriptionDisposableOpt.isDefined) { + readManySubscriptionDisposableOpt.get.dispose() + } + } + + private def shouldIgnore(statusCode: Int, subStatusCode: Int): Boolean = { + val returnValue = writeConfig.itemWriteStrategy match { + case ItemWriteStrategy.ItemAppend => Exceptions.isResourceExistsException(statusCode) + case ItemWriteStrategy.ItemPatchIfExists => Exceptions.isNotFoundExceptionCore(statusCode, subStatusCode) + case ItemWriteStrategy.ItemDelete => Exceptions.isNotFoundExceptionCore(statusCode, subStatusCode) + case ItemWriteStrategy.ItemDeleteIfNotModified => Exceptions.isNotFoundExceptionCore(statusCode, subStatusCode) || + Exceptions.isPreconditionFailedException(statusCode) + case ItemWriteStrategy.ItemOverwriteIfNotModified => + Exceptions.isResourceExistsException(statusCode) || + Exceptions.isNotFoundExceptionCore(statusCode, subStatusCode) || + Exceptions.isPreconditionFailedException(statusCode) + case _ => false + } + + returnValue + } + + private def shouldRetry(statusCode: Int, subStatusCode: Int, operationContext: OperationContext): Boolean = { + var returnValue = false + if (operationContext.attemptNumber < writeConfig.maxRetryCount) { + returnValue = writeConfig.itemWriteStrategy match { + case ItemWriteStrategy.ItemBulkUpdate => + this.shouldRetryForItemPatchBulkUpdate(statusCode, subStatusCode) + // Upsert can return 404/0 in rare cases (when due to TTL expiration there is a race condition + case ItemWriteStrategy.ItemOverwrite => + Exceptions.canBeTransientFailure(statusCode, subStatusCode) || + statusCode == 0 || // Gateway mode reports inability to connect due to PoolAcquirePendingLimitException as status code 0 + Exceptions.isNotFoundExceptionCore(statusCode, subStatusCode) + case _ => + Exceptions.canBeTransientFailure(statusCode, subStatusCode) || + statusCode == 0 // Gateway mode reports inability to connect due to PoolAcquirePendingLimitException as status code 0 + } + } + + log.logDebug(s"Should retry statusCode '$statusCode:$subStatusCode' -> $returnValue, " + + s"Context: ${operationContext.toString} $getThreadInfo") + + returnValue + } + + private def shouldRetryForItemPatchBulkUpdate(statusCode: Int, subStatusCode: Int): Boolean = { + Exceptions.canBeTransientFailure(statusCode, subStatusCode) || + statusCode == 0 || // Gateway mode reports inability to connect due to + // PoolAcquirePendingLimitException as status code 0 + Exceptions.isResourceExistsException(statusCode) || + Exceptions.isPreconditionFailedException(statusCode) + } + + private def getId(objectNode: ObjectNode) = { + val idField = objectNode.get(CosmosConstants.Properties.Id) + assume(idField != null && idField.isTextual) + idField.textValue() + } + + /** + * Don't wait for any remaining work but signal to the writer the ungraceful close + * Should not throw any exceptions + */ + override def abort(shouldThrow: Boolean): Unit = { + if (shouldThrow) { + log.logError(s"Abort, Context: ${operationContext.toString} $getThreadInfo") + // signal an exception that will be thrown for any pending work/flushAndClose if no other exception has + // been registered + captureIfFirstFailure( + new IllegalStateException(s"The Spark task was aborted, Context: ${operationContext.toString}")) + } else { + log.logWarning(s"BulkWriter aborted and commit retried, Context: ${operationContext.toString} $getThreadInfo") + } + cancelWork() + } + + private class OperationContext + ( + itemIdInput: String, + partitionKeyValueInput: PartitionKey, + eTagInput: Option[String], + val attemptNumber: Int, + val sequenceNumber: Long, + /** starts from 1 * */ + sourceItemInput: Option[ObjectNode] = None) // for patchBulkUpdate: source item refers to the original objectNode from which SDK constructs the final bulk item operation + { + private val ctxCore: OperationContextCore = OperationContextCore(itemIdInput, partitionKeyValueInput, eTagInput, sourceItemInput) + + override def equals(obj: Any): Boolean = ctxCore.equals(obj) + + override def hashCode(): Int = ctxCore.hashCode() + + override def toString: String = { + ctxCore.toString + s", attemptNumber = $attemptNumber" + } + + def itemId: String = ctxCore.itemId + + def partitionKeyValue: PartitionKey = ctxCore.partitionKeyValue + + def eTag: Option[String] = ctxCore.eTag + + def sourceItem: Option[ObjectNode] = ctxCore.sourceItem + } + + private case class OperationContextCore + ( + itemId: String, + partitionKeyValue: PartitionKey, + eTag: Option[String], + sourceItem: Option[ObjectNode] = None) // for patchBulkUpdate: source item refers to the original objectNode from which SDK constructs the final bulk item operation + { + override def productPrefix: String = "OperationContext" + } + + + private case class ReadManyOperation( + cosmosItemIdentity: CosmosItemIdentity, + objectNode: ObjectNode, + operationContext: OperationContext) +} + +private object TransactionalBulkWriter { + private val log = new DefaultDiagnostics().getLogger(this.getClass) + //scalastyle:off magic.number + private val maxDelayOn408RequestTimeoutInMs = 3000 + private val minDelayOn408RequestTimeoutInMs = 500 + private val maxItemOperationsToShowInErrorMessage = 10 + private val BULK_WRITER_REQUESTS_BOUNDED_ELASTIC_THREAD_NAME = "bulk-writer-requests-bounded-elastic" + private val BULK_WRITER_INPUT_BOUNDED_ELASTIC_THREAD_NAME = "bulk-writer-input-bounded-elastic" + private val BULK_WRITER_RESPONSES_BOUNDED_ELASTIC_THREAD_NAME = "bulk-writer-responses-bounded-elastic" + private val READ_MANY_BOUNDED_ELASTIC_THREAD_NAME = "read-many-bounded-elastic" + private val TTL_FOR_SCHEDULER_WORKER_IN_SECONDS = 60 // same as BoundedElasticScheduler.DEFAULT_TTL_SECONDS + //scalastyle:on magic.number + + // let's say the spark executor VM has 16 CPU cores. + // let's say we have a cosmos container with 1M RU which is 167 partitions + // let's say we are ingesting items of size 1KB + // let's say max request size is 1MB + // hence we want 1MB/ 1KB items per partition to be buffered + // 1024 * 167 items should get buffered on a 16 CPU core VM + // so per CPU core we want (1024 * 167 / 16) max items to be buffered + // Reduced the targeted buffer from 2MB per partition and core to 1 MB because + // we had a few customers seeing to high CPU usage with the previous setting + // Reason is that several customers use larger than 1 KB documents so we need + // to be less aggressive with the buffering + val DefaultMaxPendingOperationPerCore: Int = 1024 * 167 / 16 + + val emitFailureHandler: EmitFailureHandler = + (signalType, emitResult) => { + if (emitResult.equals(EmitResult.FAIL_NON_SERIALIZED)) { + log.logDebug(s"emitFailureHandler - Signal: ${signalType.toString}, Result: ${emitResult.toString}") + true + } else { + log.logError(s"emitFailureHandler - Signal: ${signalType.toString}, Result: ${emitResult.toString}") + false + } + } + + private val emitFailureHandlerForComplete: EmitFailureHandler = (signalType, emitResult) => { + if (emitResult.equals(EmitResult.FAIL_NON_SERIALIZED)) { + log.logWarning(s"emitFailureHandlerForComplete - Signal:$signalType, Result:$emitResult") + true + } else if (emitResult.equals(EmitResult.FAIL_CANCELLED) || + emitResult.equals(EmitResult.FAIL_TERMINATED) || + emitResult.equals(EmitResult.FAIL_OVERFLOW)) { + log.logWarning(s"emitFailureHandlerForComplete - Signal:$signalType, Result:$emitResult") + false + } else { + log.logError(s"emitFailureHandlerForComplete - Signal:$signalType, Result:$emitResult") + false + } + } + + private val bulkProcessingThresholds = new CosmosBulkExecutionThresholdsState() + + private val maxPendingOperationsPerJVM: Int = DefaultMaxPendingOperationPerCore * SparkUtils.getNumberOfHostCPUCores + + // Custom bounded elastic scheduler to consume input flux + val bulkWriterRequestsBoundedElastic: Scheduler = Schedulers.newBoundedElastic( + Schedulers.DEFAULT_BOUNDED_ELASTIC_SIZE, + Schedulers.DEFAULT_BOUNDED_ELASTIC_QUEUESIZE + 2 * maxPendingOperationsPerJVM, + BULK_WRITER_REQUESTS_BOUNDED_ELASTIC_THREAD_NAME, + TTL_FOR_SCHEDULER_WORKER_IN_SECONDS, true) + + // Custom bounded elastic scheduler to consume input flux + val bulkWriterInputBoundedElastic: Scheduler = Schedulers.newBoundedElastic( + Schedulers.DEFAULT_BOUNDED_ELASTIC_SIZE, + Schedulers.DEFAULT_BOUNDED_ELASTIC_QUEUESIZE + 2 * maxPendingOperationsPerJVM, + BULK_WRITER_INPUT_BOUNDED_ELASTIC_THREAD_NAME, + TTL_FOR_SCHEDULER_WORKER_IN_SECONDS, true) + + // Custom bounded elastic scheduler to switch off IO thread to process response. + val bulkWriterResponsesBoundedElastic: Scheduler = Schedulers.newBoundedElastic( + Schedulers.DEFAULT_BOUNDED_ELASTIC_SIZE, + Schedulers.DEFAULT_BOUNDED_ELASTIC_QUEUESIZE + maxPendingOperationsPerJVM, + BULK_WRITER_RESPONSES_BOUNDED_ELASTIC_THREAD_NAME, + TTL_FOR_SCHEDULER_WORKER_IN_SECONDS, true) + + + // Custom bounded elastic scheduler to switch off IO thread to process response. + val readManyBoundedElastic: Scheduler = Schedulers.newBoundedElastic( + 2 * Schedulers.DEFAULT_BOUNDED_ELASTIC_SIZE, + Schedulers.DEFAULT_BOUNDED_ELASTIC_QUEUESIZE + maxPendingOperationsPerJVM, + READ_MANY_BOUNDED_ELASTIC_THREAD_NAME, + TTL_FOR_SCHEDULER_WORKER_IN_SECONDS, true) + + def getThreadInfo: String = { + val t = Thread.currentThread() + val group = Option.apply(t.getThreadGroup) match { + case Some(group) => group.getName + case None => "n/a" + } + s"Thread[Name: ${t.getName}, Group: $group, IsDaemon: ${t.isDaemon} Id: ${t.getId}]" + } + + private class BulkOperationFailedException(statusCode: Int, subStatusCode: Int, message:String, cause: Throwable) + extends CosmosException(statusCode, message, null, cause) { + BridgeInternal.setSubStatusCode(this, subStatusCode) + } +} + +//scalastyle:on multiple.string.literals +//scalastyle:on null +//scalastyle:on file.size.limit diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/BulkWriterITest.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/BulkWriterITest.scala index bf0c59bf757d..5ea56b39d756 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/BulkWriterITest.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/BulkWriterITest.scala @@ -408,7 +408,7 @@ class BulkWriterITest extends IntegrationSpec with CosmosClient with AutoCleanab val containerProperties = container.read().block().getProperties val partitionKeyDefinition = containerProperties.getPartitionKeyDefinition val containerConfig = CosmosContainerConfig(container.getDatabase.getId, container.getId, None) - val writeConfig = CosmosWriteConfig(ItemWriteStrategy.ItemAppend, maxRetryCount = 5, bulkEnabled = true, Some(900)) + val writeConfig = CosmosWriteConfig(ItemWriteStrategy.ItemAppend, maxRetryCount = 5, bulkEnabled = true, bulkEnableTransactions = false, Some(900)) val bulkWriter = new BulkWriter( container, containerConfig, diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala new file mode 100644 index 000000000000..730167417feb --- /dev/null +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala @@ -0,0 +1,637 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.spark + +import com.azure.cosmos.implementation.{TestConfigurations, Utils} +import com.azure.cosmos.models.{PartitionKey, PartitionKeyBuilder} +import com.azure.cosmos.CosmosException +import com.fasterxml.jackson.databind.node.ObjectNode +import org.apache.spark.sql.types._ +import org.apache.spark.sql.{Row, SaveMode} + +import java.util.UUID +// scalastyle:off underscore.import +import scala.collection.JavaConverters._ +// scalastyle:on underscore.import + +class TransactionalBatchITest extends IntegrationSpec + with Spark + with AutoCleanableCosmosContainersWithPkAsPartitionKey { + + //scalastyle:off multiple.string.literals + //scalastyle:off magic.number + + // Helper method to get root cause of exception + private def getRootCause(t: Throwable): Throwable = { + if (t.getCause == null) t else getRootCause(t.getCause) + } + + "Transactional Batch" should "create items atomically" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) + val partitionKeyValue = UUID.randomUUID().toString + val item1Id = s"test-item1-${UUID.randomUUID()}" + val item2Id = s"test-item2-${UUID.randomUUID()}" + + // Create batch operations DataFrame with flat columns (like normal writes) + val schema = StructType(Seq( + StructField("id", StringType, nullable = false), + StructField("pk", StringType, nullable = false), + StructField("name", StringType, nullable = false) + )) + + val batchOperations = Seq( + Row(item1Id, partitionKeyValue, "Alice"), + Row(item2Id, partitionKeyValue, "Bob") + ) + + val operationsDf = spark.createDataFrame(batchOperations.asJava, schema) + + // Execute transactional batch using regular write with bulkEnableTransactions config + operationsDf.write + .format("cosmos.oltp") + .option("spark.cosmos.accountEndpoint", cosmosEndpoint) + .option("spark.cosmos.accountKey", cosmosMasterKey) + .option("spark.cosmos.database", cosmosDatabase) + .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) + .option("spark.cosmos.write.strategy", "ItemOverwrite") + .option("spark.cosmos.write.bulk.enabled", "true") + .option("spark.cosmos.write.bulk.enableTransactions", "true") + .mode(SaveMode.Append) + .save() + + // Verify items were created + val item1 = container.readItem(item1Id, new PartitionKey(partitionKeyValue), classOf[ObjectNode]).block() + item1 should not be null + item1.getItem.get("name").asText() shouldEqual "Alice" + + val item2 = container.readItem(item2Id, new PartitionKey(partitionKeyValue), classOf[ObjectNode]).block() + item2 should not be null + item2.getItem.get("name").asText() shouldEqual "Bob" + } + + it should "rollback all operations on failure" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) + val partitionKeyValue = UUID.randomUUID().toString + val item1Id = s"test-item1-${UUID.randomUUID()}" + val duplicateId = s"duplicate-${UUID.randomUUID()}" + + // First create an item that we'll try to create again (should fail) + val existingDoc = Utils.getSimpleObjectMapper.createObjectNode() + existingDoc.put("id", duplicateId) + existingDoc.put("pk", partitionKeyValue) + existingDoc.put("name", "Existing") + container.createItem(existingDoc).block() + + // Create batch with one valid create and one duplicate create (should fail entire batch) + val schema = StructType(Seq( + StructField("id", StringType, nullable = false), + StructField("pk", StringType, nullable = false), + StructField("name", StringType, nullable = false) + )) + + val batchOperations = Seq( + Row(item1Id, partitionKeyValue, "NewItem"), + Row(duplicateId, partitionKeyValue, "Duplicate") + ) + + val operationsDf = spark.createDataFrame(batchOperations.asJava, schema) + + // This should fail because we're using ItemAppend strategy (create only) and duplicateId already exists + val exception = intercept[Exception] { + operationsDf.write + .format("cosmos.oltp") + .option("spark.cosmos.accountEndpoint", cosmosEndpoint) + .option("spark.cosmos.accountKey", cosmosMasterKey) + .option("spark.cosmos.database", cosmosDatabase) + .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) + .option("spark.cosmos.write.strategy", "ItemAppend") + .option("spark.cosmos.write.bulk.enabled", "true") + .option("spark.cosmos.write.bulk.enableTransactions", "true") + .mode(SaveMode.Append) + .save() + } + + // Verify the exception message indicates transactional batch failure + // Spark wraps our exception in "Writing job aborted", check the root cause + // Transactional batch failures may show as statusCode 424 (Failed Dependency) when one operation fails + val rootCause = getRootCause(exception) + assert(rootCause.getMessage.contains("424") || rootCause.getMessage.contains("409"), + s"Expected transactional batch failure error (424 or 409), got: ${rootCause.getMessage}") + + // Verify item1 was NOT created (rolled back) + try { + container.readItem(item1Id, new PartitionKey(partitionKeyValue), classOf[ObjectNode]).block() + fail("Item1 should not exist after batch rollback") + } catch { + case e: CosmosException if e.getStatusCode == 404 => + // Expected - item doesn't exist after rollback (404 Not Found) + } + } + + it should "support simplified schema with default upsert operation" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) + val partitionKeyValue = UUID.randomUUID().toString + val item1Id = s"test-item1-${UUID.randomUUID()}" + val item2Id = s"test-item2-${UUID.randomUUID()}" + + // Create batch operations DataFrame with simplified schema + val simplifiedSchema = StructType(Seq( + StructField("id", StringType, nullable = false), + StructField("pk", StringType, nullable = false), + StructField("name", StringType, nullable = false), + StructField("age", IntegerType, nullable = false) + )) + + val batchOperations = Seq( + Row(item1Id, partitionKeyValue, "Alice", 30), + Row(item2Id, partitionKeyValue, "Bob", 25) + ) + + val operationsDf = spark.createDataFrame(batchOperations.asJava, simplifiedSchema) + + // Execute transactional batch - defaults to ItemOverwrite (upsert) + operationsDf.write + .format("cosmos.oltp") + .option("spark.cosmos.accountEndpoint", cosmosEndpoint) + .option("spark.cosmos.accountKey", cosmosMasterKey) + .option("spark.cosmos.database", cosmosDatabase) + .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) + .option("spark.cosmos.write.strategy", "ItemOverwrite") + .option("spark.cosmos.write.bulk.enabled", "true") + .option("spark.cosmos.write.bulk.enableTransactions", "true") + .mode(SaveMode.Append) + .save() + + // Verify items were created + val item1 = container.readItem(item1Id, new PartitionKey(partitionKeyValue), classOf[ObjectNode]).block() + item1 should not be null + item1.getItem.get("name").asText() shouldEqual "Alice" + item1.getItem.get("age").asInt() shouldEqual 30 + + val item2 = container.readItem(item2Id, new PartitionKey(partitionKeyValue), classOf[ObjectNode]).block() + item2 should not be null + item2.getItem.get("name").asText() shouldEqual "Bob" + item2.getItem.get("age").asInt() shouldEqual 25 + } + + it should "preserve order with simplified schema" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) + val partitionKeyValue = UUID.randomUUID().toString + val baseId = s"order-test-${UUID.randomUUID()}" + + // Create batch with multiple operations in specific order + val simplifiedSchema = StructType(Seq( + StructField("id", StringType, nullable = false), + StructField("pk", StringType, nullable = false), + StructField("sequence", IntegerType, nullable = false) + )) + + val batchOperations = (1 to 10).map { i => + Row(s"$baseId-$i", partitionKeyValue, i) + } + + val operationsDf = spark.createDataFrame(batchOperations.asJava, simplifiedSchema) + + operationsDf.write + .format("cosmos.oltp") + .option("spark.cosmos.accountEndpoint", cosmosEndpoint) + .option("spark.cosmos.accountKey", cosmosMasterKey) + .option("spark.cosmos.database", cosmosDatabase) + .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) + .option("spark.cosmos.write.strategy", "ItemOverwrite") + .option("spark.cosmos.write.bulk.enabled", "true") + .option("spark.cosmos.write.bulk.enableTransactions", "true") + .mode(SaveMode.Append) + .save() + + // Verify items were created in order + (1 to 10).foreach { i => + val item = container.readItem(s"$baseId-$i", new PartitionKey(partitionKeyValue), classOf[ObjectNode]).block() + item should not be null + item.getItem.get("sequence").asInt() shouldEqual i + } + } + + it should "support update operations with simplified schema" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) + val partitionKeyValue = UUID.randomUUID().toString + val itemId = s"test-item-${UUID.randomUUID()}" + + // First create an item + val initialDoc = Utils.getSimpleObjectMapper.createObjectNode() + initialDoc.put("id", itemId) + initialDoc.put("pk", partitionKeyValue) + initialDoc.put("name", "InitialName") + initialDoc.put("version", 1) + container.createItem(initialDoc).block() + + // Now update it using transactional batch (defaults to upsert) + val simplifiedSchema = StructType(Seq( + StructField("id", StringType, nullable = false), + StructField("pk", StringType, nullable = false), + StructField("name", StringType, nullable = false), + StructField("version", IntegerType, nullable = false) + )) + + val updateOperations = Seq( + Row(itemId, partitionKeyValue, "UpdatedName", 2) + ) + + val operationsDf = spark.createDataFrame(updateOperations.asJava, simplifiedSchema) + + operationsDf.write + .format("cosmos.oltp") + .option("spark.cosmos.accountEndpoint", cosmosEndpoint) + .option("spark.cosmos.accountKey", cosmosMasterKey) + .option("spark.cosmos.database", cosmosDatabase) + .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) + .option("spark.cosmos.write.strategy", "ItemOverwrite") + .option("spark.cosmos.write.bulk.enabled", "true") + .option("spark.cosmos.write.bulk.enableTransactions", "true") + .mode(SaveMode.Append) + .save() + + // Verify item was updated + val updatedItem = container.readItem(itemId, new PartitionKey(partitionKeyValue), classOf[ObjectNode]).block() + updatedItem should not be null + updatedItem.getItem.get("name").asText() shouldEqual "UpdatedName" + updatedItem.getItem.get("version").asInt() shouldEqual 2 + } + + it should "fail when more than 100 operations for a single partition key" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + + val partitionKeyValue = UUID.randomUUID().toString + + // Create 101 operations for the same partition key + val schema = StructType(Seq( + StructField("id", StringType, nullable = false), + StructField("pk", StringType, nullable = false), + StructField("name", StringType, nullable = false) + )) + + val batchOperations = (1 to 101).map { i => + Row(s"item-$i-${UUID.randomUUID()}", partitionKeyValue, s"Name-$i") + } + + val operationsDf = spark.createDataFrame(batchOperations.asJava, schema) + + // Should throw exception due to exceeding 100 operations per partition key limit + val exception = intercept[Exception] { + operationsDf.write + .format("cosmos.oltp") + .option("spark.cosmos.accountEndpoint", cosmosEndpoint) + .option("spark.cosmos.accountKey", cosmosMasterKey) + .option("spark.cosmos.database", cosmosDatabase) + .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) + .option("spark.cosmos.write.strategy", "ItemOverwrite") + .option("spark.cosmos.write.bulk.enabled", "true") + .option("spark.cosmos.write.bulk.enableTransactions", "true") + .mode(SaveMode.Append) + .save() + } + + // Spark wraps our exception in "Writing job aborted", check the root cause + val rootCause = getRootCause(exception) + rootCause.getMessage should include("exceeds") + rootCause.getMessage should include("101 operations") + rootCause.getMessage should include("maximum allowed limit of 100") + } + + "Transactional Batch with Hierarchical Partition Keys" should "create items atomically with PermId and SourceId" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + + // Create container with hierarchical partition keys (PermId, SourceId) + val containerName = s"test-hpk-${UUID.randomUUID()}" + val containerProperties = new com.azure.cosmos.models.CosmosContainerProperties( + containerName, + new com.azure.cosmos.models.PartitionKeyDefinition() + ) + val paths = new java.util.ArrayList[String]() + paths.add("/PermId") + paths.add("/SourceId") + containerProperties.getPartitionKeyDefinition.setPaths(paths) + containerProperties.getPartitionKeyDefinition.setKind(com.azure.cosmos.models.PartitionKind.MULTI_HASH) + containerProperties.getPartitionKeyDefinition.setVersion(com.azure.cosmos.models.PartitionKeyDefinitionVersion.V2) + cosmosClient.getDatabase(cosmosDatabase).createContainer(containerProperties).block() + val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(containerName) + + try { + val permId = "MSFT" + val sourceId = "Bloomberg" + val item1Id = s"${UUID.randomUUID()}" + val item2Id = s"${UUID.randomUUID()}" + + // Create batch operations with hierarchical partition keys + val schema = StructType(Seq( + StructField("id", StringType, nullable = false), + StructField("PermId", StringType, nullable = false), + StructField("SourceId", StringType, nullable = false), + StructField("price", org.apache.spark.sql.types.DoubleType, nullable = false) + )) + + val batchOperations = Seq( + Row(item1Id, permId, sourceId, 100.5), + Row(item2Id, permId, sourceId, 101.25) + ) + + val operationsDf = spark.createDataFrame(batchOperations.asJava, schema) + + operationsDf.write + .format("cosmos.oltp") + .option("spark.cosmos.accountEndpoint", cosmosEndpoint) + .option("spark.cosmos.accountKey", cosmosMasterKey) + .option("spark.cosmos.database", cosmosDatabase) + .option("spark.cosmos.container", containerName) + .option("spark.cosmos.write.strategy", "ItemOverwrite") + .option("spark.cosmos.write.bulk.enabled", "true") + .option("spark.cosmos.write.bulk.enableTransactions", "true") + .mode(SaveMode.Append) + .save() + + // Verify items were created + val pk = new PartitionKeyBuilder().add(permId).add(sourceId).build() + val item1 = container.readItem(item1Id, pk, classOf[ObjectNode]).block() + item1 should not be null + item1.getItem.get("price").asDouble() shouldEqual 100.5 + + val item2 = container.readItem(item2Id, pk, classOf[ObjectNode]).block() + item2 should not be null + item2.getItem.get("price").asDouble() shouldEqual 101.25 + } finally { + // Clean up container + container.delete().block() + } + } + + it should "handle temporal updates for financial instrument timelines" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + + // Create container with hierarchical partition keys (PermId, SourceId) + val containerName = s"test-hpk-temporal-${UUID.randomUUID()}" + val containerProperties = new com.azure.cosmos.models.CosmosContainerProperties( + containerName, + new com.azure.cosmos.models.PartitionKeyDefinition() + ) + val paths = new java.util.ArrayList[String]() + paths.add("/PermId") + paths.add("/SourceId") + containerProperties.getPartitionKeyDefinition.setPaths(paths) + containerProperties.getPartitionKeyDefinition.setKind(com.azure.cosmos.models.PartitionKind.MULTI_HASH) + containerProperties.getPartitionKeyDefinition.setVersion(com.azure.cosmos.models.PartitionKeyDefinitionVersion.V2) + cosmosClient.getDatabase(cosmosDatabase).createContainer(containerProperties).block() + val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(containerName) + + try { + val permId = "MSFT" + val sourceId = "Bloomberg" + val oldRecordId = "2024-01-01T00:00:00Z" + val newRecordId = "2024-06-01T00:00:00Z" + + // First, create initial record + val initialDoc = Utils.getSimpleObjectMapper.createObjectNode() + initialDoc.put("id", oldRecordId) + initialDoc.put("PermId", permId) + initialDoc.put("SourceId", sourceId) + initialDoc.put("price", 100.0) + initialDoc.putNull("valid_to") + container.createItem(initialDoc).block() + + // Now perform atomic temporal update: close old record + create new record + val schema = StructType(Seq( + StructField("id", StringType, nullable = false), + StructField("PermId", StringType, nullable = false), + StructField("SourceId", StringType, nullable = false), + StructField("price", org.apache.spark.sql.types.DoubleType, nullable = false), + StructField("valid_to", StringType, nullable = true) + )) + + val batchOperations = Seq( + // Close old record by setting valid_to (using upsert to replace) + Row(oldRecordId, permId, sourceId, 100.0, "2024-06-01T00:00:00Z"), + // Create new record with new price + Row(newRecordId, permId, sourceId, 150.0, null) + ) + + val operationsDf = spark.createDataFrame(batchOperations.asJava, schema) + + operationsDf.write + .format("cosmos.oltp") + .option("spark.cosmos.accountEndpoint", cosmosEndpoint) + .option("spark.cosmos.accountKey", cosmosMasterKey) + .option("spark.cosmos.database", cosmosDatabase) + .option("spark.cosmos.container", containerName) + .option("spark.cosmos.write.strategy", "ItemOverwrite") + .option("spark.cosmos.write.bulk.enabled", "true") + .option("spark.cosmos.write.bulk.enableTransactions", "true") + .mode(SaveMode.Append) + .save() + + // Verify old record was closed + val pk = new PartitionKeyBuilder().add(permId).add(sourceId).build() + val oldRecord = container.readItem(oldRecordId, pk, classOf[ObjectNode]).block() + oldRecord should not be null + oldRecord.getItem.get("valid_to").asText() shouldEqual "2024-06-01T00:00:00Z" + + // Verify new record was created + val newRecord = container.readItem(newRecordId, pk, classOf[ObjectNode]).block() + newRecord should not be null + newRecord.getItem.get("price").asDouble() shouldEqual 150.0 + newRecord.getItem.get("valid_to").isNull shouldBe true + } finally { + // Clean up container + container.delete().block() + } + } + + it should "handle operations across multiple PermId/SourceId combinations" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + + // Create container with hierarchical partition keys (PermId, SourceId) + val containerName = s"test-hpk-multi-${UUID.randomUUID()}" + val containerProperties = new com.azure.cosmos.models.CosmosContainerProperties( + containerName, + new com.azure.cosmos.models.PartitionKeyDefinition() + ) + val paths = new java.util.ArrayList[String]() + paths.add("/PermId") + paths.add("/SourceId") + containerProperties.getPartitionKeyDefinition.setPaths(paths) + containerProperties.getPartitionKeyDefinition.setKind(com.azure.cosmos.models.PartitionKind.MULTI_HASH) + containerProperties.getPartitionKeyDefinition.setVersion(com.azure.cosmos.models.PartitionKeyDefinitionVersion.V2) + cosmosClient.getDatabase(cosmosDatabase).createContainer(containerProperties).block() + val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(containerName) + + try { + // Create operations for different PermId/SourceId combinations + val schema = StructType(Seq( + StructField("id", StringType, nullable = false), + StructField("PermId", StringType, nullable = false), + StructField("SourceId", StringType, nullable = false), + StructField("price", org.apache.spark.sql.types.DoubleType, nullable = false) + )) + + val batchOperations = Seq( + // MSFT from Bloomberg + Row(s"${UUID.randomUUID()}", "MSFT", "Bloomberg", 100.0), + Row(s"${UUID.randomUUID()}", "MSFT", "Bloomberg", 101.0), + // MSFT from Reuters + Row(s"${UUID.randomUUID()}", "MSFT", "Reuters", 100.5), + // AAPL from Bloomberg + Row(s"${UUID.randomUUID()}", "AAPL", "Bloomberg", 150.0) + ) + + val operationsDf = spark.createDataFrame(batchOperations.asJava, schema) + + operationsDf.write + .format("cosmos.oltp") + .option("spark.cosmos.accountEndpoint", cosmosEndpoint) + .option("spark.cosmos.accountKey", cosmosMasterKey) + .option("spark.cosmos.database", cosmosDatabase) + .option("spark.cosmos.container", containerName) + .option("spark.cosmos.write.strategy", "ItemOverwrite") + .option("spark.cosmos.write.bulk.enabled", "true") + .option("spark.cosmos.write.bulk.enableTransactions", "true") + .mode(SaveMode.Append) + .save() + + // All operations should succeed since they're across different partition key combinations + // Each unique PermId/SourceId combination is treated as a separate transactional batch + } finally { + // Clean up container + container.delete().block() + } + } + + it should "fail when more than 100 operations for a single hierarchical partition key" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + + // Create container with hierarchical partition keys (PermId, SourceId) + val containerName = s"test-hpk-limit-${UUID.randomUUID()}" + val containerProperties = new com.azure.cosmos.models.CosmosContainerProperties( + containerName, + new com.azure.cosmos.models.PartitionKeyDefinition() + ) + val paths = new java.util.ArrayList[String]() + paths.add("/PermId") + paths.add("/SourceId") + containerProperties.getPartitionKeyDefinition.setPaths(paths) + containerProperties.getPartitionKeyDefinition.setKind(com.azure.cosmos.models.PartitionKind.MULTI_HASH) + containerProperties.getPartitionKeyDefinition.setVersion(com.azure.cosmos.models.PartitionKeyDefinitionVersion.V2) + cosmosClient.getDatabase(cosmosDatabase).createContainer(containerProperties).block() + val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(containerName) + + try { + val permId = "MSFT" + val sourceId = "Bloomberg" + + // Create 101 operations for the same hierarchical partition key + val schema = StructType(Seq( + StructField("id", StringType, nullable = false), + StructField("PermId", StringType, nullable = false), + StructField("SourceId", StringType, nullable = false), + StructField("price", org.apache.spark.sql.types.DoubleType, nullable = false) + )) + + val batchOperations = (1 to 101).map { i => + Row(s"${UUID.randomUUID()}", permId, sourceId, 100.0 + i) + } + + val operationsDf = spark.createDataFrame(batchOperations.asJava, schema) + + // Should throw exception due to exceeding 100 operations per partition key limit + val exception = intercept[Exception] { + operationsDf.write + .format("cosmos.oltp") + .option("spark.cosmos.accountEndpoint", cosmosEndpoint) + .option("spark.cosmos.accountKey", cosmosMasterKey) + .option("spark.cosmos.database", cosmosDatabase) + .option("spark.cosmos.container", containerName) + .option("spark.cosmos.write.strategy", "ItemOverwrite") + .option("spark.cosmos.write.bulk.enabled", "true") + .option("spark.cosmos.write.bulk.enableTransactions", "true") + .mode(SaveMode.Append) + .save() + } + + // Spark wraps our exception in "Writing job aborted", check the root cause + val rootCause = getRootCause(exception) + rootCause.getMessage should include("exceeds") + rootCause.getMessage should include("101 operations") + rootCause.getMessage should include("maximum allowed limit of 100") + } finally { + // Clean up container + container.delete().block() + } + } + + it should "handle multiple partition keys with repartitioning" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) + + // Create operations for multiple partition keys intentionally in random order + val pk1 = UUID.randomUUID().toString + val pk2 = UUID.randomUUID().toString + val pk3 = UUID.randomUUID().toString + + val schema = StructType(Seq( + StructField("id", StringType, nullable = false), + StructField("pk", StringType, nullable = false), + StructField("value", StringType, nullable = false) + )) + + // Interleave operations for different partition keys + // With repartitioning, operations should be grouped by partition key + val batchOperations = Seq( + Row(s"${pk1}_1", pk1, "value1"), + Row(s"${pk2}_1", pk2, "value1"), + Row(s"${pk3}_1", pk3, "value1"), + Row(s"${pk1}_2", pk1, "value2"), + Row(s"${pk2}_2", pk2, "value2"), + Row(s"${pk3}_2", pk3, "value2"), + Row(s"${pk1}_3", pk1, "value3"), + Row(s"${pk2}_3", pk2, "value3"), + Row(s"${pk3}_3", pk3, "value3") + ) + + val operationsDf = spark.createDataFrame(batchOperations.asJava, schema) + + // Execute transactional batch + operationsDf.write + .format("cosmos.oltp") + .option("spark.cosmos.accountEndpoint", cosmosEndpoint) + .option("spark.cosmos.accountKey", cosmosMasterKey) + .option("spark.cosmos.database", cosmosDatabase) + .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) + .option("spark.cosmos.write.strategy", "ItemOverwrite") + .option("spark.cosmos.write.bulk.enabled", "true") + .option("spark.cosmos.write.bulk.enableTransactions", "true") + .mode(SaveMode.Append) + .save() + + // Verify all operations succeeded + Seq(pk1, pk2, pk3).foreach { pk => + (1 to 3).foreach { i => + val item = container.readItem(s"${pk}_$i", new PartitionKey(pk), classOf[ObjectNode]).block() + item should not be null + item.getItem.get("value").asText() shouldEqual s"value$i" + } + } + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/TransactionalBulkExecutor.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/TransactionalBulkExecutor.java new file mode 100644 index 000000000000..9d23e55d41a8 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/TransactionalBulkExecutor.java @@ -0,0 +1,1162 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.batch; + +import com.azure.cosmos.BridgeInternal; +import com.azure.cosmos.CosmosAsyncClient; +import com.azure.cosmos.CosmosAsyncContainer; +import com.azure.cosmos.CosmosBridgeInternal; +import com.azure.cosmos.CosmosEndToEndOperationLatencyPolicyConfig; +import com.azure.cosmos.CosmosException; +import com.azure.cosmos.CosmosItemSerializer; +import com.azure.cosmos.ThrottlingRetryOptions; +import com.azure.cosmos.implementation.AsyncDocumentClient; +import com.azure.cosmos.implementation.CosmosBulkExecutionOptionsImpl; +import com.azure.cosmos.implementation.CosmosSchedulers; +import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.implementation.ImplementationBridgeHelpers; +import com.azure.cosmos.implementation.OperationType; +import com.azure.cosmos.implementation.RequestOptions; +import com.azure.cosmos.implementation.ResourceType; +import com.azure.cosmos.implementation.UUIDs; +import com.azure.cosmos.implementation.apachecommons.lang.tuple.Pair; +import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple; +import com.azure.cosmos.models.CosmosBatchOperationResult; +import com.azure.cosmos.models.CosmosBatchResponse; +import com.azure.cosmos.models.CosmosBatch; +import com.azure.cosmos.models.CosmosBatchRequestOptions; +import com.azure.cosmos.models.CosmosBatchItemRequestOptions; +import com.azure.cosmos.models.CosmosBulkItemResponse; +import com.azure.cosmos.models.CosmosBulkOperationResponse; +import com.azure.cosmos.models.CosmosItemOperation; +import com.azure.cosmos.models.CosmosItemOperationType; +import com.azure.cosmos.models.CosmosPatchOperations; +import com.azure.cosmos.models.PartitionKey; +import com.azure.cosmos.models.ModelBridgeInternal; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import reactor.core.Disposable; +import reactor.core.Exceptions; +import reactor.core.publisher.Flux; +import reactor.core.publisher.FluxProcessor; +import reactor.core.publisher.FluxSink; +import reactor.core.publisher.GroupedFlux; +import reactor.core.publisher.Mono; +import reactor.core.publisher.SignalType; +import reactor.core.publisher.Sinks; +import reactor.core.publisher.UnicastProcessor; +import reactor.core.scheduler.Scheduler; +import reactor.util.function.Tuple2; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; + +import static com.azure.core.util.FluxUtil.withContext; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; + +/** + * The Core logic of bulk execution is here. + * + * The actual execution of the flux of operations. It is done in following steps: + + * 1. Getting partition key range ID and grouping operations using that id. + * 2. For the flux of operations in a group, adding buffering based on size and a duration. + * 3. For the operation we get in after buffering, process it using a batch request and return + * a wrapper having request, response(if-any) and exception(if-any). Either response or exception will be there. + * + * 4. Any internal retry is done by adding in an intermediate sink for each grouped flux. + * 5. Any operation which failed due to partition key range gone is retried by putting it in the main sink which leads + * to re-calculation of partition key range id. + * 6. At the end and this is very essential, we close all the sinks as the sink continues to waits for more and the + * execution isn't finished even if all the operations have been executed(figured out by completion call of source) + * + * Note: Sink will move to a new interface from 3.5 and this is documentation for it: + * - https://github.com/reactor/reactor-core/blob/master/docs/asciidoc/processors.adoc + * + * For our use case, Sinks.many().unicast() will work. + */ +public final class TransactionalBulkExecutor implements Disposable { + + private final static Logger logger = LoggerFactory.getLogger(TransactionalBulkExecutor.class); + private final static AtomicLong instanceCount = new AtomicLong(0); + private static final ImplementationBridgeHelpers.CosmosAsyncClientHelper.CosmosAsyncClientAccessor clientAccessor = + ImplementationBridgeHelpers.CosmosAsyncClientHelper.getCosmosAsyncClientAccessor(); + private static final ImplementationBridgeHelpers.CosmosBatchResponseHelper.CosmosBatchResponseAccessor cosmosBatchResponseAccessor = + ImplementationBridgeHelpers.CosmosBatchResponseHelper.getCosmosBatchResponseAccessor(); + + private final CosmosAsyncContainer container; + private final int maxMicroBatchPayloadSizeInBytes; + private final AsyncDocumentClient docClientWrapper; + private final String operationContextText; + private final OperationContextAndListenerTuple operationListener; + private final ThrottlingRetryOptions throttlingRetryOptions; + private final Flux inputOperations; + + // Options for bulk execution. + private final Long maxMicroBatchIntervalInMs; + + private final TContext batchContext; + private final ConcurrentMap partitionScopeThresholds; + private final CosmosBulkExecutionOptionsImpl cosmosBulkExecutionOptions; + + // Handle gone error: + private final AtomicBoolean mainSourceCompleted; + private final AtomicBoolean isDisposed = new AtomicBoolean(false); + private final AtomicBoolean isShutdown = new AtomicBoolean(false); + private final AtomicInteger totalCount; + private final static Sinks.EmitFailureHandler serializedEmitFailureHandler = new SerializedEmitFailureHandler(); + private final static Sinks.EmitFailureHandler serializedCompleteEmitFailureHandler = + new SerializedCompleteEmitFailureHandler(); + private final Sinks.Many mainSink; + private final List> groupSinks; + private final CosmosAsyncClient cosmosClient; + private final String bulkSpanName; + private final AtomicReference scheduledFutureForFlush; + private final String identifier = "TransactionalBulkExecutor-" + instanceCount.incrementAndGet(); + private final BulkExecutorDiagnosticsTracker diagnosticsTracker; + private final CosmosItemSerializer effectiveItemSerializer; + private final Scheduler executionScheduler; + + @SuppressWarnings({"unchecked"}) + public TransactionalBulkExecutor(CosmosAsyncContainer container, + Flux inputOperations, + CosmosBulkExecutionOptionsImpl cosmosBulkOptions) { + + checkNotNull(container, "expected non-null container"); + checkNotNull(inputOperations, "expected non-null inputOperations"); + checkNotNull(cosmosBulkOptions, "expected non-null bulkOptions"); + + this.maxMicroBatchPayloadSizeInBytes = cosmosBulkOptions.getMaxMicroBatchPayloadSizeInBytes(); + this.cosmosBulkExecutionOptions = cosmosBulkOptions; + this.container = container; + this.bulkSpanName = "nonTransactionalBatch." + this.container.getId(); + this.inputOperations = inputOperations; + this.docClientWrapper = CosmosBridgeInternal.getAsyncDocumentClient(container.getDatabase()); + this.cosmosClient = ImplementationBridgeHelpers + .CosmosAsyncDatabaseHelper + .getCosmosAsyncDatabaseAccessor() + .getCosmosAsyncClient(container.getDatabase()); + this.effectiveItemSerializer = this.docClientWrapper.getEffectiveItemSerializer(cosmosBulkOptions.getCustomItemSerializer()); + + this.throttlingRetryOptions = docClientWrapper.getConnectionPolicy().getThrottlingRetryOptions(); + + // Fill the option first, to make the BulkProcessingOptions immutable, as if accessed directly, we might get + // different values when a new group is created. + maxMicroBatchIntervalInMs = cosmosBulkExecutionOptions.getMaxMicroBatchInterval().toMillis(); + batchContext = (TContext) cosmosBulkExecutionOptions.getLegacyBatchScopedContext(); + this.partitionScopeThresholds = ImplementationBridgeHelpers.CosmosBulkExecutionThresholdsStateHelper + .getBulkExecutionThresholdsAccessor() + .getPartitionScopeThresholds(cosmosBulkExecutionOptions.getThresholdsState()); + operationListener = cosmosBulkExecutionOptions.getOperationContextAndListenerTuple(); + if (operationListener != null && + operationListener.getOperationContext() != null) { + operationContextText = identifier + "[" + operationListener.getOperationContext().toString() + "]"; + } else { + operationContextText = identifier +"[n/a]"; + } + + this.diagnosticsTracker = cosmosBulkExecutionOptions.getDiagnosticsTracker(); + + // Initialize sink for handling gone error. + mainSourceCompleted = new AtomicBoolean(false); + totalCount = new AtomicInteger(0); + mainSink = Sinks.many().unicast().onBackpressureBuffer(); + groupSinks = new CopyOnWriteArrayList<>(); + + this.scheduledFutureForFlush = new AtomicReference<>(CosmosSchedulers + .BULK_EXECUTOR_FLUSH_BOUNDED_ELASTIC + .schedulePeriodically( + this::onFlush, + this.maxMicroBatchIntervalInMs, + this.maxMicroBatchIntervalInMs, + TimeUnit.MILLISECONDS)); + + Scheduler schedulerSnapshotFromOptions = cosmosBulkOptions.getSchedulerOverride(); + this.executionScheduler = schedulerSnapshotFromOptions != null + ? schedulerSnapshotFromOptions + : CosmosSchedulers.BULK_EXECUTOR_BOUNDED_ELASTIC; + + logger.debug("Instantiated BulkExecutor, Context: {}", + this.operationContextText); + } + + @Override + public void dispose() { + if (this.isDisposed.compareAndSet(false, true)) { + long totalCountSnapshot = totalCount.get(); + if (totalCountSnapshot == 0) { + completeAllSinks(); + } else { + this.shutdown(); + } + } + } + + @Override + public boolean isDisposed() { + return this.isDisposed.get(); + } + + private void cancelFlushTask(boolean initializeAggressiveFlush) { + long flushIntervalAfterDrainingIncomingFlux = Math.min( + this.maxMicroBatchIntervalInMs, + BatchRequestResponseConstants + .DEFAULT_MAX_MICRO_BATCH_INTERVAL_AFTER_DRAINING_INCOMING_FLUX_IN_MILLISECONDS); + + Disposable newFlushTask = initializeAggressiveFlush + ? CosmosSchedulers + .BULK_EXECUTOR_FLUSH_BOUNDED_ELASTIC + .schedulePeriodically( + this::onFlush, + flushIntervalAfterDrainingIncomingFlux, + flushIntervalAfterDrainingIncomingFlux, + TimeUnit.MILLISECONDS) + : null; + + Disposable scheduledFutureSnapshot = this.scheduledFutureForFlush.getAndSet(newFlushTask); + + if (scheduledFutureSnapshot != null) { + try { + scheduledFutureSnapshot.dispose(); + logDebugOrWarning("Cancelled all future scheduled tasks {}, Context: {}", getThreadInfo(), this.operationContextText); + } catch (Exception e) { + logger.warn("Failed to cancel scheduled tasks{}, Context: {}", getThreadInfo(), this.operationContextText, e); + } + } + } + + private void logInfoOrWarning(String msg, Object... args) { + if (this.diagnosticsTracker == null || !this.diagnosticsTracker.verboseLoggingAfterReEnqueueingRetriesEnabled()) { + logger.info(msg, args); + } else { + logger.warn(msg, args); + } + } + + private void logTraceOrWarning(String msg, Object... args) { + if (this.diagnosticsTracker == null || !this.diagnosticsTracker.verboseLoggingAfterReEnqueueingRetriesEnabled()) { + logger.trace(msg, args); + } else { + logger.warn(msg, args); + } + } + + private void logDebugOrWarning(String msg, Object... args) { + if (this.diagnosticsTracker == null || !this.diagnosticsTracker.verboseLoggingAfterReEnqueueingRetriesEnabled()) { + logger.debug(msg, args); + } else { + logger.warn(msg, args); + } + } + + private void shutdown() { + if (this.isShutdown.compareAndSet(false, true)) { + logDebugOrWarning("Shutting down, Context: {}", this.operationContextText); + + groupSinks.forEach(FluxSink::complete); + logger.debug("All group sinks completed, Context: {}", this.operationContextText); + + this.cancelFlushTask(false); + } + } + + public Flux> execute() { + return this + .executeCore() + .doFinally((SignalType signal) -> { + if (signal == SignalType.ON_COMPLETE) { + logDebugOrWarning("BulkExecutor.execute flux completed - # left items {}, Context: {}, {}", + this.totalCount.get(), + this.operationContextText, + getThreadInfo()); + } else { + int itemsLeftSnapshot = this.totalCount.get(); + if (itemsLeftSnapshot > 0) { + logInfoOrWarning("BulkExecutor.execute flux terminated - Signal: {} - # left items {}, Context: {}, {}", + signal, + itemsLeftSnapshot, + this.operationContextText, + getThreadInfo()); + } else { + logDebugOrWarning("BulkExecutor.execute flux terminated - Signal: {} - # left items {}, Context: {}, {}", + signal, + itemsLeftSnapshot, + this.operationContextText, + getThreadInfo()); + } + } + + this.dispose(); + }); + } + + private Flux> executeCore() { + + // The groupBy below is running into a hang if the flatMap above is + // not allowing at least a concurrency of the number of unique values + // you groupBy on. + // The groupBy is used to isolate Cosmos physical partitions + // so when there is no config override we enforce that the flatMap is using a concurrency of + // Math.max(default concurrency (256), #of partitions * 2 (to accommodate for some splits)) + // The config override can be used by the Spark connector when customers follow best practices and + // repartition the data frame to avoid that each Spark partition contains data spread across all + // physical partitions. When repartitioning the incoming data it is possible to ensure that each + // Spark partition will only target a subset of Cosmos partitions. This will improve the efficiency + // and mean fewer than #of Partitions concurrency will be needed for + // large containers. (with hundreds of physical partitions) + Integer nullableMaxConcurrentCosmosPartitions = cosmosBulkExecutionOptions.getMaxConcurrentCosmosPartitions(); + Mono maxConcurrentCosmosPartitionsMono = nullableMaxConcurrentCosmosPartitions != null ? + Mono.just(Math.max(256, nullableMaxConcurrentCosmosPartitions)) : + ImplementationBridgeHelpers + .CosmosAsyncContainerHelper + .getCosmosAsyncContainerAccessor() + .getFeedRanges(this.container, false).map(ranges -> Math.max(256, ranges.size() * 2)); + + return + maxConcurrentCosmosPartitionsMono + .subscribeOn(this.executionScheduler) + .flatMapMany(maxConcurrentCosmosPartitions -> { + + logDebugOrWarning("BulkExecutor.execute with MaxConcurrentPartitions: {}, Context: {}", + maxConcurrentCosmosPartitions, + this.operationContextText); + + return this.inputOperations + .publishOn(this.executionScheduler) + .onErrorMap(throwable -> { + logger.warn("{}: Error observed when processing inputOperations. Cause: {}, Context: {}", + getThreadInfo(), + throwable.getMessage(), + this.operationContextText, + throwable); + + return throwable; + }) + .doOnNext((CosmosItemOperation cosmosItemOperation) -> { + // Set the retry policy before starting execution. Should only happens once. + BulkExecutorUtil.setRetryPolicyForBulk( + docClientWrapper, + this.container, + cosmosItemOperation, + this.throttlingRetryOptions); + + if (cosmosItemOperation != FlushBuffersItemOperation.singleton()) { + totalCount.incrementAndGet(); + } + + logger.trace( + "SetupRetryPolicy, {}, TotalCount: {}, Context: {}, {}", + getItemOperationDiagnostics(cosmosItemOperation), + totalCount.get(), + this.operationContextText, + getThreadInfo() + ); + }) + .doOnComplete(() -> { + mainSourceCompleted.set(true); + + long totalCountSnapshot = totalCount.get(); + logDebugOrWarning("Main source completed - # left items {}, Context: {}", + totalCountSnapshot, + this.operationContextText); + if (totalCountSnapshot == 0) { + // This is needed as there can be case that onComplete was called after last element was processed + // So complete the sink here also if count is 0, if source has completed and count isn't zero, + // then the last element in the doOnNext will close it. Sink doesn't mind in case of a double close. + + completeAllSinks(); + } else { + this.cancelFlushTask(true); + this.onFlush(); + + logDebugOrWarning("Scheduled new flush operation {}, Context: {}", getThreadInfo(), this.operationContextText); + } + }) + .mergeWith(mainSink.asFlux()) + .subscribeOn(this.executionScheduler) + .flatMap( + operation -> { + logger.trace("Before Resolve PkRangeId, {}, Context: {} {}", + getItemOperationDiagnostics(operation), + this.operationContextText, + getThreadInfo()); + + // resolve partition key range id again for operations which comes in main sink due to gone retry. + return BulkExecutorUtil.resolvePartitionKeyRangeId(this.docClientWrapper, this.container, operation) + .map((String pkRangeId) -> { + PartitionScopeThresholds partitionScopeThresholds = + this.partitionScopeThresholds.computeIfAbsent( + pkRangeId, + (newPkRangeId) -> new PartitionScopeThresholds(newPkRangeId, this.cosmosBulkExecutionOptions)); + + logger.trace("Resolved PkRangeId, {}, PKRangeId: {} Context: {} {}", + getItemOperationDiagnostics(operation), + pkRangeId, + this.operationContextText, + getThreadInfo()); + + return Pair.of(partitionScopeThresholds, operation); + }); + }) + .groupBy(Pair::getKey, Pair::getValue) + .flatMap( + this::executePartitionedGroup, + maxConcurrentCosmosPartitions) + .subscribeOn(this.executionScheduler) + .doOnNext(requestAndResponse -> { + + int totalCountAfterDecrement = totalCount.decrementAndGet(); + boolean mainSourceCompletedSnapshot = mainSourceCompleted.get(); + if (totalCountAfterDecrement == 0 && mainSourceCompletedSnapshot) { + // It is possible that count is zero but there are more elements in the source. + // Count 0 also signifies that there are no pending elements in any sink. + logDebugOrWarning("All work completed, {}, TotalCount: {}, Context: {} {}", + getItemOperationDiagnostics(requestAndResponse.getOperation()), + totalCountAfterDecrement, + this.operationContextText, + getThreadInfo()); + completeAllSinks(); + } else { + if (totalCountAfterDecrement == 0) { + logDebugOrWarning( + "No Work left - but mainSource not yet completed, Context: {} {}", + this.operationContextText, + getThreadInfo()); + } + logTraceOrWarning( + "Work left - TotalCount after decrement: {}, main sink completed {}, {}, Context: {} {}", + totalCountAfterDecrement, + mainSourceCompletedSnapshot, + getItemOperationDiagnostics(requestAndResponse.getOperation()), + this.operationContextText, + getThreadInfo()); + } + }) + .doOnComplete(() -> { + int totalCountSnapshot = totalCount.get(); + boolean mainSourceCompletedSnapshot = mainSourceCompleted.get(); + if (totalCountSnapshot == 0 && mainSourceCompletedSnapshot) { + // It is possible that count is zero but there are more elements in the source. + // Count 0 also signifies that there are no pending elements in any sink. + logDebugOrWarning("DoOnComplete: All work completed, Context: {}", this.operationContextText); + completeAllSinks(); + } else { + logDebugOrWarning( + "DoOnComplete: Work left - TotalCount after decrement: {}, main sink completed {}, Context: {} {}", + totalCountSnapshot, + mainSourceCompletedSnapshot, + this.operationContextText, + getThreadInfo()); + } + }); + }); + } + + private Flux> executePartitionedGroup( + GroupedFlux partitionedGroupFluxOfInputOperations) { + + final PartitionScopeThresholds thresholds = partitionedGroupFluxOfInputOperations.key(); + + final FluxProcessor groupFluxProcessor = + UnicastProcessor.create().serialize(); + final FluxSink groupSink = groupFluxProcessor.sink(FluxSink.OverflowStrategy.BUFFER); + groupSinks.add(groupSink); + + AtomicReference currentPartitionKey = new AtomicReference<>(null); + AtomicLong currentMicroBatchSize = new AtomicLong(0); + + return partitionedGroupFluxOfInputOperations + .mergeWith(groupFluxProcessor) + .onBackpressureBuffer() + .timestamp() + .subscribeOn(this.executionScheduler) + .bufferUntil(timeStampItemOperationTuple -> { + long timestamp = timeStampItemOperationTuple.getT1(); + CosmosItemOperation itemOperation = timeStampItemOperationTuple.getT2(); + + logger.trace( + "TransactionalBulkExecutor.BufferUntil - enqueued {}, {}, Context: {} {}", + timestamp, + getItemOperationDiagnostics(itemOperation), + this.operationContextText, + getThreadInfo()); + + if (itemOperation == FlushBuffersItemOperation.singleton()) { + long currentMicroBatchSizeSnapshot = currentMicroBatchSize.get(); + if (currentMicroBatchSizeSnapshot > 0) { + logger.trace( + "TransactionalBulkExecutor - Flushing PKRange {} (batch size: {}) due to FlushItemOperation, Context: {} {}", + thresholds.getPartitionKeyRangeId(), + currentMicroBatchSizeSnapshot, + this.operationContextText, + getThreadInfo()); + + currentPartitionKey.set(null); + currentMicroBatchSize.set(0); + + return true; + } + + // avoid counting flush operations for the micro batch size calculation + return false; + } + + // For transactional batches, we flush whenever the logical partition key changes + PartitionKey itemPartitionKey = itemOperation.getPartitionKeyValue(); + PartitionKey currentPK = currentPartitionKey.get(); + long currentBatchSize = currentMicroBatchSize.get(); + + // Check if partition key has changed BEFORE incrementing batch size + boolean partitionKeyChanged = currentPK != null && !currentPK.equals(itemPartitionKey); + + if (partitionKeyChanged) { + logger.trace( + "TransactionalBulkExecutor - Partition key changed, PKRange {}, " + + "Previous PK: {}, New PK: {}, Previous batch size: {}, Context: {} {}", + thresholds.getPartitionKeyRangeId(), + currentPK, + itemPartitionKey, + currentBatchSize, + this.operationContextText, + getThreadInfo()); + + // Reset for the new batch + currentPartitionKey.set(itemPartitionKey); + currentMicroBatchSize.set(1); // Current item is first of new batch + + // Return true with cutBefore=true means: flush previous batch, start new batch with current item + return true; + } + + // First item sets the partition key + if (currentPK == null) { + currentPartitionKey.set(itemPartitionKey); + currentMicroBatchSize.set(1); + return false; + } + + // For transactional batches, we cannot split operations for the same partition key + // across multiple batches (would break atomicity). If we're about to exceed the limit, + // the operation will fail when we try to create the batch request. + // Just increment and let the batch creation fail with appropriate error. + currentMicroBatchSize.incrementAndGet(); + return false; + }, true) // cutBefore=true: when predicate is true, current element starts NEXT buffer + .flatMap( + (List> timeStampAndItemOperationTuples) -> { + List operations = new ArrayList<>(timeStampAndItemOperationTuples.size()); + PartitionKey batchPartitionKey = null; + + for (Tuple2 timeStampAndItemOperationTuple : + timeStampAndItemOperationTuples) { + + CosmosItemOperation itemOperation = timeStampAndItemOperationTuple.getT2(); + if (itemOperation == FlushBuffersItemOperation.singleton()) { + continue; + } + operations.add(itemOperation); + + // Capture the partition key for this batch (all items should have the same PK) + if (batchPartitionKey == null) { + batchPartitionKey = itemOperation.getPartitionKeyValue(); + } + } + + if (operations.isEmpty()) { + logger.trace("Empty operations list after filtering, Context: {}", this.operationContextText); + return Flux.empty(); + } + + logDebugOrWarning( + "Flushing transactional batch for PKRange {} with {} operations, PK: {}, Context: {} {}", + thresholds.getPartitionKeyRangeId(), + operations.size(), + batchPartitionKey, + this.operationContextText, + getThreadInfo()); + + return executeTransactionalBatch(operations, batchPartitionKey, thresholds, groupSink); + }, + this.cosmosBulkExecutionOptions.getMaxMicroBatchConcurrency()); + } + + private int calculateTotalSerializedLength(AtomicInteger currentTotalSerializedLength, CosmosItemOperation item) { + if (item instanceof CosmosItemOperationBase) { + return currentTotalSerializedLength.accumulateAndGet( + ((CosmosItemOperationBase) item).getSerializedLength(this.effectiveItemSerializer), + Integer::sum); + } + + return currentTotalSerializedLength.get(); + } + + private Flux> executeTransactionalBatch( + List operations, + PartitionKey partitionKey, + PartitionScopeThresholds thresholds, + FluxSink groupSink) { + + if (operations.size() == 0) { + logger.trace("Empty operations list, Context: {}", this.operationContextText); + return Flux.empty(); + } + + // Create a CosmosBatch for transactional execution + CosmosBatch cosmosBatch = CosmosBatch.createCosmosBatch(partitionKey); + + // Add all operations to the batch + for (CosmosItemOperation operation : operations) { + if (operation instanceof ItemBulkOperation) { + ItemBulkOperation itemBulkOperation = (ItemBulkOperation) operation; + CosmosBatchItemRequestOptions itemRequestOptions = itemBulkOperation.getRequestOptions() != null ? + new CosmosBatchItemRequestOptions() : null; + + switch (itemBulkOperation.getOperationType()) { + case CREATE: + cosmosBatch.createItemOperation( + itemBulkOperation.getItem(), + itemRequestOptions); + break; + case UPSERT: + cosmosBatch.upsertItemOperation( + itemBulkOperation.getItem(), + itemRequestOptions); + break; + case REPLACE: + cosmosBatch.replaceItemOperation( + itemBulkOperation.getId(), + itemBulkOperation.getItem(), + itemRequestOptions); + break; + case DELETE: + cosmosBatch.deleteItemOperation( + itemBulkOperation.getId(), + itemRequestOptions); + break; + case READ: + cosmosBatch.readItemOperation( + itemBulkOperation.getId(), + itemRequestOptions); + break; + case PATCH: + // For PATCH, the item is actually CosmosPatchOperations + Object patchItem = itemBulkOperation.getItemInternal(); + if (patchItem instanceof com.azure.cosmos.models.CosmosPatchOperations) { + com.azure.cosmos.models.CosmosBatchPatchItemRequestOptions patchRequestOptions = + itemRequestOptions != null ? new com.azure.cosmos.models.CosmosBatchPatchItemRequestOptions() : null; + cosmosBatch.patchItemOperation( + itemBulkOperation.getId(), + (com.azure.cosmos.models.CosmosPatchOperations) patchItem, + patchRequestOptions); + } + break; + default: + logger.warn("Unsupported operation type: {}", itemBulkOperation.getOperationType()); + break; + } + } + } + + CosmosBatchRequestOptions batchRequestOptions = new CosmosBatchRequestOptions(); + + // Execute the transactional batch + return Flux.just(cosmosBatch) + .publishOn(this.executionScheduler) + .flatMap((CosmosBatch batch) -> + this.executeTransactionalBatchRequest(batch, batchRequestOptions, operations, groupSink, thresholds)); + } + + private Flux> executeTransactionalBatchRequest( + CosmosBatch cosmosBatch, + CosmosBatchRequestOptions requestOptions, + List operations, + FluxSink groupSink, + PartitionScopeThresholds thresholds) { + + String batchTrackingId = UUIDs.nonBlockingRandomUUID().toString(); + logTraceOrWarning( + "Executing transactional batch - batch TrackingId: %s", + batchTrackingId); + + // Validate that transactional batch does not exceed Cosmos DB limit + if (operations.size() > BatchRequestResponseConstants.MAX_OPERATIONS_IN_DIRECT_MODE_BATCH_REQUEST) { + String errorMessage = String.format( + "Transactional batch operation failed: partition key '%s' has %d operations, " + + "which exceeds the maximum allowed limit of %d operations per transactional batch. " + + "Transactional batches require all-or-nothing execution and cannot be split across multiple requests.", + cosmosBatch.getPartitionKeyValue(), + operations.size(), + BatchRequestResponseConstants.MAX_OPERATIONS_IN_DIRECT_MODE_BATCH_REQUEST); + + logger.error(errorMessage); + + // Return error responses for all operations in this batch + return Flux.fromIterable(operations) + .map(operation -> { + Exception exception = new IllegalArgumentException(errorMessage); + TContext actualContext = this.getActualContext(operation); + return ModelBridgeInternal.createCosmosBulkOperationResponse( + operation, + exception, + actualContext); + }); + } + + // Create SinglePartitionKeyServerBatchRequest with atomic batch semantics + final SinglePartitionKeyServerBatchRequest request = SinglePartitionKeyServerBatchRequest.createBatchRequest( + cosmosBatch.getPartitionKeyValue(), + operations, + this.effectiveItemSerializer); + request.setAtomicBatch(true); + request.setShouldContinueOnError(false); + + // Set up request options + RequestOptions options = new RequestOptions(); + options.setThroughputControlGroupName(cosmosBulkExecutionOptions.getThroughputControlGroupName()); + options.setExcludedRegions(cosmosBulkExecutionOptions.getExcludedRegions()); + options.setKeywordIdentifiers(cosmosBulkExecutionOptions.getKeywordIdentifiers()); + + CosmosEndToEndOperationLatencyPolicyConfig e2eLatencyPolicySnapshot = + cosmosBulkExecutionOptions.getCosmosEndToEndLatencyPolicyConfig(); + if (e2eLatencyPolicySnapshot != null) { + options.setCosmosEndToEndLatencyPolicyConfig(e2eLatencyPolicySnapshot); + } + + Map customOptions = cosmosBulkExecutionOptions.getHeaders(); + if (customOptions != null && !customOptions.isEmpty()) { + for(Map.Entry entry : customOptions.entrySet()) { + options.setHeader(entry.getKey(), entry.getValue()); + } + } + options.setOperationContextAndListenerTuple(operationListener); + + return this.docClientWrapper + .executeBatchRequest( + BridgeInternal.getLink(this.container), + request, + options, + false, + false) + .subscribeOn(this.executionScheduler) + .flatMapMany(response -> { + + logTraceOrWarning( + "Response for transactional batch - status code %s, ActivityId: %s, batch TrackingId %s", + response.getStatusCode(), + response.getActivityId(), + batchTrackingId); + + if (diagnosticsTracker != null && response.getDiagnostics() != null) { + diagnosticsTracker.trackDiagnostics(response.getDiagnostics().getDiagnosticsContext()); + } + + return Flux + .fromIterable(response.getResults()) + .publishOn(this.executionScheduler) + .flatMap((CosmosBatchOperationResult result) -> + handleTransactionalBatchOperationResult(response, result, groupSink, thresholds)); + }) + .onErrorResume((Throwable throwable) -> { + + if (!(throwable instanceof Exception)) { + throw Exceptions.propagate(throwable); + } + + Exception exception = (Exception) throwable; + + return Flux + .fromIterable(operations) + .publishOn(this.executionScheduler) + .flatMap((CosmosItemOperation itemOperation) -> + handleTransactionalBatchExecutionException(itemOperation, exception, groupSink, thresholds)); + }); + } + + // Helper functions + private Mono> handleTransactionalBatchOperationResult( + CosmosBatchResponse response, + CosmosBatchOperationResult operationResult, + FluxSink groupSink, + PartitionScopeThresholds thresholds) { + + CosmosBulkItemResponse cosmosBulkItemResponse = ModelBridgeInternal + .createCosmosBulkItemResponse(operationResult, response); + CosmosItemOperation itemOperation = operationResult.getOperation(); + TContext actualContext = this.getActualContext(itemOperation); + + logDebugOrWarning( + "HandleTransactionalBatchOperationResult - PKRange {}, Response Status Code {}, " + + "Operation Status Code, {}, {}, Context: {} {}", + thresholds.getPartitionKeyRangeId(), + response.getStatusCode(), + operationResult.getStatusCode(), + getItemOperationDiagnostics(itemOperation), + this.operationContextText, + getThreadInfo()); + + if (!operationResult.isSuccessStatusCode()) { + + if (itemOperation instanceof ItemBulkOperation) { + + ItemBulkOperation itemBulkOperation = (ItemBulkOperation) itemOperation; + return itemBulkOperation.getRetryPolicy().shouldRetry(operationResult).flatMap( + result -> { + if (result.shouldRetry) { + logDebugOrWarning( + "HandleTransactionalBatchOperationResult - enqueue retry, PKRange {}, Response " + + "Status Code {}, Operation Status Code, {}, {}, Context: {} {}", + thresholds.getPartitionKeyRangeId(), + response.getStatusCode(), + operationResult.getStatusCode(), + getItemOperationDiagnostics(itemOperation), + this.operationContextText, + getThreadInfo()); + return this.enqueueForRetry(result.backOffTime, groupSink, itemOperation, thresholds); + } else { + // reduce log noise level for commonly expected/normal status codes + if (response.getStatusCode() == HttpConstants.StatusCodes.CONFLICT || + response.getStatusCode() == HttpConstants.StatusCodes.PRECONDITION_FAILED) { + + logDebugOrWarning( + "HandleTransactionalBatchOperationResult - Fail, PKRange {}, Response Status " + + "Code {}, Operation Status Code {}, {}, Context: {} {}", + thresholds.getPartitionKeyRangeId(), + response.getStatusCode(), + operationResult.getStatusCode(), + getItemOperationDiagnostics(itemOperation), + this.operationContextText, + getThreadInfo()); + } else { + logger.error( + "HandleTransactionalBatchOperationResult - Fail, PKRange {}, Response Status " + + "Code {}, Operation Status Code {}, {}, Context: {} {}", + thresholds.getPartitionKeyRangeId(), + response.getStatusCode(), + operationResult.getStatusCode(), + getItemOperationDiagnostics(itemOperation), + this.operationContextText, + getThreadInfo()); + } + return Mono.just(ModelBridgeInternal.createCosmosBulkOperationResponse( + itemOperation, cosmosBulkItemResponse, actualContext)); + } + }); + + } else { + throw new UnsupportedOperationException("Unknown CosmosItemOperation."); + } + } + + thresholds.recordSuccessfulOperation(); + return Mono.just(ModelBridgeInternal.createCosmosBulkOperationResponse( + itemOperation, + cosmosBulkItemResponse, + actualContext)); + } + + private TContext getActualContext(CosmosItemOperation itemOperation) { + ItemBulkOperation itemBulkOperation = null; + + if (itemOperation instanceof ItemBulkOperation) { + itemBulkOperation = (ItemBulkOperation) itemOperation; + } + + if (itemBulkOperation == null) { + return this.batchContext; + } + + TContext operationContext = itemBulkOperation.getContext(); + if (operationContext != null) { + return operationContext; + } + + return this.batchContext; + } + + private Mono> handleTransactionalBatchExecutionException( + CosmosItemOperation itemOperation, + Exception exception, + FluxSink groupSink, + PartitionScopeThresholds thresholds) { + + logDebugOrWarning( + "HandleTransactionalBatchExecutionException, PKRange {}, Error: {}, {}, Context: {} {}", + thresholds.getPartitionKeyRangeId(), + exception, + getItemOperationDiagnostics(itemOperation), + this.operationContextText, + getThreadInfo()); + + if (exception instanceof CosmosException && itemOperation instanceof ItemBulkOperation) { + CosmosException cosmosException = (CosmosException) exception; + ItemBulkOperation itemBulkOperation = (ItemBulkOperation) itemOperation; + + // First check if it failed due to split, so the operations need to go in a different pk range group. So + // add it in the mainSink. + + return itemBulkOperation.getRetryPolicy() + .shouldRetryInMainSink(cosmosException.getStatusCode(), cosmosException.getSubStatusCode(), itemBulkOperation, cosmosException) + .flatMap(shouldRetryInMainSink -> { + if (shouldRetryInMainSink) { + logDebugOrWarning( + "HandleTransactionalBatchExecutionException - Retry in main sink, PKRange {}, Error: " + + "{}, {}, Context: {} {}", + thresholds.getPartitionKeyRangeId(), + exception, + getItemOperationDiagnostics(itemOperation), + this.operationContextText, + getThreadInfo()); + // retry - but don't mark as enqueued for retry in thresholds + mainSink.emitNext(itemOperation, serializedEmitFailureHandler); + return Mono.empty(); + } else { + logDebugOrWarning( + "HandleTransactionalBatchExecutionException - Retry other, PKRange {}, Error: " + + "{}, {}, Context: {} {}", + thresholds.getPartitionKeyRangeId(), + exception, + getItemOperationDiagnostics(itemOperation), + this.operationContextText, + getThreadInfo()); + return retryOtherExceptions( + itemOperation, + exception, + groupSink, + cosmosException, + itemBulkOperation, + thresholds); + } + }); + } + + TContext actualContext = this.getActualContext(itemOperation); + return Mono.just(ModelBridgeInternal.createCosmosBulkOperationResponse(itemOperation, exception, actualContext)); + } + + private Mono> enqueueForRetry( + Duration backOffTime, + FluxSink groupSink, + CosmosItemOperation itemOperation, + PartitionScopeThresholds thresholds) { + + thresholds.recordEnqueuedRetry(); + if (backOffTime == null || backOffTime.isZero()) { + groupSink.next(itemOperation); + return Mono.empty(); + } else { + return Mono + .delay(backOffTime) + .flatMap((dummy) -> { + groupSink.next(itemOperation); + return Mono.empty(); + }); + } + } + + private Mono> retryOtherExceptions( + CosmosItemOperation itemOperation, + Exception exception, + FluxSink groupSink, + CosmosException cosmosException, + ItemBulkOperation itemBulkOperation, + PartitionScopeThresholds thresholds) { + + TContext actualContext = this.getActualContext(itemOperation); + return itemBulkOperation.getRetryPolicy().shouldRetry(cosmosException).flatMap(result -> { + if (result.shouldRetry) { + return this.enqueueForRetry(result.backOffTime, groupSink, itemBulkOperation, thresholds); + } else { + return Mono.just(ModelBridgeInternal.createCosmosBulkOperationResponse( + itemOperation, exception, actualContext)); + } + }); + } + + private Mono executeBatchRequest( + PartitionKeyRangeServerBatchRequest serverRequest, + PartitionScopeThresholds partitionScopeThresholds) { + + RequestOptions options = new RequestOptions(); + options.setThroughputControlGroupName(cosmosBulkExecutionOptions.getThroughputControlGroupName()); + options.setExcludedRegions(cosmosBulkExecutionOptions.getExcludedRegions()); + options.setKeywordIdentifiers(cosmosBulkExecutionOptions.getKeywordIdentifiers()); + + CosmosEndToEndOperationLatencyPolicyConfig e2eLatencyPolicySnapshot = + cosmosBulkExecutionOptions.getCosmosEndToEndLatencyPolicyConfig(); + if (e2eLatencyPolicySnapshot != null) { + options.setCosmosEndToEndLatencyPolicyConfig(e2eLatencyPolicySnapshot); + } + + // This logic is to handle custom bulk options which can be passed through encryption or through some other project + Map customOptions = cosmosBulkExecutionOptions.getHeaders(); + if (customOptions != null && !customOptions.isEmpty()) { + for(Map.Entry entry : customOptions.entrySet()) { + options.setHeader(entry.getKey(), entry.getValue()); + } + } + options.setOperationContextAndListenerTuple(operationListener); + + // The request options here are used for the BulkRequest exchanged with the service + // If contentResponseOnWrite is not enabled here (or at the client level) the + // service will not even send a bulk response payload - so all the + // CosmosBulItemRequestOptions are irrelevant - all payloads will be null + // Instead we should automatically enforce contentResponseOnWrite for all + // bulk requests whenever at least one of the item operations requires a content response (either + // because it is a read operation or because contentResponseOnWrite was enabled explicitly) + if (!this.docClientWrapper.isContentResponseOnWriteEnabled() && + serverRequest.getOperations().size() > 0) { + + for (CosmosItemOperation itemOperation : serverRequest.getOperations()) { + if (itemOperation instanceof ItemBulkOperation) { + + ItemBulkOperation itemBulkOperation = (ItemBulkOperation) itemOperation; + if (itemBulkOperation.getOperationType() == CosmosItemOperationType.READ || + (itemBulkOperation.getRequestOptions() != null && + itemBulkOperation.getRequestOptions().isContentResponseOnWriteEnabled() != null && + itemBulkOperation.getRequestOptions().isContentResponseOnWriteEnabled())) { + + options.setContentResponseOnWriteEnabled(true); + break; + } + } + } + } + + return withContext(context -> { + final Mono responseMono = + this.docClientWrapper + .executeBatchRequest( + BridgeInternal.getLink(this.container), + serverRequest, + options, + false, + true) // disable the staled resource exception handling as it is being handled in the BulkOperationRetryPolicy + .flatMap(cosmosBatchResponse -> { + + cosmosBatchResponseAccessor.setGlobalOpCount( + cosmosBatchResponse, partitionScopeThresholds.getTotalOperationCountSnapshot()); + + PartitionScopeThresholds.CurrentIntervalThresholds currentIntervalThresholdsSnapshot + = partitionScopeThresholds.getCurrentThresholds(); + + cosmosBatchResponseAccessor.setOpCountPerEvaluation( + cosmosBatchResponse, currentIntervalThresholdsSnapshot.currentOperationCount.get()); + cosmosBatchResponseAccessor.setRetriedOpCountPerEvaluation( + cosmosBatchResponse, currentIntervalThresholdsSnapshot.currentRetriedOperationCount.get()); + cosmosBatchResponseAccessor.setTargetMaxMicroBatchSize( + cosmosBatchResponse, partitionScopeThresholds.getTargetMicroBatchSizeSnapshot()); + + return Mono.just(cosmosBatchResponse); + }); + + return clientAccessor.getDiagnosticsProvider(this.cosmosClient) + .traceEnabledBatchResponsePublisher( + responseMono, + context, + this.bulkSpanName, + this.container.getDatabase().getId(), + this.container.getId(), + this.cosmosClient, + options.getConsistencyLevel(), + OperationType.Batch, + ResourceType.Document, + options, + this.cosmosBulkExecutionOptions.getMaxMicroBatchSize()); + }); + } + + private void completeAllSinks() { + logInfoOrWarning("Closing all sinks, Context: {}", this.operationContextText); + + logger.debug("Executor service shut down, Context: {}", this.operationContextText); + mainSink.emitComplete(serializedCompleteEmitFailureHandler); + + this.shutdown(); + } + + private void onFlush() { + try { + this.groupSinks.forEach(sink -> sink.next(FlushBuffersItemOperation.singleton())); + } catch(Throwable t) { + logger.error("Callback invocation 'onFlush' failed. Context: {}", this.operationContextText, t); + } + } + + private static String getItemOperationDiagnostics(CosmosItemOperation operation) { + + if (operation == FlushBuffersItemOperation.singleton()) { + return "ItemOperation[Type: Flush]"; + } + + return "ItemOperation[Type: " + + operation.getOperationType().toString() + + ", PK: " + + (operation.getPartitionKeyValue() != null ? operation.getPartitionKeyValue().toString() : "n/a") + + ", id: " + + operation.getId() + + "]"; + } + + private static String getThreadInfo() { + StringBuilder sb = new StringBuilder(); + Thread t = Thread.currentThread(); + sb + .append("Thread[") + .append("Name: ") + .append(t.getName()) + .append(",Group: ") + .append(t.getThreadGroup() != null ? t.getThreadGroup().getName() : "n/a") + .append(", isDaemon: ") + .append(t.isDaemon()) + .append(", Id: ") + .append(t.getId()) + .append("]"); + + return sb.toString(); + } + + private static class SerializedEmitFailureHandler implements Sinks.EmitFailureHandler { + + @Override + public boolean onEmitFailure(SignalType signalType, Sinks.EmitResult emitResult) { + if (emitResult.equals(Sinks.EmitResult.FAIL_NON_SERIALIZED)) { + logger.debug("SerializedEmitFailureHandler.onEmitFailure - Signal:{}, Result: {}", signalType, emitResult); + + return true; + } + + logger.error("SerializedEmitFailureHandler.onEmitFailure - Signal:{}, Result: {}", signalType, emitResult); + return false; + } + } + + private static class SerializedCompleteEmitFailureHandler implements Sinks.EmitFailureHandler { + + @Override + public boolean onEmitFailure(SignalType signalType, Sinks.EmitResult emitResult) { + if (emitResult.equals(Sinks.EmitResult.FAIL_NON_SERIALIZED)) { + logger.debug("SerializedCompleteEmitFailureHandler.onEmitFailure - Signal:{}, Result: {}", signalType, emitResult); + + return true; + } + + if (emitResult == Sinks.EmitResult.FAIL_CANCELLED || emitResult == Sinks.EmitResult.FAIL_TERMINATED) { + logger.debug("SerializedCompleteEmitFailureHandler.onEmitFailure - Main sink already completed, Signal:{}, Result: {}", signalType, emitResult); + return false; + } + + logger.error("SerializedCompleteEmitFailureHandler.onEmitFailure - Signal:{}, Result: {}", signalType, emitResult); + return false; + } + } +} From 9240f7e243d198e4340b46d16af06149583ebb00 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Sat, 6 Dec 2025 14:16:10 +0000 Subject: [PATCH 02/16] Add per-row operation type support for transactional batch --- .../azure/cosmos/spark/AsyncItemWriter.scala | 8 + .../com/azure/cosmos/spark/BulkWriter.scala | 6 + .../com/azure/cosmos/spark/CosmosConfig.scala | 2 +- .../azure/cosmos/spark/CosmosWriterBase.scala | 27 +- .../com/azure/cosmos/spark/PointWriter.scala | 4 + .../spark/TransactionalBulkWriter.scala | 90 ++++-- .../spark/TransactionalBatchITest.scala | 265 ++++++++++++++++++ 7 files changed, 374 insertions(+), 28 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/AsyncItemWriter.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/AsyncItemWriter.scala index e35b43692de1..b04455603676 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/AsyncItemWriter.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/AsyncItemWriter.scala @@ -14,6 +14,14 @@ private trait AsyncItemWriter { */ def scheduleWrite(partitionKeyValue: PartitionKey, objectNode: ObjectNode): Unit + /** + * Schedule a write to happen in async and return immediately with per-row operation type + * @param partitionKeyValue the partition key value + * @param objectNode the json object node + * @param operationType optional operation type (create, upsert, replace, delete) for this specific row + */ + def scheduleWrite(partitionKeyValue: PartitionKey, objectNode: ObjectNode, operationType: Option[String]): Unit + /** * Wait for all remaining work * Throws if any of the work resulted in failure diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/BulkWriter.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/BulkWriter.scala index aec77f8762e8..56e2a691857f 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/BulkWriter.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/BulkWriter.scala @@ -646,6 +646,12 @@ private class BulkWriter } override def scheduleWrite(partitionKeyValue: PartitionKey, objectNode: ObjectNode): Unit = { + scheduleWrite(partitionKeyValue, objectNode, None) + } + + override def scheduleWrite(partitionKeyValue: PartitionKey, objectNode: ObjectNode, operationType: Option[String]): Unit = { + // BulkWriter doesn't support per-row operation types - it uses global ItemWriteStrategy + // The operationType parameter is ignored here for interface compatibility Preconditions.checkState(!closed.get()) throwIfCapturedExceptionExists() diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala index 3c931fd85892..21a6a374bc8d 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala @@ -1439,7 +1439,7 @@ private[spark] object DiagnosticsConfig { private object ItemWriteStrategy extends Enumeration { type ItemWriteStrategy = Value - val ItemOverwrite, ItemAppend, ItemDelete, ItemDeleteIfNotModified, ItemOverwriteIfNotModified, ItemPatch, ItemPatchIfExists, ItemBulkUpdate = Value + val ItemOverwrite, ItemAppend, ItemDelete, ItemDeleteIfNotModified, ItemOverwriteIfNotModified, ItemPatch, ItemPatchIfExists, ItemBulkUpdate, ItemTransactionalBatch = Value } private object CosmosPatchOperationTypes extends Enumeration { diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosWriterBase.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosWriterBase.scala index 0edfbf2170c9..331a6ac3d3f7 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosWriterBase.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosWriterBase.scala @@ -62,7 +62,13 @@ private abstract class CosmosWriterBase( private val writer: AtomicReference[AsyncItemWriter] = new AtomicReference( if (cosmosWriteConfig.bulkEnabled) { - if (cosmosWriteConfig.bulkEnableTransactions) { + // Use TransactionalBulkWriter if either: + // 1. bulkEnableTransactions is explicitly set, OR + // 2. ItemWriteStrategy is ItemTransactionalBatch + val useTransactional = cosmosWriteConfig.bulkEnableTransactions || + cosmosWriteConfig.itemWriteStrategy == ItemWriteStrategy.ItemTransactionalBatch + + if (useTransactional) { new TransactionalBulkWriter( container, cosmosTargetContainerConfig, @@ -94,6 +100,23 @@ private abstract class CosmosWriterBase( override def write(internalRow: InternalRow): Unit = { val objectNode = cosmosRowConverter.fromInternalRowToObjectNode(internalRow, inputSchema) + // Extract operationType if column exists (for per-row operation support) + val operationType: Option[String] = if (inputSchema.fieldNames.contains("operationType")) { + val opTypeIndex = inputSchema.fieldIndex("operationType") + if (!internalRow.isNullAt(opTypeIndex)) { + Some(internalRow.getString(opTypeIndex)) + } else { + None + } + } else { + None + } + + // Remove operationType from objectNode if present (don't persist to Cosmos) + if (objectNode.has("operationType")) { + objectNode.remove("operationType") + } + require(objectNode.has(CosmosConstants.Properties.Id) && objectNode.get(CosmosConstants.Properties.Id).isTextual, s"${CosmosConstants.Properties.Id} is a mandatory field. " + @@ -107,7 +130,7 @@ private abstract class CosmosWriterBase( } val partitionKeyValue = PartitionKeyHelper.getPartitionKeyPath(objectNode, partitionKeyDefinition) - writer.get.scheduleWrite(partitionKeyValue, objectNode) + writer.get.scheduleWrite(partitionKeyValue, objectNode, operationType) } override def commit(): WriterCommitMessage = { diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/PointWriter.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/PointWriter.scala index 8f07bf5339d5..4f804af75687 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/PointWriter.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/PointWriter.scala @@ -76,6 +76,10 @@ private class PointWriter(container: CosmosAsyncContainer, } override def scheduleWrite(partitionKeyValue: PartitionKey, objectNode: ObjectNode): Unit = { + scheduleWrite(partitionKeyValue, objectNode, None) + } + + override def scheduleWrite(partitionKeyValue: PartitionKey, objectNode: ObjectNode, operationType: Option[String]): Unit = { checkState(!closed.get()) val etag = getETag(objectNode) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala index 7d5007c78920..afa8e545bd15 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala @@ -655,6 +655,10 @@ private class TransactionalBulkWriter } override def scheduleWrite(partitionKeyValue: PartitionKey, objectNode: ObjectNode): Unit = { + scheduleWrite(partitionKeyValue, objectNode, None) + } + + override def scheduleWrite(partitionKeyValue: PartitionKey, objectNode: ObjectNode, operationType: Option[String]): Unit = { Preconditions.checkState(!closed.get()) throwIfCapturedExceptionExists() @@ -664,7 +668,9 @@ private class TransactionalBulkWriter partitionKeyValue, getETag(objectNode), 1, - monotonicOperationCounter.incrementAndGet()) + monotonicOperationCounter.incrementAndGet(), + None, + operationType) val numberOfIntervalsWithIdenticalActiveOperationSnapshots = new AtomicLong(0) // Don't clone the activeOperations for the first iteration // to reduce perf impact before the Semaphore has been acquired @@ -739,10 +745,28 @@ private class TransactionalBulkWriter objectNode: ObjectNode, operationContext: OperationContext): Unit = { - val bulkItemOperation = writeConfig.itemWriteStrategy match { - case ItemWriteStrategy.ItemOverwrite => + // Determine the effective operation type from per-row operationType or global strategy + val effectiveOperationType = operationContext.operationType match { + case Some(opType) => opType.toLowerCase + case None => writeConfig.itemWriteStrategy match { + case ItemWriteStrategy.ItemOverwrite => "upsert" + case ItemWriteStrategy.ItemAppend => "create" + case ItemWriteStrategy.ItemDelete | ItemWriteStrategy.ItemDeleteIfNotModified => "delete" + case ItemWriteStrategy.ItemOverwriteIfNotModified => "replace" + case ItemWriteStrategy.ItemTransactionalBatch => "upsert" // Default to upsert for transactional batch + case ItemWriteStrategy.ItemPatch | ItemWriteStrategy.ItemPatchIfExists => "patch" + case _ => throw new RuntimeException(s"${writeConfig.itemWriteStrategy} not supported for transactional batch") + } + } + + val bulkItemOperation = effectiveOperationType match { + case "create" => + CosmosBulkOperations.getCreateItemOperation(objectNode, partitionKeyValue, operationContext) + + case "upsert" => CosmosBulkOperations.getUpsertItemOperation(objectNode, partitionKeyValue, operationContext) - case ItemWriteStrategy.ItemOverwriteIfNotModified => + + case "replace" => operationContext.eTag match { case Some(eTag) => CosmosBulkOperations.getReplaceItemOperation( @@ -751,29 +775,41 @@ private class TransactionalBulkWriter partitionKeyValue, new CosmosBulkItemRequestOptions().setIfMatchETag(eTag), operationContext) - case _ => CosmosBulkOperations.getCreateItemOperation(objectNode, partitionKeyValue, operationContext) + case _ => + CosmosBulkOperations.getReplaceItemOperation( + operationContext.itemId, + objectNode, + partitionKeyValue, + operationContext) } - case ItemWriteStrategy.ItemAppend => - CosmosBulkOperations.getCreateItemOperation(objectNode, partitionKeyValue, operationContext) - case ItemWriteStrategy.ItemDelete => - CosmosBulkOperations.getDeleteItemOperation(operationContext.itemId, partitionKeyValue, operationContext) - case ItemWriteStrategy.ItemDeleteIfNotModified => - CosmosBulkOperations.getDeleteItemOperation( + + case "delete" => + operationContext.eTag match { + case Some(eTag) => + CosmosBulkOperations.getDeleteItemOperation( + operationContext.itemId, + partitionKeyValue, + new CosmosBulkItemRequestOptions().setIfMatchETag(eTag), + operationContext) + case _ => + CosmosBulkOperations.getDeleteItemOperation( + operationContext.itemId, + partitionKeyValue, + operationContext) + } + + case "patch" => + getPatchItemOperation( operationContext.itemId, partitionKeyValue, - operationContext.eTag match { - case Some(eTag) => new CosmosBulkItemRequestOptions().setIfMatchETag(eTag) - case _ => new CosmosBulkItemRequestOptions() - }, + partitionKeyDefinition, + objectNode, operationContext) - case ItemWriteStrategy.ItemPatch | ItemWriteStrategy.ItemPatchIfExists => getPatchItemOperation( - operationContext.itemId, - partitionKeyValue, - partitionKeyDefinition, - objectNode, - operationContext) + case _ => - throw new RuntimeException(s"${writeConfig.itemWriteStrategy} not supported") + throw new IllegalArgumentException( + s"Unsupported operationType '$effectiveOperationType'. " + + s"Supported types for transactional batch: create, upsert, replace, delete") } this.emitBulkInput(bulkItemOperation) @@ -1363,9 +1399,10 @@ private class TransactionalBulkWriter val attemptNumber: Int, val sequenceNumber: Long, /** starts from 1 * */ - sourceItemInput: Option[ObjectNode] = None) // for patchBulkUpdate: source item refers to the original objectNode from which SDK constructs the final bulk item operation + sourceItemInput: Option[ObjectNode] = None, // for patchBulkUpdate: source item refers to the original objectNode from which SDK constructs the final bulk item operation + operationTypeInput: Option[String] = None) // per-row operation type (create, upsert, replace, delete) { - private val ctxCore: OperationContextCore = OperationContextCore(itemIdInput, partitionKeyValueInput, eTagInput, sourceItemInput) + private val ctxCore: OperationContextCore = OperationContextCore(itemIdInput, partitionKeyValueInput, eTagInput, sourceItemInput, operationTypeInput) override def equals(obj: Any): Boolean = ctxCore.equals(obj) @@ -1382,6 +1419,8 @@ private class TransactionalBulkWriter def eTag: Option[String] = ctxCore.eTag def sourceItem: Option[ObjectNode] = ctxCore.sourceItem + + def operationType: Option[String] = ctxCore.operationType } private case class OperationContextCore @@ -1389,7 +1428,8 @@ private class TransactionalBulkWriter itemId: String, partitionKeyValue: PartitionKey, eTag: Option[String], - sourceItem: Option[ObjectNode] = None) // for patchBulkUpdate: source item refers to the original objectNode from which SDK constructs the final bulk item operation + sourceItem: Option[ObjectNode] = None, // for patchBulkUpdate: source item refers to the original objectNode from which SDK constructs the final bulk item operation + operationType: Option[String] = None) // per-row operation type (create, upsert, replace, delete) { override def productPrefix: String = "OperationContext" } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala index 730167417feb..a696077722f8 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala @@ -634,4 +634,269 @@ class TransactionalBatchITest extends IntegrationSpec } } } + + "Transactional batch with per-row operation types" should "support mixed operations in same partition key" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) + + // Pre-create an item for replace operation + val itemForReplace = Utils.getSimpleObjectMapper.createObjectNode() + itemForReplace.put("id", "item-to-replace") + itemForReplace.put("pk", "mixed-ops") + itemForReplace.put("prop", "original-value") + container.createItem(itemForReplace).block() + + // Pre-create item for delete operation + val itemForDelete = Utils.getSimpleObjectMapper.createObjectNode() + itemForDelete.put("id", "item-to-delete") + itemForDelete.put("pk", "mixed-ops") + itemForDelete.put("prop", "to-be-deleted") + container.createItem(itemForDelete).block() + + // Create DataFrame with mixed operations + val schema = StructType(Seq( + StructField("id", StringType, nullable = false), + StructField("pk", StringType, nullable = false), + StructField("operationType", StringType, nullable = false), + StructField("prop", StringType, nullable = false) + )) + + val mixedOpsItems = Seq( + Row("item-new-create", "mixed-ops", "create", "created-value"), + Row("item-new-upsert", "mixed-ops", "upsert", "upserted-value"), + Row("item-to-replace", "mixed-ops", "replace", "replaced-value"), + Row("item-to-delete", "mixed-ops", "delete", "delete-value") + ) + + val mixedOpsDf = spark.createDataFrame(mixedOpsItems.asJava, schema) + + val writeCfg = Map( + "spark.cosmos.accountEndpoint" -> cosmosEndpoint, + "spark.cosmos.accountKey" -> cosmosMasterKey, + "spark.cosmos.database" -> cosmosDatabase, + "spark.cosmos.container" -> cosmosContainersWithPkAsPartitionKey, + "spark.cosmos.write.strategy" -> "ItemTransactionalBatch", + "spark.cosmos.write.maxRetryCount" -> "3", + "spark.cosmos.write.bulk.enabled" -> "true" + ) + + // Write with mixed operations + mixedOpsDf.write.format("cosmos.oltp").mode(SaveMode.Append).options(writeCfg).save() + + // Verify all operations succeeded + // 1. Check created item + val createdItem = container.readItem("item-new-create", new PartitionKey("mixed-ops"), classOf[ObjectNode]).block().getItem + createdItem.get("prop").asText() shouldEqual "created-value" + + // 2. Check upserted item + val upsertedItem = container.readItem("item-new-upsert", new PartitionKey("mixed-ops"), classOf[ObjectNode]).block().getItem + upsertedItem.get("prop").asText() shouldEqual "upserted-value" + + // 3. Check replaced item + val replacedItem = container.readItem("item-to-replace", new PartitionKey("mixed-ops"), classOf[ObjectNode]).block().getItem + replacedItem.get("prop").asText() shouldEqual "replaced-value" + + // 4. Verify deleted item no longer exists + try { + container.readItem("item-to-delete", new PartitionKey("mixed-ops"), classOf[ObjectNode]).block() + fail("Item should have been deleted") + } catch { + case e: CosmosException => e.getStatusCode shouldEqual 404 + } + } + + it should "rollback all mixed operations when one fails" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) + + // Pre-create an item + val existingItem = Utils.getSimpleObjectMapper.createObjectNode() + existingItem.put("id", "existing-item") + existingItem.put("pk", "rollback-test") + existingItem.put("prop", "original") + container.createItem(existingItem).block() + + // Create DataFrame with mixed operations where create will fail (duplicate) + val schema = StructType(Seq( + StructField("id", StringType, nullable = false), + StructField("pk", StringType, nullable = false), + StructField("operationType", StringType, nullable = false), + StructField("prop", StringType, nullable = false) + )) + + val rollbackItems = Seq( + Row("new-item-1", "rollback-test", "create", "new-value-1"), + Row("existing-item", "rollback-test", "create", "duplicate-will-fail"), + Row("new-item-2", "rollback-test", "upsert", "new-value-2") + ) + + val rollbackDf = spark.createDataFrame(rollbackItems.asJava, schema) + + val writeCfg = Map( + "spark.cosmos.accountEndpoint" -> cosmosEndpoint, + "spark.cosmos.accountKey" -> cosmosMasterKey, + "spark.cosmos.database" -> cosmosDatabase, + "spark.cosmos.container" -> cosmosContainersWithPkAsPartitionKey, + "spark.cosmos.write.strategy" -> "ItemTransactionalBatch", + "spark.cosmos.write.maxRetryCount" -> "3", + "spark.cosmos.write.bulk.enabled" -> "true" + ) + + // Attempt write - should fail atomically + try { + rollbackDf.write.format("cosmos.oltp").mode(SaveMode.Append).options(writeCfg).save() + fail("Transactional batch should have failed due to duplicate create") + } catch { + case e: Exception => + // Expected failure (409 Conflict or 424 Failed Dependency) + println(s"Expected failure: ${e.getMessage}") + } + + // Verify NO operations were committed (atomic rollback) + // 1. Original item should be unchanged + val originalItem = container.readItem("existing-item", new PartitionKey("rollback-test"), classOf[ObjectNode]).block().getItem + originalItem.get("prop").asText() shouldEqual "original" + + // 2. new-item-1 should NOT exist (rollback) + try { + container.readItem("new-item-1", new PartitionKey("rollback-test"), classOf[ObjectNode]).block() + fail("new-item-1 should not exist due to rollback") + } catch { + case e: CosmosException => e.getStatusCode shouldEqual 404 + } + + // 3. new-item-2 should NOT exist (rollback) + try { + container.readItem("new-item-2", new PartitionKey("rollback-test"), classOf[ObjectNode]).block() + fail("new-item-2 should not exist due to rollback") + } catch { + case e: CosmosException => e.getStatusCode shouldEqual 404 + } + } + + it should "support delete operations in transactional batch" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) + + // Pre-create items to delete + val itemToDelete1 = Utils.getSimpleObjectMapper.createObjectNode() + itemToDelete1.put("id", "delete-item-1") + itemToDelete1.put("pk", "delete-test") + itemToDelete1.put("prop", "delete-me-1") + container.createItem(itemToDelete1).block() + + val itemToDelete2 = Utils.getSimpleObjectMapper.createObjectNode() + itemToDelete2.put("id", "delete-item-2") + itemToDelete2.put("pk", "delete-test") + itemToDelete2.put("prop", "delete-me-2") + container.createItem(itemToDelete2).block() + + // Create DataFrame with create + delete operations + val schema = StructType(Seq( + StructField("id", StringType, nullable = false), + StructField("pk", StringType, nullable = false), + StructField("operationType", StringType, nullable = false), + StructField("prop", StringType, nullable = false) + )) + + val deleteItems = Seq( + Row("new-item", "delete-test", "create", "new-value"), + Row("delete-item-1", "delete-test", "delete", "ignored"), + Row("delete-item-2", "delete-test", "delete", "ignored") + ) + + val deleteDf = spark.createDataFrame(deleteItems.asJava, schema) + + val writeCfg = Map( + "spark.cosmos.accountEndpoint" -> cosmosEndpoint, + "spark.cosmos.accountKey" -> cosmosMasterKey, + "spark.cosmos.database" -> cosmosDatabase, + "spark.cosmos.container" -> cosmosContainersWithPkAsPartitionKey, + "spark.cosmos.write.strategy" -> "ItemTransactionalBatch", + "spark.cosmos.write.maxRetryCount" -> "3", + "spark.cosmos.write.bulk.enabled" -> "true" + ) + + // Execute transactional batch + deleteDf.write.format("cosmos.oltp").mode(SaveMode.Append).options(writeCfg).save() + + // Verify new item was created + val newItem = container.readItem("new-item", new PartitionKey("delete-test"), classOf[ObjectNode]).block().getItem + newItem.get("prop").asText() shouldEqual "new-value" + + // Verify items were deleted + try { + container.readItem("delete-item-1", new PartitionKey("delete-test"), classOf[ObjectNode]).block() + fail("delete-item-1 should have been deleted") + } catch { + case e: CosmosException => e.getStatusCode shouldEqual 404 + } + + try { + container.readItem("delete-item-2", new PartitionKey("delete-test"), classOf[ObjectNode]).block() + fail("delete-item-2 should have been deleted") + } catch { + case e: CosmosException => e.getStatusCode shouldEqual 404 + } + } + + it should "default to global strategy when operationType column absent" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) + + // Create DataFrame WITHOUT operationType column + val schema = StructType(Seq( + StructField("id", StringType, nullable = false), + StructField("pk", StringType, nullable = false), + StructField("prop", StringType, nullable = false) + )) + + val backwardCompatItems = Seq( + Row("compat-item-1", "compat-test", "value-1"), + Row("compat-item-2", "compat-test", "value-2") + ) + + val compatDf = spark.createDataFrame(backwardCompatItems.asJava, schema) + + val writeCfg = Map( + "spark.cosmos.accountEndpoint" -> cosmosEndpoint, + "spark.cosmos.accountKey" -> cosmosMasterKey, + "spark.cosmos.database" -> cosmosDatabase, + "spark.cosmos.container" -> cosmosContainersWithPkAsPartitionKey, + "spark.cosmos.write.strategy" -> "ItemTransactionalBatch", + "spark.cosmos.write.maxRetryCount" -> "3", + "spark.cosmos.write.bulk.enabled" -> "true" + ) + + // Write without operationType column (should default to upsert) + compatDf.write.format("cosmos.oltp").mode(SaveMode.Append).options(writeCfg).save() + + // Verify items were created (upsert default for ItemTransactionalBatch) + val item1 = container.readItem("compat-item-1", new PartitionKey("compat-test"), classOf[ObjectNode]).block().getItem + item1.get("prop").asText() shouldEqual "value-1" + + val item2 = container.readItem("compat-item-2", new PartitionKey("compat-test"), classOf[ObjectNode]).block().getItem + item2.get("prop").asText() shouldEqual "value-2" + + // Update items - should succeed with upsert + val updateItems = Seq( + Row("compat-item-1", "compat-test", "updated-1"), + Row("compat-item-2", "compat-test", "updated-2") + ) + + val updateDf = spark.createDataFrame(updateItems.asJava, schema) + + updateDf.write.format("cosmos.oltp").mode(SaveMode.Append).options(writeCfg).save() + + // Verify items were updated + val updatedItem1 = container.readItem("compat-item-1", new PartitionKey("compat-test"), classOf[ObjectNode]).block().getItem + updatedItem1.get("prop").asText() shouldEqual "updated-1" + + val updatedItem2 = container.readItem("compat-item-2", new PartitionKey("compat-test"), classOf[ObjectNode]).block().getItem + updatedItem2.get("prop").asText() shouldEqual "updated-2" + } } From 4d4110229bad42e5939b464911884187cc945428 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Mon, 8 Dec 2025 12:17:19 +0000 Subject: [PATCH 03/16] Simplify config: remove bulkEnableTransactions, use ItemTransactionalBatch only --- .../cosmos/spark/ItemsWriterBuilder.scala | 4 +- .../com/azure/cosmos/spark/CosmosConfig.scala | 11 ----- .../azure/cosmos/spark/CosmosWriterBase.scala | 7 +--- .../azure/cosmos/spark/BulkWriterITest.scala | 2 +- .../spark/TransactionalBatchITest.scala | 42 +++++++------------ 5 files changed, 21 insertions(+), 45 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3-5/src/main/scala/com/azure/cosmos/spark/ItemsWriterBuilder.scala b/sdk/cosmos/azure-cosmos-spark_3-5/src/main/scala/com/azure/cosmos/spark/ItemsWriterBuilder.scala index 1ace9707297f..4c5c6f782bb6 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5/src/main/scala/com/azure/cosmos/spark/ItemsWriterBuilder.scala +++ b/sdk/cosmos/azure-cosmos-spark_3-5/src/main/scala/com/azure/cosmos/spark/ItemsWriterBuilder.scala @@ -87,7 +87,7 @@ private class ItemsWriterBuilder override def supportedCustomMetrics(): Array[CustomMetric] = supportedCosmosMetrics override def requiredDistribution(): Distribution = { - if (writeConfig.bulkEnabled && writeConfig.bulkEnableTransactions) { + if (writeConfig.bulkEnabled && writeConfig.itemWriteStrategy == ItemWriteStrategy.ItemTransactionalBatch) { // For transactional writes, partition by all partition key columns val partitionKeyPaths = getPartitionKeyColumnNames() if (partitionKeyPaths.nonEmpty) { @@ -103,7 +103,7 @@ private class ItemsWriterBuilder } override def requiredOrdering(): Array[SortOrder] = { - if (writeConfig.bulkEnabled && writeConfig.bulkEnableTransactions) { + if (writeConfig.bulkEnabled && writeConfig.itemWriteStrategy == ItemWriteStrategy.ItemTransactionalBatch) { // For transactional writes, order by all partition key columns (ascending) val partitionKeyPaths = getPartitionKeyColumnNames() if (partitionKeyPaths.nonEmpty) { diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala index 21a6a374bc8d..d566074911c0 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala @@ -113,7 +113,6 @@ private[spark] object CosmosConfigNames { val ClientTelemetryEnabled = "spark.cosmos.clientTelemetry.enabled" // keep this to avoid breaking changes val ClientTelemetryEndpoint = "spark.cosmos.clientTelemetry.endpoint" // keep this to avoid breaking changes val WriteBulkEnabled = "spark.cosmos.write.bulk.enabled" - val WriteBulkEnableTransactions = "spark.cosmos.write.bulk.enableTransactions" val WriteBulkMaxPendingOperations = "spark.cosmos.write.bulk.maxPendingOperations" val WriteBulkMaxBatchSize = "spark.cosmos.write.bulk.maxBatchSize" val WriteBulkMinTargetBatchSize = "spark.cosmos.write.bulk.minTargetBatchSize" @@ -243,7 +242,6 @@ private[spark] object CosmosConfigNames { ClientTelemetryEnabled, ClientTelemetryEndpoint, WriteBulkEnabled, - WriteBulkEnableTransactions, WriteBulkMaxPendingOperations, WriteBulkMaxConcurrentPartitions, WriteBulkPayloadSizeInBytes, @@ -1464,7 +1462,6 @@ private case class CosmosPatchConfigs(columnConfigsMap: TrieMap[String, CosmosPa private case class CosmosWriteConfig(itemWriteStrategy: ItemWriteStrategy, maxRetryCount: Int, bulkEnabled: Boolean, - bulkEnableTransactions: Boolean = false, bulkMaxPendingOperations: Option[Int] = None, pointMaxConcurrency: Option[Int] = None, maxConcurrentCosmosPartitions: Option[Int] = None, @@ -1489,12 +1486,6 @@ private object CosmosWriteConfig { parseFromStringFunction = bulkEnabledAsString => bulkEnabledAsString.toBoolean, helpMessage = "Cosmos DB Item Write bulk enabled") - private val bulkEnableTransactions = CosmosConfigEntry[Boolean](key = CosmosConfigNames.WriteBulkEnableTransactions, - defaultValue = Option.apply(false), - mandatory = false, - parseFromStringFunction = enableTransactionsAsString => enableTransactionsAsString.toBoolean, - helpMessage = "Cosmos DB Item Write enable transactional batch - requires bulk write to be enabled and Spark 3.5+") - private val microBatchPayloadSizeInBytes = CosmosConfigEntry[Int](key = CosmosConfigNames.WriteBulkPayloadSizeInBytes, defaultValue = Option.apply(BatchRequestResponseConstants.DEFAULT_MAX_DIRECT_MODE_BATCH_REQUEST_BODY_SIZE_IN_BYTES), mandatory = false, @@ -1767,7 +1758,6 @@ private object CosmosWriteConfig { val itemWriteStrategyOpt = CosmosConfigEntry.parse(cfg, itemWriteStrategy) val maxRetryCountOpt = CosmosConfigEntry.parse(cfg, maxRetryCount) val bulkEnabledOpt = CosmosConfigEntry.parse(cfg, bulkEnabled) - val bulkEnableTransactionsOpt = CosmosConfigEntry.parse(cfg, bulkEnableTransactions) var patchConfigsOpt = Option.empty[CosmosPatchConfigs] val throughputControlConfigOpt = CosmosThroughputControlConfig.parseThroughputControlConfig(cfg) val microBatchPayloadSizeInBytesOpt = CosmosConfigEntry.parse(cfg, microBatchPayloadSizeInBytes) @@ -1798,7 +1788,6 @@ private object CosmosWriteConfig { itemWriteStrategyOpt.get, maxRetryCountOpt.get, bulkEnabled = bulkEnabledOpt.get, - bulkEnableTransactions = bulkEnableTransactionsOpt.getOrElse(false), bulkMaxPendingOperations = CosmosConfigEntry.parse(cfg, bulkMaxPendingOperations), pointMaxConcurrency = CosmosConfigEntry.parse(cfg, pointWriteConcurrency), maxConcurrentCosmosPartitions = CosmosConfigEntry.parse(cfg, bulkMaxConcurrentPartitions), diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosWriterBase.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosWriterBase.scala index 331a6ac3d3f7..b3f3ea43bf1f 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosWriterBase.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosWriterBase.scala @@ -62,11 +62,8 @@ private abstract class CosmosWriterBase( private val writer: AtomicReference[AsyncItemWriter] = new AtomicReference( if (cosmosWriteConfig.bulkEnabled) { - // Use TransactionalBulkWriter if either: - // 1. bulkEnableTransactions is explicitly set, OR - // 2. ItemWriteStrategy is ItemTransactionalBatch - val useTransactional = cosmosWriteConfig.bulkEnableTransactions || - cosmosWriteConfig.itemWriteStrategy == ItemWriteStrategy.ItemTransactionalBatch + // Use TransactionalBulkWriter if ItemWriteStrategy is ItemTransactionalBatch + val useTransactional = cosmosWriteConfig.itemWriteStrategy == ItemWriteStrategy.ItemTransactionalBatch if (useTransactional) { new TransactionalBulkWriter( diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/BulkWriterITest.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/BulkWriterITest.scala index 5ea56b39d756..9ca7f459cf15 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/BulkWriterITest.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/BulkWriterITest.scala @@ -408,7 +408,7 @@ class BulkWriterITest extends IntegrationSpec with CosmosClient with AutoCleanab val containerProperties = container.read().block().getProperties val partitionKeyDefinition = containerProperties.getPartitionKeyDefinition val containerConfig = CosmosContainerConfig(container.getDatabase.getId, container.getId, None) - val writeConfig = CosmosWriteConfig(ItemWriteStrategy.ItemAppend, maxRetryCount = 5, bulkEnabled = true, bulkEnableTransactions = false, Some(900)) + val writeConfig = CosmosWriteConfig(ItemWriteStrategy.ItemAppend, maxRetryCount = 5, bulkEnabled = true, bulkMaxPendingOperations = Some(900)) val bulkWriter = new BulkWriter( container, containerConfig, diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala index a696077722f8..22d5727c22cc 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala @@ -48,16 +48,15 @@ class TransactionalBatchITest extends IntegrationSpec val operationsDf = spark.createDataFrame(batchOperations.asJava, schema) - // Execute transactional batch using regular write with bulkEnableTransactions config + // Execute transactional batch using ItemTransactionalBatch write strategy operationsDf.write .format("cosmos.oltp") .option("spark.cosmos.accountEndpoint", cosmosEndpoint) .option("spark.cosmos.accountKey", cosmosMasterKey) .option("spark.cosmos.database", cosmosDatabase) .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) - .option("spark.cosmos.write.strategy", "ItemOverwrite") + .option("spark.cosmos.write.strategy", "ItemTransactionalBatch") .option("spark.cosmos.write.bulk.enabled", "true") - .option("spark.cosmos.write.bulk.enableTransactions", "true") .mode(SaveMode.Append) .save() @@ -90,17 +89,18 @@ class TransactionalBatchITest extends IntegrationSpec val schema = StructType(Seq( StructField("id", StringType, nullable = false), StructField("pk", StringType, nullable = false), + StructField("operationType", StringType, nullable = false), StructField("name", StringType, nullable = false) )) val batchOperations = Seq( - Row(item1Id, partitionKeyValue, "NewItem"), - Row(duplicateId, partitionKeyValue, "Duplicate") + Row(item1Id, partitionKeyValue, "create", "NewItem"), + Row(duplicateId, partitionKeyValue, "create", "Duplicate") ) val operationsDf = spark.createDataFrame(batchOperations.asJava, schema) - // This should fail because we're using ItemAppend strategy (create only) and duplicateId already exists + // This should fail because we're using create operations and duplicateId already exists val exception = intercept[Exception] { operationsDf.write .format("cosmos.oltp") @@ -108,9 +108,8 @@ class TransactionalBatchITest extends IntegrationSpec .option("spark.cosmos.accountKey", cosmosMasterKey) .option("spark.cosmos.database", cosmosDatabase) .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) - .option("spark.cosmos.write.strategy", "ItemAppend") + .option("spark.cosmos.write.strategy", "ItemTransactionalBatch") .option("spark.cosmos.write.bulk.enabled", "true") - .option("spark.cosmos.write.bulk.enableTransactions", "true") .mode(SaveMode.Append) .save() } @@ -162,9 +161,8 @@ class TransactionalBatchITest extends IntegrationSpec .option("spark.cosmos.accountKey", cosmosMasterKey) .option("spark.cosmos.database", cosmosDatabase) .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) - .option("spark.cosmos.write.strategy", "ItemOverwrite") + .option("spark.cosmos.write.strategy", "ItemTransactionalBatch") .option("spark.cosmos.write.bulk.enabled", "true") - .option("spark.cosmos.write.bulk.enableTransactions", "true") .mode(SaveMode.Append) .save() @@ -206,9 +204,8 @@ class TransactionalBatchITest extends IntegrationSpec .option("spark.cosmos.accountKey", cosmosMasterKey) .option("spark.cosmos.database", cosmosDatabase) .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) - .option("spark.cosmos.write.strategy", "ItemOverwrite") + .option("spark.cosmos.write.strategy", "ItemTransactionalBatch") .option("spark.cosmos.write.bulk.enabled", "true") - .option("spark.cosmos.write.bulk.enableTransactions", "true") .mode(SaveMode.Append) .save() @@ -255,9 +252,8 @@ class TransactionalBatchITest extends IntegrationSpec .option("spark.cosmos.accountKey", cosmosMasterKey) .option("spark.cosmos.database", cosmosDatabase) .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) - .option("spark.cosmos.write.strategy", "ItemOverwrite") + .option("spark.cosmos.write.strategy", "ItemTransactionalBatch") .option("spark.cosmos.write.bulk.enabled", "true") - .option("spark.cosmos.write.bulk.enableTransactions", "true") .mode(SaveMode.Append) .save() @@ -295,9 +291,8 @@ class TransactionalBatchITest extends IntegrationSpec .option("spark.cosmos.accountKey", cosmosMasterKey) .option("spark.cosmos.database", cosmosDatabase) .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) - .option("spark.cosmos.write.strategy", "ItemOverwrite") + .option("spark.cosmos.write.strategy", "ItemTransactionalBatch") .option("spark.cosmos.write.bulk.enabled", "true") - .option("spark.cosmos.write.bulk.enableTransactions", "true") .mode(SaveMode.Append) .save() } @@ -355,9 +350,8 @@ class TransactionalBatchITest extends IntegrationSpec .option("spark.cosmos.accountKey", cosmosMasterKey) .option("spark.cosmos.database", cosmosDatabase) .option("spark.cosmos.container", containerName) - .option("spark.cosmos.write.strategy", "ItemOverwrite") + .option("spark.cosmos.write.strategy", "ItemTransactionalBatch") .option("spark.cosmos.write.bulk.enabled", "true") - .option("spark.cosmos.write.bulk.enableTransactions", "true") .mode(SaveMode.Append) .save() @@ -434,9 +428,8 @@ class TransactionalBatchITest extends IntegrationSpec .option("spark.cosmos.accountKey", cosmosMasterKey) .option("spark.cosmos.database", cosmosDatabase) .option("spark.cosmos.container", containerName) - .option("spark.cosmos.write.strategy", "ItemOverwrite") + .option("spark.cosmos.write.strategy", "ItemTransactionalBatch") .option("spark.cosmos.write.bulk.enabled", "true") - .option("spark.cosmos.write.bulk.enableTransactions", "true") .mode(SaveMode.Append) .save() @@ -503,9 +496,8 @@ class TransactionalBatchITest extends IntegrationSpec .option("spark.cosmos.accountKey", cosmosMasterKey) .option("spark.cosmos.database", cosmosDatabase) .option("spark.cosmos.container", containerName) - .option("spark.cosmos.write.strategy", "ItemOverwrite") + .option("spark.cosmos.write.strategy", "ItemTransactionalBatch") .option("spark.cosmos.write.bulk.enabled", "true") - .option("spark.cosmos.write.bulk.enableTransactions", "true") .mode(SaveMode.Append) .save() @@ -562,9 +554,8 @@ class TransactionalBatchITest extends IntegrationSpec .option("spark.cosmos.accountKey", cosmosMasterKey) .option("spark.cosmos.database", cosmosDatabase) .option("spark.cosmos.container", containerName) - .option("spark.cosmos.write.strategy", "ItemOverwrite") + .option("spark.cosmos.write.strategy", "ItemTransactionalBatch") .option("spark.cosmos.write.bulk.enabled", "true") - .option("spark.cosmos.write.bulk.enableTransactions", "true") .mode(SaveMode.Append) .save() } @@ -619,9 +610,8 @@ class TransactionalBatchITest extends IntegrationSpec .option("spark.cosmos.accountKey", cosmosMasterKey) .option("spark.cosmos.database", cosmosDatabase) .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) - .option("spark.cosmos.write.strategy", "ItemOverwrite") + .option("spark.cosmos.write.strategy", "ItemTransactionalBatch") .option("spark.cosmos.write.bulk.enabled", "true") - .option("spark.cosmos.write.bulk.enableTransactions", "true") .mode(SaveMode.Append) .save() From d73b3747098fb79186346bfc6f742a7a2e779e69 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Tue, 9 Dec 2025 18:16:26 +0000 Subject: [PATCH 04/16] refactored to use CosmosBatch, removed operationType per row, and added docs to ingestion.md --- .../cosmos/spark/ItemsWriterBuilder.scala | 4 +- .../docs/configuration-reference.md | 2 +- .../docs/scenarios/Ingestion.md | 55 ++ .../com/azure/cosmos/spark/CosmosConfig.scala | 14 +- .../azure/cosmos/spark/CosmosWriterBase.scala | 15 +- .../spark/TransactionalBulkWriter.scala | 250 ++--- .../spark/TransactionalBatchITest.scala | 537 ++++------ .../batch/TransactionalBulkExecutor.java | 918 +++--------------- 8 files changed, 531 insertions(+), 1264 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3-5/src/main/scala/com/azure/cosmos/spark/ItemsWriterBuilder.scala b/sdk/cosmos/azure-cosmos-spark_3-5/src/main/scala/com/azure/cosmos/spark/ItemsWriterBuilder.scala index 4c5c6f782bb6..a7260319bdb2 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5/src/main/scala/com/azure/cosmos/spark/ItemsWriterBuilder.scala +++ b/sdk/cosmos/azure-cosmos-spark_3-5/src/main/scala/com/azure/cosmos/spark/ItemsWriterBuilder.scala @@ -87,7 +87,7 @@ private class ItemsWriterBuilder override def supportedCustomMetrics(): Array[CustomMetric] = supportedCosmosMetrics override def requiredDistribution(): Distribution = { - if (writeConfig.bulkEnabled && writeConfig.itemWriteStrategy == ItemWriteStrategy.ItemTransactionalBatch) { + if (writeConfig.bulkEnabled && writeConfig.bulkTransactional) { // For transactional writes, partition by all partition key columns val partitionKeyPaths = getPartitionKeyColumnNames() if (partitionKeyPaths.nonEmpty) { @@ -103,7 +103,7 @@ private class ItemsWriterBuilder } override def requiredOrdering(): Array[SortOrder] = { - if (writeConfig.bulkEnabled && writeConfig.itemWriteStrategy == ItemWriteStrategy.ItemTransactionalBatch) { + if (writeConfig.bulkEnabled && writeConfig.bulkTransactional) { // For transactional writes, order by all partition key columns (ascending) val partitionKeyPaths = getPartitionKeyColumnNames() if (partitionKeyPaths.nonEmpty) { diff --git a/sdk/cosmos/azure-cosmos-spark_3/docs/configuration-reference.md b/sdk/cosmos/azure-cosmos-spark_3/docs/configuration-reference.md index 33952bfc381d..e25cb78ab9df 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/docs/configuration-reference.md +++ b/sdk/cosmos/azure-cosmos-spark_3/docs/configuration-reference.md @@ -67,7 +67,7 @@ | `spark.cosmos.write.point.maxConcurrency` | None | Cosmos DB Item Write Max concurrency. If not specified it will be determined based on the Spark executor VM Size | | `spark.cosmos.write.bulk.maxPendingOperations` | None | Cosmos DB Item Write bulk mode maximum pending operations. Defines a limit of bulk operations being processed concurrently. If not specified it will be determined based on the Spark executor VM Size. If the volume of data is large for the provisioned throughput on the destination container, this setting can be adjusted by following the estimation of `1000 x Cores` | | `spark.cosmos.write.bulk.enabled` | `true` | Cosmos DB Item Write bulk enabled | -| `spark.cosmos.write.bulk.targetedPayloadSizeInBytes` | `220201` | When the targeted payload size is reached for buffered documents, the request is sent to the backend. The default value is optimized for small documents <= 10 KB - when documents often exceed 110 KB, it can help to increase this value to up to about `1500000` (should still be smaller than 2 MB). | +| `spark.cosmos.write.bulk.transactional` | `false` | Enable transactional batch mode for bulk writes. When enabled, all operations for the same partition key are executed atomically (all succeed or all fail). Requires ordering and clustering by partition key columns. Only supports upsert operations. Cannot exceed 100 operations or 2MB per partition key. See [Transactional Batch documentation](https://learn.microsoft.com/azure/cosmos-db/nosql/transactional-batch) for details. |\n| `spark.cosmos.write.bulk.targetedPayloadSizeInBytes` | `220201` | When the targeted payload size is reached for buffered documents, the request is sent to the backend. The default value is optimized for small documents <= 10 KB - when documents often exceed 110 KB, it can help to increase this value to up to about `1500000` (should still be smaller than 2 MB). | | `spark.cosmos.write.bulk.initialBatchSize` | `100` | Cosmos DB initial bulk micro batch size - a micro batch will be flushed to the backend when the number of documents enqueued exceeds this size - or the target payload size is met. The micro batch size is getting automatically tuned based on the throttling rate. By default the initial micro batch size is 100. Reduce this when you want to avoid that the first few requests consume too many RUs. | | `spark.cosmos.write.bulk.maxBatchSize` | `100` | Cosmos DB max. bulk micro batch size - a micro batch will be flushed to the backend when the number of documents enqueued exceeds this size - or the target payload size is met. The micro batch size is getting automatically tuned based on the throttling rate. By default the max. micro batch size is 100. Use this setting only when migrating Spark 2.4 workloads - for other scenarios relying on the auto-tuning combined with throughput control will result in better experience. | | `spark.cosmos.write.flush.noProgress.maxIntervalInSeconds` | `180` | The time interval in seconds that write operations will wait when no progress can be made for bulk writes before forcing a retry. The retry will reinitialize the bulk write process - so, any delays on the retry can be sure to be actual service issues. The default value of 3 min should be sufficient to prevent false negatives when there is a short service-side write unavailability - like for partition splits or merges. Increase it only if you regularly see these transient errors to exceed a time period of 180 seconds. | diff --git a/sdk/cosmos/azure-cosmos-spark_3/docs/scenarios/Ingestion.md b/sdk/cosmos/azure-cosmos-spark_3/docs/scenarios/Ingestion.md index 0825a513458f..7056ffa1d883 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/docs/scenarios/Ingestion.md +++ b/sdk/cosmos/azure-cosmos-spark_3/docs/scenarios/Ingestion.md @@ -79,6 +79,61 @@ When your container has a "unique key constraint policy" any 409 "Conflict" (ind - For `ItemOverwrite` a 409 - Conflict due to unique key violation will result in an error - and the Spark job will fail. *NOTE: Conflicts due to pk+id being identical to another document won't even result in a 409 - because with Upsert the existing document would simply be updated.* - For `ItemAppend` like conflicts on pk+id any unique key policy constraint violation will be ignored. +### Transactional batch writes + +For scenarios requiring atomic all-or-nothing semantics within a partition, you can enable transactional batch writes using the `spark.cosmos.write.bulk.transactional` configuration. When enabled, all operations for a single partition key value either succeed or fail together. + +#### Configuration + +**Python example:** +```python +df.write \ + .format("cosmos.oltp") \ + .option("spark.cosmos.accountEndpoint", cosmosEndpoint) \ + .option("spark.cosmos.accountKey", cosmosKey) \ + .option("spark.cosmos.database", cosmosDatabase) \ + .option("spark.cosmos.container", cosmosContainer) \ + .option("spark.cosmos.write.bulk.enabled", "true") \ + .option("spark.cosmos.write.bulk.transactional", "true") \ + .mode("Append") \ + .save() +``` + +**Scala example:** +```scala +df.write + .format("cosmos.oltp") + .option("spark.cosmos.accountEndpoint", cosmosEndpoint) + .option("spark.cosmos.accountKey", cosmosKey) + .option("spark.cosmos.database", cosmosDatabase) + .option("spark.cosmos.container", cosmosContainer) + .option("spark.cosmos.write.bulk.enabled", "true") + .option("spark.cosmos.write.bulk.transactional", "true") + .mode(SaveMode.Append) + .save() +``` + +#### Characteristics and limitations + +- **Atomic semantics**: All operations for the same partition key succeed or all fail (rollback) +- **Operation type**: Only upsert operations are supported (equivalent to `ItemOverwrite` write strategy) +- **Partition grouping**: Spark automatically partitions and orders data by partition key columns +- **Size limits**: Maximum 100 operations per transaction; maximum 2MB total payload per transaction +- **Partition key requirement**: All operations in a transaction must share the same partition key value +- **Bulk mode required**: Must have `spark.cosmos.write.bulk.enabled=true` (enabled by default) + +#### Use cases + +Transactional batch writes are ideal for: +- Financial transactions requiring consistency across multiple documents +- Order processing where order header and line items must be committed together +- Multi-document updates that must be atomic (e.g., inventory adjustments) +- Any scenario where partial success would leave data in an inconsistent state + +#### Error handling + +If any operation in a transaction fails (e.g., insufficient RUs, document too large, transaction exceeds 100 operations), the entire transaction is rolled back and no documents are modified. The Spark task will fail and retry according to Spark's retry policy. + ## Preparation Below are a couple of tips/best-practices that can help you to prepare for a data migration into a Cosmos DB container. diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala index d566074911c0..96e26eaf0f1a 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala @@ -113,6 +113,7 @@ private[spark] object CosmosConfigNames { val ClientTelemetryEnabled = "spark.cosmos.clientTelemetry.enabled" // keep this to avoid breaking changes val ClientTelemetryEndpoint = "spark.cosmos.clientTelemetry.endpoint" // keep this to avoid breaking changes val WriteBulkEnabled = "spark.cosmos.write.bulk.enabled" + val WriteBulkTransactional = "spark.cosmos.write.bulk.transactional" val WriteBulkMaxPendingOperations = "spark.cosmos.write.bulk.maxPendingOperations" val WriteBulkMaxBatchSize = "spark.cosmos.write.bulk.maxBatchSize" val WriteBulkMinTargetBatchSize = "spark.cosmos.write.bulk.minTargetBatchSize" @@ -242,6 +243,7 @@ private[spark] object CosmosConfigNames { ClientTelemetryEnabled, ClientTelemetryEndpoint, WriteBulkEnabled, + WriteBulkTransactional, WriteBulkMaxPendingOperations, WriteBulkMaxConcurrentPartitions, WriteBulkPayloadSizeInBytes, @@ -1437,7 +1439,7 @@ private[spark] object DiagnosticsConfig { private object ItemWriteStrategy extends Enumeration { type ItemWriteStrategy = Value - val ItemOverwrite, ItemAppend, ItemDelete, ItemDeleteIfNotModified, ItemOverwriteIfNotModified, ItemPatch, ItemPatchIfExists, ItemBulkUpdate, ItemTransactionalBatch = Value + val ItemOverwrite, ItemAppend, ItemDelete, ItemDeleteIfNotModified, ItemOverwriteIfNotModified, ItemPatch, ItemPatchIfExists, ItemBulkUpdate = Value } private object CosmosPatchOperationTypes extends Enumeration { @@ -1462,6 +1464,7 @@ private case class CosmosPatchConfigs(columnConfigsMap: TrieMap[String, CosmosPa private case class CosmosWriteConfig(itemWriteStrategy: ItemWriteStrategy, maxRetryCount: Int, bulkEnabled: Boolean, + bulkTransactional: Boolean = false, bulkMaxPendingOperations: Option[Int] = None, pointMaxConcurrency: Option[Int] = None, maxConcurrentCosmosPartitions: Option[Int] = None, @@ -1486,6 +1489,12 @@ private object CosmosWriteConfig { parseFromStringFunction = bulkEnabledAsString => bulkEnabledAsString.toBoolean, helpMessage = "Cosmos DB Item Write bulk enabled") + private val bulkTransactional = CosmosConfigEntry[Boolean](key = CosmosConfigNames.WriteBulkTransactional, + defaultValue = Option.apply(false), + mandatory = false, + parseFromStringFunction = bulkTransactionalAsString => bulkTransactionalAsString.toBoolean, + helpMessage = "Cosmos DB Item Write bulk transactional batch mode enabled") + private val microBatchPayloadSizeInBytes = CosmosConfigEntry[Int](key = CosmosConfigNames.WriteBulkPayloadSizeInBytes, defaultValue = Option.apply(BatchRequestResponseConstants.DEFAULT_MAX_DIRECT_MODE_BATCH_REQUEST_BODY_SIZE_IN_BYTES), mandatory = false, @@ -1758,6 +1767,7 @@ private object CosmosWriteConfig { val itemWriteStrategyOpt = CosmosConfigEntry.parse(cfg, itemWriteStrategy) val maxRetryCountOpt = CosmosConfigEntry.parse(cfg, maxRetryCount) val bulkEnabledOpt = CosmosConfigEntry.parse(cfg, bulkEnabled) + val bulkTransactionalOpt = CosmosConfigEntry.parse(cfg, bulkTransactional) var patchConfigsOpt = Option.empty[CosmosPatchConfigs] val throughputControlConfigOpt = CosmosThroughputControlConfig.parseThroughputControlConfig(cfg) val microBatchPayloadSizeInBytesOpt = CosmosConfigEntry.parse(cfg, microBatchPayloadSizeInBytes) @@ -1768,6 +1778,7 @@ private object CosmosWriteConfig { .parse(cfg, writeOnRetryCommitInterceptor).flatten assert(bulkEnabledOpt.isDefined, s"Parameter '${CosmosConfigNames.WriteBulkEnabled}' is missing.") + assert(bulkTransactionalOpt.isDefined, s"Parameter '${CosmosConfigNames.WriteBulkTransactional}' is missing.") // parsing above already validated this assert(itemWriteStrategyOpt.isDefined, s"Parameter '${CosmosConfigNames.WriteStrategy}' is missing.") @@ -1788,6 +1799,7 @@ private object CosmosWriteConfig { itemWriteStrategyOpt.get, maxRetryCountOpt.get, bulkEnabled = bulkEnabledOpt.get, + bulkTransactional = bulkTransactionalOpt.get, bulkMaxPendingOperations = CosmosConfigEntry.parse(cfg, bulkMaxPendingOperations), pointMaxConcurrency = CosmosConfigEntry.parse(cfg, pointWriteConcurrency), maxConcurrentCosmosPartitions = CosmosConfigEntry.parse(cfg, bulkMaxConcurrentPartitions), diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosWriterBase.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosWriterBase.scala index b3f3ea43bf1f..ece29491dd0a 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosWriterBase.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosWriterBase.scala @@ -62,8 +62,8 @@ private abstract class CosmosWriterBase( private val writer: AtomicReference[AsyncItemWriter] = new AtomicReference( if (cosmosWriteConfig.bulkEnabled) { - // Use TransactionalBulkWriter if ItemWriteStrategy is ItemTransactionalBatch - val useTransactional = cosmosWriteConfig.itemWriteStrategy == ItemWriteStrategy.ItemTransactionalBatch + // Use TransactionalBulkWriter if bulkTransactional config is enabled + val useTransactional = cosmosWriteConfig.bulkTransactional if (useTransactional) { new TransactionalBulkWriter( @@ -127,7 +127,16 @@ private abstract class CosmosWriterBase( } val partitionKeyValue = PartitionKeyHelper.getPartitionKeyPath(objectNode, partitionKeyDefinition) - writer.get.scheduleWrite(partitionKeyValue, objectNode, operationType) + + // Call the appropriate scheduleWrite method based on whether operationType is specified + operationType match { + case Some(opType) => + // Per-row operation type specified - use 3-parameter method + writer.get.scheduleWrite(partitionKeyValue, objectNode, Some(opType)) + case None => + // No per-row operation type - use 2-parameter method (global strategy) + writer.get.scheduleWrite(partitionKeyValue, objectNode) + } } override def commit(): WriterCommitMessage = { diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala index afa8e545bd15..b62c1d8d57ec 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala @@ -77,6 +77,13 @@ private class TransactionalBulkWriter // multiplied by 2 to leave space for partition splits during ingestion case None => 2 * ContainerFeedRangesCache.getFeedRanges(container, containerConfig.feedRangeRefreshIntervalInSecondsOpt).block().size } + // Validate write strategy for transactional batches + if (writeConfig.itemWriteStrategy != ItemWriteStrategy.ItemOverwrite) { + throw new IllegalArgumentException( + s"Transactional batches only support ItemOverwrite (upsert) write strategy. " + + s"Requested strategy: ${writeConfig.itemWriteStrategy}") + } + log.logInfo( s"BulkWriter instantiated (Host CPU count: $cpuCount, maxPendingOperations: $maxPendingOperations, " + s"maxConcurrentPartitions: $maxConcurrentPartitions ...") @@ -102,9 +109,16 @@ private class TransactionalBulkWriter private val pendingReadManyRetries = java.util.concurrent.ConcurrentHashMap.newKeySet[ReadManyOperation]().asScala private val activeTasks = new AtomicInteger(0) private val errorCaptureFirstException = new AtomicReference[Throwable]() - private val bulkInputEmitter: Sinks.Many[CosmosItemOperation] = Sinks.many().unicast().onBackpressureBuffer() + private val bulkInputEmitter: Sinks.Many[CosmosBatch] = Sinks.many().unicast().onBackpressureBuffer() + + // Partition key buffering for batch construction + // Buffer holds tuples of (ObjectNode, OperationContext, effectiveOperationType) + private var currentPartitionKey: PartitionKey = null + private val currentBatchOperations = new mutable.ListBuffer[(ObjectNode, OperationContext, String)]() + private val batchConstructionLock = new Object() private val activeBulkWriteOperations =java.util.concurrent.ConcurrentHashMap.newKeySet[CosmosItemOperation]().asScala + private val operationContextMap = new java.util.concurrent.ConcurrentHashMap[CosmosItemOperation, OperationContext]().asScala private val activeReadManyOperations = java.util.concurrent.ConcurrentHashMap.newKeySet[ReadManyOperation]().asScala private val semaphore = new Semaphore(maxPendingOperations) @@ -383,16 +397,11 @@ private class TransactionalBulkWriter .get .patchBulkUpdateItem(resultMap.get(readManyOperation.cosmosItemIdentity), cosmosPatchBulkUpdateOperations) - // create bulk item operation + // Add operation to batch - ReadMany results become batch operations val etagOpt = Option.apply(rootNode.get(CosmosConstants.Properties.ETag)) - val bulkItemOperation = etagOpt match { + val (effectiveOperationType, operationContext) = etagOpt match { case Some(etag) => - CosmosBulkOperations.getReplaceItemOperation( - readManyOperation.cosmosItemIdentity.getId, - rootNode, - readManyOperation.cosmosItemIdentity.getPartitionKey, - new CosmosBulkItemRequestOptions().setIfMatchETag(etag.asText()), - new OperationContext( + ("replace", new OperationContext( readManyOperation.operationContext.itemId, readManyOperation.operationContext.partitionKeyValue, Some(etag.asText()), @@ -400,20 +409,24 @@ private class TransactionalBulkWriter monotonicOperationCounter.incrementAndGet(), Some(readManyOperation.objectNode) )) - case None => CosmosBulkOperations.getCreateItemOperation( - rootNode, - readManyOperation.cosmosItemIdentity.getPartitionKey, - new OperationContext( - readManyOperation.operationContext.itemId, - readManyOperation.operationContext.partitionKeyValue, - eTagInput = None, - readManyOperation.operationContext.attemptNumber, - monotonicOperationCounter.incrementAndGet(), - Some(readManyOperation.objectNode) + case None => + ("create", new OperationContext( + readManyOperation.operationContext.itemId, + readManyOperation.operationContext.partitionKeyValue, + eTagInput = None, + readManyOperation.operationContext.attemptNumber, + monotonicOperationCounter.incrementAndGet(), + Some(readManyOperation.objectNode) )) } - this.emitBulkInput(bulkItemOperation) + // Add to batch buffer - will be flushed on PK change or 100-op limit + addOperationToBatch( + readManyOperation.cosmosItemIdentity.getPartitionKey, + rootNode, + operationContext, + effectiveOperationType + ) } }) .onErrorResume(throwable => { @@ -600,7 +613,10 @@ private class TransactionalBulkWriter shouldSkipTaskCompletion.set(true) } if (pendingRetriesFound || itemOperationFound) { - val context = itemOperation.getContext[OperationContext] + // Look up context from our external map since ItemBatchOperation.getContext() returns null + val context = operationContextMap.remove(itemOperation).getOrElse( + throw new IllegalStateException(s"Cannot find context for operation: ${itemOperation.getId}") + ) val itemResponse = resp.getResponse if (resp.getException != null) { @@ -623,6 +639,8 @@ private class TransactionalBulkWriter // no error case outputMetricsPublisher.trackWriteOperation(1, None) totalSuccessfulIngestionMetrics.getAndIncrement() + log.logTrace(s"Successfully processed operation: ${itemOperation.getOperationType} " + + s"for item ${itemOperation.getId}, PK: ${itemOperation.getPartitionKeyValue}") } } } @@ -655,10 +673,6 @@ private class TransactionalBulkWriter } override def scheduleWrite(partitionKeyValue: PartitionKey, objectNode: ObjectNode): Unit = { - scheduleWrite(partitionKeyValue, objectNode, None) - } - - override def scheduleWrite(partitionKeyValue: PartitionKey, objectNode: ObjectNode, operationType: Option[String]): Unit = { Preconditions.checkState(!closed.get()) throwIfCapturedExceptionExists() @@ -669,8 +683,7 @@ private class TransactionalBulkWriter getETag(objectNode), 1, monotonicOperationCounter.incrementAndGet(), - None, - operationType) + None) val numberOfIntervalsWithIdenticalActiveOperationSnapshots = new AtomicLong(0) // Don't clone the activeOperations for the first iteration // to reduce perf impact before the Semaphore has been acquired @@ -713,6 +726,18 @@ private class TransactionalBulkWriter scheduleWriteInternal(partitionKeyValue, objectNode, operationContext) } + /** + * Per-row operation type is not supported for transactional batches. + * All operations in a transactional batch must be upserts (ItemOverwrite strategy). + * This method is implemented to satisfy the AsyncItemWriter interface but throws an exception. + */ + override def scheduleWrite(partitionKeyValue: PartitionKey, objectNode: ObjectNode, operationType: Option[String]): Unit = { + throw new UnsupportedOperationException( + "Per-row operation types are not supported for transactional batches. " + + "All operations in a transactional batch must use ItemOverwrite (upsert) strategy. " + + "Use scheduleWrite(partitionKeyValue, objectNode) instead.") + } + private def scheduleWriteInternal(partitionKeyValue: PartitionKey, objectNode: ObjectNode, operationContext: OperationContext): Unit = { @@ -745,81 +770,86 @@ private class TransactionalBulkWriter objectNode: ObjectNode, operationContext: OperationContext): Unit = { - // Determine the effective operation type from per-row operationType or global strategy - val effectiveOperationType = operationContext.operationType match { - case Some(opType) => opType.toLowerCase - case None => writeConfig.itemWriteStrategy match { - case ItemWriteStrategy.ItemOverwrite => "upsert" - case ItemWriteStrategy.ItemAppend => "create" - case ItemWriteStrategy.ItemDelete | ItemWriteStrategy.ItemDeleteIfNotModified => "delete" - case ItemWriteStrategy.ItemOverwriteIfNotModified => "replace" - case ItemWriteStrategy.ItemTransactionalBatch => "upsert" // Default to upsert for transactional batch - case ItemWriteStrategy.ItemPatch | ItemWriteStrategy.ItemPatchIfExists => "patch" - case _ => throw new RuntimeException(s"${writeConfig.itemWriteStrategy} not supported for transactional batch") + // For transactional batches, only upsert is supported + // This is validated in the constructor + val effectiveOperationType = "upsert" + + // Buffer operation for batch construction + this.addOperationToBatch(partitionKeyValue, objectNode, operationContext, effectiveOperationType) + } + + private[this] def addOperationToBatch(partitionKey: PartitionKey, + objectNode: ObjectNode, + operationContext: OperationContext, + effectiveOperationType: String): Unit = { + + batchConstructionLock.synchronized { + // Check if partition key changed - flush existing batch if needed + if (currentPartitionKey != null && !currentPartitionKey.equals(partitionKey)) { + flushCurrentBatch() + } + + // Initialize partition key if this is the first operation + if (currentPartitionKey == null) { + currentPartitionKey = partitionKey } + + // Add operation data to buffer + currentBatchOperations += ((objectNode, operationContext, effectiveOperationType)) + + // Note: We don't auto-flush at 100 operations in transactional mode because: + // 1. Auto-splitting would break transactional atomicity + // 2. The service will reject batches with >100 operations with a clear error + // 3. Users should ensure their transactional batches respect the 100 operation limit } + } - val bulkItemOperation = effectiveOperationType match { - case "create" => - CosmosBulkOperations.getCreateItemOperation(objectNode, partitionKeyValue, operationContext) - - case "upsert" => - CosmosBulkOperations.getUpsertItemOperation(objectNode, partitionKeyValue, operationContext) - - case "replace" => - operationContext.eTag match { - case Some(eTag) => - CosmosBulkOperations.getReplaceItemOperation( - operationContext.itemId, - objectNode, - partitionKeyValue, - new CosmosBulkItemRequestOptions().setIfMatchETag(eTag), - operationContext) - case _ => - CosmosBulkOperations.getReplaceItemOperation( - operationContext.itemId, - objectNode, - partitionKeyValue, - operationContext) - } - - case "delete" => - operationContext.eTag match { - case Some(eTag) => - CosmosBulkOperations.getDeleteItemOperation( - operationContext.itemId, - partitionKeyValue, - new CosmosBulkItemRequestOptions().setIfMatchETag(eTag), - operationContext) - case _ => - CosmosBulkOperations.getDeleteItemOperation( - operationContext.itemId, - partitionKeyValue, - operationContext) - } - - case "patch" => - getPatchItemOperation( - operationContext.itemId, - partitionKeyValue, - partitionKeyDefinition, - objectNode, - operationContext) + private[this] def flushCurrentBatch(): Unit = { + // Must be called within batchConstructionLock.synchronized + if (currentBatchOperations.nonEmpty && currentPartitionKey != null) { + // Build the batch using the builder API + val batch = CosmosBatch.createCosmosBatch(currentPartitionKey) + + // Build the batch and keep track of context for each operation + // For transactional batches, all operations are upserts + val contextList = new scala.collection.mutable.ListBuffer[OperationContext]() + + currentBatchOperations.foreach { case (objectNode, operationContext, _) => + batch.upsertItemOperation(objectNode) + contextList += operationContext + } + + // After building the batch, get the operations and track them with their contexts + val batchOperations = batch.getOperations() - case _ => - throw new IllegalArgumentException( - s"Unsupported operationType '$effectiveOperationType'. " + - s"Supported types for transactional batch: create, upsert, replace, delete") + // Map each operation to its context and add to tracking + for (i <- 0 until batchOperations.size()) { + val operation = batchOperations.get(i) + val context = contextList(i) + operationContextMap.put(operation, context) + activeBulkWriteOperations.add(operation) + } + + // Emit the batch + bulkInputEmitter.emitNext(batch, emitFailureHandler) + + // Clear the buffer + currentBatchOperations.clear() + currentPartitionKey = null } + } - this.emitBulkInput(bulkItemOperation) + private[this] def getPatchOperationsForBatch(itemId: String, objectNode: ObjectNode): CosmosPatchOperations = { + assert(writeConfig.patchConfigs.isDefined) + assert(cosmosPatchHelperOpt.isDefined) + val cosmosPatchHelper = cosmosPatchHelperOpt.get + cosmosPatchHelper.createCosmosPatchOperations(itemId, partitionKeyDefinition, objectNode) } - private[this] def emitBulkInput(bulkItemOperation: CosmosItemOperation): Unit = { - activeBulkWriteOperations.add(bulkItemOperation) - - // Simply emit the operation - TransactionalBulkExecutor will handle batching by partition key - bulkInputEmitter.emitNext(bulkItemOperation, emitFailureHandler) + private[this] def finalFlushBatch(): Unit = { + batchConstructionLock.synchronized { + flushCurrentBatch() + } } private[this] def getPatchItemOperation(itemId: String, @@ -1100,6 +1130,9 @@ private class TransactionalBulkWriter this.synchronized { try { if (!closed.get()) { + // Flush any remaining batched operations before closing + finalFlushBatch() + log.logInfo(s"flushAndClose invoked $getThreadInfo") log.logInfo(s"completed so far ${totalSuccessfulIngestionMetrics.get()}, " + s"pending bulkWrite asks ${activeBulkWriteOperations.size}, pending readMany tasks ${activeReadManyOperations.size} $getThreadInfo") @@ -1158,18 +1191,13 @@ private class TransactionalBulkWriter } activeOperationsSnapshot.foreach(operation => { - if (activeBulkWriteOperations.contains(operation)) { - // re-validating whether the operation is still active - if so, just re-enqueue another retry - // this is harmless - because all bulkItemOperations from Spark connector are always idempotent - - // For FAIL_NON_SERIALIZED, will keep retry, while for other errors, use the default behavior - // TransactionalBulkExecutor will handle batching by partition key - bulkInputEmitter.emitNext(operation, emitFailureHandler) - log.logWarning(s"Re-enqueued a retry for pending active write task '${operation.getOperationType} " - + s"(${operation.getPartitionKeyValue}/${operation.getId})' " - + s"- Attempt: ${numberOfIntervalsWithIdenticalActiveOperationSnapshots.get} - " - + s"Context: ${operationContext.toString} $getThreadInfo") - } + // Note: In transactional batch mode, we don't retry individual operations + // because batches are constructed and executed atomically. The executor + // handles retries at the batch level. Individual operation retry would require + // reconstructing batches, which could lead to incorrect retry behavior. + // Failed batches will be retried by Spark's task retry mechanism. + log.logTrace(s"Skipping individual operation retry in transactional batch mode - " + + s"Context: ${operationContext.toString} $getThreadInfo") }) activeReadManyOperationsSnapshot.foreach(operation => { @@ -1399,10 +1427,9 @@ private class TransactionalBulkWriter val attemptNumber: Int, val sequenceNumber: Long, /** starts from 1 * */ - sourceItemInput: Option[ObjectNode] = None, // for patchBulkUpdate: source item refers to the original objectNode from which SDK constructs the final bulk item operation - operationTypeInput: Option[String] = None) // per-row operation type (create, upsert, replace, delete) + sourceItemInput: Option[ObjectNode] = None) // for patchBulkUpdate: source item refers to the original objectNode from which SDK constructs the final bulk item operation { - private val ctxCore: OperationContextCore = OperationContextCore(itemIdInput, partitionKeyValueInput, eTagInput, sourceItemInput, operationTypeInput) + private val ctxCore: OperationContextCore = OperationContextCore(itemIdInput, partitionKeyValueInput, eTagInput, sourceItemInput) override def equals(obj: Any): Boolean = ctxCore.equals(obj) @@ -1419,8 +1446,6 @@ private class TransactionalBulkWriter def eTag: Option[String] = ctxCore.eTag def sourceItem: Option[ObjectNode] = ctxCore.sourceItem - - def operationType: Option[String] = ctxCore.operationType } private case class OperationContextCore @@ -1428,8 +1453,7 @@ private class TransactionalBulkWriter itemId: String, partitionKeyValue: PartitionKey, eTag: Option[String], - sourceItem: Option[ObjectNode] = None, // for patchBulkUpdate: source item refers to the original objectNode from which SDK constructs the final bulk item operation - operationType: Option[String] = None) // per-row operation type (create, upsert, replace, delete) + sourceItem: Option[ObjectNode] = None) // for patchBulkUpdate: source item refers to the original objectNode from which SDK constructs the final bulk item operation { override def productPrefix: String = "OperationContext" } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala index 22d5727c22cc..4027a2398f78 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala @@ -48,14 +48,14 @@ class TransactionalBatchITest extends IntegrationSpec val operationsDf = spark.createDataFrame(batchOperations.asJava, schema) - // Execute transactional batch using ItemTransactionalBatch write strategy + // Execute transactional batch using bulk transactional mode operationsDf.write .format("cosmos.oltp") .option("spark.cosmos.accountEndpoint", cosmosEndpoint) .option("spark.cosmos.accountKey", cosmosMasterKey) .option("spark.cosmos.database", cosmosDatabase) .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) - .option("spark.cosmos.write.strategy", "ItemTransactionalBatch") + .option("spark.cosmos.write.bulk.transactional", "true") .option("spark.cosmos.write.bulk.enabled", "true") .mode(SaveMode.Append) .save() @@ -70,65 +70,186 @@ class TransactionalBatchITest extends IntegrationSpec item2.getItem.get("name").asText() shouldEqual "Bob" } - it should "rollback all operations on failure" in { + it should "rollback all operations on batch size limit failure" in { val cosmosEndpoint = TestConfigurations.HOST val cosmosMasterKey = TestConfigurations.MASTER_KEY val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) val partitionKeyValue = UUID.randomUUID().toString - val item1Id = s"test-item1-${UUID.randomUUID()}" - val duplicateId = s"duplicate-${UUID.randomUUID()}" - // First create an item that we'll try to create again (should fail) - val existingDoc = Utils.getSimpleObjectMapper.createObjectNode() - existingDoc.put("id", duplicateId) - existingDoc.put("pk", partitionKeyValue) - existingDoc.put("name", "Existing") - container.createItem(existingDoc).block() + // Create a batch that exceeds the 2MB size limit for transactional batches + // Each item is ~20KB, so 101 items will exceed the limit + val schema = StructType(Seq( + StructField("id", StringType, nullable = false), + StructField("pk", StringType, nullable = false), + StructField("largeField", StringType, nullable = false) + )) + + val largeString = "x" * 20000 // 20KB string + val batchOperations = (1 to 101).map { i => + Row(s"item-$i-${UUID.randomUUID()}", partitionKeyValue, largeString) + } + + val operationsDf = spark.createDataFrame(batchOperations.asJava, schema) + + // This should fail because batch exceeds size limit + val exception = intercept[Exception] { + operationsDf.write + .format("cosmos.oltp") + .option("spark.cosmos.accountEndpoint", cosmosEndpoint) + .option("spark.cosmos.accountKey", cosmosMasterKey) + .option("spark.cosmos.database", cosmosDatabase) + .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) + .option("spark.cosmos.write.bulk.transactional", "true") + .option("spark.cosmos.write.bulk.enabled", "true") + .mode(SaveMode.Append) + .save() + } + + // Verify the exception indicates batch failure (either size or count limit) + val rootCause = getRootCause(exception) + assert(rootCause.getMessage.contains("more operations") || rootCause.getMessage.contains("size") || rootCause.getMessage.contains("limit"), + s"Expected batch limit error, got: ${rootCause.getMessage}") + + // Verify NO items were created (rolled back) - atomic rollback ensures nothing persisted + val itemsCreated = try { + val countList = container.queryItems( + s"SELECT VALUE COUNT(1) FROM c WHERE c.pk = '$partitionKeyValue'", + classOf[Long] + ).collectList().block() + + if (countList.isEmpty) { + 0 + } else { + countList.get(0).toInt + } + } catch { + case _: Exception => 0 + } + assert(itemsCreated == 0, s"No items should exist due to batch rollback, but found $itemsCreated") + } + + it should "rollback all operations when one has blank ID" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + val partitionKeyValue = "rollback-blank-id-test" + + // Create a batch where the 5th item has a blank ID (validation error) + val batchOperations = Seq( + ("id1", partitionKeyValue, "value1"), + ("id2", partitionKeyValue, "value2"), + ("id3", partitionKeyValue, "value3"), + ("id4", partitionKeyValue, "value4"), + ("", partitionKeyValue, "blank-id"), // Blank ID will cause error + ("id6", partitionKeyValue, "value6") + ) + + val schema = StructType(Seq( + StructField("id", StringType), + StructField("pk", StringType), + StructField("value", StringType) + )) + + val operationsDf = spark.createDataFrame( + batchOperations.map(op => Row(op._1, op._2, op._3)).asJava, + schema + ) + + // Should throw exception due to blank ID + val exception = intercept[Exception] { + operationsDf.write + .format("cosmos.oltp") + .option("spark.cosmos.accountEndpoint", cosmosEndpoint) + .option("spark.cosmos.accountKey", cosmosMasterKey) + .option("spark.cosmos.database", cosmosDatabase) + .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) + .option("spark.cosmos.write.bulk.transactional", "true") + .mode(SaveMode.Append) + .save() + } + + // Verify the first item (id1) was NOT created - proving rollback worked + val cosmosContainer = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) + val queryResult = cosmosContainer + .queryItems(s"SELECT * FROM c WHERE c.pk = '$partitionKeyValue'", classOf[ObjectNode]) + .collectList() + .block() + + // Should have zero items - all operations rolled back + queryResult.size() shouldBe 0 + } + + it should "reject unsupported write strategies" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + val partitionKeyValue = UUID.randomUUID().toString + val item1Id = s"test-item1-${UUID.randomUUID()}" - // Create batch with one valid create and one duplicate create (should fail entire batch) + // Create DataFrame with simple schema val schema = StructType(Seq( StructField("id", StringType, nullable = false), StructField("pk", StringType, nullable = false), - StructField("operationType", StringType, nullable = false), StructField("name", StringType, nullable = false) )) val batchOperations = Seq( - Row(item1Id, partitionKeyValue, "create", "NewItem"), - Row(duplicateId, partitionKeyValue, "create", "Duplicate") + Row(item1Id, partitionKeyValue, "TestItem") ) val operationsDf = spark.createDataFrame(batchOperations.asJava, schema) - // This should fail because we're using create operations and duplicateId already exists - val exception = intercept[Exception] { + // Test ItemAppend (create) - should fail + val appendException = intercept[Exception] { operationsDf.write .format("cosmos.oltp") .option("spark.cosmos.accountEndpoint", cosmosEndpoint) .option("spark.cosmos.accountKey", cosmosMasterKey) .option("spark.cosmos.database", cosmosDatabase) .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) - .option("spark.cosmos.write.strategy", "ItemTransactionalBatch") + .option("spark.cosmos.write.bulk.transactional", "true") .option("spark.cosmos.write.bulk.enabled", "true") + .option("spark.cosmos.write.strategy", "ItemAppend") .mode(SaveMode.Append) .save() } + val appendRootCause = getRootCause(appendException) + assert(appendRootCause.getMessage.contains("Transactional batches only support ItemOverwrite"), + s"Expected ItemAppend rejection, got: ${appendRootCause.getMessage}") - // Verify the exception message indicates transactional batch failure - // Spark wraps our exception in "Writing job aborted", check the root cause - // Transactional batch failures may show as statusCode 424 (Failed Dependency) when one operation fails - val rootCause = getRootCause(exception) - assert(rootCause.getMessage.contains("424") || rootCause.getMessage.contains("409"), - s"Expected transactional batch failure error (424 or 409), got: ${rootCause.getMessage}") + // Test ItemDelete - should fail + val deleteException = intercept[Exception] { + operationsDf.write + .format("cosmos.oltp") + .option("spark.cosmos.accountEndpoint", cosmosEndpoint) + .option("spark.cosmos.accountKey", cosmosMasterKey) + .option("spark.cosmos.database", cosmosDatabase) + .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) + .option("spark.cosmos.write.bulk.transactional", "true") + .option("spark.cosmos.write.bulk.enabled", "true") + .option("spark.cosmos.write.strategy", "ItemDelete") + .mode(SaveMode.Append) + .save() + } + val deleteRootCause = getRootCause(deleteException) + assert(deleteRootCause.getMessage.contains("Transactional batches only support ItemOverwrite"), + s"Expected ItemDelete rejection, got: ${deleteRootCause.getMessage}") - // Verify item1 was NOT created (rolled back) - try { - container.readItem(item1Id, new PartitionKey(partitionKeyValue), classOf[ObjectNode]).block() - fail("Item1 should not exist after batch rollback") - } catch { - case e: CosmosException if e.getStatusCode == 404 => - // Expected - item doesn't exist after rollback (404 Not Found) + // Test ItemOverwriteIfNotModified - should fail + val replaceException = intercept[Exception] { + operationsDf.write + .format("cosmos.oltp") + .option("spark.cosmos.accountEndpoint", cosmosEndpoint) + .option("spark.cosmos.accountKey", cosmosMasterKey) + .option("spark.cosmos.database", cosmosDatabase) + .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) + .option("spark.cosmos.write.bulk.transactional", "true") + .option("spark.cosmos.write.bulk.enabled", "true") + .option("spark.cosmos.write.strategy", "ItemOverwriteIfNotModified") + .mode(SaveMode.Append) + .save() } + val replaceRootCause = getRootCause(replaceException) + assert(replaceRootCause.getMessage.contains("Transactional batches only support ItemOverwrite"), + s"Expected ItemOverwriteIfNotModified rejection, got: ${replaceRootCause.getMessage}") } it should "support simplified schema with default upsert operation" in { @@ -161,7 +282,7 @@ class TransactionalBatchITest extends IntegrationSpec .option("spark.cosmos.accountKey", cosmosMasterKey) .option("spark.cosmos.database", cosmosDatabase) .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) - .option("spark.cosmos.write.strategy", "ItemTransactionalBatch") + .option("spark.cosmos.write.bulk.transactional", "true") .option("spark.cosmos.write.bulk.enabled", "true") .mode(SaveMode.Append) .save() @@ -204,7 +325,7 @@ class TransactionalBatchITest extends IntegrationSpec .option("spark.cosmos.accountKey", cosmosMasterKey) .option("spark.cosmos.database", cosmosDatabase) .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) - .option("spark.cosmos.write.strategy", "ItemTransactionalBatch") + .option("spark.cosmos.write.bulk.transactional", "true") .option("spark.cosmos.write.bulk.enabled", "true") .mode(SaveMode.Append) .save() @@ -252,7 +373,7 @@ class TransactionalBatchITest extends IntegrationSpec .option("spark.cosmos.accountKey", cosmosMasterKey) .option("spark.cosmos.database", cosmosDatabase) .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) - .option("spark.cosmos.write.strategy", "ItemTransactionalBatch") + .option("spark.cosmos.write.bulk.transactional", "true") .option("spark.cosmos.write.bulk.enabled", "true") .mode(SaveMode.Append) .save() @@ -291,7 +412,7 @@ class TransactionalBatchITest extends IntegrationSpec .option("spark.cosmos.accountKey", cosmosMasterKey) .option("spark.cosmos.database", cosmosDatabase) .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) - .option("spark.cosmos.write.strategy", "ItemTransactionalBatch") + .option("spark.cosmos.write.bulk.transactional", "true") .option("spark.cosmos.write.bulk.enabled", "true") .mode(SaveMode.Append) .save() @@ -299,9 +420,8 @@ class TransactionalBatchITest extends IntegrationSpec // Spark wraps our exception in "Writing job aborted", check the root cause val rootCause = getRootCause(exception) - rootCause.getMessage should include("exceeds") - rootCause.getMessage should include("101 operations") - rootCause.getMessage should include("maximum allowed limit of 100") + // Cosmos DB rejects batches with > 100 operations + rootCause.getMessage should include("Batch request has more operations than what is supported") } "Transactional Batch with Hierarchical Partition Keys" should "create items atomically with PermId and SourceId" in { @@ -350,7 +470,7 @@ class TransactionalBatchITest extends IntegrationSpec .option("spark.cosmos.accountKey", cosmosMasterKey) .option("spark.cosmos.database", cosmosDatabase) .option("spark.cosmos.container", containerName) - .option("spark.cosmos.write.strategy", "ItemTransactionalBatch") + .option("spark.cosmos.write.bulk.transactional", "true") .option("spark.cosmos.write.bulk.enabled", "true") .mode(SaveMode.Append) .save() @@ -428,7 +548,7 @@ class TransactionalBatchITest extends IntegrationSpec .option("spark.cosmos.accountKey", cosmosMasterKey) .option("spark.cosmos.database", cosmosDatabase) .option("spark.cosmos.container", containerName) - .option("spark.cosmos.write.strategy", "ItemTransactionalBatch") + .option("spark.cosmos.write.bulk.transactional", "true") .option("spark.cosmos.write.bulk.enabled", "true") .mode(SaveMode.Append) .save() @@ -496,7 +616,7 @@ class TransactionalBatchITest extends IntegrationSpec .option("spark.cosmos.accountKey", cosmosMasterKey) .option("spark.cosmos.database", cosmosDatabase) .option("spark.cosmos.container", containerName) - .option("spark.cosmos.write.strategy", "ItemTransactionalBatch") + .option("spark.cosmos.write.bulk.transactional", "true") .option("spark.cosmos.write.bulk.enabled", "true") .mode(SaveMode.Append) .save() @@ -528,46 +648,47 @@ class TransactionalBatchITest extends IntegrationSpec cosmosClient.getDatabase(cosmosDatabase).createContainer(containerProperties).block() val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(containerName) - try { - val permId = "MSFT" - val sourceId = "Bloomberg" + val permId = "MSFT" + val sourceId = "Bloomberg" - // Create 101 operations for the same hierarchical partition key - val schema = StructType(Seq( - StructField("id", StringType, nullable = false), - StructField("PermId", StringType, nullable = false), - StructField("SourceId", StringType, nullable = false), - StructField("price", org.apache.spark.sql.types.DoubleType, nullable = false) - )) + // Create 101 operations for the same hierarchical partition key + val schema = StructType(Seq( + StructField("id", StringType, nullable = false), + StructField("PermId", StringType, nullable = false), + StructField("SourceId", StringType, nullable = false), + StructField("price", org.apache.spark.sql.types.DoubleType, nullable = false) + )) - val batchOperations = (1 to 101).map { i => - Row(s"${UUID.randomUUID()}", permId, sourceId, 100.0 + i) - } + val batchOperations = (1 to 101).map { i => + Row(s"${UUID.randomUUID()}", permId, sourceId, 100.0 + i) + } - val operationsDf = spark.createDataFrame(batchOperations.asJava, schema) + val operationsDf = spark.createDataFrame(batchOperations.asJava, schema) - // Should throw exception due to exceeding 100 operations per partition key limit - val exception = intercept[Exception] { - operationsDf.write - .format("cosmos.oltp") - .option("spark.cosmos.accountEndpoint", cosmosEndpoint) - .option("spark.cosmos.accountKey", cosmosMasterKey) - .option("spark.cosmos.database", cosmosDatabase) - .option("spark.cosmos.container", containerName) - .option("spark.cosmos.write.strategy", "ItemTransactionalBatch") - .option("spark.cosmos.write.bulk.enabled", "true") - .mode(SaveMode.Append) - .save() - } + // Should throw exception due to exceeding 100 operations per partition key limit + val exception = intercept[Exception] { + operationsDf.write + .format("cosmos.oltp") + .option("spark.cosmos.accountEndpoint", cosmosEndpoint) + .option("spark.cosmos.accountKey", cosmosMasterKey) + .option("spark.cosmos.database", cosmosDatabase) + .option("spark.cosmos.container", containerName) + .option("spark.cosmos.write.bulk.transactional", "true") + .option("spark.cosmos.write.bulk.enabled", "true") + .mode(SaveMode.Append) + .save() + } - // Spark wraps our exception in "Writing job aborted", check the root cause - val rootCause = getRootCause(exception) - rootCause.getMessage should include("exceeds") - rootCause.getMessage should include("101 operations") - rootCause.getMessage should include("maximum allowed limit of 100") - } finally { - // Clean up container + // Spark wraps our exception in "Writing job aborted", check the root cause + val rootCause = getRootCause(exception) + // Cosmos DB rejects batches with > 100 operations + rootCause.getMessage should include("Batch request has more operations than what is supported") + + // Clean up container after assertions pass + try { container.delete().block() + } catch { + case e: Exception => // Ignore cleanup failures } } @@ -610,7 +731,7 @@ class TransactionalBatchITest extends IntegrationSpec .option("spark.cosmos.accountKey", cosmosMasterKey) .option("spark.cosmos.database", cosmosDatabase) .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) - .option("spark.cosmos.write.strategy", "ItemTransactionalBatch") + .option("spark.cosmos.write.bulk.transactional", "true") .option("spark.cosmos.write.bulk.enabled", "true") .mode(SaveMode.Append) .save() @@ -625,268 +746,4 @@ class TransactionalBatchITest extends IntegrationSpec } } - "Transactional batch with per-row operation types" should "support mixed operations in same partition key" in { - val cosmosEndpoint = TestConfigurations.HOST - val cosmosMasterKey = TestConfigurations.MASTER_KEY - val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) - - // Pre-create an item for replace operation - val itemForReplace = Utils.getSimpleObjectMapper.createObjectNode() - itemForReplace.put("id", "item-to-replace") - itemForReplace.put("pk", "mixed-ops") - itemForReplace.put("prop", "original-value") - container.createItem(itemForReplace).block() - - // Pre-create item for delete operation - val itemForDelete = Utils.getSimpleObjectMapper.createObjectNode() - itemForDelete.put("id", "item-to-delete") - itemForDelete.put("pk", "mixed-ops") - itemForDelete.put("prop", "to-be-deleted") - container.createItem(itemForDelete).block() - - // Create DataFrame with mixed operations - val schema = StructType(Seq( - StructField("id", StringType, nullable = false), - StructField("pk", StringType, nullable = false), - StructField("operationType", StringType, nullable = false), - StructField("prop", StringType, nullable = false) - )) - - val mixedOpsItems = Seq( - Row("item-new-create", "mixed-ops", "create", "created-value"), - Row("item-new-upsert", "mixed-ops", "upsert", "upserted-value"), - Row("item-to-replace", "mixed-ops", "replace", "replaced-value"), - Row("item-to-delete", "mixed-ops", "delete", "delete-value") - ) - - val mixedOpsDf = spark.createDataFrame(mixedOpsItems.asJava, schema) - - val writeCfg = Map( - "spark.cosmos.accountEndpoint" -> cosmosEndpoint, - "spark.cosmos.accountKey" -> cosmosMasterKey, - "spark.cosmos.database" -> cosmosDatabase, - "spark.cosmos.container" -> cosmosContainersWithPkAsPartitionKey, - "spark.cosmos.write.strategy" -> "ItemTransactionalBatch", - "spark.cosmos.write.maxRetryCount" -> "3", - "spark.cosmos.write.bulk.enabled" -> "true" - ) - - // Write with mixed operations - mixedOpsDf.write.format("cosmos.oltp").mode(SaveMode.Append).options(writeCfg).save() - - // Verify all operations succeeded - // 1. Check created item - val createdItem = container.readItem("item-new-create", new PartitionKey("mixed-ops"), classOf[ObjectNode]).block().getItem - createdItem.get("prop").asText() shouldEqual "created-value" - - // 2. Check upserted item - val upsertedItem = container.readItem("item-new-upsert", new PartitionKey("mixed-ops"), classOf[ObjectNode]).block().getItem - upsertedItem.get("prop").asText() shouldEqual "upserted-value" - - // 3. Check replaced item - val replacedItem = container.readItem("item-to-replace", new PartitionKey("mixed-ops"), classOf[ObjectNode]).block().getItem - replacedItem.get("prop").asText() shouldEqual "replaced-value" - - // 4. Verify deleted item no longer exists - try { - container.readItem("item-to-delete", new PartitionKey("mixed-ops"), classOf[ObjectNode]).block() - fail("Item should have been deleted") - } catch { - case e: CosmosException => e.getStatusCode shouldEqual 404 - } - } - - it should "rollback all mixed operations when one fails" in { - val cosmosEndpoint = TestConfigurations.HOST - val cosmosMasterKey = TestConfigurations.MASTER_KEY - val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) - - // Pre-create an item - val existingItem = Utils.getSimpleObjectMapper.createObjectNode() - existingItem.put("id", "existing-item") - existingItem.put("pk", "rollback-test") - existingItem.put("prop", "original") - container.createItem(existingItem).block() - - // Create DataFrame with mixed operations where create will fail (duplicate) - val schema = StructType(Seq( - StructField("id", StringType, nullable = false), - StructField("pk", StringType, nullable = false), - StructField("operationType", StringType, nullable = false), - StructField("prop", StringType, nullable = false) - )) - - val rollbackItems = Seq( - Row("new-item-1", "rollback-test", "create", "new-value-1"), - Row("existing-item", "rollback-test", "create", "duplicate-will-fail"), - Row("new-item-2", "rollback-test", "upsert", "new-value-2") - ) - - val rollbackDf = spark.createDataFrame(rollbackItems.asJava, schema) - - val writeCfg = Map( - "spark.cosmos.accountEndpoint" -> cosmosEndpoint, - "spark.cosmos.accountKey" -> cosmosMasterKey, - "spark.cosmos.database" -> cosmosDatabase, - "spark.cosmos.container" -> cosmosContainersWithPkAsPartitionKey, - "spark.cosmos.write.strategy" -> "ItemTransactionalBatch", - "spark.cosmos.write.maxRetryCount" -> "3", - "spark.cosmos.write.bulk.enabled" -> "true" - ) - - // Attempt write - should fail atomically - try { - rollbackDf.write.format("cosmos.oltp").mode(SaveMode.Append).options(writeCfg).save() - fail("Transactional batch should have failed due to duplicate create") - } catch { - case e: Exception => - // Expected failure (409 Conflict or 424 Failed Dependency) - println(s"Expected failure: ${e.getMessage}") - } - - // Verify NO operations were committed (atomic rollback) - // 1. Original item should be unchanged - val originalItem = container.readItem("existing-item", new PartitionKey("rollback-test"), classOf[ObjectNode]).block().getItem - originalItem.get("prop").asText() shouldEqual "original" - - // 2. new-item-1 should NOT exist (rollback) - try { - container.readItem("new-item-1", new PartitionKey("rollback-test"), classOf[ObjectNode]).block() - fail("new-item-1 should not exist due to rollback") - } catch { - case e: CosmosException => e.getStatusCode shouldEqual 404 - } - - // 3. new-item-2 should NOT exist (rollback) - try { - container.readItem("new-item-2", new PartitionKey("rollback-test"), classOf[ObjectNode]).block() - fail("new-item-2 should not exist due to rollback") - } catch { - case e: CosmosException => e.getStatusCode shouldEqual 404 - } - } - - it should "support delete operations in transactional batch" in { - val cosmosEndpoint = TestConfigurations.HOST - val cosmosMasterKey = TestConfigurations.MASTER_KEY - val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) - - // Pre-create items to delete - val itemToDelete1 = Utils.getSimpleObjectMapper.createObjectNode() - itemToDelete1.put("id", "delete-item-1") - itemToDelete1.put("pk", "delete-test") - itemToDelete1.put("prop", "delete-me-1") - container.createItem(itemToDelete1).block() - - val itemToDelete2 = Utils.getSimpleObjectMapper.createObjectNode() - itemToDelete2.put("id", "delete-item-2") - itemToDelete2.put("pk", "delete-test") - itemToDelete2.put("prop", "delete-me-2") - container.createItem(itemToDelete2).block() - - // Create DataFrame with create + delete operations - val schema = StructType(Seq( - StructField("id", StringType, nullable = false), - StructField("pk", StringType, nullable = false), - StructField("operationType", StringType, nullable = false), - StructField("prop", StringType, nullable = false) - )) - - val deleteItems = Seq( - Row("new-item", "delete-test", "create", "new-value"), - Row("delete-item-1", "delete-test", "delete", "ignored"), - Row("delete-item-2", "delete-test", "delete", "ignored") - ) - - val deleteDf = spark.createDataFrame(deleteItems.asJava, schema) - - val writeCfg = Map( - "spark.cosmos.accountEndpoint" -> cosmosEndpoint, - "spark.cosmos.accountKey" -> cosmosMasterKey, - "spark.cosmos.database" -> cosmosDatabase, - "spark.cosmos.container" -> cosmosContainersWithPkAsPartitionKey, - "spark.cosmos.write.strategy" -> "ItemTransactionalBatch", - "spark.cosmos.write.maxRetryCount" -> "3", - "spark.cosmos.write.bulk.enabled" -> "true" - ) - - // Execute transactional batch - deleteDf.write.format("cosmos.oltp").mode(SaveMode.Append).options(writeCfg).save() - - // Verify new item was created - val newItem = container.readItem("new-item", new PartitionKey("delete-test"), classOf[ObjectNode]).block().getItem - newItem.get("prop").asText() shouldEqual "new-value" - - // Verify items were deleted - try { - container.readItem("delete-item-1", new PartitionKey("delete-test"), classOf[ObjectNode]).block() - fail("delete-item-1 should have been deleted") - } catch { - case e: CosmosException => e.getStatusCode shouldEqual 404 - } - - try { - container.readItem("delete-item-2", new PartitionKey("delete-test"), classOf[ObjectNode]).block() - fail("delete-item-2 should have been deleted") - } catch { - case e: CosmosException => e.getStatusCode shouldEqual 404 - } - } - - it should "default to global strategy when operationType column absent" in { - val cosmosEndpoint = TestConfigurations.HOST - val cosmosMasterKey = TestConfigurations.MASTER_KEY - val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) - - // Create DataFrame WITHOUT operationType column - val schema = StructType(Seq( - StructField("id", StringType, nullable = false), - StructField("pk", StringType, nullable = false), - StructField("prop", StringType, nullable = false) - )) - - val backwardCompatItems = Seq( - Row("compat-item-1", "compat-test", "value-1"), - Row("compat-item-2", "compat-test", "value-2") - ) - - val compatDf = spark.createDataFrame(backwardCompatItems.asJava, schema) - - val writeCfg = Map( - "spark.cosmos.accountEndpoint" -> cosmosEndpoint, - "spark.cosmos.accountKey" -> cosmosMasterKey, - "spark.cosmos.database" -> cosmosDatabase, - "spark.cosmos.container" -> cosmosContainersWithPkAsPartitionKey, - "spark.cosmos.write.strategy" -> "ItemTransactionalBatch", - "spark.cosmos.write.maxRetryCount" -> "3", - "spark.cosmos.write.bulk.enabled" -> "true" - ) - - // Write without operationType column (should default to upsert) - compatDf.write.format("cosmos.oltp").mode(SaveMode.Append).options(writeCfg).save() - - // Verify items were created (upsert default for ItemTransactionalBatch) - val item1 = container.readItem("compat-item-1", new PartitionKey("compat-test"), classOf[ObjectNode]).block().getItem - item1.get("prop").asText() shouldEqual "value-1" - - val item2 = container.readItem("compat-item-2", new PartitionKey("compat-test"), classOf[ObjectNode]).block().getItem - item2.get("prop").asText() shouldEqual "value-2" - - // Update items - should succeed with upsert - val updateItems = Seq( - Row("compat-item-1", "compat-test", "updated-1"), - Row("compat-item-2", "compat-test", "updated-2") - ) - - val updateDf = spark.createDataFrame(updateItems.asJava, schema) - - updateDf.write.format("cosmos.oltp").mode(SaveMode.Append).options(writeCfg).save() - - // Verify items were updated - val updatedItem1 = container.readItem("compat-item-1", new PartitionKey("compat-test"), classOf[ObjectNode]).block().getItem - updatedItem1.get("prop").asText() shouldEqual "updated-1" - - val updatedItem2 = container.readItem("compat-item-2", new PartitionKey("compat-test"), classOf[ObjectNode]).block().getItem - updatedItem2.get("prop").asText() shouldEqual "updated-2" - } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/TransactionalBulkExecutor.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/TransactionalBulkExecutor.java index 9d23e55d41a8..f67cf44109ff 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/TransactionalBulkExecutor.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/TransactionalBulkExecutor.java @@ -100,13 +100,9 @@ public final class TransactionalBulkExecutor implements Disposable { private final String operationContextText; private final OperationContextAndListenerTuple operationListener; private final ThrottlingRetryOptions throttlingRetryOptions; - private final Flux inputOperations; - - // Options for bulk execution. - private final Long maxMicroBatchIntervalInMs; + private final Flux inputBatches; private final TContext batchContext; - private final ConcurrentMap partitionScopeThresholds; private final CosmosBulkExecutionOptionsImpl cosmosBulkExecutionOptions; // Handle gone error: @@ -114,14 +110,8 @@ public final class TransactionalBulkExecutor implements Disposable { private final AtomicBoolean isDisposed = new AtomicBoolean(false); private final AtomicBoolean isShutdown = new AtomicBoolean(false); private final AtomicInteger totalCount; - private final static Sinks.EmitFailureHandler serializedEmitFailureHandler = new SerializedEmitFailureHandler(); - private final static Sinks.EmitFailureHandler serializedCompleteEmitFailureHandler = - new SerializedCompleteEmitFailureHandler(); - private final Sinks.Many mainSink; - private final List> groupSinks; private final CosmosAsyncClient cosmosClient; private final String bulkSpanName; - private final AtomicReference scheduledFutureForFlush; private final String identifier = "TransactionalBulkExecutor-" + instanceCount.incrementAndGet(); private final BulkExecutorDiagnosticsTracker diagnosticsTracker; private final CosmosItemSerializer effectiveItemSerializer; @@ -129,7 +119,7 @@ public final class TransactionalBulkExecutor implements Disposable { @SuppressWarnings({"unchecked"}) public TransactionalBulkExecutor(CosmosAsyncContainer container, - Flux inputOperations, + Flux inputOperations, CosmosBulkExecutionOptionsImpl cosmosBulkOptions) { checkNotNull(container, "expected non-null container"); @@ -139,8 +129,8 @@ public TransactionalBulkExecutor(CosmosAsyncContainer container, this.maxMicroBatchPayloadSizeInBytes = cosmosBulkOptions.getMaxMicroBatchPayloadSizeInBytes(); this.cosmosBulkExecutionOptions = cosmosBulkOptions; this.container = container; - this.bulkSpanName = "nonTransactionalBatch." + this.container.getId(); - this.inputOperations = inputOperations; + this.bulkSpanName = "transactionalBatch." + this.container.getId(); + this.inputBatches = inputOperations; this.docClientWrapper = CosmosBridgeInternal.getAsyncDocumentClient(container.getDatabase()); this.cosmosClient = ImplementationBridgeHelpers .CosmosAsyncDatabaseHelper @@ -152,11 +142,7 @@ public TransactionalBulkExecutor(CosmosAsyncContainer container, // Fill the option first, to make the BulkProcessingOptions immutable, as if accessed directly, we might get // different values when a new group is created. - maxMicroBatchIntervalInMs = cosmosBulkExecutionOptions.getMaxMicroBatchInterval().toMillis(); batchContext = (TContext) cosmosBulkExecutionOptions.getLegacyBatchScopedContext(); - this.partitionScopeThresholds = ImplementationBridgeHelpers.CosmosBulkExecutionThresholdsStateHelper - .getBulkExecutionThresholdsAccessor() - .getPartitionScopeThresholds(cosmosBulkExecutionOptions.getThresholdsState()); operationListener = cosmosBulkExecutionOptions.getOperationContextAndListenerTuple(); if (operationListener != null && operationListener.getOperationContext() != null) { @@ -167,19 +153,11 @@ public TransactionalBulkExecutor(CosmosAsyncContainer container, this.diagnosticsTracker = cosmosBulkExecutionOptions.getDiagnosticsTracker(); - // Initialize sink for handling gone error. + // For transactional batches, no sinks needed - batches are pre-constructed mainSourceCompleted = new AtomicBoolean(false); totalCount = new AtomicInteger(0); - mainSink = Sinks.many().unicast().onBackpressureBuffer(); - groupSinks = new CopyOnWriteArrayList<>(); - this.scheduledFutureForFlush = new AtomicReference<>(CosmosSchedulers - .BULK_EXECUTOR_FLUSH_BOUNDED_ELASTIC - .schedulePeriodically( - this::onFlush, - this.maxMicroBatchIntervalInMs, - this.maxMicroBatchIntervalInMs, - TimeUnit.MILLISECONDS)); + // No periodic flush scheduler needed - batches arrive pre-constructed from writer Scheduler schedulerSnapshotFromOptions = cosmosBulkOptions.getSchedulerOverride(); this.executionScheduler = schedulerSnapshotFromOptions != null @@ -207,34 +185,6 @@ public boolean isDisposed() { return this.isDisposed.get(); } - private void cancelFlushTask(boolean initializeAggressiveFlush) { - long flushIntervalAfterDrainingIncomingFlux = Math.min( - this.maxMicroBatchIntervalInMs, - BatchRequestResponseConstants - .DEFAULT_MAX_MICRO_BATCH_INTERVAL_AFTER_DRAINING_INCOMING_FLUX_IN_MILLISECONDS); - - Disposable newFlushTask = initializeAggressiveFlush - ? CosmosSchedulers - .BULK_EXECUTOR_FLUSH_BOUNDED_ELASTIC - .schedulePeriodically( - this::onFlush, - flushIntervalAfterDrainingIncomingFlux, - flushIntervalAfterDrainingIncomingFlux, - TimeUnit.MILLISECONDS) - : null; - - Disposable scheduledFutureSnapshot = this.scheduledFutureForFlush.getAndSet(newFlushTask); - - if (scheduledFutureSnapshot != null) { - try { - scheduledFutureSnapshot.dispose(); - logDebugOrWarning("Cancelled all future scheduled tasks {}, Context: {}", getThreadInfo(), this.operationContextText); - } catch (Exception e) { - logger.warn("Failed to cancel scheduled tasks{}, Context: {}", getThreadInfo(), this.operationContextText, e); - } - } - } - private void logInfoOrWarning(String msg, Object... args) { if (this.diagnosticsTracker == null || !this.diagnosticsTracker.verboseLoggingAfterReEnqueueingRetriesEnabled()) { logger.info(msg, args); @@ -262,11 +212,7 @@ private void logDebugOrWarning(String msg, Object... args) { private void shutdown() { if (this.isShutdown.compareAndSet(false, true)) { logDebugOrWarning("Shutting down, Context: {}", this.operationContextText); - - groupSinks.forEach(FluxSink::complete); - logger.debug("All group sinks completed, Context: {}", this.operationContextText); - - this.cancelFlushTask(false); + logger.debug("Shutdown complete, Context: {}", this.operationContextText); } } @@ -302,419 +248,54 @@ public Flux> execute() { private Flux> executeCore() { - // The groupBy below is running into a hang if the flatMap above is - // not allowing at least a concurrency of the number of unique values - // you groupBy on. - // The groupBy is used to isolate Cosmos physical partitions - // so when there is no config override we enforce that the flatMap is using a concurrency of - // Math.max(default concurrency (256), #of partitions * 2 (to accommodate for some splits)) - // The config override can be used by the Spark connector when customers follow best practices and - // repartition the data frame to avoid that each Spark partition contains data spread across all - // physical partitions. When repartitioning the incoming data it is possible to ensure that each - // Spark partition will only target a subset of Cosmos partitions. This will improve the efficiency - // and mean fewer than #of Partitions concurrency will be needed for - // large containers. (with hundreds of physical partitions) + // For transactional batches, batches are pre-constructed by the writer with proper partition key grouping. + // We just need to execute each pre-built batch. Integer nullableMaxConcurrentCosmosPartitions = cosmosBulkExecutionOptions.getMaxConcurrentCosmosPartitions(); - Mono maxConcurrentCosmosPartitionsMono = nullableMaxConcurrentCosmosPartitions != null ? - Mono.just(Math.max(256, nullableMaxConcurrentCosmosPartitions)) : - ImplementationBridgeHelpers - .CosmosAsyncContainerHelper - .getCosmosAsyncContainerAccessor() - .getFeedRanges(this.container, false).map(ranges -> Math.max(256, ranges.size() * 2)); - - return - maxConcurrentCosmosPartitionsMono - .subscribeOn(this.executionScheduler) - .flatMapMany(maxConcurrentCosmosPartitions -> { - - logDebugOrWarning("BulkExecutor.execute with MaxConcurrentPartitions: {}, Context: {}", - maxConcurrentCosmosPartitions, - this.operationContextText); - - return this.inputOperations - .publishOn(this.executionScheduler) - .onErrorMap(throwable -> { - logger.warn("{}: Error observed when processing inputOperations. Cause: {}, Context: {}", - getThreadInfo(), - throwable.getMessage(), - this.operationContextText, - throwable); - - return throwable; - }) - .doOnNext((CosmosItemOperation cosmosItemOperation) -> { - // Set the retry policy before starting execution. Should only happens once. - BulkExecutorUtil.setRetryPolicyForBulk( - docClientWrapper, - this.container, - cosmosItemOperation, - this.throttlingRetryOptions); - - if (cosmosItemOperation != FlushBuffersItemOperation.singleton()) { - totalCount.incrementAndGet(); - } - - logger.trace( - "SetupRetryPolicy, {}, TotalCount: {}, Context: {}, {}", - getItemOperationDiagnostics(cosmosItemOperation), - totalCount.get(), - this.operationContextText, - getThreadInfo() - ); - }) - .doOnComplete(() -> { - mainSourceCompleted.set(true); - - long totalCountSnapshot = totalCount.get(); - logDebugOrWarning("Main source completed - # left items {}, Context: {}", - totalCountSnapshot, - this.operationContextText); - if (totalCountSnapshot == 0) { - // This is needed as there can be case that onComplete was called after last element was processed - // So complete the sink here also if count is 0, if source has completed and count isn't zero, - // then the last element in the doOnNext will close it. Sink doesn't mind in case of a double close. - - completeAllSinks(); - } else { - this.cancelFlushTask(true); - this.onFlush(); - - logDebugOrWarning("Scheduled new flush operation {}, Context: {}", getThreadInfo(), this.operationContextText); - } - }) - .mergeWith(mainSink.asFlux()) - .subscribeOn(this.executionScheduler) - .flatMap( - operation -> { - logger.trace("Before Resolve PkRangeId, {}, Context: {} {}", - getItemOperationDiagnostics(operation), - this.operationContextText, - getThreadInfo()); - - // resolve partition key range id again for operations which comes in main sink due to gone retry. - return BulkExecutorUtil.resolvePartitionKeyRangeId(this.docClientWrapper, this.container, operation) - .map((String pkRangeId) -> { - PartitionScopeThresholds partitionScopeThresholds = - this.partitionScopeThresholds.computeIfAbsent( - pkRangeId, - (newPkRangeId) -> new PartitionScopeThresholds(newPkRangeId, this.cosmosBulkExecutionOptions)); - - logger.trace("Resolved PkRangeId, {}, PKRangeId: {} Context: {} {}", - getItemOperationDiagnostics(operation), - pkRangeId, - this.operationContextText, - getThreadInfo()); - - return Pair.of(partitionScopeThresholds, operation); - }); - }) - .groupBy(Pair::getKey, Pair::getValue) - .flatMap( - this::executePartitionedGroup, - maxConcurrentCosmosPartitions) - .subscribeOn(this.executionScheduler) - .doOnNext(requestAndResponse -> { - - int totalCountAfterDecrement = totalCount.decrementAndGet(); - boolean mainSourceCompletedSnapshot = mainSourceCompleted.get(); - if (totalCountAfterDecrement == 0 && mainSourceCompletedSnapshot) { - // It is possible that count is zero but there are more elements in the source. - // Count 0 also signifies that there are no pending elements in any sink. - logDebugOrWarning("All work completed, {}, TotalCount: {}, Context: {} {}", - getItemOperationDiagnostics(requestAndResponse.getOperation()), - totalCountAfterDecrement, - this.operationContextText, - getThreadInfo()); - completeAllSinks(); - } else { - if (totalCountAfterDecrement == 0) { - logDebugOrWarning( - "No Work left - but mainSource not yet completed, Context: {} {}", - this.operationContextText, - getThreadInfo()); - } - logTraceOrWarning( - "Work left - TotalCount after decrement: {}, main sink completed {}, {}, Context: {} {}", - totalCountAfterDecrement, - mainSourceCompletedSnapshot, - getItemOperationDiagnostics(requestAndResponse.getOperation()), - this.operationContextText, - getThreadInfo()); - } - }) - .doOnComplete(() -> { - int totalCountSnapshot = totalCount.get(); - boolean mainSourceCompletedSnapshot = mainSourceCompleted.get(); - if (totalCountSnapshot == 0 && mainSourceCompletedSnapshot) { - // It is possible that count is zero but there are more elements in the source. - // Count 0 also signifies that there are no pending elements in any sink. - logDebugOrWarning("DoOnComplete: All work completed, Context: {}", this.operationContextText); - completeAllSinks(); - } else { - logDebugOrWarning( - "DoOnComplete: Work left - TotalCount after decrement: {}, main sink completed {}, Context: {} {}", - totalCountSnapshot, - mainSourceCompletedSnapshot, - this.operationContextText, - getThreadInfo()); - } - }); - }); - } - - private Flux> executePartitionedGroup( - GroupedFlux partitionedGroupFluxOfInputOperations) { - - final PartitionScopeThresholds thresholds = partitionedGroupFluxOfInputOperations.key(); + int maxConcurrentBatches = nullableMaxConcurrentCosmosPartitions != null ? + Math.max(256, nullableMaxConcurrentCosmosPartitions) : 256; - final FluxProcessor groupFluxProcessor = - UnicastProcessor.create().serialize(); - final FluxSink groupSink = groupFluxProcessor.sink(FluxSink.OverflowStrategy.BUFFER); - groupSinks.add(groupSink); - - AtomicReference currentPartitionKey = new AtomicReference<>(null); - AtomicLong currentMicroBatchSize = new AtomicLong(0); - - return partitionedGroupFluxOfInputOperations - .mergeWith(groupFluxProcessor) - .onBackpressureBuffer() - .timestamp() - .subscribeOn(this.executionScheduler) - .bufferUntil(timeStampItemOperationTuple -> { - long timestamp = timeStampItemOperationTuple.getT1(); - CosmosItemOperation itemOperation = timeStampItemOperationTuple.getT2(); + logDebugOrWarning("TransactionalBulkExecutor.executeCore with MaxConcurrentBatches: {}, Context: {}", + maxConcurrentBatches, + this.operationContextText); - logger.trace( - "TransactionalBulkExecutor.BufferUntil - enqueued {}, {}, Context: {} {}", - timestamp, - getItemOperationDiagnostics(itemOperation), + return this.inputBatches + .publishOn(this.executionScheduler) + .onErrorMap(throwable -> { + logger.warn("{}: Error observed when processing input batches. Cause: {}, Context: {}", + getThreadInfo(), + throwable.getMessage(), this.operationContextText, - getThreadInfo()); - - if (itemOperation == FlushBuffersItemOperation.singleton()) { - long currentMicroBatchSizeSnapshot = currentMicroBatchSize.get(); - if (currentMicroBatchSizeSnapshot > 0) { - logger.trace( - "TransactionalBulkExecutor - Flushing PKRange {} (batch size: {}) due to FlushItemOperation, Context: {} {}", - thresholds.getPartitionKeyRangeId(), - currentMicroBatchSizeSnapshot, - this.operationContextText, - getThreadInfo()); - - currentPartitionKey.set(null); - currentMicroBatchSize.set(0); - - return true; - } - - // avoid counting flush operations for the micro batch size calculation - return false; - } - - // For transactional batches, we flush whenever the logical partition key changes - PartitionKey itemPartitionKey = itemOperation.getPartitionKeyValue(); - PartitionKey currentPK = currentPartitionKey.get(); - long currentBatchSize = currentMicroBatchSize.get(); - - // Check if partition key has changed BEFORE incrementing batch size - boolean partitionKeyChanged = currentPK != null && !currentPK.equals(itemPartitionKey); - - if (partitionKeyChanged) { - logger.trace( - "TransactionalBulkExecutor - Partition key changed, PKRange {}, " + - "Previous PK: {}, New PK: {}, Previous batch size: {}, Context: {} {}", - thresholds.getPartitionKeyRangeId(), - currentPK, - itemPartitionKey, - currentBatchSize, - this.operationContextText, - getThreadInfo()); - - // Reset for the new batch - currentPartitionKey.set(itemPartitionKey); - currentMicroBatchSize.set(1); // Current item is first of new batch - - // Return true with cutBefore=true means: flush previous batch, start new batch with current item - return true; - } - - // First item sets the partition key - if (currentPK == null) { - currentPartitionKey.set(itemPartitionKey); - currentMicroBatchSize.set(1); - return false; - } - - // For transactional batches, we cannot split operations for the same partition key - // across multiple batches (would break atomicity). If we're about to exceed the limit, - // the operation will fail when we try to create the batch request. - // Just increment and let the batch creation fail with appropriate error. - currentMicroBatchSize.incrementAndGet(); - return false; - }, true) // cutBefore=true: when predicate is true, current element starts NEXT buffer + throwable); + return throwable; + }) .flatMap( - (List> timeStampAndItemOperationTuples) -> { - List operations = new ArrayList<>(timeStampAndItemOperationTuples.size()); - PartitionKey batchPartitionKey = null; - - for (Tuple2 timeStampAndItemOperationTuple : - timeStampAndItemOperationTuples) { - - CosmosItemOperation itemOperation = timeStampAndItemOperationTuple.getT2(); - if (itemOperation == FlushBuffersItemOperation.singleton()) { - continue; - } - operations.add(itemOperation); - - // Capture the partition key for this batch (all items should have the same PK) - if (batchPartitionKey == null) { - batchPartitionKey = itemOperation.getPartitionKeyValue(); - } - } - - if (operations.isEmpty()) { - logger.trace("Empty operations list after filtering, Context: {}", this.operationContextText); - return Flux.empty(); - } - - logDebugOrWarning( - "Flushing transactional batch for PKRange {} with {} operations, PK: {}, Context: {} {}", - thresholds.getPartitionKeyRangeId(), - operations.size(), - batchPartitionKey, - this.operationContextText, - getThreadInfo()); - - return executeTransactionalBatch(operations, batchPartitionKey, thresholds, groupSink); - }, - this.cosmosBulkExecutionOptions.getMaxMicroBatchConcurrency()); - } - - private int calculateTotalSerializedLength(AtomicInteger currentTotalSerializedLength, CosmosItemOperation item) { - if (item instanceof CosmosItemOperationBase) { - return currentTotalSerializedLength.accumulateAndGet( - ((CosmosItemOperationBase) item).getSerializedLength(this.effectiveItemSerializer), - Integer::sum); - } - - return currentTotalSerializedLength.get(); + cosmosBatch -> executeTransactionalBatch(cosmosBatch), + maxConcurrentBatches) + .subscribeOn(this.executionScheduler); } - private Flux> executeTransactionalBatch( - List operations, - PartitionKey partitionKey, - PartitionScopeThresholds thresholds, - FluxSink groupSink) { - - if (operations.size() == 0) { - logger.trace("Empty operations list, Context: {}", this.operationContextText); - return Flux.empty(); - } - - // Create a CosmosBatch for transactional execution - CosmosBatch cosmosBatch = CosmosBatch.createCosmosBatch(partitionKey); + private Flux> executeTransactionalBatch(CosmosBatch cosmosBatch) { + // Extract operations from the pre-built batch + List operations = cosmosBatch.getOperations(); - // Add all operations to the batch - for (CosmosItemOperation operation : operations) { - if (operation instanceof ItemBulkOperation) { - ItemBulkOperation itemBulkOperation = (ItemBulkOperation) operation; - CosmosBatchItemRequestOptions itemRequestOptions = itemBulkOperation.getRequestOptions() != null ? - new CosmosBatchItemRequestOptions() : null; - - switch (itemBulkOperation.getOperationType()) { - case CREATE: - cosmosBatch.createItemOperation( - itemBulkOperation.getItem(), - itemRequestOptions); - break; - case UPSERT: - cosmosBatch.upsertItemOperation( - itemBulkOperation.getItem(), - itemRequestOptions); - break; - case REPLACE: - cosmosBatch.replaceItemOperation( - itemBulkOperation.getId(), - itemBulkOperation.getItem(), - itemRequestOptions); - break; - case DELETE: - cosmosBatch.deleteItemOperation( - itemBulkOperation.getId(), - itemRequestOptions); - break; - case READ: - cosmosBatch.readItemOperation( - itemBulkOperation.getId(), - itemRequestOptions); - break; - case PATCH: - // For PATCH, the item is actually CosmosPatchOperations - Object patchItem = itemBulkOperation.getItemInternal(); - if (patchItem instanceof com.azure.cosmos.models.CosmosPatchOperations) { - com.azure.cosmos.models.CosmosBatchPatchItemRequestOptions patchRequestOptions = - itemRequestOptions != null ? new com.azure.cosmos.models.CosmosBatchPatchItemRequestOptions() : null; - cosmosBatch.patchItemOperation( - itemBulkOperation.getId(), - (com.azure.cosmos.models.CosmosPatchOperations) patchItem, - patchRequestOptions); - } - break; - default: - logger.warn("Unsupported operation type: {}", itemBulkOperation.getOperationType()); - break; - } - } + if (operations == null || operations.isEmpty()) { + logger.trace("Empty batch received, Context: {}", this.operationContextText); + return Flux.empty(); } - CosmosBatchRequestOptions batchRequestOptions = new CosmosBatchRequestOptions(); - - // Execute the transactional batch - return Flux.just(cosmosBatch) - .publishOn(this.executionScheduler) - .flatMap((CosmosBatch batch) -> - this.executeTransactionalBatchRequest(batch, batchRequestOptions, operations, groupSink, thresholds)); - } - - private Flux> executeTransactionalBatchRequest( - CosmosBatch cosmosBatch, - CosmosBatchRequestOptions requestOptions, - List operations, - FluxSink groupSink, - PartitionScopeThresholds thresholds) { - String batchTrackingId = UUIDs.nonBlockingRandomUUID().toString(); - logTraceOrWarning( - "Executing transactional batch - batch TrackingId: %s", - batchTrackingId); - - // Validate that transactional batch does not exceed Cosmos DB limit - if (operations.size() > BatchRequestResponseConstants.MAX_OPERATIONS_IN_DIRECT_MODE_BATCH_REQUEST) { - String errorMessage = String.format( - "Transactional batch operation failed: partition key '%s' has %d operations, " + - "which exceeds the maximum allowed limit of %d operations per transactional batch. " + - "Transactional batches require all-or-nothing execution and cannot be split across multiple requests.", - cosmosBatch.getPartitionKeyValue(), - operations.size(), - BatchRequestResponseConstants.MAX_OPERATIONS_IN_DIRECT_MODE_BATCH_REQUEST); - - logger.error(errorMessage); - - // Return error responses for all operations in this batch - return Flux.fromIterable(operations) - .map(operation -> { - Exception exception = new IllegalArgumentException(errorMessage); - TContext actualContext = this.getActualContext(operation); - return ModelBridgeInternal.createCosmosBulkOperationResponse( - operation, - exception, - actualContext); - }); - } + PartitionKey partitionKey = cosmosBatch.getPartitionKeyValue(); + + logDebugOrWarning( + "Executing transactional batch - {} operations, PK: {}, TrackingId: {}, Context: {}", + operations.size(), + partitionKey, + batchTrackingId, + this.operationContextText); // Create SinglePartitionKeyServerBatchRequest with atomic batch semantics final SinglePartitionKeyServerBatchRequest request = SinglePartitionKeyServerBatchRequest.createBatchRequest( - cosmosBatch.getPartitionKeyValue(), + partitionKey, operations, this.effectiveItemSerializer); request.setAtomicBatch(true); @@ -748,120 +329,87 @@ private Flux> executeTransactionalBatchReq false, false) .subscribeOn(this.executionScheduler) - .flatMapMany(response -> { - - logTraceOrWarning( - "Response for transactional batch - status code %s, ActivityId: %s, batch TrackingId %s", - response.getStatusCode(), - response.getActivityId(), - batchTrackingId); + .flatMapMany(cosmosBatchResponse -> { + List results = cosmosBatchResponse.getResults(); + + if (results == null || results.size() != operations.size()) { + String errorMessage = String.format( + "Transactional batch response mismatch: expected %d results, got %d", + operations.size(), + results != null ? results.size() : 0); + logger.error("{}, Context: {}", errorMessage, this.operationContextText); - if (diagnosticsTracker != null && response.getDiagnostics() != null) { - diagnosticsTracker.trackDiagnostics(response.getDiagnostics().getDiagnosticsContext()); - } - - return Flux - .fromIterable(response.getResults()) - .publishOn(this.executionScheduler) - .flatMap((CosmosBatchOperationResult result) -> - handleTransactionalBatchOperationResult(response, result, groupSink, thresholds)); - }) - .onErrorResume((Throwable throwable) -> { - - if (!(throwable instanceof Exception)) { - throw Exceptions.propagate(throwable); + // Return error responses for all operations + return Flux.fromIterable(operations) + .map(operation -> { + Exception exception = new IllegalStateException(errorMessage); + TContext actualContext = this.getActualContext(operation); + return ModelBridgeInternal.createCosmosBulkOperationResponse( + operation, + exception, + actualContext); + }); } - - Exception exception = (Exception) throwable; - - return Flux - .fromIterable(operations) - .publishOn(this.executionScheduler) - .flatMap((CosmosItemOperation itemOperation) -> - handleTransactionalBatchExecutionException(itemOperation, exception, groupSink, thresholds)); - }); - } - - // Helper functions - private Mono> handleTransactionalBatchOperationResult( - CosmosBatchResponse response, - CosmosBatchOperationResult operationResult, - FluxSink groupSink, - PartitionScopeThresholds thresholds) { - - CosmosBulkItemResponse cosmosBulkItemResponse = ModelBridgeInternal - .createCosmosBulkItemResponse(operationResult, response); - CosmosItemOperation itemOperation = operationResult.getOperation(); - TContext actualContext = this.getActualContext(itemOperation); - - logDebugOrWarning( - "HandleTransactionalBatchOperationResult - PKRange {}, Response Status Code {}, " + - "Operation Status Code, {}, {}, Context: {} {}", - thresholds.getPartitionKeyRangeId(), - response.getStatusCode(), - operationResult.getStatusCode(), - getItemOperationDiagnostics(itemOperation), - this.operationContextText, - getThreadInfo()); - - if (!operationResult.isSuccessStatusCode()) { - - if (itemOperation instanceof ItemBulkOperation) { - - ItemBulkOperation itemBulkOperation = (ItemBulkOperation) itemOperation; - return itemBulkOperation.getRetryPolicy().shouldRetry(operationResult).flatMap( - result -> { - if (result.shouldRetry) { + + // Map results back to operations + return Flux.range(0, operations.size()) + .map(index -> { + CosmosItemOperation operation = operations.get(index); + CosmosBatchOperationResult operationResult = results.get(index); + TContext actualContext = this.getActualContext(operation); + + CosmosBulkItemResponse cosmosBulkItemResponse = ModelBridgeInternal.createCosmosBulkItemResponse( + operationResult, + cosmosBatchResponse); + + if (!operationResult.isSuccessStatusCode()) { logDebugOrWarning( - "HandleTransactionalBatchOperationResult - enqueue retry, PKRange {}, Response " + - "Status Code {}, Operation Status Code, {}, {}, Context: {} {}", - thresholds.getPartitionKeyRangeId(), - response.getStatusCode(), + "Transactional batch operation failed - PK: {}, Status: {}, Operation: {}, Context: {}", + String.valueOf(partitionKey), operationResult.getStatusCode(), - getItemOperationDiagnostics(itemOperation), - this.operationContextText, - getThreadInfo()); - return this.enqueueForRetry(result.backOffTime, groupSink, itemOperation, thresholds); - } else { - // reduce log noise level for commonly expected/normal status codes - if (response.getStatusCode() == HttpConstants.StatusCodes.CONFLICT || - response.getStatusCode() == HttpConstants.StatusCodes.PRECONDITION_FAILED) { - - logDebugOrWarning( - "HandleTransactionalBatchOperationResult - Fail, PKRange {}, Response Status " + - "Code {}, Operation Status Code {}, {}, Context: {} {}", - thresholds.getPartitionKeyRangeId(), - response.getStatusCode(), - operationResult.getStatusCode(), - getItemOperationDiagnostics(itemOperation), - this.operationContextText, - getThreadInfo()); - } else { - logger.error( - "HandleTransactionalBatchOperationResult - Fail, PKRange {}, Response Status " + - "Code {}, Operation Status Code {}, {}, Context: {} {}", - thresholds.getPartitionKeyRangeId(), - response.getStatusCode(), - operationResult.getStatusCode(), - getItemOperationDiagnostics(itemOperation), - this.operationContextText, - getThreadInfo()); - } - return Mono.just(ModelBridgeInternal.createCosmosBulkOperationResponse( - itemOperation, cosmosBulkItemResponse, actualContext)); + getItemOperationDiagnostics(operation), + this.operationContextText); } + + return ModelBridgeInternal.createCosmosBulkOperationResponse( + operation, + cosmosBulkItemResponse, + actualContext); }); + }) + .onErrorResume(throwable -> { + logger.error( + "Transactional batch execution failed - PK: {}, Error: {}, Context: {}", + String.valueOf(partitionKey), + throwable.getMessage(), + this.operationContextText, + throwable); + + // Convert Throwable to Exception if needed + Exception exception = throwable instanceof Exception ? + (Exception) throwable : + new RuntimeException(throwable); + + // Return error responses for all operations in the batch + return Flux.fromIterable(operations) + .map(operation -> { + TContext actualContext = this.getActualContext(operation); + return ModelBridgeInternal.createCosmosBulkOperationResponse( + operation, + exception, + actualContext); + }); + }); + } - } else { - throw new UnsupportedOperationException("Unknown CosmosItemOperation."); - } + private int calculateTotalSerializedLength(AtomicInteger currentTotalSerializedLength, CosmosItemOperation item) { + if (item instanceof CosmosItemOperationBase) { + return currentTotalSerializedLength.accumulateAndGet( + ((CosmosItemOperationBase) item).getSerializedLength(this.effectiveItemSerializer), + Integer::sum); } - thresholds.recordSuccessfulOperation(); - return Mono.just(ModelBridgeInternal.createCosmosBulkOperationResponse( - itemOperation, - cosmosBulkItemResponse, - actualContext)); + return currentTotalSerializedLength.get(); } private TContext getActualContext(CosmosItemOperation itemOperation) { @@ -883,215 +431,12 @@ private TContext getActualContext(CosmosItemOperation itemOperation) { return this.batchContext; } - private Mono> handleTransactionalBatchExecutionException( - CosmosItemOperation itemOperation, - Exception exception, - FluxSink groupSink, - PartitionScopeThresholds thresholds) { - - logDebugOrWarning( - "HandleTransactionalBatchExecutionException, PKRange {}, Error: {}, {}, Context: {} {}", - thresholds.getPartitionKeyRangeId(), - exception, - getItemOperationDiagnostics(itemOperation), - this.operationContextText, - getThreadInfo()); - - if (exception instanceof CosmosException && itemOperation instanceof ItemBulkOperation) { - CosmosException cosmosException = (CosmosException) exception; - ItemBulkOperation itemBulkOperation = (ItemBulkOperation) itemOperation; - - // First check if it failed due to split, so the operations need to go in a different pk range group. So - // add it in the mainSink. - - return itemBulkOperation.getRetryPolicy() - .shouldRetryInMainSink(cosmosException.getStatusCode(), cosmosException.getSubStatusCode(), itemBulkOperation, cosmosException) - .flatMap(shouldRetryInMainSink -> { - if (shouldRetryInMainSink) { - logDebugOrWarning( - "HandleTransactionalBatchExecutionException - Retry in main sink, PKRange {}, Error: " + - "{}, {}, Context: {} {}", - thresholds.getPartitionKeyRangeId(), - exception, - getItemOperationDiagnostics(itemOperation), - this.operationContextText, - getThreadInfo()); - // retry - but don't mark as enqueued for retry in thresholds - mainSink.emitNext(itemOperation, serializedEmitFailureHandler); - return Mono.empty(); - } else { - logDebugOrWarning( - "HandleTransactionalBatchExecutionException - Retry other, PKRange {}, Error: " + - "{}, {}, Context: {} {}", - thresholds.getPartitionKeyRangeId(), - exception, - getItemOperationDiagnostics(itemOperation), - this.operationContextText, - getThreadInfo()); - return retryOtherExceptions( - itemOperation, - exception, - groupSink, - cosmosException, - itemBulkOperation, - thresholds); - } - }); - } - - TContext actualContext = this.getActualContext(itemOperation); - return Mono.just(ModelBridgeInternal.createCosmosBulkOperationResponse(itemOperation, exception, actualContext)); - } - - private Mono> enqueueForRetry( - Duration backOffTime, - FluxSink groupSink, - CosmosItemOperation itemOperation, - PartitionScopeThresholds thresholds) { - - thresholds.recordEnqueuedRetry(); - if (backOffTime == null || backOffTime.isZero()) { - groupSink.next(itemOperation); - return Mono.empty(); - } else { - return Mono - .delay(backOffTime) - .flatMap((dummy) -> { - groupSink.next(itemOperation); - return Mono.empty(); - }); - } - } - - private Mono> retryOtherExceptions( - CosmosItemOperation itemOperation, - Exception exception, - FluxSink groupSink, - CosmosException cosmosException, - ItemBulkOperation itemBulkOperation, - PartitionScopeThresholds thresholds) { - - TContext actualContext = this.getActualContext(itemOperation); - return itemBulkOperation.getRetryPolicy().shouldRetry(cosmosException).flatMap(result -> { - if (result.shouldRetry) { - return this.enqueueForRetry(result.backOffTime, groupSink, itemBulkOperation, thresholds); - } else { - return Mono.just(ModelBridgeInternal.createCosmosBulkOperationResponse( - itemOperation, exception, actualContext)); - } - }); - } - - private Mono executeBatchRequest( - PartitionKeyRangeServerBatchRequest serverRequest, - PartitionScopeThresholds partitionScopeThresholds) { - - RequestOptions options = new RequestOptions(); - options.setThroughputControlGroupName(cosmosBulkExecutionOptions.getThroughputControlGroupName()); - options.setExcludedRegions(cosmosBulkExecutionOptions.getExcludedRegions()); - options.setKeywordIdentifiers(cosmosBulkExecutionOptions.getKeywordIdentifiers()); - - CosmosEndToEndOperationLatencyPolicyConfig e2eLatencyPolicySnapshot = - cosmosBulkExecutionOptions.getCosmosEndToEndLatencyPolicyConfig(); - if (e2eLatencyPolicySnapshot != null) { - options.setCosmosEndToEndLatencyPolicyConfig(e2eLatencyPolicySnapshot); - } - - // This logic is to handle custom bulk options which can be passed through encryption or through some other project - Map customOptions = cosmosBulkExecutionOptions.getHeaders(); - if (customOptions != null && !customOptions.isEmpty()) { - for(Map.Entry entry : customOptions.entrySet()) { - options.setHeader(entry.getKey(), entry.getValue()); - } - } - options.setOperationContextAndListenerTuple(operationListener); - - // The request options here are used for the BulkRequest exchanged with the service - // If contentResponseOnWrite is not enabled here (or at the client level) the - // service will not even send a bulk response payload - so all the - // CosmosBulItemRequestOptions are irrelevant - all payloads will be null - // Instead we should automatically enforce contentResponseOnWrite for all - // bulk requests whenever at least one of the item operations requires a content response (either - // because it is a read operation or because contentResponseOnWrite was enabled explicitly) - if (!this.docClientWrapper.isContentResponseOnWriteEnabled() && - serverRequest.getOperations().size() > 0) { - - for (CosmosItemOperation itemOperation : serverRequest.getOperations()) { - if (itemOperation instanceof ItemBulkOperation) { - - ItemBulkOperation itemBulkOperation = (ItemBulkOperation) itemOperation; - if (itemBulkOperation.getOperationType() == CosmosItemOperationType.READ || - (itemBulkOperation.getRequestOptions() != null && - itemBulkOperation.getRequestOptions().isContentResponseOnWriteEnabled() != null && - itemBulkOperation.getRequestOptions().isContentResponseOnWriteEnabled())) { - - options.setContentResponseOnWriteEnabled(true); - break; - } - } - } - } - - return withContext(context -> { - final Mono responseMono = - this.docClientWrapper - .executeBatchRequest( - BridgeInternal.getLink(this.container), - serverRequest, - options, - false, - true) // disable the staled resource exception handling as it is being handled in the BulkOperationRetryPolicy - .flatMap(cosmosBatchResponse -> { - - cosmosBatchResponseAccessor.setGlobalOpCount( - cosmosBatchResponse, partitionScopeThresholds.getTotalOperationCountSnapshot()); - - PartitionScopeThresholds.CurrentIntervalThresholds currentIntervalThresholdsSnapshot - = partitionScopeThresholds.getCurrentThresholds(); - - cosmosBatchResponseAccessor.setOpCountPerEvaluation( - cosmosBatchResponse, currentIntervalThresholdsSnapshot.currentOperationCount.get()); - cosmosBatchResponseAccessor.setRetriedOpCountPerEvaluation( - cosmosBatchResponse, currentIntervalThresholdsSnapshot.currentRetriedOperationCount.get()); - cosmosBatchResponseAccessor.setTargetMaxMicroBatchSize( - cosmosBatchResponse, partitionScopeThresholds.getTargetMicroBatchSizeSnapshot()); - - return Mono.just(cosmosBatchResponse); - }); - - return clientAccessor.getDiagnosticsProvider(this.cosmosClient) - .traceEnabledBatchResponsePublisher( - responseMono, - context, - this.bulkSpanName, - this.container.getDatabase().getId(), - this.container.getId(), - this.cosmosClient, - options.getConsistencyLevel(), - OperationType.Batch, - ResourceType.Document, - options, - this.cosmosBulkExecutionOptions.getMaxMicroBatchSize()); - }); - } - private void completeAllSinks() { - logInfoOrWarning("Closing all sinks, Context: {}", this.operationContextText); - + logInfoOrWarning("Completing execution, Context: {}", this.operationContextText); logger.debug("Executor service shut down, Context: {}", this.operationContextText); - mainSink.emitComplete(serializedCompleteEmitFailureHandler); - this.shutdown(); } - private void onFlush() { - try { - this.groupSinks.forEach(sink -> sink.next(FlushBuffersItemOperation.singleton())); - } catch(Throwable t) { - logger.error("Callback invocation 'onFlush' failed. Context: {}", this.operationContextText, t); - } - } - private static String getItemOperationDiagnostics(CosmosItemOperation operation) { if (operation == FlushBuffersItemOperation.singleton()) { @@ -1124,39 +469,4 @@ private static String getThreadInfo() { return sb.toString(); } - - private static class SerializedEmitFailureHandler implements Sinks.EmitFailureHandler { - - @Override - public boolean onEmitFailure(SignalType signalType, Sinks.EmitResult emitResult) { - if (emitResult.equals(Sinks.EmitResult.FAIL_NON_SERIALIZED)) { - logger.debug("SerializedEmitFailureHandler.onEmitFailure - Signal:{}, Result: {}", signalType, emitResult); - - return true; - } - - logger.error("SerializedEmitFailureHandler.onEmitFailure - Signal:{}, Result: {}", signalType, emitResult); - return false; - } - } - - private static class SerializedCompleteEmitFailureHandler implements Sinks.EmitFailureHandler { - - @Override - public boolean onEmitFailure(SignalType signalType, Sinks.EmitResult emitResult) { - if (emitResult.equals(Sinks.EmitResult.FAIL_NON_SERIALIZED)) { - logger.debug("SerializedCompleteEmitFailureHandler.onEmitFailure - Signal:{}, Result: {}", signalType, emitResult); - - return true; - } - - if (emitResult == Sinks.EmitResult.FAIL_CANCELLED || emitResult == Sinks.EmitResult.FAIL_TERMINATED) { - logger.debug("SerializedCompleteEmitFailureHandler.onEmitFailure - Main sink already completed, Signal:{}, Result: {}", signalType, emitResult); - return false; - } - - logger.error("SerializedCompleteEmitFailureHandler.onEmitFailure - Signal:{}, Result: {}", signalType, emitResult); - return false; - } - } } From 2b0f09dac16c9049a2b46be9c9dc51055b54dba8 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Wed, 10 Dec 2025 13:57:54 +0000 Subject: [PATCH 05/16] Address PR feedback: fix resource leak, improve logging, clarify docs --- .../cosmos/spark/ItemsWriterBuilder.scala | 95 ++++++++++++------- .../docs/configuration-reference.md | 2 +- .../docs/scenarios/Ingestion.md | 20 +++- .../com/azure/cosmos/spark/CosmosConfig.scala | 7 +- .../batch/TransactionalBulkExecutor.java | 2 +- 5 files changed, 82 insertions(+), 44 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3-5/src/main/scala/com/azure/cosmos/spark/ItemsWriterBuilder.scala b/sdk/cosmos/azure-cosmos-spark_3-5/src/main/scala/com/azure/cosmos/spark/ItemsWriterBuilder.scala index a7260319bdb2..77e67a2a5f35 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5/src/main/scala/com/azure/cosmos/spark/ItemsWriterBuilder.scala +++ b/sdk/cosmos/azure-cosmos-spark_3-5/src/main/scala/com/azure/cosmos/spark/ItemsWriterBuilder.scala @@ -59,18 +59,21 @@ private class ItemsWriterBuilder ) } + // Extract userConfig conversion to avoid repeated calls + private[this] val userConfigMap = userConfig.asCaseSensitiveMap().asScala.toMap + private[this] val writeConfig = CosmosWriteConfig.parseWriteConfig( - userConfig.asCaseSensitiveMap().asScala.toMap, + userConfigMap, inputSchema ) private[this] val containerConfig = CosmosContainerConfig.parseCosmosContainerConfig( - userConfig.asCaseSensitiveMap().asScala.toMap + userConfigMap ) override def toBatch(): BatchWrite = new ItemsBatchWriter( - userConfig.asCaseSensitiveMap().asScala.toMap, + userConfigMap, inputSchema, cosmosClientStateHandles, diagnosticsConfig, @@ -78,7 +81,7 @@ private class ItemsWriterBuilder override def toStreaming: StreamingWrite = new ItemsBatchWriter( - userConfig.asCaseSensitiveMap().asScala.toMap, + userConfigMap, inputSchema, cosmosClientStateHandles, diagnosticsConfig, @@ -88,6 +91,7 @@ private class ItemsWriterBuilder override def requiredDistribution(): Distribution = { if (writeConfig.bulkEnabled && writeConfig.bulkTransactional) { + log.logInfo("Transactional batch mode enabled - configuring data distribution by partition key columns") // For transactional writes, partition by all partition key columns val partitionKeyPaths = getPartitionKeyColumnNames() if (partitionKeyPaths.nonEmpty) { @@ -125,38 +129,20 @@ private class ItemsWriterBuilder private def getPartitionKeyColumnNames(): Seq[String] = { try { - // Need to create a temporary container client to get partition key definition - val clientCacheItem = CosmosClientCache( - CosmosClientConfiguration( - userConfig.asCaseSensitiveMap().asScala.toMap, - ReadConsistencyStrategy.EVENTUAL, - sparkEnvironmentInfo - ), - Some(cosmosClientStateHandles.value.cosmosClientMetadataCaches), - "ItemsWriterBuilder-PKLookup" - ) - - val container = ThroughputControlHelper.getContainer( - userConfig.asCaseSensitiveMap().asScala.toMap, - containerConfig, - clientCacheItem, - None - ) - - val containerProperties = SparkBridgeInternal.getContainerPropertiesFromCollectionCache(container) - val partitionKeyDefinition = containerProperties.getPartitionKeyDefinition - - // Release the client - clientCacheItem.close() - - if (partitionKeyDefinition != null && partitionKeyDefinition.getPaths != null) { - val paths = partitionKeyDefinition.getPaths.asScala - paths.map(path => { - // Remove leading '/' from partition key path (e.g., "/pk" -> "pk") - if (path.startsWith("/")) path.substring(1) else path - }).toSeq - } else { - Seq.empty[String] + // Use loan pattern to ensure client is properly closed + using(createClientForPartitionKeyLookup()) { clientCacheItem => + val container = ThroughputControlHelper.getContainer( + userConfigMap, + containerConfig, + clientCacheItem, + None + ) + + // Simplified retrieval using SparkBridgeInternal directly + val containerProperties = SparkBridgeInternal.getContainerPropertiesFromCollectionCache(container) + val partitionKeyDefinition = containerProperties.getPartitionKeyDefinition + + extractPartitionKeyPaths(partitionKeyDefinition) } } catch { case ex: Exception => @@ -164,5 +150,42 @@ private class ItemsWriterBuilder Seq.empty[String] } } + + private def createClientForPartitionKeyLookup(): CosmosClientCacheItem = { + CosmosClientCache( + CosmosClientConfiguration( + userConfigMap, + ReadConsistencyStrategy.EVENTUAL, + sparkEnvironmentInfo + ), + Some(cosmosClientStateHandles.value.cosmosClientMetadataCaches), + "ItemsWriterBuilder-PKLookup" + ) + } + + private def extractPartitionKeyPaths(partitionKeyDefinition: com.azure.cosmos.models.PartitionKeyDefinition): Seq[String] = { + if (partitionKeyDefinition != null && partitionKeyDefinition.getPaths != null) { + val paths = partitionKeyDefinition.getPaths.asScala + if (paths.isEmpty) { + log.logError("Partition key definition has 0 columns - this should not happen for modern containers") + } + paths.map(path => { + // Remove leading '/' from partition key path (e.g., "/pk" -> "pk") + if (path.startsWith("/")) path.substring(1) else path + }).toSeq + } else { + log.logError("Partition key definition is null - this should not happen for modern containers") + Seq.empty[String] + } + } + + // Scala loan pattern to ensure resources are properly cleaned up + private def using[A <: { def close(): Unit }, B](resource: A)(f: A => B): B = { + try { + f(resource) + } finally { + if (resource != null) resource.close() + } + } } } diff --git a/sdk/cosmos/azure-cosmos-spark_3/docs/configuration-reference.md b/sdk/cosmos/azure-cosmos-spark_3/docs/configuration-reference.md index e25cb78ab9df..53357f1e22f8 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/docs/configuration-reference.md +++ b/sdk/cosmos/azure-cosmos-spark_3/docs/configuration-reference.md @@ -67,7 +67,7 @@ | `spark.cosmos.write.point.maxConcurrency` | None | Cosmos DB Item Write Max concurrency. If not specified it will be determined based on the Spark executor VM Size | | `spark.cosmos.write.bulk.maxPendingOperations` | None | Cosmos DB Item Write bulk mode maximum pending operations. Defines a limit of bulk operations being processed concurrently. If not specified it will be determined based on the Spark executor VM Size. If the volume of data is large for the provisioned throughput on the destination container, this setting can be adjusted by following the estimation of `1000 x Cores` | | `spark.cosmos.write.bulk.enabled` | `true` | Cosmos DB Item Write bulk enabled | -| `spark.cosmos.write.bulk.transactional` | `false` | Enable transactional batch mode for bulk writes. When enabled, all operations for the same partition key are executed atomically (all succeed or all fail). Requires ordering and clustering by partition key columns. Only supports upsert operations. Cannot exceed 100 operations or 2MB per partition key. See [Transactional Batch documentation](https://learn.microsoft.com/azure/cosmos-db/nosql/transactional-batch) for details. |\n| `spark.cosmos.write.bulk.targetedPayloadSizeInBytes` | `220201` | When the targeted payload size is reached for buffered documents, the request is sent to the backend. The default value is optimized for small documents <= 10 KB - when documents often exceed 110 KB, it can help to increase this value to up to about `1500000` (should still be smaller than 2 MB). | +| `spark.cosmos.write.bulk.transactional` | `false` | Enable transactional batch mode for bulk writes. When enabled, all operations for the same partition key are executed atomically (all succeed or all fail). Requires ordering and clustering by partition key columns. Only supports upsert operations. Cannot exceed 100 operations or 2MB per partition key. **Note**: For containers using hierarchical partition keys (HPK), transactional scope applies only to **logical partitions** (complete partition key paths), not partial top-level keys. See [Transactional Batch documentation](https://learn.microsoft.com/azure/cosmos-db/nosql/transactional-batch) for details. |\n| `spark.cosmos.write.bulk.targetedPayloadSizeInBytes` | `220201` | When the targeted payload size is reached for buffered documents, the request is sent to the backend. The default value is optimized for small documents <= 10 KB - when documents often exceed 110 KB, it can help to increase this value to up to about `1500000` (should still be smaller than 2 MB). | | `spark.cosmos.write.bulk.initialBatchSize` | `100` | Cosmos DB initial bulk micro batch size - a micro batch will be flushed to the backend when the number of documents enqueued exceeds this size - or the target payload size is met. The micro batch size is getting automatically tuned based on the throttling rate. By default the initial micro batch size is 100. Reduce this when you want to avoid that the first few requests consume too many RUs. | | `spark.cosmos.write.bulk.maxBatchSize` | `100` | Cosmos DB max. bulk micro batch size - a micro batch will be flushed to the backend when the number of documents enqueued exceeds this size - or the target payload size is met. The micro batch size is getting automatically tuned based on the throttling rate. By default the max. micro batch size is 100. Use this setting only when migrating Spark 2.4 workloads - for other scenarios relying on the auto-tuning combined with throughput control will result in better experience. | | `spark.cosmos.write.flush.noProgress.maxIntervalInSeconds` | `180` | The time interval in seconds that write operations will wait when no progress can be made for bulk writes before forcing a retry. The retry will reinitialize the bulk write process - so, any delays on the retry can be sure to be actual service issues. The default value of 3 min should be sufficient to prevent false negatives when there is a short service-side write unavailability - like for partition splits or merges. Increase it only if you regularly see these transient errors to exceed a time period of 180 seconds. | diff --git a/sdk/cosmos/azure-cosmos-spark_3/docs/scenarios/Ingestion.md b/sdk/cosmos/azure-cosmos-spark_3/docs/scenarios/Ingestion.md index 7056ffa1d883..b7a206aec6d6 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/docs/scenarios/Ingestion.md +++ b/sdk/cosmos/azure-cosmos-spark_3/docs/scenarios/Ingestion.md @@ -118,21 +118,35 @@ df.write - **Atomic semantics**: All operations for the same partition key succeed or all fail (rollback) - **Operation type**: Only upsert operations are supported (equivalent to `ItemOverwrite` write strategy) - **Partition grouping**: Spark automatically partitions and orders data by partition key columns -- **Size limits**: Maximum 100 operations per transaction; maximum 2MB total payload per transaction +- **Size limits**: Maximum 100 operations per transaction; maximum 2MB total payload per transaction. These limits are **enforced service-side by Azure Cosmos DB** and cannot be configured client-side. - **Partition key requirement**: All operations in a transaction must share the same partition key value +- **Hierarchical partition keys**: For containers using hierarchical partition keys (HPK), transactional scope applies only to **logical partitions** (complete partition key paths), not partial top-level keys - **Bulk mode required**: Must have `spark.cosmos.write.bulk.enabled=true` (enabled by default) #### Use cases Transactional batch writes are ideal for: -- Financial transactions requiring consistency across multiple documents +- Financial transactions requiring consistency across multiple documents (e.g., double-entry bookkeeping, account balance updates) - Order processing where order header and line items must be committed together - Multi-document updates that must be atomic (e.g., inventory adjustments) +- **Temporal data scenarios** where historical versioning must be consistent (e.g., financial instrument versioning, audit trails, snapshot consistency) - Any scenario where partial success would leave data in an inconsistent state #### Error handling -If any operation in a transaction fails (e.g., insufficient RUs, document too large, transaction exceeds 100 operations), the entire transaction is rolled back and no documents are modified. The Spark task will fail and retry according to Spark's retry policy. +Errors during transactional batch execution fall into two categories: + +**Transient errors** (automatically retried): +- Insufficient RUs (429 throttling) - the entire transaction is retried according to the SDK's retry policy +- Temporary network issues or service unavailability - automatic retry with exponential backoff + +**Non-transient errors** (Spark job fails): +- Transaction size limit exceeded (> 100 operations or > 2MB payload) +- Document validation errors (e.g., blank ID, schema violations) +- Unique key constraint violations +- Any service-side rejection that cannot be resolved by retry + +When any operation in a transaction fails with a non-transient error, the entire transaction is rolled back and no documents are modified. The Spark task will fail and may retry according to Spark's retry policy. ## Preparation Below are a couple of tips/best-practices that can help you to prepare for a data migration into a Cosmos DB container. diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala index 96e26eaf0f1a..448f7d2c91f1 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala @@ -1490,10 +1490,12 @@ private object CosmosWriteConfig { helpMessage = "Cosmos DB Item Write bulk enabled") private val bulkTransactional = CosmosConfigEntry[Boolean](key = CosmosConfigNames.WriteBulkTransactional, - defaultValue = Option.apply(false), mandatory = false, + defaultValue = Option.apply(false), parseFromStringFunction = bulkTransactionalAsString => bulkTransactionalAsString.toBoolean, - helpMessage = "Cosmos DB Item Write bulk transactional batch mode enabled") + helpMessage = "Cosmos DB Item Write bulk transactional batch mode enabled - requires bulk write to be enabled. " + + "Spark 3.5+ provides automatic distribution/ordering for transactional batch. " + + "On Spark 3.3/3.4, transactional batch is supported but requires manual sorting for optimal performance.") private val microBatchPayloadSizeInBytes = CosmosConfigEntry[Int](key = CosmosConfigNames.WriteBulkPayloadSizeInBytes, defaultValue = Option.apply(BatchRequestResponseConstants.DEFAULT_MAX_DIRECT_MODE_BATCH_REQUEST_BODY_SIZE_IN_BYTES), @@ -1778,7 +1780,6 @@ private object CosmosWriteConfig { .parse(cfg, writeOnRetryCommitInterceptor).flatten assert(bulkEnabledOpt.isDefined, s"Parameter '${CosmosConfigNames.WriteBulkEnabled}' is missing.") - assert(bulkTransactionalOpt.isDefined, s"Parameter '${CosmosConfigNames.WriteBulkTransactional}' is missing.") // parsing above already validated this assert(itemWriteStrategyOpt.isDefined, s"Parameter '${CosmosConfigNames.WriteStrategy}' is missing.") diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/TransactionalBulkExecutor.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/TransactionalBulkExecutor.java index f67cf44109ff..fe7d4cd50418 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/TransactionalBulkExecutor.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/TransactionalBulkExecutor.java @@ -164,7 +164,7 @@ public TransactionalBulkExecutor(CosmosAsyncContainer container, ? schedulerSnapshotFromOptions : CosmosSchedulers.BULK_EXECUTOR_BOUNDED_ELASTIC; - logger.debug("Instantiated BulkExecutor, Context: {}", + logger.info("Instantiated BulkExecutor, Context: {}", this.operationContextText); } From 5dbae894a6c9a2bc144f757a5bd2a1aa2faaf472 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Wed, 10 Dec 2025 15:12:14 +0000 Subject: [PATCH 06/16] Remove read many parts --- .../spark/TransactionalBulkWriter.scala | 510 ++---------------- 1 file changed, 33 insertions(+), 477 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala index b62c1d8d57ec..3292eb3b69d9 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala @@ -7,7 +7,7 @@ import com.azure.cosmos.implementation.apachecommons.lang.StringUtils import com.azure.cosmos.implementation.batch.{BatchRequestResponseConstants, BulkExecutorDiagnosticsTracker, ItemBulkOperation, TransactionalBulkExecutor} import com.azure.cosmos.implementation.{CosmosDaemonThreadFactory, CosmosBulkExecutionOptionsImpl, ImplementationBridgeHelpers, UUIDs} import com.azure.cosmos.models._ -import com.azure.cosmos.spark.TransactionalBulkWriter.{BulkOperationFailedException, bulkWriterInputBoundedElastic, bulkWriterRequestsBoundedElastic, bulkWriterResponsesBoundedElastic, getThreadInfo, readManyBoundedElastic} +import com.azure.cosmos.spark.TransactionalBulkWriter.{BulkOperationFailedException, bulkWriterInputBoundedElastic, bulkWriterRequestsBoundedElastic, bulkWriterResponsesBoundedElastic, getThreadInfo} import com.azure.cosmos.spark.diagnostics.DefaultDiagnostics import com.azure.cosmos.{BridgeInternal, CosmosAsyncContainer, CosmosDiagnosticsContext, CosmosEndToEndOperationLatencyPolicyConfigBuilder, CosmosException} import reactor.core.Scannable @@ -88,25 +88,12 @@ private class TransactionalBulkWriter s"BulkWriter instantiated (Host CPU count: $cpuCount, maxPendingOperations: $maxPendingOperations, " + s"maxConcurrentPartitions: $maxConcurrentPartitions ...") - // Artificial operation used to signale to the bufferUntil operator that - // the buffer should be flushed. A timer-based scheduler will publish this - // dummy operation for every batchIntervalInMs ms. This operation - // is filtered out and will never be flushed to the backend - private val readManyFlushOperationSingleton = ReadManyOperation( - new CosmosItemIdentity( - new PartitionKey("ReadManyOperation.FlushSingleton"), - "ReadManyOperation.FlushSingleton" - ), - null, - null - ) private val closed = new AtomicBoolean(false) private val lock = new ReentrantLock private val pendingTasksCompleted = lock.newCondition private val pendingRetries = new AtomicLong(0) private val pendingBulkWriteRetries = java.util.concurrent.ConcurrentHashMap.newKeySet[CosmosItemOperation]().asScala - private val pendingReadManyRetries = java.util.concurrent.ConcurrentHashMap.newKeySet[ReadManyOperation]().asScala private val activeTasks = new AtomicInteger(0) private val errorCaptureFirstException = new AtomicReference[Throwable]() private val bulkInputEmitter: Sinks.Many[CosmosBatch] = Sinks.many().unicast().onBackpressureBuffer() @@ -117,9 +104,14 @@ private class TransactionalBulkWriter private val currentBatchOperations = new mutable.ListBuffer[(ObjectNode, OperationContext, String)]() private val batchConstructionLock = new Object() + private val cosmosPatchHelperOpt = writeConfig.itemWriteStrategy match { + case ItemWriteStrategy.ItemPatch | ItemWriteStrategy.ItemPatchIfExists => + Some(new CosmosPatchHelper(diagnosticsConfig, writeConfig.patchConfigs.get)) + case _ => None + } + private val activeBulkWriteOperations =java.util.concurrent.ConcurrentHashMap.newKeySet[CosmosItemOperation]().asScala private val operationContextMap = new java.util.concurrent.ConcurrentHashMap[CosmosItemOperation, OperationContext]().asScala - private val activeReadManyOperations = java.util.concurrent.ConcurrentHashMap.newKeySet[ReadManyOperation]().asScala private val semaphore = new Semaphore(maxPendingOperations) private val totalScheduledMetrics = new AtomicLong(0) @@ -185,45 +177,9 @@ private class TransactionalBulkWriter } private val operationContext = initializeOperationContext() - private val cosmosPatchHelperOpt = writeConfig.itemWriteStrategy match { - case ItemWriteStrategy.ItemPatch | ItemWriteStrategy.ItemPatchIfExists | ItemWriteStrategy.ItemBulkUpdate => - Some(new CosmosPatchHelper(diagnosticsConfig, writeConfig.patchConfigs.get)) - case _ => None - } - private val readManyInputEmitterOpt: Option[Sinks.Many[ReadManyOperation]] = { - writeConfig.itemWriteStrategy match { - case ItemWriteStrategy.ItemBulkUpdate => Some(Sinks.many().unicast().onBackpressureBuffer()) - case _ => None - } - } - private val batchIntervalInMs = cosmosBulkExecutionOptionsImpl - .getMaxMicroBatchInterval - .toMillis - - private[this] val flushExecutorHolder: Option[(ScheduledThreadPoolExecutor, ScheduledFuture[_])] = { - writeConfig.itemWriteStrategy match { - case ItemWriteStrategy.ItemBulkUpdate => - val executor = new ScheduledThreadPoolExecutor( - 1, - new CosmosDaemonThreadFactory( - "BulkWriterReadManyFlush" + UUIDs.nonBlockingRandomUUID() - )) - executor.setExecuteExistingDelayedTasksAfterShutdownPolicy(false) - executor.setRemoveOnCancelPolicy(true) - - val future:ScheduledFuture[_] = executor.scheduleWithFixedDelay( - () => this.onFlushReadMany(), - batchIntervalInMs, - batchIntervalInMs, - TimeUnit.MILLISECONDS) - - Some(executor, future) - case _ => None - } - } - private def initializeOperationContext(): SparkTaskContext = { + private def initializeOperationContext(): SparkTaskContext = { val taskContext = TaskContext.get val diagnosticsContext: DiagnosticsContext = DiagnosticsContext(UUIDs.nonBlockingRandomUUID(), "BulkWriter") @@ -252,266 +208,8 @@ private class TransactionalBulkWriter } } - private def onFlushReadMany(): Unit = { - if (this.readManyInputEmitterOpt.isEmpty) { - throw new IllegalStateException("Callback onFlushReadMany should only be scheduled for bulk update.") - } - try { - this.readManyInputEmitterOpt.get.tryEmitNext(readManyFlushOperationSingleton) match { - case EmitResult.OK => log.logInfo("onFlushReadMany Successfully emitted flush") - case faultEmitResult => - log.logError(s"Callback invocation 'onFlush' failed with result: $faultEmitResult.") - } - } - catch { - case t: Throwable => - log.logError("Callback invocation 'onFlush' failed.", t) - } - } - - private val readManySubscriptionDisposableOpt: Option[Disposable] = { - writeConfig.itemWriteStrategy match { - case ItemWriteStrategy.ItemBulkUpdate => Some(createReadManySubscriptionDisposable()) - case _ => None - } - } - - private def createReadManySubscriptionDisposable(): Disposable = { - log.logTrace(s"readManySubscriptionDisposable, Context: ${operationContext.toString} $getThreadInfo") - // We start from using the bulk batch size and interval and concurrency - // If in the future, there is a need to separate the configuration, can re-consider - val bulkBatchSize = writeConfig.maxMicroBatchSize match { - case Some(customMaxMicroBatchSize) => Math.min( - BatchRequestResponseConstants.MAX_OPERATIONS_IN_DIRECT_MODE_BATCH_REQUEST, - Math.max(1, customMaxMicroBatchSize)) - case None => BatchRequestResponseConstants.MAX_OPERATIONS_IN_DIRECT_MODE_BATCH_REQUEST - } - - val batchConcurrency = cosmosBulkExecutionOptionsImpl.getMaxMicroBatchConcurrency - - val firstRecordTimeStamp = new AtomicLong(-1) - val currentMicroBatchSize = new AtomicLong(0) - - readManyInputEmitterOpt - .get - .asFlux() - .publishOn(readManyBoundedElastic) - .timestamp - .bufferUntil(timestampReadManyOperationTuple => { - val timestamp = timestampReadManyOperationTuple.getT1 - val readManyOperation = timestampReadManyOperationTuple.getT2 - - if (readManyOperation eq readManyFlushOperationSingleton) { - log.logTrace(s"FlushSingletonReceived, Context: ${operationContext.toString}") - val currentMicroBatchSizeSnapshot = currentMicroBatchSize.get - if (currentMicroBatchSizeSnapshot > 0) { - firstRecordTimeStamp.set(-1) - currentMicroBatchSize.set(0) - log.logTrace(s"FlushSingletonReceived - flushing batch, Context: ${operationContext.toString}") - true - } else { - // avoid counting flush operations for the micro batch size calculation - log.logTrace(s"FlushSingletonReceived - empty buffer, nothing to flush, Context: ${operationContext.toString}") - false - } - } else { - - firstRecordTimeStamp.compareAndSet(-1, timestamp) - val age = timestamp - firstRecordTimeStamp.get - val batchSize = currentMicroBatchSize.incrementAndGet - if (batchSize >= bulkBatchSize || age >= batchIntervalInMs) { - log.logTrace(s"BatchIntervalExpired - flushing batch, Context: ${operationContext.toString}") - firstRecordTimeStamp.set(-1) - currentMicroBatchSize.set(0) - true - } else { - false - } - } - }) - .subscribeOn(readManyBoundedElastic) - .asScala - .flatMap(timestampReadManyOperationTuples => { - val readManyOperations = timestampReadManyOperationTuples - .filter(candidate => !candidate.getT2.equals(readManyFlushOperationSingleton)) - .map(tuple => tuple.getT2) - - if (readManyOperations.isEmpty) { - Mono.empty() - } else { - val cosmosIdentitySet = readManyOperations.map(option => option.cosmosItemIdentity).toSet - - // for each batch, use readMany to read items from cosmosdb - val requestOptions = new CosmosReadManyRequestOptions() - val requestOptionsImpl = ImplementationBridgeHelpers - .CosmosReadManyRequestOptionsHelper - .getCosmosReadManyRequestOptionsAccessor.getImpl(requestOptions) - ThroughputControlHelper.populateThroughputControlGroupName( - requestOptionsImpl, - writeConfig.throughputControlConfig) - ImplementationBridgeHelpers - .CosmosAsyncContainerHelper - .getCosmosAsyncContainerAccessor - .readMany(container, cosmosIdentitySet.toList.asJava, requestOptions, classOf[ObjectNode]) - .switchIfEmpty( - // For Java SDK, empty pages will not be returned (this can happen when all the items does not exists yet) - // create a fake empty page response - Mono.just( - ImplementationBridgeHelpers - .FeedResponseHelper - .getFeedResponseAccessor - .createFeedResponse(new util.ArrayList[ObjectNode](), null, null))) - .doOnNext(feedResponse => { - // Tracking the bytes read as part of client-side patch (readMany + replace) as bytes written as well - // to have a way to indicate the additional work happening here - outputMetricsPublisher.trackWriteOperation( - 0, - Option.apply(feedResponse.getCosmosDiagnostics) match { - case Some(diagnostics) => Option.apply(diagnostics.getDiagnosticsContext) - case None => None - }) - val resultMap = new TrieMap[CosmosItemIdentity, ObjectNode]() - for (itemNode: ObjectNode <- feedResponse.getResults.asScala) { - resultMap += ( - new CosmosItemIdentity( - PartitionKeyHelper.getPartitionKeyPath(itemNode, partitionKeyDefinition), - itemNode.get(CosmosConstants.Properties.Id).asText()) -> itemNode) - } - - // It is possible that multiple cosmosPatchBulkUpdateOperations were targeting for the same item - // Currently, we are still creating one bulk item operation for each cosmosPatchBulkUpdateOperations - // for easier exception and semaphore handling - // However a consequences of it could be, the generated bulk item operation will fail due to conflicts or pre-condition failure - // If this turns out to be a problem, we can do more optimization here: merge multiple cosmosPatchBulkUpdateOperations into one bulkItemOperation - // But even the above approach can only work within the same batch but not for the whole spark partition processing. - for (readManyOperation <- readManyOperations) { - val cosmosPatchBulkUpdateOperations = - cosmosPatchHelperOpt - .get - .createCosmosPatchBulkUpdateOperations(readManyOperation.objectNode) - - val rootNode = - cosmosPatchHelperOpt - .get - .patchBulkUpdateItem(resultMap.get(readManyOperation.cosmosItemIdentity), cosmosPatchBulkUpdateOperations) - - // Add operation to batch - ReadMany results become batch operations - val etagOpt = Option.apply(rootNode.get(CosmosConstants.Properties.ETag)) - val (effectiveOperationType, operationContext) = etagOpt match { - case Some(etag) => - ("replace", new OperationContext( - readManyOperation.operationContext.itemId, - readManyOperation.operationContext.partitionKeyValue, - Some(etag.asText()), - readManyOperation.operationContext.attemptNumber, - monotonicOperationCounter.incrementAndGet(), - Some(readManyOperation.objectNode) - )) - case None => - ("create", new OperationContext( - readManyOperation.operationContext.itemId, - readManyOperation.operationContext.partitionKeyValue, - eTagInput = None, - readManyOperation.operationContext.attemptNumber, - monotonicOperationCounter.incrementAndGet(), - Some(readManyOperation.objectNode) - )) - } - - // Add to batch buffer - will be flushed on PK change or 100-op limit - addOperationToBatch( - readManyOperation.cosmosItemIdentity.getPartitionKey, - rootNode, - operationContext, - effectiveOperationType - ) - } - }) - .onErrorResume(throwable => { - for (readManyOperation <- readManyOperations) { - handleReadManyExceptions(throwable, readManyOperation) - } - - Mono.empty() - }) - .doFinally(_ => { - for (readManyOperation <- readManyOperations) { - val activeReadManyOperationFound = activeReadManyOperations.remove(readManyOperation) - // for ItemBulkUpdate strategy, each active task includes two stages: ReadMany + BulkWrite - // so we are not going to make task complete here - if (!activeReadManyOperationFound) { - // can't find the read-many operation in list of active operations! - logInfoOrWarning(s"Cannot find active read-many for '" - + s"${readManyOperation.cosmosItemIdentity.getPartitionKey}/" - + s"${readManyOperation.cosmosItemIdentity.getId}'. This can happen when " - + s"retries get re-enqueued.") - - if (pendingReadManyRetries.remove(readManyOperation)) { - pendingRetries.decrementAndGet() - } - } - } - }) - .`then`(Mono.empty()) - } - }, batchConcurrency) - .subscribe() - } - - private def handleReadManyExceptions(throwable: Throwable, ReadManyOperation: ReadManyOperation): Unit = { - throwable match { - case e: CosmosException => - outputMetricsPublisher.trackWriteOperation( - 0, - Option.apply(e.getDiagnostics) match { - case Some(diagnostics) => Option.apply(diagnostics.getDiagnosticsContext) - case None => None - }) - val requestOperationContext = ReadManyOperation.operationContext - if (shouldRetry(e.getStatusCode, e.getSubStatusCode, requestOperationContext)) { - log.logInfo(s"for itemId=[${requestOperationContext.itemId}], partitionKeyValue=[${requestOperationContext.partitionKeyValue}], " + - s"encountered status code '${e.getStatusCode}:${e.getSubStatusCode}' in read many, will retry! " + - s"attemptNumber=${requestOperationContext.attemptNumber}, exceptionMessage=${e.getMessage}, " + - s"Context: {${operationContext.toString}} $getThreadInfo") - - // the task will be re-queued at the beginning of the flow, so mark it complete here - markTaskCompletion() - - this.scheduleRetry( - trackPendingRetryAction = () => pendingReadManyRetries.add(ReadManyOperation), - clearPendingRetryAction = () => pendingReadManyRetries.remove(ReadManyOperation), - ReadManyOperation.cosmosItemIdentity.getPartitionKey, - ReadManyOperation.objectNode, - ReadManyOperation.operationContext, - e.getStatusCode) - - } else { - // Non-retryable exception or has exceeded the max retry count - val requestOperationContext = ReadManyOperation.operationContext - log.logError(s"for itemId=[${requestOperationContext.itemId}], partitionKeyValue=[${requestOperationContext.partitionKeyValue}], " + - s"encountered status code '${e.getStatusCode}:${e.getSubStatusCode}', all retries exhausted! " + - s"attemptNumber=${requestOperationContext.attemptNumber}, exceptionMessage=${e.getMessage}, " + - s"Context: {${operationContext.toString} $getThreadInfo") - - val message = s"All retries exhausted for readMany - " + - s"statusCode=[${e.getStatusCode}:${e.getSubStatusCode}] " + - s"itemId=[${requestOperationContext.itemId}], partitionKeyValue=[${requestOperationContext.partitionKeyValue}]" - - val exceptionToBeThrown = new BulkOperationFailedException(e.getStatusCode, e.getSubStatusCode, message, e) - captureIfFirstFailure(exceptionToBeThrown) - cancelWork() - markTaskCompletion() - } - case _ => // handle non cosmos exceptions - log.logError(s"Unexpected failure code path in Bulk ingestion readMany stage, " + - s"Context: ${operationContext.toString} $getThreadInfo", throwable) - captureIfFirstFailure(throwable) - cancelWork() - markTaskCompletion() - } - } private def scheduleRetry( trackPendingRetryAction: () => Boolean, @@ -692,15 +390,12 @@ private class TransactionalBulkWriter // allow staleness for ten additional minutes - which is perfectly fine var activeBulkWriteOperationsSnapshot = mutable.Set.empty[CosmosItemOperation] var pendingBulkWriteRetriesSnapshot = mutable.Set.empty[CosmosItemOperation] - var activeReadManyOperationsSnapshot = mutable.Set.empty[ReadManyOperation] - var pendingReadManyRetriesSnapshot = mutable.Set.empty[ReadManyOperation] log.logTrace( s"Before TryAcquire ${totalScheduledMetrics.get}, Context: ${operationContext.toString} $getThreadInfo") while (!semaphore.tryAcquire(activeTasksSemaphoreTimeout, TimeUnit.MINUTES)) { log.logDebug(s"Not able to acquire semaphore, Context: ${operationContext.toString} $getThreadInfo") - if (subscriptionDisposable.isDisposed || - (readManySubscriptionDisposableOpt.isDefined && readManySubscriptionDisposableOpt.get.isDisposed)) { + if (subscriptionDisposable.isDisposed) { captureIfFirstFailure( new IllegalStateException("Can't accept any new work - BulkWriter has been disposed already")) } @@ -709,15 +404,11 @@ private class TransactionalBulkWriter "Semaphore acquisition", activeBulkWriteOperationsSnapshot, pendingBulkWriteRetriesSnapshot, - activeReadManyOperationsSnapshot, - pendingReadManyRetriesSnapshot, numberOfIntervalsWithIdenticalActiveOperationSnapshots, allowRetryOnNewBulkWriterInstance = false) activeBulkWriteOperationsSnapshot = activeBulkWriteOperations.clone() pendingBulkWriteRetriesSnapshot = pendingBulkWriteRetries.clone() - activeReadManyOperationsSnapshot = activeReadManyOperations.clone() - pendingReadManyRetriesSnapshot = pendingReadManyRetries.clone() } val cnt = totalScheduledMetrics.getAndIncrement() @@ -747,23 +438,7 @@ private class TransactionalBulkWriter s"Context: ${operationContext.toString} $getThreadInfo") } - // The handling will make sure that during retry: - // For itemBulkUpdate -> the retry will go through readMany stage -> bulkWrite stage. - // For other strategies -> the retry will only go through bulk write stage - writeConfig.itemWriteStrategy match { - case ItemWriteStrategy.ItemBulkUpdate => scheduleReadManyInternal(partitionKeyValue, objectNode, operationContext) - case _ => scheduleBulkWriteInternal(partitionKeyValue, objectNode, operationContext) - } - } - - private def scheduleReadManyInternal(partitionKeyValue: PartitionKey, - objectNode: ObjectNode, - operationContext: OperationContext): Unit = { - - // For FAIL_NON_SERIALIZED, will keep retry, while for other errors, use the default behavior - val readManyOperation = ReadManyOperation(new CosmosItemIdentity(partitionKeyValue, operationContext.itemId), objectNode, operationContext) - activeReadManyOperations.add(readManyOperation) - readManyInputEmitterOpt.get.emitNext(readManyOperation, emitFailureHandler) + scheduleBulkWriteInternal(partitionKeyValue, objectNode, operationContext) } private def scheduleBulkWriteInternal(partitionKeyValue: PartitionKey, @@ -921,9 +596,7 @@ private class TransactionalBulkWriter s"attemptNumber=${context.attemptNumber}, exceptionMessage=$exceptionMessage, " + s"Context: {${operationContext.toString}} $getThreadInfo") - // If the write strategy is patchBulkUpdate, the OperationContext.sourceItem will not be the original objectNode, - // It is computed through read item from cosmosdb, and then patch the item locally. - // During retry, it is important to use the original objectNode (for example for preCondition failure, it requires to go through the readMany step again) + // During retry, use the original objectNode to ensure consistency val sourceItem = itemOperation match { case _: ItemBulkOperation[ObjectNode, OperationContext] => context.sourceItem match { @@ -975,8 +648,7 @@ private class TransactionalBulkWriter } private[this] def getActiveOperationsLog( - activeOperationsSnapshot: mutable.Set[CosmosItemOperation], - activeReadManyOperationsSnapshot: mutable.Set[ReadManyOperation]): String = { + activeOperationsSnapshot: mutable.Set[CosmosItemOperation]): String = { val sb = new StringBuilder() activeOperationsSnapshot @@ -992,20 +664,6 @@ private class TransactionalBulkWriter sb.append(s"${ctx.partitionKeyValue}/${ctx.itemId}/${ctx.eTag}(${ctx.attemptNumber})") }) - // add readMany snapshot logs - activeReadManyOperationsSnapshot - .take(TransactionalBulkWriter.maxItemOperationsToShowInErrorMessage - activeOperationsSnapshot.size) - .foreach(readManyOperation => { - if (sb.nonEmpty) { - sb.append(", ") - } - - sb.append("ReadMany") - sb.append("->") - val ctx = readManyOperation.operationContext - sb.append(s"${ctx.partitionKeyValue}/${ctx.itemId}/${ctx.eTag}(${ctx.attemptNumber})") - }) - sb.toString() } @@ -1029,39 +687,18 @@ private class TransactionalBulkWriter } } - private[this] def sameReadManyOperations - ( - snapshot: mutable.Set[ReadManyOperation], - current: mutable.Set[ReadManyOperation] - ): Boolean = { - - if (snapshot.size != current.size) { - false - } else { - snapshot.forall(snapshotOperation => { - current.exists( - currentOperation => snapshotOperation.cosmosItemIdentity == currentOperation.cosmosItemIdentity - && Objects.equals(snapshotOperation.objectNode, currentOperation.objectNode) - ) - }) - } - } - private[this] def throwIfProgressStaled ( operationName: String, activeOperationsSnapshot: mutable.Set[CosmosItemOperation], pendingRetriesSnapshot: mutable.Set[CosmosItemOperation], - activeReadManyOperationsSnapshot: mutable.Set[ReadManyOperation], - pendingReadManyOperationsSnapshot: mutable.Set[ReadManyOperation], numberOfIntervalsWithIdenticalActiveOperationSnapshots: AtomicLong, allowRetryOnNewBulkWriterInstance: Boolean ): Unit = { - val operationsLog = getActiveOperationsLog(activeOperationsSnapshot, activeReadManyOperationsSnapshot) + val operationsLog = getActiveOperationsLog(activeOperationsSnapshot) - if (sameBulkWriteOperations(pendingRetriesSnapshot ++ activeOperationsSnapshot , activeBulkWriteOperations ++ pendingBulkWriteRetries) - && sameReadManyOperations(pendingReadManyOperationsSnapshot ++ activeReadManyOperationsSnapshot , activeReadManyOperations ++ pendingReadManyRetries)) { + if (sameBulkWriteOperations(pendingRetriesSnapshot ++ activeOperationsSnapshot , activeBulkWriteOperations ++ pendingBulkWriteRetries)) { numberOfIntervalsWithIdenticalActiveOperationSnapshots.incrementAndGet() log.logWarning( @@ -1079,8 +716,7 @@ private class TransactionalBulkWriter val secondsWithoutProgress = numberOfIntervalsWithIdenticalActiveOperationSnapshots.get * writeConfig.flushCloseIntervalInSeconds - val maxNoProgressIntervalInSeconds = if (commitAttempt == 1 && allowRetryOnNewBulkWriterInstance - && this.activeReadManyOperations.isEmpty && this.pendingReadManyRetries.isEmpty) { + val maxNoProgressIntervalInSeconds = if (commitAttempt == 1 && allowRetryOnNewBulkWriterInstance) { writeConfig.maxInitialNoProgressIntervalInSeconds } else { writeConfig.maxRetryNoProgressIntervalInSeconds @@ -1089,32 +725,23 @@ private class TransactionalBulkWriter if (maxAllowedIntervalWithoutAnyProgressExceeded) { - val exception = if (activeReadManyOperationsSnapshot.isEmpty) { - val retriableRemainingOperations = if (allowRetryOnNewBulkWriterInstance) { - Some( - (pendingRetriesSnapshot ++ activeOperationsSnapshot) - .toList - .sortBy(op => op.getContext[OperationContext].sequenceNumber) - ) - } else { - None - } - - new BulkWriterNoProgressException( - s"Stale bulk ingestion identified in $operationName - the following active operations have not been " + - s"completed (first ${TransactionalBulkWriter.maxItemOperationsToShowInErrorMessage} shown) or progressed after " + - s"$maxNoProgressIntervalInSeconds seconds: $operationsLog", - commitAttempt, - retriableRemainingOperations) + val retriableRemainingOperations = if (allowRetryOnNewBulkWriterInstance) { + Some( + (pendingRetriesSnapshot ++ activeOperationsSnapshot) + .toList + .sortBy(op => op.getContext[OperationContext].sequenceNumber) + ) } else { - new BulkWriterNoProgressException( - s"Stale bulk ingestion as well as readMany operations identified in $operationName - the following active operations have not been " + - s"completed (first ${TransactionalBulkWriter.maxItemOperationsToShowInErrorMessage} shown) or progressed after " + - s"${maxNoProgressIntervalInSeconds} : $operationsLog", - commitAttempt, - None) + None } + val exception = new BulkWriterNoProgressException( + s"Stale bulk ingestion identified in $operationName - the following active operations have not been " + + s"completed (first ${TransactionalBulkWriter.maxItemOperationsToShowInErrorMessage} shown) or progressed after " + + s"$maxNoProgressIntervalInSeconds seconds: $operationsLog", + commitAttempt, + retriableRemainingOperations) + captureIfFirstFailure(exception) cancelWork() @@ -1135,7 +762,7 @@ private class TransactionalBulkWriter log.logInfo(s"flushAndClose invoked $getThreadInfo") log.logInfo(s"completed so far ${totalSuccessfulIngestionMetrics.get()}, " + - s"pending bulkWrite asks ${activeBulkWriteOperations.size}, pending readMany tasks ${activeReadManyOperations.size} $getThreadInfo") + s"pending bulkWrite asks ${activeBulkWriteOperations.size} $getThreadInfo") // error handling, if there is any error and the subscription is cancelled // the remaining tasks will not be processed hence we never reach activeTasks = 0 @@ -1152,17 +779,13 @@ private class TransactionalBulkWriter s"Waiting for pending activeTasks $activeTasksSnapshot and/or pendingRetries " + s"$pendingRetriesSnapshot, Context: ${operationContext.toString} $getThreadInfo") val activeOperationsSnapshot = activeBulkWriteOperations.clone() - val activeReadManyOperationsSnapshot = activeReadManyOperations.clone() val pendingOperationsSnapshot = pendingBulkWriteRetries.clone() - val pendingReadManyOperationsSnapshot = pendingReadManyRetries.clone() val awaitCompleted = pendingTasksCompleted.await(writeConfig.flushCloseIntervalInSeconds, TimeUnit.SECONDS) if (!awaitCompleted) { throwIfProgressStaled( "FlushAndClose", activeOperationsSnapshot, pendingOperationsSnapshot, - activeReadManyOperationsSnapshot, - pendingReadManyOperationsSnapshot, numberOfIntervalsWithIdenticalActiveOperationSnapshots, allowRetryOnNewBulkWriterInstance = true ) @@ -1176,7 +799,7 @@ private class TransactionalBulkWriter log.logWarning(s"Starting to re-enqueue retries. Enabling verbose logs. " + s"Number of intervals with identical pending operations: " + s"$numberOfIntervalsWithIdenticalActiveOperationSnapshots Active Bulk Operations: " - + s"$activeOperationsSnapshot, Active Read-Many Operations: $activeReadManyOperationsSnapshot, " + + s"$activeOperationsSnapshot, " + s"PendingRetries: $pendingRetriesSnapshot, Buffered tasks: $buffered " + s"Attempt: ${numberOfIntervalsWithIdenticalActiveOperationSnapshots.get} - " + s"Context: ${operationContext.toString} $getThreadInfo") @@ -1184,7 +807,7 @@ private class TransactionalBulkWriter log.logWarning(s"Reattempting to re-enqueue retries. Enabling verbose logs. " + s"Number of intervals with identical pending operations: " + s"$numberOfIntervalsWithIdenticalActiveOperationSnapshots Active Bulk Operations: " - + s"$activeOperationsSnapshot, Active Read-Many Operations: $activeReadManyOperationsSnapshot, " + + s"$activeOperationsSnapshot, " + s"PendingRetries: $pendingRetriesSnapshot, Buffered tasks: $buffered " + s"Attempt: ${numberOfIntervalsWithIdenticalActiveOperationSnapshots.get} - " + s"Context: ${operationContext.toString} $getThreadInfo") @@ -1199,20 +822,6 @@ private class TransactionalBulkWriter log.logTrace(s"Skipping individual operation retry in transactional batch mode - " + s"Context: ${operationContext.toString} $getThreadInfo") }) - - activeReadManyOperationsSnapshot.foreach(operation => { - if (activeReadManyOperations.contains(operation)) { - // re-validating whether the operation is still active - if so, just re-enqueue another retry - // this is harmless - because all bulkItemOperations from Spark connector are always idempotent - - // For FAIL_NON_SERIALIZED, will keep retry, while for other errors, use the default behavior - readManyInputEmitterOpt.get.emitNext(operation, emitFailureHandler) - log.logWarning(s"Re-enqueued a retry for pending active read-many task '" - + s"(${operation.cosmosItemIdentity.getPartitionKey}/${operation.cosmosItemIdentity.getId})' " - + s"- Attempt: ${numberOfIntervalsWithIdenticalActiveOperationSnapshots.get} - " - + s"Context: ${operationContext.toString} $getThreadInfo") - } - }) } } @@ -1240,17 +849,11 @@ private class TransactionalBulkWriter semaphore.release(Math.max(0, activeTasks.get())) bulkInputEmitter.emitComplete(TransactionalBulkWriter.emitFailureHandlerForComplete) - // complete readManyInputEmitter - if (readManyInputEmitterOpt.isDefined) { - readManyInputEmitterOpt.get.emitComplete(TransactionalBulkWriter.emitFailureHandlerForComplete) - } - throwIfCapturedExceptionExists() assume(activeTasks.get() <= 0) assume(activeBulkWriteOperations.isEmpty) - assume(activeReadManyOperations.isEmpty) assume(semaphore.availablePermits() >= maxPendingOperations) if (totalScheduledMetrics.get() != totalSuccessfulIngestionMetrics.get) { @@ -1266,37 +869,6 @@ private class TransactionalBulkWriter } } finally { subscriptionDisposable.dispose() - readManySubscriptionDisposableOpt match { - case Some(readManySubscriptionDisposable) => - - - readManySubscriptionDisposable.dispose() - case _ => - } - - flushExecutorHolder match { - case Some(executorAndFutureTuple) => - val executor: ScheduledThreadPoolExecutor = executorAndFutureTuple._1 - val future: ScheduledFuture[_] = executorAndFutureTuple._2 - - try { - future.cancel(true) - log.logDebug(s"Cancelled all future scheduled tasks $getThreadInfo, Context: ${operationContext.toString}") - } catch { - case e: Exception => - log.logWarning(s"Failed to cancel scheduled tasks $getThreadInfo, Context: ${operationContext.toString}", e) - } - - try { - log.logDebug(s"Shutting down the executor service, Context: ${operationContext.toString}") - executor.shutdownNow - log.logDebug(s"Successfully shut down the executor service, Context: ${operationContext.toString}") - } catch { - case e: Exception => - log.logWarning(s"Failed to shut down the executor service, Context: ${operationContext.toString}", e) - } - case _ => - } closed.set(true) } } @@ -1340,12 +912,9 @@ private class TransactionalBulkWriter private def cancelWork(): Unit = { logInfoOrWarning(s"cancelling remaining unprocessed tasks ${activeTasks.get} " + - s"[bulkWrite tasks ${activeBulkWriteOperations.size}, readMany tasks ${activeReadManyOperations.size} ]" + + s"[bulkWrite tasks ${activeBulkWriteOperations.size}]" + s"Context: ${operationContext.toString}") subscriptionDisposable.dispose() - if (readManySubscriptionDisposableOpt.isDefined) { - readManySubscriptionDisposableOpt.get.dispose() - } } private def shouldIgnore(statusCode: Int, subStatusCode: Int): Boolean = { @@ -1458,11 +1027,6 @@ private class TransactionalBulkWriter override def productPrefix: String = "OperationContext" } - - private case class ReadManyOperation( - cosmosItemIdentity: CosmosItemIdentity, - objectNode: ObjectNode, - operationContext: OperationContext) } private object TransactionalBulkWriter { @@ -1474,7 +1038,6 @@ private object TransactionalBulkWriter { private val BULK_WRITER_REQUESTS_BOUNDED_ELASTIC_THREAD_NAME = "bulk-writer-requests-bounded-elastic" private val BULK_WRITER_INPUT_BOUNDED_ELASTIC_THREAD_NAME = "bulk-writer-input-bounded-elastic" private val BULK_WRITER_RESPONSES_BOUNDED_ELASTIC_THREAD_NAME = "bulk-writer-responses-bounded-elastic" - private val READ_MANY_BOUNDED_ELASTIC_THREAD_NAME = "read-many-bounded-elastic" private val TTL_FOR_SCHEDULER_WORKER_IN_SECONDS = 60 // same as BoundedElasticScheduler.DEFAULT_TTL_SECONDS //scalastyle:on magic.number @@ -1543,13 +1106,6 @@ private object TransactionalBulkWriter { TTL_FOR_SCHEDULER_WORKER_IN_SECONDS, true) - // Custom bounded elastic scheduler to switch off IO thread to process response. - val readManyBoundedElastic: Scheduler = Schedulers.newBoundedElastic( - 2 * Schedulers.DEFAULT_BOUNDED_ELASTIC_SIZE, - Schedulers.DEFAULT_BOUNDED_ELASTIC_QUEUESIZE + maxPendingOperationsPerJVM, - READ_MANY_BOUNDED_ELASTIC_THREAD_NAME, - TTL_FOR_SCHEDULER_WORKER_IN_SECONDS, true) - def getThreadInfo: String = { val t = Thread.currentThread() val group = Option.apply(t.getThreadGroup) match { From 28846369dfdcf46a6e3e5b3480b86264c5566458 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Wed, 10 Dec 2025 15:42:47 +0000 Subject: [PATCH 07/16] Remove item patch parts --- .../spark/TransactionalBulkWriter.scala | 81 +++---------------- 1 file changed, 10 insertions(+), 71 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala index 3292eb3b69d9..7bc065d73733 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala @@ -104,12 +104,6 @@ private class TransactionalBulkWriter private val currentBatchOperations = new mutable.ListBuffer[(ObjectNode, OperationContext, String)]() private val batchConstructionLock = new Object() - private val cosmosPatchHelperOpt = writeConfig.itemWriteStrategy match { - case ItemWriteStrategy.ItemPatch | ItemWriteStrategy.ItemPatchIfExists => - Some(new CosmosPatchHelper(diagnosticsConfig, writeConfig.patchConfigs.get)) - case _ => None - } - private val activeBulkWriteOperations =java.util.concurrent.ConcurrentHashMap.newKeySet[CosmosItemOperation]().asScala private val operationContextMap = new java.util.concurrent.ConcurrentHashMap[CosmosItemOperation, OperationContext]().asScala private val semaphore = new Semaphore(maxPendingOperations) @@ -514,40 +508,12 @@ private class TransactionalBulkWriter } } - private[this] def getPatchOperationsForBatch(itemId: String, objectNode: ObjectNode): CosmosPatchOperations = { - assert(writeConfig.patchConfigs.isDefined) - assert(cosmosPatchHelperOpt.isDefined) - val cosmosPatchHelper = cosmosPatchHelperOpt.get - cosmosPatchHelper.createCosmosPatchOperations(itemId, partitionKeyDefinition, objectNode) - } - private[this] def finalFlushBatch(): Unit = { batchConstructionLock.synchronized { flushCurrentBatch() } } - private[this] def getPatchItemOperation(itemId: String, - partitionKey: PartitionKey, - partitionKeyDefinition: PartitionKeyDefinition, - objectNode: ObjectNode, - context: OperationContext): CosmosItemOperation = { - - assert(writeConfig.patchConfigs.isDefined) - assert(cosmosPatchHelperOpt.isDefined) - val patchConfigs = writeConfig.patchConfigs.get - val cosmosPatchHelper = cosmosPatchHelperOpt.get - - val cosmosPatchOperations = cosmosPatchHelper.createCosmosPatchOperations(itemId, partitionKeyDefinition, objectNode) - - val requestOptions = new CosmosBulkPatchItemRequestOptions() - if (patchConfigs.filter.isDefined && !StringUtils.isEmpty(patchConfigs.filter.get)) { - requestOptions.setFilterPredicate(patchConfigs.filter.get) - } - - CosmosBulkOperations.getPatchItemOperation(itemId, partitionKey, cosmosPatchOperations, requestOptions, context) - } - //scalastyle:off method.length //scalastyle:off cyclomatic.complexity private[this] def handleNonSuccessfulStatusCode @@ -918,38 +884,19 @@ private class TransactionalBulkWriter } private def shouldIgnore(statusCode: Int, subStatusCode: Int): Boolean = { - val returnValue = writeConfig.itemWriteStrategy match { - case ItemWriteStrategy.ItemAppend => Exceptions.isResourceExistsException(statusCode) - case ItemWriteStrategy.ItemPatchIfExists => Exceptions.isNotFoundExceptionCore(statusCode, subStatusCode) - case ItemWriteStrategy.ItemDelete => Exceptions.isNotFoundExceptionCore(statusCode, subStatusCode) - case ItemWriteStrategy.ItemDeleteIfNotModified => Exceptions.isNotFoundExceptionCore(statusCode, subStatusCode) || - Exceptions.isPreconditionFailedException(statusCode) - case ItemWriteStrategy.ItemOverwriteIfNotModified => - Exceptions.isResourceExistsException(statusCode) || - Exceptions.isNotFoundExceptionCore(statusCode, subStatusCode) || - Exceptions.isPreconditionFailedException(statusCode) - case _ => false - } - - returnValue + // Transactional batches only support ItemOverwrite - no errors to ignore + false } private def shouldRetry(statusCode: Int, subStatusCode: Int, operationContext: OperationContext): Boolean = { - var returnValue = false - if (operationContext.attemptNumber < writeConfig.maxRetryCount) { - returnValue = writeConfig.itemWriteStrategy match { - case ItemWriteStrategy.ItemBulkUpdate => - this.shouldRetryForItemPatchBulkUpdate(statusCode, subStatusCode) - // Upsert can return 404/0 in rare cases (when due to TTL expiration there is a race condition - case ItemWriteStrategy.ItemOverwrite => - Exceptions.canBeTransientFailure(statusCode, subStatusCode) || - statusCode == 0 || // Gateway mode reports inability to connect due to PoolAcquirePendingLimitException as status code 0 - Exceptions.isNotFoundExceptionCore(statusCode, subStatusCode) - case _ => - Exceptions.canBeTransientFailure(statusCode, subStatusCode) || - statusCode == 0 // Gateway mode reports inability to connect due to PoolAcquirePendingLimitException as status code 0 - } - } + var returnValue = false + if (operationContext.attemptNumber < writeConfig.maxRetryCount) { + // Transactional batches only support ItemOverwrite (upsert) + // Upsert can return 404/0 in rare cases (when due to TTL expiration there is a race condition) + returnValue = Exceptions.canBeTransientFailure(statusCode, subStatusCode) || + statusCode == 0 || // Gateway mode reports inability to connect due to PoolAcquirePendingLimitException as status code 0 + Exceptions.isNotFoundExceptionCore(statusCode, subStatusCode) + } log.logDebug(s"Should retry statusCode '$statusCode:$subStatusCode' -> $returnValue, " + s"Context: ${operationContext.toString} $getThreadInfo") @@ -957,14 +904,6 @@ private class TransactionalBulkWriter returnValue } - private def shouldRetryForItemPatchBulkUpdate(statusCode: Int, subStatusCode: Int): Boolean = { - Exceptions.canBeTransientFailure(statusCode, subStatusCode) || - statusCode == 0 || // Gateway mode reports inability to connect due to - // PoolAcquirePendingLimitException as status code 0 - Exceptions.isResourceExistsException(statusCode) || - Exceptions.isPreconditionFailedException(statusCode) - } - private def getId(objectNode: ObjectNode) = { val idField = objectNode.get(CosmosConstants.Properties.Id) assume(idField != null && idField.isTextual) From 0f91e741c666a31fe73ee906e3a6e186c3480b73 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Wed, 10 Dec 2025 16:32:21 +0000 Subject: [PATCH 08/16] Revert operation type changes --- .../azure/cosmos/spark/AsyncItemWriter.scala | 8 ------ .../com/azure/cosmos/spark/BulkWriter.scala | 6 ---- .../azure/cosmos/spark/CosmosWriterBase.scala | 28 +------------------ .../com/azure/cosmos/spark/PointWriter.scala | 4 --- .../spark/TransactionalBulkWriter.scala | 12 -------- 5 files changed, 1 insertion(+), 57 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/AsyncItemWriter.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/AsyncItemWriter.scala index b04455603676..e35b43692de1 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/AsyncItemWriter.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/AsyncItemWriter.scala @@ -14,14 +14,6 @@ private trait AsyncItemWriter { */ def scheduleWrite(partitionKeyValue: PartitionKey, objectNode: ObjectNode): Unit - /** - * Schedule a write to happen in async and return immediately with per-row operation type - * @param partitionKeyValue the partition key value - * @param objectNode the json object node - * @param operationType optional operation type (create, upsert, replace, delete) for this specific row - */ - def scheduleWrite(partitionKeyValue: PartitionKey, objectNode: ObjectNode, operationType: Option[String]): Unit - /** * Wait for all remaining work * Throws if any of the work resulted in failure diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/BulkWriter.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/BulkWriter.scala index 56e2a691857f..aec77f8762e8 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/BulkWriter.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/BulkWriter.scala @@ -646,12 +646,6 @@ private class BulkWriter } override def scheduleWrite(partitionKeyValue: PartitionKey, objectNode: ObjectNode): Unit = { - scheduleWrite(partitionKeyValue, objectNode, None) - } - - override def scheduleWrite(partitionKeyValue: PartitionKey, objectNode: ObjectNode, operationType: Option[String]): Unit = { - // BulkWriter doesn't support per-row operation types - it uses global ItemWriteStrategy - // The operationType parameter is ignored here for interface compatibility Preconditions.checkState(!closed.get()) throwIfCapturedExceptionExists() diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosWriterBase.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosWriterBase.scala index ece29491dd0a..c798864c729a 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosWriterBase.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosWriterBase.scala @@ -97,23 +97,6 @@ private abstract class CosmosWriterBase( override def write(internalRow: InternalRow): Unit = { val objectNode = cosmosRowConverter.fromInternalRowToObjectNode(internalRow, inputSchema) - // Extract operationType if column exists (for per-row operation support) - val operationType: Option[String] = if (inputSchema.fieldNames.contains("operationType")) { - val opTypeIndex = inputSchema.fieldIndex("operationType") - if (!internalRow.isNullAt(opTypeIndex)) { - Some(internalRow.getString(opTypeIndex)) - } else { - None - } - } else { - None - } - - // Remove operationType from objectNode if present (don't persist to Cosmos) - if (objectNode.has("operationType")) { - objectNode.remove("operationType") - } - require(objectNode.has(CosmosConstants.Properties.Id) && objectNode.get(CosmosConstants.Properties.Id).isTextual, s"${CosmosConstants.Properties.Id} is a mandatory field. " + @@ -127,16 +110,7 @@ private abstract class CosmosWriterBase( } val partitionKeyValue = PartitionKeyHelper.getPartitionKeyPath(objectNode, partitionKeyDefinition) - - // Call the appropriate scheduleWrite method based on whether operationType is specified - operationType match { - case Some(opType) => - // Per-row operation type specified - use 3-parameter method - writer.get.scheduleWrite(partitionKeyValue, objectNode, Some(opType)) - case None => - // No per-row operation type - use 2-parameter method (global strategy) - writer.get.scheduleWrite(partitionKeyValue, objectNode) - } + writer.get.scheduleWrite(partitionKeyValue, objectNode) } override def commit(): WriterCommitMessage = { diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/PointWriter.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/PointWriter.scala index 4f804af75687..8f07bf5339d5 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/PointWriter.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/PointWriter.scala @@ -76,10 +76,6 @@ private class PointWriter(container: CosmosAsyncContainer, } override def scheduleWrite(partitionKeyValue: PartitionKey, objectNode: ObjectNode): Unit = { - scheduleWrite(partitionKeyValue, objectNode, None) - } - - override def scheduleWrite(partitionKeyValue: PartitionKey, objectNode: ObjectNode, operationType: Option[String]): Unit = { checkState(!closed.get()) val etag = getETag(objectNode) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala index 7bc065d73733..43c7fc678af5 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala @@ -411,18 +411,6 @@ private class TransactionalBulkWriter scheduleWriteInternal(partitionKeyValue, objectNode, operationContext) } - /** - * Per-row operation type is not supported for transactional batches. - * All operations in a transactional batch must be upserts (ItemOverwrite strategy). - * This method is implemented to satisfy the AsyncItemWriter interface but throws an exception. - */ - override def scheduleWrite(partitionKeyValue: PartitionKey, objectNode: ObjectNode, operationType: Option[String]): Unit = { - throw new UnsupportedOperationException( - "Per-row operation types are not supported for transactional batches. " + - "All operations in a transactional batch must use ItemOverwrite (upsert) strategy. " + - "Use scheduleWrite(partitionKeyValue, objectNode) instead.") - } - private def scheduleWriteInternal(partitionKeyValue: PartitionKey, objectNode: ObjectNode, operationContext: OperationContext): Unit = { From f7379addb89fd30d062fdba4ad0ee354dd35f492 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Wed, 10 Dec 2025 17:20:16 +0000 Subject: [PATCH 09/16] Refactor semaphore to limit batches not operations --- .../spark/TransactionalBulkWriter.scala | 111 +++++++++++------- .../spark/TransactionalBatchITest.scala | 60 ++++++++++ 2 files changed, 131 insertions(+), 40 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala index 43c7fc678af5..6321c664c5d8 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala @@ -63,10 +63,16 @@ private class TransactionalBulkWriter private val verboseLoggingAfterReEnqueueingRetriesEnabled = new AtomicBoolean(false) private val cpuCount = SparkUtils.getNumberOfHostCPUCores - // each bulk writer allows up to maxPendingOperations being buffered - // there is one bulk writer per spark task/partition - // and default config will create one executor per core on the executor host - // so multiplying by cpuCount in the default config is too aggressive + + // NOTE: The public API config property is "maxPendingOperations" for backward compatibility, + // but internally TransactionalBulkWriter limits concurrent *batches* not individual operations. + // This is semantically correct because transactional batches are atomic units. + // We convert maxPendingOperations to maxPendingBatches by assuming ~50 operations per batch. + private val maxPendingBatches = { + val maxOps = writeConfig.bulkMaxPendingOperations.getOrElse(DefaultMaxPendingOperationPerCore) + // Assume average batch size of 50 operations - limit concurrent batches accordingly + Math.max(1, maxOps / 50) + } private val maxPendingOperations = writeConfig.bulkMaxPendingOperations .getOrElse(DefaultMaxPendingOperationPerCore) private val maxConcurrentPartitions = writeConfig.maxConcurrentCosmosPartitions match { @@ -85,8 +91,8 @@ private class TransactionalBulkWriter } log.logInfo( - s"BulkWriter instantiated (Host CPU count: $cpuCount, maxPendingOperations: $maxPendingOperations, " + - s"maxConcurrentPartitions: $maxConcurrentPartitions ...") + s"TransactionalBulkWriter instantiated (Host CPU count: $cpuCount, maxPendingBatches: $maxPendingBatches, " + + s"maxPendingOperations: $maxPendingOperations, maxConcurrentPartitions: $maxConcurrentPartitions ...") private val closed = new AtomicBoolean(false) @@ -106,7 +112,11 @@ private class TransactionalBulkWriter private val activeBulkWriteOperations =java.util.concurrent.ConcurrentHashMap.newKeySet[CosmosItemOperation]().asScala private val operationContextMap = new java.util.concurrent.ConcurrentHashMap[CosmosItemOperation, OperationContext]().asScala - private val semaphore = new Semaphore(maxPendingOperations) + // Semaphore limits number of outstanding batches (not individual operations) + private val semaphore = new Semaphore(maxPendingBatches) + private val activeBatches = new AtomicInteger(0) + // Map each batch's first operation to the batch size for semaphore release tracking + private val batchSizeMap = new java.util.concurrent.ConcurrentHashMap[CosmosItemOperation, Integer]() private val totalScheduledMetrics = new AtomicLong(0) private val totalSuccessfulIngestionMetrics = new AtomicLong(0) @@ -288,10 +298,15 @@ private class TransactionalBulkWriter resp => { val isGettingRetried = new AtomicBoolean(false) val shouldSkipTaskCompletion = new AtomicBoolean(false) + val isFirstOperationInBatch = new AtomicBoolean(false) try { val itemOperation = resp.getOperation val itemOperationFound = activeBulkWriteOperations.remove(itemOperation) val pendingRetriesFound = pendingBulkWriteRetries.remove(itemOperation) + + // Check if this is the first operation in a batch (used for semaphore release) + val batchSizeOption = Option(batchSizeMap.remove(itemOperation)) + isFirstOperationInBatch.set(batchSizeOption.isDefined) if (pendingRetriesFound) { pendingRetries.decrementAndGet() @@ -337,8 +352,12 @@ private class TransactionalBulkWriter } } finally { - if (!isGettingRetried.get) { + // Release semaphore when we process the first operation of a batch + // This indicates the entire batch has been processed + if (!isGettingRetried.get && isFirstOperationInBatch.get) { + activeBatches.decrementAndGet() semaphore.release() + log.logTrace(s"Released semaphore for completed batch, activeBatches: ${activeBatches.get} $getThreadInfo") } } @@ -368,7 +387,6 @@ private class TransactionalBulkWriter Preconditions.checkState(!closed.get()) throwIfCapturedExceptionExists() - val activeTasksSemaphoreTimeout = 10 val operationContext = new OperationContext( getId(objectNode), partitionKeyValue, @@ -376,34 +394,6 @@ private class TransactionalBulkWriter 1, monotonicOperationCounter.incrementAndGet(), None) - val numberOfIntervalsWithIdenticalActiveOperationSnapshots = new AtomicLong(0) - // Don't clone the activeOperations for the first iteration - // to reduce perf impact before the Semaphore has been acquired - // this means if the semaphore can't be acquired within 10 minutes - // the first attempt will always assume it wasn't stale - so effectively we - // allow staleness for ten additional minutes - which is perfectly fine - var activeBulkWriteOperationsSnapshot = mutable.Set.empty[CosmosItemOperation] - var pendingBulkWriteRetriesSnapshot = mutable.Set.empty[CosmosItemOperation] - - log.logTrace( - s"Before TryAcquire ${totalScheduledMetrics.get}, Context: ${operationContext.toString} $getThreadInfo") - while (!semaphore.tryAcquire(activeTasksSemaphoreTimeout, TimeUnit.MINUTES)) { - log.logDebug(s"Not able to acquire semaphore, Context: ${operationContext.toString} $getThreadInfo") - if (subscriptionDisposable.isDisposed) { - captureIfFirstFailure( - new IllegalStateException("Can't accept any new work - BulkWriter has been disposed already")) - } - - throwIfProgressStaled( - "Semaphore acquisition", - activeBulkWriteOperationsSnapshot, - pendingBulkWriteRetriesSnapshot, - numberOfIntervalsWithIdenticalActiveOperationSnapshots, - allowRetryOnNewBulkWriterInstance = false) - - activeBulkWriteOperationsSnapshot = activeBulkWriteOperations.clone() - pendingBulkWriteRetriesSnapshot = pendingBulkWriteRetries.clone() - } val cnt = totalScheduledMetrics.getAndIncrement() log.logTrace(s"total scheduled $cnt, Context: ${operationContext.toString} $getThreadInfo") @@ -464,6 +454,36 @@ private class TransactionalBulkWriter private[this] def flushCurrentBatch(): Unit = { // Must be called within batchConstructionLock.synchronized if (currentBatchOperations.nonEmpty && currentPartitionKey != null) { + // Acquire semaphore before emitting batch - this limits concurrent batches + val activeTasksSemaphoreTimeout = 10 + val numberOfIntervalsWithIdenticalActiveOperationSnapshots = new AtomicLong(0) + var activeBulkWriteOperationsSnapshot = mutable.Set.empty[CosmosItemOperation] + var pendingBulkWriteRetriesSnapshot = mutable.Set.empty[CosmosItemOperation] + + val dummyContext = new OperationContext("batch-flush", currentPartitionKey, null, 1, 0, None) + + log.logTrace(s"Before acquiring semaphore for batch emission, activeBatches: ${activeBatches.get} $getThreadInfo") + while (!semaphore.tryAcquire(activeTasksSemaphoreTimeout, TimeUnit.MINUTES)) { + log.logDebug(s"Not able to acquire semaphore for batch, activeBatches: ${activeBatches.get} $getThreadInfo") + if (subscriptionDisposable.isDisposed) { + captureIfFirstFailure( + new IllegalStateException("Can't accept any new work - BulkWriter has been disposed already")) + } + + throwIfProgressStaled( + "Batch semaphore acquisition", + activeBulkWriteOperationsSnapshot, + pendingBulkWriteRetriesSnapshot, + numberOfIntervalsWithIdenticalActiveOperationSnapshots, + allowRetryOnNewBulkWriterInstance = false) + + activeBulkWriteOperationsSnapshot = activeBulkWriteOperations.clone() + pendingBulkWriteRetriesSnapshot = pendingBulkWriteRetries.clone() + } + + activeBatches.incrementAndGet() + log.logTrace(s"Acquired semaphore for batch emission, activeBatches: ${activeBatches.get} $getThreadInfo") + // Build the batch using the builder API val batch = CosmosBatch.createCosmosBatch(currentPartitionKey) @@ -478,14 +498,20 @@ private class TransactionalBulkWriter // After building the batch, get the operations and track them with their contexts val batchOperations = batch.getOperations() + val batchSize = batchOperations.size() // Map each operation to its context and add to tracking - for (i <- 0 until batchOperations.size()) { + for (i <- 0 until batchSize) { val operation = batchOperations.get(i) val context = contextList(i) operationContextMap.put(operation, context) activeBulkWriteOperations.add(operation) } + + // Track batch size using first operation as key - needed for semaphore release + if (batchSize > 0) { + batchSizeMap.put(batchOperations.get(0), batchSize) + } // Emit the batch bulkInputEmitter.emitNext(batch, emitFailureHandler) @@ -800,7 +826,12 @@ private class TransactionalBulkWriter } logInfoOrWarning(s"invoking bulkInputEmitter.onComplete(), Context: ${operationContext.toString} $getThreadInfo") - semaphore.release(Math.max(0, activeTasks.get())) + // Release any remaining batch permits + val remainingBatches = activeBatches.get() + if (remainingBatches > 0) { + semaphore.release(remainingBatches) + log.logDebug(s"Released $remainingBatches batch permits during cleanup") + } bulkInputEmitter.emitComplete(TransactionalBulkWriter.emitFailureHandlerForComplete) throwIfCapturedExceptionExists() @@ -808,7 +839,7 @@ private class TransactionalBulkWriter assume(activeTasks.get() <= 0) assume(activeBulkWriteOperations.isEmpty) - assume(semaphore.availablePermits() >= maxPendingOperations) + assume(semaphore.availablePermits() >= maxPendingBatches) if (totalScheduledMetrics.get() != totalSuccessfulIngestionMetrics.get) { log.logWarning(s"flushAndClose completed with no error but inconsistent total success and " + diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala index 4027a2398f78..9fc55d512566 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala @@ -746,4 +746,64 @@ class TransactionalBatchITest extends IntegrationSpec } } + it should "enforce batch-level backpressure with small maxPendingOperations" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) + + // Use a small maxPendingOperations value to force batch-level limiting + // With maxPendingOperations=50, maxPendingBatches = 50/50 = 1 + // This means only 1 batch should be in-flight at a time + val maxPendingOperations = 50 + + // Create 200 operations across 4 batches (50 operations per partition key = 1 batch each) + // This will test that the semaphore properly limits concurrent batches + val partitionKeys = (1 to 4).map(_ => UUID.randomUUID().toString) + + val schema = StructType(Seq( + StructField("id", StringType, nullable = false), + StructField("pk", StringType, nullable = false), + StructField("value", IntegerType, nullable = false) + )) + + // Create 50 operations per partition key (forms 1 batch per PK) + val allOperations = partitionKeys.flatMap { pk => + (1 to 50).map { i => + Row(s"item-$i-${UUID.randomUUID()}", pk, i) + } + } + + val operationsDf = spark.createDataFrame(allOperations.asJava, schema) + + // Execute with very small maxPendingOperations to force batch limiting + operationsDf.write + .format("cosmos.oltp") + .option("spark.cosmos.accountEndpoint", cosmosEndpoint) + .option("spark.cosmos.accountKey", cosmosMasterKey) + .option("spark.cosmos.database", cosmosDatabase) + .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) + .option("spark.cosmos.write.bulk.transactional", "true") + .option("spark.cosmos.write.bulk.enabled", "true") + .option("spark.cosmos.write.bulk.maxPendingOperations", maxPendingOperations.toString) + .mode(SaveMode.Append) + .save() + + // Verify all items were created successfully + partitionKeys.foreach { pk => + val queryResult = container + .queryItems(s"SELECT VALUE COUNT(1) FROM c WHERE c.pk = '$pk'", classOf[Long]) + .collectList() + .block() + + val count = if (queryResult.isEmpty) 0L else queryResult.get(0) + assert(count == 50, s"Expected 50 items for partition key $pk, but found $count") + } + + // If we get here without deadlock or timeout, batch-level backpressure is working + // The test verifies: + // 1. Operations complete successfully even with tight batch limit + // 2. No deadlocks occur from semaphore management + // 3. All batches are properly tracked and released + } + } From d14d40ed6b96953465ce3f0619159ac73e9ac5fa Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Wed, 10 Dec 2025 18:20:35 +0000 Subject: [PATCH 10/16] Add batch size validation for transactional batches. Refactored calculateTotalSerializedLength to validate 2MB payload limit and return HTTP 413 errors with clear messages when exceeded. --- .../batch/TransactionalBulkExecutor.java | 39 +++++++++++++++---- 1 file changed, 32 insertions(+), 7 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/TransactionalBulkExecutor.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/TransactionalBulkExecutor.java index fe7d4cd50418..e9402f124e6d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/TransactionalBulkExecutor.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/TransactionalBulkExecutor.java @@ -286,6 +286,31 @@ private Flux> executeTransactionalBatch(Co String batchTrackingId = UUIDs.nonBlockingRandomUUID().toString(); PartitionKey partitionKey = cosmosBatch.getPartitionKeyValue(); + // Validate batch payload size before execution + int totalSerializedLength = this.calculateTotalSerializedLength(operations); + if (totalSerializedLength > this.maxMicroBatchPayloadSizeInBytes) { + String errorMessage = String.format( + "Transactional batch exceeds maximum payload size: %d bytes (limit: %d bytes), PK: %s", + totalSerializedLength, + this.maxMicroBatchPayloadSizeInBytes, + partitionKey); + logger.error("{}, Context: {}", errorMessage, this.operationContextText); + + // Return error responses for all operations in the batch + CosmosException payloadSizeException = BridgeInternal.createCosmosException( + HttpConstants.StatusCodes.REQUEST_ENTITY_TOO_LARGE, + errorMessage); + + return Flux.fromIterable(operations) + .map(operation -> { + TContext actualContext = this.getActualContext(operation); + return ModelBridgeInternal.createCosmosBulkOperationResponse( + operation, + payloadSizeException, + actualContext); + }); + } + logDebugOrWarning( "Executing transactional batch - {} operations, PK: {}, TrackingId: {}, Context: {}", operations.size(), @@ -402,14 +427,14 @@ private Flux> executeTransactionalBatch(Co }); } - private int calculateTotalSerializedLength(AtomicInteger currentTotalSerializedLength, CosmosItemOperation item) { - if (item instanceof CosmosItemOperationBase) { - return currentTotalSerializedLength.accumulateAndGet( - ((CosmosItemOperationBase) item).getSerializedLength(this.effectiveItemSerializer), - Integer::sum); + private int calculateTotalSerializedLength(List operations) { + int totalLength = 0; + for (CosmosItemOperation operation : operations) { + if (operation instanceof CosmosItemOperationBase) { + totalLength += ((CosmosItemOperationBase) operation).getSerializedLength(this.effectiveItemSerializer); + } } - - return currentTotalSerializedLength.get(); + return totalLength; } private TContext getActualContext(CosmosItemOperation itemOperation) { From 57cc0eb7e0e993188de88674eadc85b2312ff7f8 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Wed, 10 Dec 2025 19:26:54 +0000 Subject: [PATCH 11/16] Fix transactional batch diagnostics prefix, simplify validation, rename test fields --- .../spark/TransactionalBulkWriter.scala | 11 ++-- .../spark/TransactionalBatchITest.scala | 52 +++++++++---------- 2 files changed, 30 insertions(+), 33 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala index 6321c664c5d8..79e751d7be98 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala @@ -84,11 +84,8 @@ private class TransactionalBulkWriter case None => 2 * ContainerFeedRangesCache.getFeedRanges(container, containerConfig.feedRangeRefreshIntervalInSecondsOpt).block().size } // Validate write strategy for transactional batches - if (writeConfig.itemWriteStrategy != ItemWriteStrategy.ItemOverwrite) { - throw new IllegalArgumentException( - s"Transactional batches only support ItemOverwrite (upsert) write strategy. " + - s"Requested strategy: ${writeConfig.itemWriteStrategy}") - } + require(writeConfig.itemWriteStrategy == ItemWriteStrategy.ItemOverwrite, + s"Transactional batches only support ItemOverwrite (upsert) write strategy. Requested: ${writeConfig.itemWriteStrategy}") log.logInfo( s"TransactionalBulkWriter instantiated (Host CPU count: $cpuCount, maxPendingBatches: $maxPendingBatches, " + @@ -186,7 +183,7 @@ private class TransactionalBulkWriter private def initializeOperationContext(): SparkTaskContext = { val taskContext = TaskContext.get - val diagnosticsContext: DiagnosticsContext = DiagnosticsContext(UUIDs.nonBlockingRandomUUID(), "BulkWriter") + val diagnosticsContext: DiagnosticsContext = DiagnosticsContext(UUIDs.nonBlockingRandomUUID(), "transactional-BulkWriter") if (taskContext != null) { val taskDiagnosticsContext = SparkTaskContext(diagnosticsContext.correlationActivityId, @@ -406,7 +403,7 @@ private class TransactionalBulkWriter operationContext: OperationContext): Unit = { activeTasks.incrementAndGet() if (operationContext.attemptNumber > 1) { - logInfoOrWarning(s"bulk scheduleWrite attemptCnt: ${operationContext.attemptNumber}, " + + logInfoOrWarning(s"TransactionalBulkWriter scheduleWrite attemptCnt: ${operationContext.attemptNumber}, " + s"Context: ${operationContext.toString} $getThreadInfo") } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala index 9fc55d512566..8a96f340f451 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala @@ -435,8 +435,8 @@ class TransactionalBatchITest extends IntegrationSpec new com.azure.cosmos.models.PartitionKeyDefinition() ) val paths = new java.util.ArrayList[String]() - paths.add("/PermId") - paths.add("/SourceId") + paths.add("/testPrimaryKey") + paths.add("/testSecondaryKey") containerProperties.getPartitionKeyDefinition.setPaths(paths) containerProperties.getPartitionKeyDefinition.setKind(com.azure.cosmos.models.PartitionKind.MULTI_HASH) containerProperties.getPartitionKeyDefinition.setVersion(com.azure.cosmos.models.PartitionKeyDefinitionVersion.V2) @@ -452,9 +452,9 @@ class TransactionalBatchITest extends IntegrationSpec // Create batch operations with hierarchical partition keys val schema = StructType(Seq( StructField("id", StringType, nullable = false), - StructField("PermId", StringType, nullable = false), - StructField("SourceId", StringType, nullable = false), - StructField("price", org.apache.spark.sql.types.DoubleType, nullable = false) + StructField("testPrimaryKey", StringType, nullable = false), + StructField("testSecondaryKey", StringType, nullable = false), + StructField("testPrice", org.apache.spark.sql.types.DoubleType, nullable = false) )) val batchOperations = Seq( @@ -479,11 +479,11 @@ class TransactionalBatchITest extends IntegrationSpec val pk = new PartitionKeyBuilder().add(permId).add(sourceId).build() val item1 = container.readItem(item1Id, pk, classOf[ObjectNode]).block() item1 should not be null - item1.getItem.get("price").asDouble() shouldEqual 100.5 + item1.getItem.get("testPrice").asDouble() shouldEqual 100.5 val item2 = container.readItem(item2Id, pk, classOf[ObjectNode]).block() item2 should not be null - item2.getItem.get("price").asDouble() shouldEqual 101.25 + item2.getItem.get("testPrice").asDouble() shouldEqual 101.25 } finally { // Clean up container container.delete().block() @@ -501,8 +501,8 @@ class TransactionalBatchITest extends IntegrationSpec new com.azure.cosmos.models.PartitionKeyDefinition() ) val paths = new java.util.ArrayList[String]() - paths.add("/PermId") - paths.add("/SourceId") + paths.add("/testPrimaryKey") + paths.add("/testSecondaryKey") containerProperties.getPartitionKeyDefinition.setPaths(paths) containerProperties.getPartitionKeyDefinition.setKind(com.azure.cosmos.models.PartitionKind.MULTI_HASH) containerProperties.getPartitionKeyDefinition.setVersion(com.azure.cosmos.models.PartitionKeyDefinitionVersion.V2) @@ -518,18 +518,18 @@ class TransactionalBatchITest extends IntegrationSpec // First, create initial record val initialDoc = Utils.getSimpleObjectMapper.createObjectNode() initialDoc.put("id", oldRecordId) - initialDoc.put("PermId", permId) - initialDoc.put("SourceId", sourceId) - initialDoc.put("price", 100.0) + initialDoc.put("testPrimaryKey", permId) + initialDoc.put("testSecondaryKey", sourceId) + initialDoc.put("testPrice", 100.0) initialDoc.putNull("valid_to") container.createItem(initialDoc).block() // Now perform atomic temporal update: close old record + create new record val schema = StructType(Seq( StructField("id", StringType, nullable = false), - StructField("PermId", StringType, nullable = false), - StructField("SourceId", StringType, nullable = false), - StructField("price", org.apache.spark.sql.types.DoubleType, nullable = false), + StructField("testPrimaryKey", StringType, nullable = false), + StructField("testSecondaryKey", StringType, nullable = false), + StructField("testPrice", org.apache.spark.sql.types.DoubleType, nullable = false), StructField("valid_to", StringType, nullable = true) )) @@ -562,7 +562,7 @@ class TransactionalBatchITest extends IntegrationSpec // Verify new record was created val newRecord = container.readItem(newRecordId, pk, classOf[ObjectNode]).block() newRecord should not be null - newRecord.getItem.get("price").asDouble() shouldEqual 150.0 + newRecord.getItem.get("testPrice").asDouble() shouldEqual 150.0 newRecord.getItem.get("valid_to").isNull shouldBe true } finally { // Clean up container @@ -581,8 +581,8 @@ class TransactionalBatchITest extends IntegrationSpec new com.azure.cosmos.models.PartitionKeyDefinition() ) val paths = new java.util.ArrayList[String]() - paths.add("/PermId") - paths.add("/SourceId") + paths.add("/testPrimaryKey") + paths.add("/testSecondaryKey") containerProperties.getPartitionKeyDefinition.setPaths(paths) containerProperties.getPartitionKeyDefinition.setKind(com.azure.cosmos.models.PartitionKind.MULTI_HASH) containerProperties.getPartitionKeyDefinition.setVersion(com.azure.cosmos.models.PartitionKeyDefinitionVersion.V2) @@ -593,9 +593,9 @@ class TransactionalBatchITest extends IntegrationSpec // Create operations for different PermId/SourceId combinations val schema = StructType(Seq( StructField("id", StringType, nullable = false), - StructField("PermId", StringType, nullable = false), - StructField("SourceId", StringType, nullable = false), - StructField("price", org.apache.spark.sql.types.DoubleType, nullable = false) + StructField("testPrimaryKey", StringType, nullable = false), + StructField("testSecondaryKey", StringType, nullable = false), + StructField("testPrice", org.apache.spark.sql.types.DoubleType, nullable = false) )) val batchOperations = Seq( @@ -640,8 +640,8 @@ class TransactionalBatchITest extends IntegrationSpec new com.azure.cosmos.models.PartitionKeyDefinition() ) val paths = new java.util.ArrayList[String]() - paths.add("/PermId") - paths.add("/SourceId") + paths.add("/testPrimaryKey") + paths.add("/testSecondaryKey") containerProperties.getPartitionKeyDefinition.setPaths(paths) containerProperties.getPartitionKeyDefinition.setKind(com.azure.cosmos.models.PartitionKind.MULTI_HASH) containerProperties.getPartitionKeyDefinition.setVersion(com.azure.cosmos.models.PartitionKeyDefinitionVersion.V2) @@ -654,9 +654,9 @@ class TransactionalBatchITest extends IntegrationSpec // Create 101 operations for the same hierarchical partition key val schema = StructType(Seq( StructField("id", StringType, nullable = false), - StructField("PermId", StringType, nullable = false), - StructField("SourceId", StringType, nullable = false), - StructField("price", org.apache.spark.sql.types.DoubleType, nullable = false) + StructField("testPrimaryKey", StringType, nullable = false), + StructField("testSecondaryKey", StringType, nullable = false), + StructField("testPrice", org.apache.spark.sql.types.DoubleType, nullable = false) )) val batchOperations = (1 to 101).map { i => From f989c4e305561eaf43197a861fca694c667f74e0 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Wed, 10 Dec 2025 20:23:26 +0000 Subject: [PATCH 12/16] batch-level semantics, remove client validation, update test titles --- .../spark/TransactionalBulkWriter.scala | 10 +++-- .../spark/TransactionalBatchITest.scala | 12 +++--- .../batch/TransactionalBulkExecutor.java | 37 ------------------- 3 files changed, 12 insertions(+), 47 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala index 79e751d7be98..91fc34ee3c6b 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala @@ -1037,26 +1037,28 @@ private object TransactionalBulkWriter { private val bulkProcessingThresholds = new CosmosBulkExecutionThresholdsState() - private val maxPendingOperationsPerJVM: Int = DefaultMaxPendingOperationPerCore * SparkUtils.getNumberOfHostCPUCores + // For batch-level backpressure: use a conservative estimate of concurrent batches across the JVM + // Each batch can have up to 100 operations, so this sizing provides headroom for the scheduler queues + private val maxPendingBatchesPerJVM: Int = DefaultMaxPendingOperationPerCore * SparkUtils.getNumberOfHostCPUCores / 50 // Custom bounded elastic scheduler to consume input flux val bulkWriterRequestsBoundedElastic: Scheduler = Schedulers.newBoundedElastic( Schedulers.DEFAULT_BOUNDED_ELASTIC_SIZE, - Schedulers.DEFAULT_BOUNDED_ELASTIC_QUEUESIZE + 2 * maxPendingOperationsPerJVM, + Schedulers.DEFAULT_BOUNDED_ELASTIC_QUEUESIZE + 2 * maxPendingBatchesPerJVM, BULK_WRITER_REQUESTS_BOUNDED_ELASTIC_THREAD_NAME, TTL_FOR_SCHEDULER_WORKER_IN_SECONDS, true) // Custom bounded elastic scheduler to consume input flux val bulkWriterInputBoundedElastic: Scheduler = Schedulers.newBoundedElastic( Schedulers.DEFAULT_BOUNDED_ELASTIC_SIZE, - Schedulers.DEFAULT_BOUNDED_ELASTIC_QUEUESIZE + 2 * maxPendingOperationsPerJVM, + Schedulers.DEFAULT_BOUNDED_ELASTIC_QUEUESIZE + 2 * maxPendingBatchesPerJVM, BULK_WRITER_INPUT_BOUNDED_ELASTIC_THREAD_NAME, TTL_FOR_SCHEDULER_WORKER_IN_SECONDS, true) // Custom bounded elastic scheduler to switch off IO thread to process response. val bulkWriterResponsesBoundedElastic: Scheduler = Schedulers.newBoundedElastic( Schedulers.DEFAULT_BOUNDED_ELASTIC_SIZE, - Schedulers.DEFAULT_BOUNDED_ELASTIC_QUEUESIZE + maxPendingOperationsPerJVM, + Schedulers.DEFAULT_BOUNDED_ELASTIC_QUEUESIZE + maxPendingBatchesPerJVM, BULK_WRITER_RESPONSES_BOUNDED_ELASTIC_THREAD_NAME, TTL_FOR_SCHEDULER_WORKER_IN_SECONDS, true) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala index 8a96f340f451..9130207a9c8e 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala @@ -424,11 +424,11 @@ class TransactionalBatchITest extends IntegrationSpec rootCause.getMessage should include("Batch request has more operations than what is supported") } - "Transactional Batch with Hierarchical Partition Keys" should "create items atomically with PermId and SourceId" in { + "Transactional Batch with Hierarchical Partition Keys" should "create items atomically with testPrimaryKey and testSecondaryKey" in { val cosmosEndpoint = TestConfigurations.HOST val cosmosMasterKey = TestConfigurations.MASTER_KEY - // Create container with hierarchical partition keys (PermId, SourceId) + // Create container with hierarchical partition keys (testPrimaryKey, testSecondaryKey) val containerName = s"test-hpk-${UUID.randomUUID()}" val containerProperties = new com.azure.cosmos.models.CosmosContainerProperties( containerName, @@ -494,7 +494,7 @@ class TransactionalBatchITest extends IntegrationSpec val cosmosEndpoint = TestConfigurations.HOST val cosmosMasterKey = TestConfigurations.MASTER_KEY - // Create container with hierarchical partition keys (PermId, SourceId) + // Create container with hierarchical partition keys (testPrimaryKey, testSecondaryKey) val containerName = s"test-hpk-temporal-${UUID.randomUUID()}" val containerProperties = new com.azure.cosmos.models.CosmosContainerProperties( containerName, @@ -570,11 +570,11 @@ class TransactionalBatchITest extends IntegrationSpec } } - it should "handle operations across multiple PermId/SourceId combinations" in { + it should "handle operations across multiple testPrimaryKey/testSecondaryKey combinations" in { val cosmosEndpoint = TestConfigurations.HOST val cosmosMasterKey = TestConfigurations.MASTER_KEY - // Create container with hierarchical partition keys (PermId, SourceId) + // Create container with hierarchical partition keys (testPrimaryKey, testSecondaryKey) val containerName = s"test-hpk-multi-${UUID.randomUUID()}" val containerProperties = new com.azure.cosmos.models.CosmosContainerProperties( containerName, @@ -633,7 +633,7 @@ class TransactionalBatchITest extends IntegrationSpec val cosmosEndpoint = TestConfigurations.HOST val cosmosMasterKey = TestConfigurations.MASTER_KEY - // Create container with hierarchical partition keys (PermId, SourceId) + // Create container with hierarchical partition keys (testPrimaryKey, testSecondaryKey) val containerName = s"test-hpk-limit-${UUID.randomUUID()}" val containerProperties = new com.azure.cosmos.models.CosmosContainerProperties( containerName, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/TransactionalBulkExecutor.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/TransactionalBulkExecutor.java index e9402f124e6d..8a4ee1e7c7f3 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/TransactionalBulkExecutor.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/TransactionalBulkExecutor.java @@ -95,7 +95,6 @@ public final class TransactionalBulkExecutor implements Disposable { ImplementationBridgeHelpers.CosmosBatchResponseHelper.getCosmosBatchResponseAccessor(); private final CosmosAsyncContainer container; - private final int maxMicroBatchPayloadSizeInBytes; private final AsyncDocumentClient docClientWrapper; private final String operationContextText; private final OperationContextAndListenerTuple operationListener; @@ -126,7 +125,6 @@ public TransactionalBulkExecutor(CosmosAsyncContainer container, checkNotNull(inputOperations, "expected non-null inputOperations"); checkNotNull(cosmosBulkOptions, "expected non-null bulkOptions"); - this.maxMicroBatchPayloadSizeInBytes = cosmosBulkOptions.getMaxMicroBatchPayloadSizeInBytes(); this.cosmosBulkExecutionOptions = cosmosBulkOptions; this.container = container; this.bulkSpanName = "transactionalBatch." + this.container.getId(); @@ -286,31 +284,6 @@ private Flux> executeTransactionalBatch(Co String batchTrackingId = UUIDs.nonBlockingRandomUUID().toString(); PartitionKey partitionKey = cosmosBatch.getPartitionKeyValue(); - // Validate batch payload size before execution - int totalSerializedLength = this.calculateTotalSerializedLength(operations); - if (totalSerializedLength > this.maxMicroBatchPayloadSizeInBytes) { - String errorMessage = String.format( - "Transactional batch exceeds maximum payload size: %d bytes (limit: %d bytes), PK: %s", - totalSerializedLength, - this.maxMicroBatchPayloadSizeInBytes, - partitionKey); - logger.error("{}, Context: {}", errorMessage, this.operationContextText); - - // Return error responses for all operations in the batch - CosmosException payloadSizeException = BridgeInternal.createCosmosException( - HttpConstants.StatusCodes.REQUEST_ENTITY_TOO_LARGE, - errorMessage); - - return Flux.fromIterable(operations) - .map(operation -> { - TContext actualContext = this.getActualContext(operation); - return ModelBridgeInternal.createCosmosBulkOperationResponse( - operation, - payloadSizeException, - actualContext); - }); - } - logDebugOrWarning( "Executing transactional batch - {} operations, PK: {}, TrackingId: {}, Context: {}", operations.size(), @@ -427,16 +400,6 @@ private Flux> executeTransactionalBatch(Co }); } - private int calculateTotalSerializedLength(List operations) { - int totalLength = 0; - for (CosmosItemOperation operation : operations) { - if (operation instanceof CosmosItemOperationBase) { - totalLength += ((CosmosItemOperationBase) operation).getSerializedLength(this.effectiveItemSerializer); - } - } - return totalLength; - } - private TContext getActualContext(CosmosItemOperation itemOperation) { ItemBulkOperation itemBulkOperation = null; From 3773412def5ed43c4b45d4a5e5d884835660fb9a Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Fri, 19 Dec 2025 11:26:58 +0000 Subject: [PATCH 13/16] Use standard Loan pattern for resource management --- .../cosmos/spark/ItemsWriterBuilder.scala | 42 ++++++++----------- 1 file changed, 18 insertions(+), 24 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3-5/src/main/scala/com/azure/cosmos/spark/ItemsWriterBuilder.scala b/sdk/cosmos/azure-cosmos-spark_3-5/src/main/scala/com/azure/cosmos/spark/ItemsWriterBuilder.scala index 77e67a2a5f35..ea759335091b 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5/src/main/scala/com/azure/cosmos/spark/ItemsWriterBuilder.scala +++ b/sdk/cosmos/azure-cosmos-spark_3-5/src/main/scala/com/azure/cosmos/spark/ItemsWriterBuilder.scala @@ -129,21 +129,24 @@ private class ItemsWriterBuilder private def getPartitionKeyColumnNames(): Seq[String] = { try { - // Use loan pattern to ensure client is properly closed - using(createClientForPartitionKeyLookup()) { clientCacheItem => - val container = ThroughputControlHelper.getContainer( - userConfigMap, - containerConfig, - clientCacheItem, - None - ) - - // Simplified retrieval using SparkBridgeInternal directly - val containerProperties = SparkBridgeInternal.getContainerPropertiesFromCollectionCache(container) - val partitionKeyDefinition = containerProperties.getPartitionKeyDefinition - - extractPartitionKeyPaths(partitionKeyDefinition) - } + Loan( + List[Option[CosmosClientCacheItem]]( + Some(createClientForPartitionKeyLookup()) + )) + .to(clientCacheItems => { + val container = ThroughputControlHelper.getContainer( + userConfigMap, + containerConfig, + clientCacheItems(0).get, + None + ) + + // Simplified retrieval using SparkBridgeInternal directly + val containerProperties = SparkBridgeInternal.getContainerPropertiesFromCollectionCache(container) + val partitionKeyDefinition = containerProperties.getPartitionKeyDefinition + + extractPartitionKeyPaths(partitionKeyDefinition) + }) } catch { case ex: Exception => log.logWarning(s"Failed to get partition key definition for transactional writes: ${ex.getMessage}") @@ -178,14 +181,5 @@ private class ItemsWriterBuilder Seq.empty[String] } } - - // Scala loan pattern to ensure resources are properly cleaned up - private def using[A <: { def close(): Unit }, B](resource: A)(f: A => B): B = { - try { - f(resource) - } finally { - if (resource != null) resource.close() - } - } } } From ab76a79a714b932e054e5f4e1b65225661ba87c1 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Fri, 19 Dec 2025 14:05:40 +0000 Subject: [PATCH 14/16] refactor error handling to be at batch level and add tests --- .../spark/TransactionalBulkWriter.scala | 392 +++++++++++++----- .../spark/TransactionalBatchITest.scala | 193 ++++++++- 2 files changed, 487 insertions(+), 98 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala index 91fc34ee3c6b..5a90ba5e1b97 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala @@ -107,13 +107,26 @@ private class TransactionalBulkWriter private val currentBatchOperations = new mutable.ListBuffer[(ObjectNode, OperationContext, String)]() private val batchConstructionLock = new Object() - private val activeBulkWriteOperations =java.util.concurrent.ConcurrentHashMap.newKeySet[CosmosItemOperation]().asScala - private val operationContextMap = new java.util.concurrent.ConcurrentHashMap[CosmosItemOperation, OperationContext]().asScala - // Semaphore limits number of outstanding batches (not individual operations) + // Batch-level tracking - BatchTracker holds all information about a batch + private case class BatchTracker( + batchId: String, + operations: List[CosmosItemOperation], + contexts: List[OperationContext], + sourceData: List[(ObjectNode, OperationContext, String)], + attemptNumber: Int + ) + + // Map batch ID to BatchTracker + private val activeBatches = new java.util.concurrent.ConcurrentHashMap[String, BatchTracker]().asScala + // Map batch ID to collected responses for that batch + private val batchResponseCollector = new java.util.concurrent.ConcurrentHashMap[String, java.util.concurrent.ConcurrentLinkedQueue[CosmosBulkOperationResponse[Object]]]().asScala + // Map individual operation to its batch ID for response grouping + private val operationToBatchId = new java.util.concurrent.ConcurrentHashMap[CosmosItemOperation, String]().asScala + // Semaphore limits number of outstanding batches private val semaphore = new Semaphore(maxPendingBatches) - private val activeBatches = new AtomicInteger(0) - // Map each batch's first operation to the batch size for semaphore release tracking - private val batchSizeMap = new java.util.concurrent.ConcurrentHashMap[CosmosItemOperation, Integer]() + private val activeBatchCount = new AtomicInteger(0) + // Track batches being retried + private val pendingBatchRetries = java.util.concurrent.ConcurrentHashMap.newKeySet[String]().asScala private val totalScheduledMetrics = new AtomicLong(0) private val totalSuccessfulIngestionMetrics = new AtomicLong(0) @@ -212,6 +225,210 @@ private class TransactionalBulkWriter + private def processBatchResponses( + batchId: String, + batchTracker: BatchTracker, + responseQueue: java.util.concurrent.ConcurrentLinkedQueue[CosmosBulkOperationResponse[Object]]): Unit = { + + var hasRetryableError = false + var hasFatalError = false + var fatalException: Throwable = null + val successCount = new AtomicInteger(0) + + // Analyze all responses to make batch-level decision + val responses = responseQueue.asScala.toList + responses.zipWithIndex.foreach { case (resp, idx) => + val context = batchTracker.contexts(idx) + val itemResponse = resp.getResponse + + if (resp.getException != null) { + Option(resp.getException) match { + case Some(cosmosException: CosmosException) => + val statusCode = cosmosException.getStatusCode + val subStatusCode = cosmosException.getSubStatusCode + + if (shouldIgnore(statusCode, subStatusCode)) { + successCount.incrementAndGet() + } else if (shouldRetry(statusCode, subStatusCode, context)) { + hasRetryableError = true + log.logDebug(s"Batch $batchId has retryable error: $statusCode:$subStatusCode for item ${context.itemId}") + } else { + hasFatalError = true + fatalException = cosmosException + log.logError(s"Batch $batchId has fatal error: $statusCode:$subStatusCode for item ${context.itemId}") + } + case _ => + hasFatalError = true + fatalException = resp.getException + log.logError(s"Batch $batchId has unexpected exception for item ${context.itemId}", resp.getException) + } + } else if (Option(itemResponse).isEmpty || !itemResponse.isSuccessStatusCode) { + val statusCode = Option(itemResponse).map(_.getStatusCode).getOrElse(CosmosConstants.StatusCodes.Timeout) + val subStatusCode = Option(itemResponse).map(_.getSubStatusCode).getOrElse(0) + + if (shouldIgnore(statusCode, subStatusCode)) { + successCount.incrementAndGet() + } else if (shouldRetry(statusCode, subStatusCode, context)) { + hasRetryableError = true + log.logDebug(s"Batch $batchId has retryable error: $statusCode:$subStatusCode for item ${context.itemId}") + } else { + hasFatalError = true + fatalException = new BulkOperationFailedException(statusCode, subStatusCode, + s"Fatal error in batch $batchId for item ${context.itemId}", null) + log.logError(s"Batch $batchId has fatal error: $statusCode:$subStatusCode for item ${context.itemId}") + } + } else { + // Success + successCount.incrementAndGet() + } + } + + // Make batch-level decision + if (hasFatalError) { + // Fatal error - fail entire batch + log.logError(s"Batch $batchId failed with fatal error, cannot retry") + captureIfFirstFailure(fatalException) + cancelWork() + cleanupBatch(batchId, batchTracker.operations.size) + } else if (hasRetryableError) { + // Retry entire batch atomically + log.logWarning(s"Batch $batchId has retryable errors, will retry entire batch atomically") + scheduleBatchRetry(batchId, batchTracker) + cleanupBatch(batchId, batchTracker.operations.size) + } else { + // All operations succeeded + log.logDebug(s"Batch $batchId completed successfully with ${successCount.get} operations") + outputMetricsPublisher.trackWriteOperation(successCount.get, None) + totalSuccessfulIngestionMetrics.addAndGet(successCount.get) + cleanupBatch(batchId, batchTracker.operations.size) + } + } + + private def cleanupBatch(batchId: String, operationCount: Int): Unit = { + // Clean up batch tracking + activeBatches.remove(batchId) + batchResponseCollector.remove(batchId) + pendingBatchRetries.remove(batchId) + + // Release semaphore + activeBatchCount.decrementAndGet() + semaphore.release() + log.logTrace(s"Released semaphore for batch $batchId, activeBatchCount: ${activeBatchCount.get} $getThreadInfo") + + // Mark task completion for all operations in batch + (1 to operationCount).foreach(_ => markTaskCompletion()) + } + + private def scheduleBatchRetry(batchId: String, batchTracker: BatchTracker): Unit = { + this.pendingRetries.incrementAndGet() + pendingBatchRetries.add(batchId) + + // Determine if this is a timeout to apply delay + val hasTimeout = batchResponseCollector.get(batchId).exists { queue => + queue.asScala.exists { resp => + val statusCode = Option(resp.getResponse).map(_.getStatusCode) + .orElse(Option(resp.getException).collect { case ce: CosmosException => ce.getStatusCode }) + .getOrElse(CosmosConstants.StatusCodes.Timeout) + Exceptions.isTimeout(statusCode) + } + } + + val deferredRetryMono = SMono.defer(() => { + log.logWarning(s"Retrying batch $batchId with ${batchTracker.operations.size} operations, " + + s"attempt ${batchTracker.attemptNumber + 1}") + + retryBatch(batchTracker) + this.pendingRetries.decrementAndGet() + SMono.empty + }) + + if (hasTimeout) { + deferredRetryMono + .delaySubscription( + Duration( + TransactionalBulkWriter.minDelayOn408RequestTimeoutInMs + + scala.util.Random.nextInt( + TransactionalBulkWriter.maxDelayOn408RequestTimeoutInMs - TransactionalBulkWriter.minDelayOn408RequestTimeoutInMs), + TimeUnit.MILLISECONDS), + Schedulers.boundedElastic()) + .subscribeOn(Schedulers.boundedElastic()) + .subscribe() + } else { + deferredRetryMono + .subscribeOn(Schedulers.boundedElastic()) + .subscribe() + } + } + + private def retryBatch(batchTracker: BatchTracker): Unit = { + // Acquire semaphore for retry + val activeTasksSemaphoreTimeout = 10 + val numberOfIntervalsWithIdenticalActiveOperationSnapshots = new AtomicLong(0) + var activeBatchesSnapshot = mutable.Set.empty[String] + var pendingBatchRetriesSnapshot = mutable.Set.empty[String] + + while (!semaphore.tryAcquire(activeTasksSemaphoreTimeout, TimeUnit.MINUTES)) { + if (subscriptionDisposable.isDisposed) { + captureIfFirstFailure( + new IllegalStateException("Can't accept any new work - BulkWriter has been disposed already")) + } + + throwIfProgressStaled( + "Batch retry semaphore acquisition", + activeBatchesSnapshot.flatMap(id => activeBatches.get(id).map(_.operations).getOrElse(List.empty)).to[mutable.Set], + pendingBatchRetriesSnapshot.flatMap(id => activeBatches.get(id).map(_.operations).getOrElse(List.empty)).to[mutable.Set], + numberOfIntervalsWithIdenticalActiveOperationSnapshots, + allowRetryOnNewBulkWriterInstance = false) + + activeBatchesSnapshot = activeBatches.keySet.to[mutable.Set] + pendingBatchRetriesSnapshot = pendingBatchRetries.to[mutable.Set] + } + + activeBatchCount.incrementAndGet() + + // Generate new batch ID for retry + val newBatchId = java.util.UUID.randomUUID().toString + + // Recreate batch with incremented attempt numbers + val partitionKey = batchTracker.operations.head.getPartitionKeyValue + val batch = CosmosBatch.createCosmosBatch(partitionKey) + val updatedContexts = new scala.collection.mutable.ListBuffer[OperationContext]() + + batchTracker.sourceData.foreach { case (objectNode, context, _) => + batch.upsertItemOperation(objectNode) + val updatedContext = new OperationContext( + context.itemId, + context.partitionKeyValue, + context.eTag, + batchTracker.attemptNumber + 1, + context.sequenceNumber, + context.sourceItem + ) + updatedContexts += updatedContext + } + + // Create new BatchTracker for retry + val batchOperations = batch.getOperations() + val operationsList = batchOperations.asScala.toList + val newBatchTracker = BatchTracker( + batchId = newBatchId, + operations = operationsList, + contexts = updatedContexts.toList, + sourceData = batchTracker.sourceData, + attemptNumber = batchTracker.attemptNumber + 1 + ) + + // Register new batch + activeBatches.put(newBatchId, newBatchTracker) + batchResponseCollector.put(newBatchId, new java.util.concurrent.ConcurrentLinkedQueue[CosmosBulkOperationResponse[Object]]()) + operationsList.foreach(op => operationToBatchId.put(op, newBatchId)) + + log.logDebug(s"Created retry batch $newBatchId (original: ${batchTracker.batchId}) with ${operationsList.size} operations") + + // Emit retry batch + bulkInputEmitter.emitNext(batch, emitFailureHandler) + } + private def scheduleRetry( trackPendingRetryAction: () => Boolean, clearPendingRetryAction: () => Boolean, @@ -293,73 +510,46 @@ private class TransactionalBulkWriter bulkOperationResponseFlux.subscribe( resp => { - val isGettingRetried = new AtomicBoolean(false) - val shouldSkipTaskCompletion = new AtomicBoolean(false) - val isFirstOperationInBatch = new AtomicBoolean(false) try { val itemOperation = resp.getOperation - val itemOperationFound = activeBulkWriteOperations.remove(itemOperation) - val pendingRetriesFound = pendingBulkWriteRetries.remove(itemOperation) - // Check if this is the first operation in a batch (used for semaphore release) - val batchSizeOption = Option(batchSizeMap.remove(itemOperation)) - isFirstOperationInBatch.set(batchSizeOption.isDefined) - - if (pendingRetriesFound) { - pendingRetries.decrementAndGet() - } - - if (!itemOperationFound) { - // can't find the item operation in list of active operations! - logInfoOrWarning(s"Cannot find active operation for '${itemOperation.getOperationType} " + + // Find which batch this operation belongs to + val batchIdOption = operationToBatchId.get(itemOperation) + + if (batchIdOption.isEmpty) { + logInfoOrWarning(s"Cannot find batch ID for operation '${itemOperation.getOperationType} " + s"${itemOperation.getPartitionKeyValue}/${itemOperation.getId}'. This can happen when " + s"retries get re-enqueued.") - shouldSkipTaskCompletion.set(true) - } - if (pendingRetriesFound || itemOperationFound) { - // Look up context from our external map since ItemBatchOperation.getContext() returns null - val context = operationContextMap.remove(itemOperation).getOrElse( - throw new IllegalStateException(s"Cannot find context for operation: ${itemOperation.getId}") - ) - val itemResponse = resp.getResponse - - if (resp.getException != null) { - Option(resp.getException) match { - case Some(cosmosException: CosmosException) => - handleNonSuccessfulStatusCode( - context, itemOperation, itemResponse, isGettingRetried, Some(cosmosException)) - case _ => - log.logWarning( - s"unexpected failure: itemId=[${context.itemId}], partitionKeyValue=[" + - s"${context.partitionKeyValue}], encountered , attemptNumber=${context.attemptNumber}, " + - s"exceptionMessage=${resp.getException.getMessage}, " + - s"Context: ${operationContext.toString} $getThreadInfo", resp.getException) - captureIfFirstFailure(resp.getException) - cancelWork() - } - } else if (Option(itemResponse).isEmpty || !itemResponse.isSuccessStatusCode) { - handleNonSuccessfulStatusCode(context, itemOperation, itemResponse, isGettingRetried, None) - } else { - // no error case - outputMetricsPublisher.trackWriteOperation(1, None) - totalSuccessfulIngestionMetrics.getAndIncrement() - log.logTrace(s"Successfully processed operation: ${itemOperation.getOperationType} " + - s"for item ${itemOperation.getId}, PK: ${itemOperation.getPartitionKeyValue}") + } else { + val batchId = batchIdOption.get + + // Add response to batch collector + batchResponseCollector.get(batchId) match { + case Some(responseQueue) => + responseQueue.add(resp) + + // Check if we've received all responses for this batch + val batchTrackerOption = activeBatches.get(batchId) + if (batchTrackerOption.isDefined) { + val batchTracker = batchTrackerOption.get + val expectedCount = batchTracker.operations.size + val receivedCount = responseQueue.size() + + if (receivedCount == expectedCount) { + // All responses received - process batch + processBatchResponses(batchId, batchTracker, responseQueue) + } + } + case None => + log.logWarning(s"Response collector not found for batch $batchId") } } } - finally { - // Release semaphore when we process the first operation of a batch - // This indicates the entire batch has been processed - if (!isGettingRetried.get && isFirstOperationInBatch.get) { - activeBatches.decrementAndGet() - semaphore.release() - log.logTrace(s"Released semaphore for completed batch, activeBatches: ${activeBatches.get} $getThreadInfo") - } - } - - if (!shouldSkipTaskCompletion.get) { - markTaskCompletion() + catch { + case ex: Exception => + log.logError(s"Error processing bulk operation response, Context: ${operationContext.toString} $getThreadInfo", ex) + captureIfFirstFailure(ex) + cancelWork() } }, errorConsumer = Option.apply( @@ -454,14 +644,12 @@ private class TransactionalBulkWriter // Acquire semaphore before emitting batch - this limits concurrent batches val activeTasksSemaphoreTimeout = 10 val numberOfIntervalsWithIdenticalActiveOperationSnapshots = new AtomicLong(0) - var activeBulkWriteOperationsSnapshot = mutable.Set.empty[CosmosItemOperation] - var pendingBulkWriteRetriesSnapshot = mutable.Set.empty[CosmosItemOperation] + var activeBatchesSnapshot = mutable.Set.empty[String] + var pendingBatchRetriesSnapshot = mutable.Set.empty[String] - val dummyContext = new OperationContext("batch-flush", currentPartitionKey, null, 1, 0, None) - - log.logTrace(s"Before acquiring semaphore for batch emission, activeBatches: ${activeBatches.get} $getThreadInfo") + log.logTrace(s"Before acquiring semaphore for batch emission, activeBatchCount: ${activeBatchCount.get} $getThreadInfo") while (!semaphore.tryAcquire(activeTasksSemaphoreTimeout, TimeUnit.MINUTES)) { - log.logDebug(s"Not able to acquire semaphore for batch, activeBatches: ${activeBatches.get} $getThreadInfo") + log.logDebug(s"Not able to acquire semaphore for batch, activeBatchCount: ${activeBatchCount.get} $getThreadInfo") if (subscriptionDisposable.isDisposed) { captureIfFirstFailure( new IllegalStateException("Can't accept any new work - BulkWriter has been disposed already")) @@ -469,17 +657,20 @@ private class TransactionalBulkWriter throwIfProgressStaled( "Batch semaphore acquisition", - activeBulkWriteOperationsSnapshot, - pendingBulkWriteRetriesSnapshot, + activeBatchesSnapshot.flatMap(batchId => activeBatches.get(batchId).map(_.operations).getOrElse(List.empty)).to[mutable.Set], + pendingBatchRetriesSnapshot.flatMap(batchId => activeBatches.get(batchId).map(_.operations).getOrElse(List.empty)).to[mutable.Set], numberOfIntervalsWithIdenticalActiveOperationSnapshots, allowRetryOnNewBulkWriterInstance = false) - activeBulkWriteOperationsSnapshot = activeBulkWriteOperations.clone() - pendingBulkWriteRetriesSnapshot = pendingBulkWriteRetries.clone() + activeBatchesSnapshot = activeBatches.keySet.to[mutable.Set] + pendingBatchRetriesSnapshot = pendingBatchRetries.to[mutable.Set] } - activeBatches.incrementAndGet() - log.logTrace(s"Acquired semaphore for batch emission, activeBatches: ${activeBatches.get} $getThreadInfo") + activeBatchCount.incrementAndGet() + log.logTrace(s"Acquired semaphore for batch emission, activeBatchCount: ${activeBatchCount.get} $getThreadInfo") + + // Generate unique batch ID + val batchId = java.util.UUID.randomUUID().toString // Build the batch using the builder API val batch = CosmosBatch.createCosmosBatch(currentPartitionKey) @@ -493,22 +684,25 @@ private class TransactionalBulkWriter contextList += operationContext } - // After building the batch, get the operations and track them with their contexts + // After building the batch, get the operations and create BatchTracker val batchOperations = batch.getOperations() - val batchSize = batchOperations.size() + val operationsList = batchOperations.asScala.toList - // Map each operation to its context and add to tracking - for (i <- 0 until batchSize) { - val operation = batchOperations.get(i) - val context = contextList(i) - operationContextMap.put(operation, context) - activeBulkWriteOperations.add(operation) - } + // Create BatchTracker with all batch information + val batchTracker = BatchTracker( + batchId = batchId, + operations = operationsList, + contexts = contextList.toList, + sourceData = currentBatchOperations.toList, + attemptNumber = 1 + ) - // Track batch size using first operation as key - needed for semaphore release - if (batchSize > 0) { - batchSizeMap.put(batchOperations.get(0), batchSize) - } + // Register batch and map operations to batch ID + activeBatches.put(batchId, batchTracker) + batchResponseCollector.put(batchId, new java.util.concurrent.ConcurrentLinkedQueue[CosmosBulkOperationResponse[Object]]()) + operationsList.foreach(op => operationToBatchId.put(op, batchId)) + + log.logDebug(s"Created batch $batchId with ${operationsList.size} operations") // Emit the batch bulkInputEmitter.emitNext(batch, emitFailureHandler) @@ -675,7 +869,9 @@ private class TransactionalBulkWriter val operationsLog = getActiveOperationsLog(activeOperationsSnapshot) - if (sameBulkWriteOperations(pendingRetriesSnapshot ++ activeOperationsSnapshot , activeBulkWriteOperations ++ pendingBulkWriteRetries)) { + val currentActiveOperations = activeBatches.values.flatMap(_.operations).to[mutable.Set] + val currentPendingRetries = pendingBatchRetries.flatMap(batchId => activeBatches.get(batchId).map(_.operations).getOrElse(List.empty)).to[mutable.Set] + if (sameBulkWriteOperations(pendingRetriesSnapshot ++ activeOperationsSnapshot, currentActiveOperations ++ currentPendingRetries)) { numberOfIntervalsWithIdenticalActiveOperationSnapshots.incrementAndGet() log.logWarning( @@ -738,8 +934,9 @@ private class TransactionalBulkWriter finalFlushBatch() log.logInfo(s"flushAndClose invoked $getThreadInfo") + val totalActiveOperations = activeBatches.values.map(_.operations.size).sum log.logInfo(s"completed so far ${totalSuccessfulIngestionMetrics.get()}, " + - s"pending bulkWrite asks ${activeBulkWriteOperations.size} $getThreadInfo") + s"pending bulkWrite tasks (operations across ${activeBatches.size} batches): $totalActiveOperations $getThreadInfo") // error handling, if there is any error and the subscription is cancelled // the remaining tasks will not be processed hence we never reach activeTasks = 0 @@ -755,8 +952,8 @@ private class TransactionalBulkWriter logInfoOrWarning( s"Waiting for pending activeTasks $activeTasksSnapshot and/or pendingRetries " + s"$pendingRetriesSnapshot, Context: ${operationContext.toString} $getThreadInfo") - val activeOperationsSnapshot = activeBulkWriteOperations.clone() - val pendingOperationsSnapshot = pendingBulkWriteRetries.clone() + val activeOperationsSnapshot = activeBatches.values.flatMap(_.operations).to[mutable.Set] + val pendingOperationsSnapshot = pendingBatchRetries.flatMap(batchId => activeBatches.get(batchId).map(_.operations).getOrElse(List.empty)).to[mutable.Set] val awaitCompleted = pendingTasksCompleted.await(writeConfig.flushCloseIntervalInSeconds, TimeUnit.SECONDS) if (!awaitCompleted) { throwIfProgressStaled( @@ -824,7 +1021,7 @@ private class TransactionalBulkWriter logInfoOrWarning(s"invoking bulkInputEmitter.onComplete(), Context: ${operationContext.toString} $getThreadInfo") // Release any remaining batch permits - val remainingBatches = activeBatches.get() + val remainingBatches = activeBatchCount.get() if (remainingBatches > 0) { semaphore.release(remainingBatches) log.logDebug(s"Released $remainingBatches batch permits during cleanup") @@ -835,7 +1032,7 @@ private class TransactionalBulkWriter assume(activeTasks.get() <= 0) - assume(activeBulkWriteOperations.isEmpty) + assume(activeBatches.isEmpty) assume(semaphore.availablePermits() >= maxPendingBatches) if (totalScheduledMetrics.get() != totalSuccessfulIngestionMetrics.get) { @@ -893,8 +1090,9 @@ private class TransactionalBulkWriter } private def cancelWork(): Unit = { + val totalActiveOperations = activeBatches.values.map(_.operations.size).sum logInfoOrWarning(s"cancelling remaining unprocessed tasks ${activeTasks.get} " + - s"[bulkWrite tasks ${activeBulkWriteOperations.size}]" + + s"[bulkWrite tasks (operations across ${activeBatches.size} batches): $totalActiveOperations] " + s"Context: ${operationContext.toString}") subscriptionDisposable.dispose() } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala index 9130207a9c8e..3b9feae2fdbd 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala @@ -4,11 +4,14 @@ package com.azure.cosmos.spark import com.azure.cosmos.implementation.{TestConfigurations, Utils} import com.azure.cosmos.models.{PartitionKey, PartitionKeyBuilder} -import com.azure.cosmos.CosmosException +import com.azure.cosmos.{CosmosAsyncClient, CosmosException} +import com.azure.cosmos.test.faultinjection._ import com.fasterxml.jackson.databind.node.ObjectNode import org.apache.spark.sql.types._ import org.apache.spark.sql.{Row, SaveMode} +import java.time.Duration + import java.util.UUID // scalastyle:off underscore.import import scala.collection.JavaConverters._ @@ -806,4 +809,192 @@ class TransactionalBatchITest extends IntegrationSpec // 3. All batches are properly tracked and released } + "Transactional Batch" should "handle multiple batches for same partition key" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) + + // Create two sets of operations for the same partition key + // This tests that multiple successful batches can be written to the same partition key + val partitionKeyValue = UUID.randomUUID().toString + + val schema = StructType(Seq( + StructField("id", StringType, nullable = false), + StructField("pk", StringType, nullable = false), + StructField("counter", IntegerType, nullable = false) + )) + + // First, write initial items that will later be updated transactionally + val initialItems = (1 to 10).map { i => + Row(s"item-$i", partitionKeyValue, 0) + } + + val initialDf = spark.createDataFrame(initialItems.asJava, schema) + initialDf.write + .format("cosmos.oltp") + .option("spark.cosmos.accountEndpoint", cosmosEndpoint) + .option("spark.cosmos.accountKey", cosmosMasterKey) + .option("spark.cosmos.database", cosmosDatabase) + .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) + .option("spark.cosmos.write.bulk.transactional", "true") + .option("spark.cosmos.write.bulk.enabled", "true") + .mode(SaveMode.Append) + .save() + + // Now update items 1-5 to counter=1, then items 6-10 to counter=2 + // Both are atomic batches for the same partition key + // If retries from the first batch interleave with the second batch, + // atomicity would be violated + val batch1 = (1 to 5).map { i => + Row(s"item-$i", partitionKeyValue, 1) + } + + val batch2 = (6 to 10).map { i => + Row(s"item-$i", partitionKeyValue, 2) + } + + val allUpdates = batch1 ++ batch2 + val updatesDf = spark.createDataFrame(allUpdates.asJava, schema) + + // Delete existing items first since Overwrite mode is not supported in transactional mode + (1 to 10).foreach { i => + container.deleteItem(s"item-$i", new PartitionKey(partitionKeyValue), null).block() + } + + updatesDf.write + .format("cosmos.oltp") + .option("spark.cosmos.accountEndpoint", cosmosEndpoint) + .option("spark.cosmos.accountKey", cosmosMasterKey) + .option("spark.cosmos.database", cosmosDatabase) + .option("spark.cosmos.container", cosmosContainersWithPkAsPartitionKey) + .option("spark.cosmos.write.bulk.transactional", "true") + .option("spark.cosmos.write.bulk.enabled", "true") + .option("spark.cosmos.write.bulk.maxPendingOperations", "10") // Force separate batches + .mode(SaveMode.Append) + .save() + + // Verify the final state: all items should have their expected counter values + // If interleaving occurred, some updates might have been lost or inconsistent + (1 to 5).foreach { i => + val item = container.readItem( + s"item-$i", + new PartitionKey(partitionKeyValue), + classOf[ObjectNode] + ).block() + + assert(item != null, s"Item item-$i should exist") + assert(item.getItem.get("counter").asInt() == 1, + s"Item item-$i should have counter=1, but got ${item.getItem.get("counter").asInt()}") + } + + (6 to 10).foreach { i => + val item = container.readItem( + s"item-$i", + new PartitionKey(partitionKeyValue), + classOf[ObjectNode] + ).block() + + assert(item != null, s"Item item-$i should exist") + assert(item.getItem.get("counter").asInt() == 2, + s"Item item-$i should have counter=2, but got ${item.getItem.get("counter").asInt()}") + } + } + + it should "handle batch-level retries for retriable errors without interleaving" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + + val partitionKeyValue = UUID.randomUUID().toString + + val schema = StructType(Seq( + StructField("id", StringType, nullable = false), + StructField("pk", StringType, nullable = false), + StructField("counter", IntegerType, nullable = false) + )) + + // Configuration for Spark connector - must match exactly for cache lookup + val cfg = Map( + "spark.cosmos.accountEndpoint" -> cosmosEndpoint, + "spark.cosmos.accountKey" -> cosmosMasterKey, + "spark.cosmos.database" -> cosmosDatabase, + "spark.cosmos.container" -> cosmosContainersWithPkAsPartitionKey + ) + + // Create initial items with counter = 0 + // This FIRST write ensures Spark creates the client and caches it + val initialItems = (1 to 10).map { i => + Row(s"item-$i", partitionKeyValue, 0) + } + + val initialDf = spark.createDataFrame(initialItems.asJava, schema) + initialDf.write + .format("cosmos.oltp") + .options(cfg) + .option("spark.cosmos.write.bulk.transactional", "true") + .option("spark.cosmos.write.bulk.enabled", "true") + .mode(SaveMode.Append) + .save() + + // NOW get the actual client that Spark created and cached + val clientFromCache = udf.CosmosAsyncClientCache + .getCosmosClientFromCache(cfg) + .getClient + .asInstanceOf[CosmosAsyncClient] + + val container = clientFromCache.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) + + // Configure fault injection to inject retriable 429 TOO_MANY_REQUEST errors on BATCH_ITEM operations + // 429 errors are retriable and should trigger batch-level retry without aborting the job + // This tests that transactional batches handle retries at the batch level + val faultInjectionRule = new FaultInjectionRuleBuilder("batch-retry-interleaving-test-" + UUID.randomUUID()) + .condition( + new FaultInjectionConditionBuilder() + .operationType(FaultInjectionOperationType.BATCH_ITEM) + .build() + ) + .result( + FaultInjectionResultBuilders + .getResultBuilder(FaultInjectionServerErrorType.TOO_MANY_REQUEST) + .times(1) // Inject 429 on first batch operation, then allow retries to succeed + .build() + ) + .duration(Duration.ofMinutes(5)) + .build() + + // Configure the fault injection rule on the container + CosmosFaultInjectionHelper.configureFaultInjectionRules(container, java.util.Collections.singletonList(faultInjectionRule)).block() + + try { + // Now write just ONE more item using NEW ID to avoid conflicts with the initial write + // This is the absolute simplest test case to verify batch-level retry handling + // With maxPendingOperations=1, this single item becomes its own batch + val singleItem = Seq(Row("item-11", partitionKeyValue, 1)) + + val updatesDf = spark.createDataFrame(singleItem.asJava, schema) + + updatesDf.write + .format("cosmos.oltp") + .options(cfg) + .option("spark.cosmos.write.bulk.transactional", "true") + .option("spark.cosmos.write.bulk.enabled", "true") + .option("spark.cosmos.write.bulk.maxPendingOperations", "1") // Ensure minimal batch size + .mode(SaveMode.Append) + .save() + + // Verify that fault injection triggered and was retried successfully + // This confirms that retriable errors (429) trigger batch-level retries + val hitCount = faultInjectionRule.getHitCount + assert(hitCount > 0, s"Fault injection should have tracked BATCH_ITEM operations, but hit count was $hitCount") + + // Verify the single item was written correctly despite the injected 429 error and retry + val item11 = container.readItem("item-11", new PartitionKey(partitionKeyValue), classOf[ObjectNode]).block() + assert(item11 != null, "Item item-11 should exist") + assert(item11.getItem.get("counter").asInt() == 1, + s"Item item-11 should have counter=1, but got ${item11.getItem.get("counter").asInt()}") + } finally { + // Clean up: disable the fault injection rule + faultInjectionRule.disable() + } + } + } From 54be55a6a70fd86426d6b7d4d8b640966a052841 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Mon, 12 Jan 2026 15:22:08 +0000 Subject: [PATCH 15/16] Fix Scala 2.12/2.13 cross-compilation --- .../spark/TransactionalBulkWriter.scala | 42 +++++++++---------- 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala index 5a90ba5e1b97..070025988060 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala @@ -364,8 +364,8 @@ private class TransactionalBulkWriter // Acquire semaphore for retry val activeTasksSemaphoreTimeout = 10 val numberOfIntervalsWithIdenticalActiveOperationSnapshots = new AtomicLong(0) - var activeBatchesSnapshot = mutable.Set.empty[String] - var pendingBatchRetriesSnapshot = mutable.Set.empty[String] + var activeBatchesSnapshot = Set.empty[String] + var pendingBatchRetriesSnapshot = Set.empty[String] while (!semaphore.tryAcquire(activeTasksSemaphoreTimeout, TimeUnit.MINUTES)) { if (subscriptionDisposable.isDisposed) { @@ -375,13 +375,13 @@ private class TransactionalBulkWriter throwIfProgressStaled( "Batch retry semaphore acquisition", - activeBatchesSnapshot.flatMap(id => activeBatches.get(id).map(_.operations).getOrElse(List.empty)).to[mutable.Set], - pendingBatchRetriesSnapshot.flatMap(id => activeBatches.get(id).map(_.operations).getOrElse(List.empty)).to[mutable.Set], + activeBatchesSnapshot.flatMap(id => activeBatches.get(id).map(_.operations).getOrElse(List.empty)).toSet, + pendingBatchRetriesSnapshot.flatMap(id => activeBatches.get(id).map(_.operations).getOrElse(List.empty)).toSet, numberOfIntervalsWithIdenticalActiveOperationSnapshots, allowRetryOnNewBulkWriterInstance = false) - activeBatchesSnapshot = activeBatches.keySet.to[mutable.Set] - pendingBatchRetriesSnapshot = pendingBatchRetries.to[mutable.Set] + activeBatchesSnapshot = activeBatches.keySet.toSet + pendingBatchRetriesSnapshot = pendingBatchRetries.toSet } activeBatchCount.incrementAndGet() @@ -644,8 +644,8 @@ private class TransactionalBulkWriter // Acquire semaphore before emitting batch - this limits concurrent batches val activeTasksSemaphoreTimeout = 10 val numberOfIntervalsWithIdenticalActiveOperationSnapshots = new AtomicLong(0) - var activeBatchesSnapshot = mutable.Set.empty[String] - var pendingBatchRetriesSnapshot = mutable.Set.empty[String] + var activeBatchesSnapshot = Set.empty[String] + var pendingBatchRetriesSnapshot = Set.empty[String] log.logTrace(s"Before acquiring semaphore for batch emission, activeBatchCount: ${activeBatchCount.get} $getThreadInfo") while (!semaphore.tryAcquire(activeTasksSemaphoreTimeout, TimeUnit.MINUTES)) { @@ -657,13 +657,13 @@ private class TransactionalBulkWriter throwIfProgressStaled( "Batch semaphore acquisition", - activeBatchesSnapshot.flatMap(batchId => activeBatches.get(batchId).map(_.operations).getOrElse(List.empty)).to[mutable.Set], - pendingBatchRetriesSnapshot.flatMap(batchId => activeBatches.get(batchId).map(_.operations).getOrElse(List.empty)).to[mutable.Set], + activeBatchesSnapshot.flatMap(batchId => activeBatches.get(batchId).map(_.operations).getOrElse(List.empty)).toSet, + pendingBatchRetriesSnapshot.flatMap(batchId => activeBatches.get(batchId).map(_.operations).getOrElse(List.empty)).toSet, numberOfIntervalsWithIdenticalActiveOperationSnapshots, allowRetryOnNewBulkWriterInstance = false) - activeBatchesSnapshot = activeBatches.keySet.to[mutable.Set] - pendingBatchRetriesSnapshot = pendingBatchRetries.to[mutable.Set] + activeBatchesSnapshot = activeBatches.keySet.toSet + pendingBatchRetriesSnapshot = pendingBatchRetries.toSet } activeBatchCount.incrementAndGet() @@ -819,7 +819,7 @@ private class TransactionalBulkWriter } private[this] def getActiveOperationsLog( - activeOperationsSnapshot: mutable.Set[CosmosItemOperation]): String = { + activeOperationsSnapshot: Set[CosmosItemOperation]): String = { val sb = new StringBuilder() activeOperationsSnapshot @@ -840,8 +840,8 @@ private class TransactionalBulkWriter private[this] def sameBulkWriteOperations ( - snapshot: mutable.Set[CosmosItemOperation], - current: mutable.Set[CosmosItemOperation] + snapshot: Set[CosmosItemOperation], + current: Set[CosmosItemOperation] ): Boolean = { if (snapshot.size != current.size) { @@ -861,16 +861,16 @@ private class TransactionalBulkWriter private[this] def throwIfProgressStaled ( operationName: String, - activeOperationsSnapshot: mutable.Set[CosmosItemOperation], - pendingRetriesSnapshot: mutable.Set[CosmosItemOperation], + activeOperationsSnapshot: Set[CosmosItemOperation], + pendingRetriesSnapshot: Set[CosmosItemOperation], numberOfIntervalsWithIdenticalActiveOperationSnapshots: AtomicLong, allowRetryOnNewBulkWriterInstance: Boolean ): Unit = { val operationsLog = getActiveOperationsLog(activeOperationsSnapshot) - val currentActiveOperations = activeBatches.values.flatMap(_.operations).to[mutable.Set] - val currentPendingRetries = pendingBatchRetries.flatMap(batchId => activeBatches.get(batchId).map(_.operations).getOrElse(List.empty)).to[mutable.Set] + val currentActiveOperations = activeBatches.values.flatMap(_.operations).toSet + val currentPendingRetries = pendingBatchRetries.flatMap(batchId => activeBatches.get(batchId).map(_.operations).getOrElse(List.empty)).toSet if (sameBulkWriteOperations(pendingRetriesSnapshot ++ activeOperationsSnapshot, currentActiveOperations ++ currentPendingRetries)) { numberOfIntervalsWithIdenticalActiveOperationSnapshots.incrementAndGet() @@ -952,8 +952,8 @@ private class TransactionalBulkWriter logInfoOrWarning( s"Waiting for pending activeTasks $activeTasksSnapshot and/or pendingRetries " + s"$pendingRetriesSnapshot, Context: ${operationContext.toString} $getThreadInfo") - val activeOperationsSnapshot = activeBatches.values.flatMap(_.operations).to[mutable.Set] - val pendingOperationsSnapshot = pendingBatchRetries.flatMap(batchId => activeBatches.get(batchId).map(_.operations).getOrElse(List.empty)).to[mutable.Set] + val activeOperationsSnapshot = activeBatches.values.flatMap(_.operations).toSet + val pendingOperationsSnapshot = pendingBatchRetries.flatMap(batchId => activeBatches.get(batchId).map(_.operations).getOrElse(List.empty)).toSet val awaitCompleted = pendingTasksCompleted.await(writeConfig.flushCloseIntervalInSeconds, TimeUnit.SECONDS) if (!awaitCompleted) { throwIfProgressStaled( From ee7de52f2b5fd49152f47805f5d2f6d003f262ca Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Mon, 12 Jan 2026 15:38:36 +0000 Subject: [PATCH 16/16] fix broken link since doc refresh --- sdk/cosmos/azure-cosmos-spark_3/docs/configuration-reference.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/docs/configuration-reference.md b/sdk/cosmos/azure-cosmos-spark_3/docs/configuration-reference.md index 53357f1e22f8..5d9907342ba7 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/docs/configuration-reference.md +++ b/sdk/cosmos/azure-cosmos-spark_3/docs/configuration-reference.md @@ -67,7 +67,7 @@ | `spark.cosmos.write.point.maxConcurrency` | None | Cosmos DB Item Write Max concurrency. If not specified it will be determined based on the Spark executor VM Size | | `spark.cosmos.write.bulk.maxPendingOperations` | None | Cosmos DB Item Write bulk mode maximum pending operations. Defines a limit of bulk operations being processed concurrently. If not specified it will be determined based on the Spark executor VM Size. If the volume of data is large for the provisioned throughput on the destination container, this setting can be adjusted by following the estimation of `1000 x Cores` | | `spark.cosmos.write.bulk.enabled` | `true` | Cosmos DB Item Write bulk enabled | -| `spark.cosmos.write.bulk.transactional` | `false` | Enable transactional batch mode for bulk writes. When enabled, all operations for the same partition key are executed atomically (all succeed or all fail). Requires ordering and clustering by partition key columns. Only supports upsert operations. Cannot exceed 100 operations or 2MB per partition key. **Note**: For containers using hierarchical partition keys (HPK), transactional scope applies only to **logical partitions** (complete partition key paths), not partial top-level keys. See [Transactional Batch documentation](https://learn.microsoft.com/azure/cosmos-db/nosql/transactional-batch) for details. |\n| `spark.cosmos.write.bulk.targetedPayloadSizeInBytes` | `220201` | When the targeted payload size is reached for buffered documents, the request is sent to the backend. The default value is optimized for small documents <= 10 KB - when documents often exceed 110 KB, it can help to increase this value to up to about `1500000` (should still be smaller than 2 MB). | +| `spark.cosmos.write.bulk.transactional` | `false` | Enable transactional batch mode for bulk writes. When enabled, all operations for the same partition key are executed atomically (all succeed or all fail). Requires ordering and clustering by partition key columns. Only supports upsert operations. Cannot exceed 100 operations or 2MB per partition key. **Note**: For containers using hierarchical partition keys (HPK), transactional scope applies only to **logical partitions** (complete partition key paths), not partial top-level keys. See [Transactional Batch documentation](https://learn.microsoft.com/azure/cosmos-db/transactional-batch) for details. |\n| `spark.cosmos.write.bulk.targetedPayloadSizeInBytes` | `220201` | When the targeted payload size is reached for buffered documents, the request is sent to the backend. The default value is optimized for small documents <= 10 KB - when documents often exceed 110 KB, it can help to increase this value to up to about `1500000` (should still be smaller than 2 MB). | | `spark.cosmos.write.bulk.initialBatchSize` | `100` | Cosmos DB initial bulk micro batch size - a micro batch will be flushed to the backend when the number of documents enqueued exceeds this size - or the target payload size is met. The micro batch size is getting automatically tuned based on the throttling rate. By default the initial micro batch size is 100. Reduce this when you want to avoid that the first few requests consume too many RUs. | | `spark.cosmos.write.bulk.maxBatchSize` | `100` | Cosmos DB max. bulk micro batch size - a micro batch will be flushed to the backend when the number of documents enqueued exceeds this size - or the target payload size is met. The micro batch size is getting automatically tuned based on the throttling rate. By default the max. micro batch size is 100. Use this setting only when migrating Spark 2.4 workloads - for other scenarios relying on the auto-tuning combined with throughput control will result in better experience. | | `spark.cosmos.write.flush.noProgress.maxIntervalInSeconds` | `180` | The time interval in seconds that write operations will wait when no progress can be made for bulk writes before forcing a retry. The retry will reinitialize the bulk write process - so, any delays on the retry can be sure to be actual service issues. The default value of 3 min should be sufficient to prevent false negatives when there is a short service-side write unavailability - like for partition splits or merges. Increase it only if you regularly see these transient errors to exceed a time period of 180 seconds. |