Fixing hangs when bulk ingesting via Cosmos Spark into containers with >255 physical partitions (Azure#26017)

* Fixing hangs when bulk ingesting via Cosmos Spark into containers with >255 physical partitions

* Update BulkExecutor.java

* Reacting to code review feedback

* Reacting to CR feedback

Author: FabianMeiswinkel

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

@@ -59,6 +59,7 @@ private[spark] object CosmosConfigNames {
   val DiagnosticsMode = "spark.cosmos.diagnostics"
   val WriteBulkEnabled = "spark.cosmos.write.bulk.enabled"
   val WriteBulkMaxPendingOperations = "spark.cosmos.write.bulk.maxPendingOperations"
+  val WriteBulkMaxConcurrentPartitions = "spark.cosmos.write.bulk.maxConcurrentCosmosPartitions"
   val WritePointMaxConcurrency = "spark.cosmos.write.point.maxConcurrency"
   val WriteStrategy = "spark.cosmos.write.strategy"
   val WriteMaxRetryCount = "spark.cosmos.write.maxRetryCount"
@@ -106,6 +107,7 @@ private[spark] object CosmosConfigNames {
     DiagnosticsMode,
     WriteBulkEnabled,
     WriteBulkMaxPendingOperations,
+    WriteBulkMaxConcurrentPartitions,
     WritePointMaxConcurrency,
     WriteStrategy,
     WriteMaxRetryCount,
@@ -458,8 +460,9 @@ private object ItemWriteStrategy extends Enumeration {
 private case class CosmosWriteConfig(itemWriteStrategy: ItemWriteStrategy,
                                      maxRetryCount: Int,
                                      bulkEnabled: Boolean,
-                                     bulkMaxPendingOperations: Option[Int] = Option.empty,
-                                     pointMaxConcurrency: Option[Int] = Option.empty)
+                                     bulkMaxPendingOperations: Option[Int] = None,
+                                     pointMaxConcurrency: Option[Int] = None,
+                                     maxConcurrentCosmosPartitions: Option[Int] = None)
 
 private object CosmosWriteConfig {
   private val bulkEnabled = CosmosConfigEntry[Boolean](key = CosmosConfigNames.WriteBulkEnabled,
@@ -476,6 +479,18 @@ private object CosmosWriteConfig {
     helpMessage = s"Cosmos DB Item Write Max Pending Operations." +
       s" If not specified it will be determined based on the Spark executor VM Size")
 
+  private val bulkMaxConcurrentPartitions = CosmosConfigEntry[Int](
+    key = CosmosConfigNames.WriteBulkMaxConcurrentPartitions,
+    mandatory = false,
+    parseFromStringFunction = bulkMaxConcurrencyAsString => bulkMaxConcurrencyAsString.toInt,
+    helpMessage = s"Cosmos DB Item Write Max Concurrent Cosmos Partitions." +
+      s" If not specified it will be determined based on the number of the container's physical partitions -" +
+      s" which would indicate every Spark partition is expected to have data from all Cosmos physical partitions." +
+      s" If specified it indicates from at most how many Cosmos Physical Partitions each Spark partition contains" +
+      s" data. So this config can be used to make bulk processing more efficient when input data in Spark has been" +
+      s" repartitioned to balance to how many Cosmos partitions each Spark partition needs to write. This is mainly" +
+      s" useful for very large containers (with hundreds of physical partitions).")
+
   private val pointWriteConcurrency = CosmosConfigEntry[Int](key = CosmosConfigNames.WritePointMaxConcurrency,
     mandatory = false,
     parseFromStringFunction = bulkMaxConcurrencyAsString => bulkMaxConcurrencyAsString.toInt,
@@ -517,7 +532,8 @@ private object CosmosWriteConfig {
       maxRetryCountOpt.get,
       bulkEnabled = bulkEnabledOpt.get,
       bulkMaxPendingOperations = CosmosConfigEntry.parse(cfg, bulkMaxPendingOperations),
-      pointMaxConcurrency = CosmosConfigEntry.parse(cfg, pointWriteConcurrency))
+      pointMaxConcurrency = CosmosConfigEntry.parse(cfg, pointWriteConcurrency),
+      maxConcurrentCosmosPartitions = CosmosConfigEntry.parse(cfg, bulkMaxConcurrentPartitions))
   }
 }
 

sdk/cosmos/azure-cosmos-spark_3-1_2-12/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala

@@ -7,6 +7,7 @@
 import com.azure.cosmos.implementation.Configs;
 import com.azure.cosmos.implementation.Constants;
 import com.azure.cosmos.implementation.CosmosPagedFluxOptions;
+import com.azure.cosmos.implementation.CosmosSchedulers;
 import com.azure.cosmos.implementation.Document;
 import com.azure.cosmos.implementation.DocumentCollection;
 import com.azure.cosmos.implementation.HttpConstants;
@@ -865,7 +866,7 @@ public <TContext> Flux<CosmosBulkOperationResponse<TContext>> executeBulkOperati
         return Flux.deferContextual(context -> {
             final BulkExecutor<TContext> executor = new BulkExecutor<>(this, operations, cosmosBulkExecutionOptions);
 
-            return executor.execute();
+            return executor.execute().publishOn(CosmosSchedulers.BULK_EXECUTOR_BOUNDED_ELASTIC);
         });
     }
 

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java

@@ -9,6 +9,7 @@
 public class CosmosSchedulers {
     private final static String COSMOS_PARALLEL_THREAD_NAME =  "cosmos-parallel";
     private final static String TRANSPORT_RESPONSE_BOUNDED_ELASTIC_THREAD_NAME = "transport-response-bounded-elastic";
+    private final static String BULK_EXECUTOR_BOUNDED_ELASTIC_THREAD_NAME = "bulk-executor-bounded-elastic";
     private final static int TTL_FOR_SCHEDULER_WORKER_IN_SECONDS = 60; // same as BoundedElasticScheduler.DEFAULT_TTL_SECONDS
 
     // Using a custom parallel scheduler to be able to schedule retries etc.
@@ -27,4 +28,13 @@ public class CosmosSchedulers {
         TTL_FOR_SCHEDULER_WORKER_IN_SECONDS,
         true
     );
+
+    // Custom bounded elastic scheduler process bulk execution tasks
+    public final static Scheduler BULK_EXECUTOR_BOUNDED_ELASTIC = Schedulers.newBoundedElastic(
+        2 * Schedulers.DEFAULT_BOUNDED_ELASTIC_SIZE,
+        Schedulers.DEFAULT_BOUNDED_ELASTIC_QUEUESIZE,
+        BULK_EXECUTOR_BOUNDED_ELASTIC_THREAD_NAME,
+        TTL_FOR_SCHEDULER_WORKER_IN_SECONDS,
+        true
+    );
 }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosSchedulers.java

@@ -288,6 +288,11 @@ CosmosBulkExecutionOptions setTargetedMicroBatchRetryRate(
 
             int getMaxMicroBatchConcurrency(CosmosBulkExecutionOptions options);
 
+            Integer getMaxConcurrentCosmosPartitions(CosmosBulkExecutionOptions options);
+
+            CosmosBulkExecutionOptions setMaxConcurrentCosmosPartitions(
+                CosmosBulkExecutionOptions options, int mxConcurrentCosmosPartitions);
+
             Duration getMaxMicroBatchInterval(CosmosBulkExecutionOptions options);
         }
     }