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..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 @@ -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,9 +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( + userConfigMap, + inputSchema + ) + + private[this] val containerConfig = CosmosContainerConfig.parseCosmosContainerConfig( + userConfigMap + ) + override def toBatch(): BatchWrite = new ItemsBatchWriter( - userConfig.asCaseSensitiveMap().asScala.toMap, + userConfigMap, inputSchema, cosmosClientStateHandles, diagnosticsConfig, @@ -66,12 +81,105 @@ private class ItemsWriterBuilder override def toStreaming: StreamingWrite = new ItemsBatchWriter( - userConfig.asCaseSensitiveMap().asScala.toMap, + userConfigMap, inputSchema, cosmosClientStateHandles, diagnosticsConfig, sparkEnvironmentInfo) override def supportedCustomMetrics(): Array[CustomMetric] = supportedCosmosMetrics + + 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) { + // 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.bulkTransactional) { + // 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 { + 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}") + 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] + } + } } } 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..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.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. | 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..b7a206aec6d6 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,75 @@ 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. 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 (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 + +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 b2137fd09dcd..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 @@ -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, @@ -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,14 @@ private object CosmosWriteConfig { parseFromStringFunction = bulkEnabledAsString => bulkEnabledAsString.toBoolean, helpMessage = "Cosmos DB Item Write bulk enabled") + private val bulkTransactional = CosmosConfigEntry[Boolean](key = CosmosConfigNames.WriteBulkTransactional, + mandatory = false, + defaultValue = Option.apply(false), + parseFromStringFunction = bulkTransactionalAsString => bulkTransactionalAsString.toBoolean, + 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), mandatory = false, @@ -1758,6 +1769,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) @@ -1788,6 +1800,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 726a6fe461ac..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 @@ -62,14 +62,28 @@ private abstract class CosmosWriterBase( private val writer: AtomicReference[AsyncItemWriter] = new AtomicReference( if (cosmosWriteConfig.bulkEnabled) { - new BulkWriter( - container, - cosmosTargetContainerConfig, - partitionKeyDefinition, - cosmosWriteConfig, - diagnosticsConfig, - getOutputMetricsPublisher(), - commitAttempt.getAndIncrement()) + // Use TransactionalBulkWriter if bulkTransactional config is enabled + val useTransactional = cosmosWriteConfig.bulkTransactional + + if (useTransactional) { + 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..070025988060 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala @@ -0,0 +1,1281 @@ +// 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} +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 + + // 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 { + // 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 + } + // Validate write strategy for transactional batches + 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, " + + s"maxPendingOperations: $maxPendingOperations, maxConcurrentPartitions: $maxConcurrentPartitions ...") + + + 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 activeTasks = new AtomicInteger(0) + private val errorCaptureFirstException = new AtomicReference[Throwable]() + 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() + + // 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 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) + + 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 def initializeOperationContext(): SparkTaskContext = { + val taskContext = TaskContext.get + + val diagnosticsContext: DiagnosticsContext = DiagnosticsContext(UUIDs.nonBlockingRandomUUID(), "transactional-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 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 = Set.empty[String] + var pendingBatchRetriesSnapshot = 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)).toSet, + pendingBatchRetriesSnapshot.flatMap(id => activeBatches.get(id).map(_.operations).getOrElse(List.empty)).toSet, + numberOfIntervalsWithIdenticalActiveOperationSnapshots, + allowRetryOnNewBulkWriterInstance = false) + + activeBatchesSnapshot = activeBatches.keySet.toSet + pendingBatchRetriesSnapshot = pendingBatchRetries.toSet + } + + 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, + 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 => { + try { + val itemOperation = resp.getOperation + + // 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.") + } 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") + } + } + } + catch { + case ex: Exception => + log.logError(s"Error processing bulk operation response, Context: ${operationContext.toString} $getThreadInfo", ex) + captureIfFirstFailure(ex) + cancelWork() + } + }, + 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 operationContext = new OperationContext( + getId(objectNode), + partitionKeyValue, + getETag(objectNode), + 1, + monotonicOperationCounter.incrementAndGet(), + None) + + 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"TransactionalBulkWriter scheduleWrite attemptCnt: ${operationContext.attemptNumber}, " + + s"Context: ${operationContext.toString} $getThreadInfo") + } + + scheduleBulkWriteInternal(partitionKeyValue, objectNode, operationContext) + } + + private def scheduleBulkWriteInternal(partitionKeyValue: PartitionKey, + objectNode: ObjectNode, + operationContext: OperationContext): Unit = { + + // 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 + } + } + + 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 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)) { + 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")) + } + + throwIfProgressStaled( + "Batch semaphore acquisition", + 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.toSet + pendingBatchRetriesSnapshot = pendingBatchRetries.toSet + } + + 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) + + // 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 create BatchTracker + val batchOperations = batch.getOperations() + val operationsList = batchOperations.asScala.toList + + // Create BatchTracker with all batch information + val batchTracker = BatchTracker( + batchId = batchId, + operations = operationsList, + contexts = contextList.toList, + sourceData = currentBatchOperations.toList, + attemptNumber = 1 + ) + + // 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) + + // Clear the buffer + currentBatchOperations.clear() + currentPartitionKey = null + } + } + + private[this] def finalFlushBatch(): Unit = { + batchConstructionLock.synchronized { + flushCurrentBatch() + } + } + + //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") + + // During retry, use the original objectNode to ensure consistency + 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: Set[CosmosItemOperation]): 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})") + }) + + sb.toString() + } + + private[this] def sameBulkWriteOperations + ( + snapshot: Set[CosmosItemOperation], + current: 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 throwIfProgressStaled + ( + operationName: String, + activeOperationsSnapshot: Set[CosmosItemOperation], + pendingRetriesSnapshot: Set[CosmosItemOperation], + numberOfIntervalsWithIdenticalActiveOperationSnapshots: AtomicLong, + allowRetryOnNewBulkWriterInstance: Boolean + ): Unit = { + + val operationsLog = getActiveOperationsLog(activeOperationsSnapshot) + + 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() + 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) { + writeConfig.maxInitialNoProgressIntervalInSeconds + } else { + writeConfig.maxRetryNoProgressIntervalInSeconds + } + val maxAllowedIntervalWithoutAnyProgressExceeded = secondsWithoutProgress >= maxNoProgressIntervalInSeconds + + if (maxAllowedIntervalWithoutAnyProgressExceeded) { + + val retriableRemainingOperations = if (allowRetryOnNewBulkWriterInstance) { + Some( + (pendingRetriesSnapshot ++ activeOperationsSnapshot) + .toList + .sortBy(op => op.getContext[OperationContext].sequenceNumber) + ) + } else { + 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() + } + + 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()) { + // Flush any remaining batched operations before closing + 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 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 + // 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 = 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( + "FlushAndClose", + activeOperationsSnapshot, + pendingOperationsSnapshot, + 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, " + + 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, " + + s"PendingRetries: $pendingRetriesSnapshot, Buffered tasks: $buffered " + + s"Attempt: ${numberOfIntervalsWithIdenticalActiveOperationSnapshots.get} - " + + s"Context: ${operationContext.toString} $getThreadInfo") + } + + activeOperationsSnapshot.foreach(operation => { + // 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") + }) + } + } + + 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") + // Release any remaining batch permits + val remainingBatches = activeBatchCount.get() + if (remainingBatches > 0) { + semaphore.release(remainingBatches) + log.logDebug(s"Released $remainingBatches batch permits during cleanup") + } + bulkInputEmitter.emitComplete(TransactionalBulkWriter.emitFailureHandlerForComplete) + + throwIfCapturedExceptionExists() + + assume(activeTasks.get() <= 0) + + assume(activeBatches.isEmpty) + assume(semaphore.availablePermits() >= maxPendingBatches) + + 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() + 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 = { + val totalActiveOperations = activeBatches.values.map(_.operations.size).sum + logInfoOrWarning(s"cancelling remaining unprocessed tasks ${activeTasks.get} " + + s"[bulkWrite tasks (operations across ${activeBatches.size} batches): $totalActiveOperations] " + + s"Context: ${operationContext.toString}") + subscriptionDisposable.dispose() + } + + private def shouldIgnore(statusCode: Int, subStatusCode: Int): Boolean = { + // 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) { + // 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") + + returnValue + } + + 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 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 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() + + // 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 * 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 * 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 + maxPendingBatchesPerJVM, + BULK_WRITER_RESPONSES_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..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, 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 new file mode 100644 index 000000000000..3b9feae2fdbd --- /dev/null +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala @@ -0,0 +1,1000 @@ +// 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.{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._ +// 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 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.bulk.transactional", "true") + .option("spark.cosmos.write.bulk.enabled", "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 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 + + // 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 DataFrame with simple schema + 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, "TestItem") + ) + + val operationsDf = spark.createDataFrame(batchOperations.asJava, schema) + + // 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.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}") + + // 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}") + + // 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 { + 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.bulk.transactional", "true") + .option("spark.cosmos.write.bulk.enabled", "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.bulk.transactional", "true") + .option("spark.cosmos.write.bulk.enabled", "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.bulk.transactional", "true") + .option("spark.cosmos.write.bulk.enabled", "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.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) + // 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 testPrimaryKey and testSecondaryKey" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + + // Create container with hierarchical partition keys (testPrimaryKey, testSecondaryKey) + 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("/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) + 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("testPrimaryKey", StringType, nullable = false), + StructField("testSecondaryKey", StringType, nullable = false), + StructField("testPrice", 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.bulk.transactional", "true") + .option("spark.cosmos.write.bulk.enabled", "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("testPrice").asDouble() shouldEqual 100.5 + + val item2 = container.readItem(item2Id, pk, classOf[ObjectNode]).block() + item2 should not be null + item2.getItem.get("testPrice").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 (testPrimaryKey, testSecondaryKey) + 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("/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) + 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("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("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) + )) + + 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.bulk.transactional", "true") + .option("spark.cosmos.write.bulk.enabled", "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("testPrice").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 testPrimaryKey/testSecondaryKey combinations" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + + // 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, + new com.azure.cosmos.models.PartitionKeyDefinition() + ) + val paths = new java.util.ArrayList[String]() + 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) + 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("testPrimaryKey", StringType, nullable = false), + StructField("testSecondaryKey", StringType, nullable = false), + StructField("testPrice", 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.bulk.transactional", "true") + .option("spark.cosmos.write.bulk.enabled", "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 (testPrimaryKey, testSecondaryKey) + 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("/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) + cosmosClient.getDatabase(cosmosDatabase).createContainer(containerProperties).block() + val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(containerName) + + 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("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 => + 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.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) + // 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 + } + } + + 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.bulk.transactional", "true") + .option("spark.cosmos.write.bulk.enabled", "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" + } + } + } + + 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 + } + + "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() + } + } + +} 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..8a4ee1e7c7f3 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/TransactionalBulkExecutor.java @@ -0,0 +1,460 @@ +// 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 AsyncDocumentClient docClientWrapper; + private final String operationContextText; + private final OperationContextAndListenerTuple operationListener; + private final ThrottlingRetryOptions throttlingRetryOptions; + private final Flux inputBatches; + + private final TContext batchContext; + 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 CosmosAsyncClient cosmosClient; + private final String bulkSpanName; + 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.cosmosBulkExecutionOptions = cosmosBulkOptions; + this.container = container; + this.bulkSpanName = "transactionalBatch." + this.container.getId(); + this.inputBatches = 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. + batchContext = (TContext) cosmosBulkExecutionOptions.getLegacyBatchScopedContext(); + operationListener = cosmosBulkExecutionOptions.getOperationContextAndListenerTuple(); + if (operationListener != null && + operationListener.getOperationContext() != null) { + operationContextText = identifier + "[" + operationListener.getOperationContext().toString() + "]"; + } else { + operationContextText = identifier +"[n/a]"; + } + + this.diagnosticsTracker = cosmosBulkExecutionOptions.getDiagnosticsTracker(); + + // For transactional batches, no sinks needed - batches are pre-constructed + mainSourceCompleted = new AtomicBoolean(false); + totalCount = new AtomicInteger(0); + + // No periodic flush scheduler needed - batches arrive pre-constructed from writer + + Scheduler schedulerSnapshotFromOptions = cosmosBulkOptions.getSchedulerOverride(); + this.executionScheduler = schedulerSnapshotFromOptions != null + ? schedulerSnapshotFromOptions + : CosmosSchedulers.BULK_EXECUTOR_BOUNDED_ELASTIC; + + logger.info("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 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); + logger.debug("Shutdown complete, Context: {}", this.operationContextText); + } + } + + 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() { + + // 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(); + int maxConcurrentBatches = nullableMaxConcurrentCosmosPartitions != null ? + Math.max(256, nullableMaxConcurrentCosmosPartitions) : 256; + + logDebugOrWarning("TransactionalBulkExecutor.executeCore with MaxConcurrentBatches: {}, Context: {}", + maxConcurrentBatches, + this.operationContextText); + + return this.inputBatches + .publishOn(this.executionScheduler) + .onErrorMap(throwable -> { + logger.warn("{}: Error observed when processing input batches. Cause: {}, Context: {}", + getThreadInfo(), + throwable.getMessage(), + this.operationContextText, + throwable); + return throwable; + }) + .flatMap( + cosmosBatch -> executeTransactionalBatch(cosmosBatch), + maxConcurrentBatches) + .subscribeOn(this.executionScheduler); + } + + private Flux> executeTransactionalBatch(CosmosBatch cosmosBatch) { + // Extract operations from the pre-built batch + List operations = cosmosBatch.getOperations(); + + if (operations == null || operations.isEmpty()) { + logger.trace("Empty batch received, Context: {}", this.operationContextText); + return Flux.empty(); + } + + String batchTrackingId = UUIDs.nonBlockingRandomUUID().toString(); + 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( + partitionKey, + 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(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); + + // 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); + }); + } + + // 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( + "Transactional batch operation failed - PK: {}, Status: {}, Operation: {}, Context: {}", + String.valueOf(partitionKey), + operationResult.getStatusCode(), + 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); + }); + }); + } + + 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 void completeAllSinks() { + logInfoOrWarning("Completing execution, Context: {}", this.operationContextText); + logger.debug("Executor service shut down, Context: {}", this.operationContextText); + this.shutdown(); + } + + 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(); + } +}