Cosmos Spark: Adding more robust retry policy for transient errors (Azure#26029)

* Cosmos Spark: Adding more robust retry policy for transient errors

* Fixing StyleCop violations

* Reacting to CR feedback

Author: FabianMeiswinkel

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

@@ -40,14 +40,16 @@ private[spark] object ContainerFeedRangesCache {
   }
 
   private[this] def refreshFeedRanges(key: String, container: CosmosAsyncContainer): SMono[List[FeedRange]] = {
-    container
-      .getFeedRanges
-      .map[List[FeedRange]](javaList => {
-        val scalaList = javaList.asScala.toList
-        cache.put(key, CachedFeedRanges(scalaList, Instant.now))
-        scalaList
-      })
-      .asScala
+
+    TransientErrorsRetryPolicy.executeWithRetry(() =>
+      container
+        .getFeedRanges
+        .map[List[FeedRange]](javaList => {
+          val scalaList = javaList.asScala.toList
+          cache.put(key, CachedFeedRanges(scalaList, Instant.now))
+          scalaList
+        })
+        .asScala)
   }
 
   private case class CachedFeedRanges(feedRanges: List[FeedRange], retrievedAt: Instant)

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

@@ -112,6 +112,12 @@ class CosmosCatalog
   override def listNamespaces(): Array[Array[String]] = {
     logDebug("catalog:listNamespaces")
 
+    TransientErrorsRetryPolicy.executeWithRetry(() => listNamespacesImpl())
+  }
+
+  private[this] def listNamespacesImpl(): Array[Array[String]] = {
+    logDebug("catalog:listNamespaces")
+
     Loan(CosmosClientCache(
       CosmosClientConfiguration(config, readConfig.forceEventualConsistency),
       None,
@@ -153,6 +159,12 @@ class CosmosCatalog
   override def loadNamespaceMetadata(
                                       namespace: Array[String]): util.Map[String, String] = {
 
+    TransientErrorsRetryPolicy.executeWithRetry(() => loadNamespaceMetadataImpl(namespace))
+  }
+
+  private[this] def loadNamespaceMetadataImpl(
+                                      namespace: Array[String]): util.Map[String, String] = {
+
     checkNamespace(namespace)
 
     Loan(CosmosClientCache(
@@ -187,6 +199,12 @@ class CosmosCatalog
   @throws(classOf[NamespaceAlreadyExistsException])
   override def createNamespace(namespace: Array[String],
                                metadata: util.Map[String, String]): Unit = {
+    TransientErrorsRetryPolicy.executeWithRetry(() => createNamespaceImpl(namespace, metadata))
+  }
+
+  @throws(classOf[NamespaceAlreadyExistsException])
+  private[this] def createNamespaceImpl(namespace: Array[String],
+                               metadata: util.Map[String, String]): Unit = {
     checkNamespace(namespace)
     val throughputPropertiesOpt =
       CosmosThroughputProperties.tryGetThroughputProperties(
@@ -234,6 +252,11 @@ class CosmosCatalog
    */
   @throws(classOf[NoSuchNamespaceException])
   override def dropNamespace(namespace: Array[String]): Boolean = {
+    TransientErrorsRetryPolicy.executeWithRetry(() => dropNamespaceImpl(namespace))
+  }
+
+  @throws(classOf[NoSuchNamespaceException])
+  private[this] def dropNamespaceImpl(namespace: Array[String]): Boolean = {
     checkNamespace(namespace)
     try {
       Loan(CosmosClientCache(
@@ -256,6 +279,10 @@ class CosmosCatalog
   }
 
   override def listTables(namespace: Array[String]): Array[Identifier] = {
+    TransientErrorsRetryPolicy.executeWithRetry(() => listTablesImpl(namespace))
+  }
+
+  private[this] def listTablesImpl(namespace: Array[String]): Array[Identifier] = {
     checkNamespace(namespace)
     val databaseName = toCosmosDatabaseName(namespace.head)
 
@@ -290,6 +317,10 @@ class CosmosCatalog
   }
 
   override def loadTable(ident: Identifier): Table = {
+    TransientErrorsRetryPolicy.executeWithRetry(() => loadTableImpl(ident))
+  }
+
+  private[this] def loadTableImpl(ident: Identifier): Table = {
     checkNamespace(ident.namespace())
     val databaseName = toCosmosDatabaseName(ident.namespace().head)
     val containerName = toCosmosContainerName(ident.name())
@@ -328,6 +359,15 @@ class CosmosCatalog
                            schema: StructType,
                            partitions: Array[Transform],
                            properties: util.Map[String, String]): Table = {
+
+    TransientErrorsRetryPolicy.executeWithRetry(() =>
+      createTableImpl(ident, schema, partitions, properties))
+  }
+
+  private[this] def createTableImpl(ident: Identifier,
+                           schema: StructType,
+                           partitions: Array[Transform],
+                           properties: util.Map[String, String]): Table = {
     checkNamespace(ident.namespace())
 
     val databaseName = toCosmosDatabaseName(ident.namespace().head)
@@ -348,6 +388,10 @@ class CosmosCatalog
   }
 
   override def dropTable(ident: Identifier): Boolean = {
+    TransientErrorsRetryPolicy.executeWithRetry(() => dropTableImpl(ident))
+  }
+
+  private[this] def dropTableImpl(ident: Identifier): Boolean = {
     checkNamespace(ident.namespace())
 
     val databaseName = toCosmosDatabaseName(ident.namespace().head)

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

@@ -14,6 +14,7 @@ private object CosmosConstants {
   val currentName: String =
     CoreUtils.getProperties(propertiesFileName).get("name")
   val userAgentSuffix = s"SparkConnector/$currentName/$currentVersion"
+  val initialMaxRetryIntervalForTransientFailuresInMs = 100
   val maxRetryIntervalForTransientFailuresInMs = 5000
   val maxRetryCountForTransientFailures = 100
   val defaultDirectRequestTimeoutInSeconds = 10L

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

@@ -30,6 +30,26 @@ private object CosmosPartitionPlanner extends BasicLoggingTrait {
   val DefaultPartitionSizeInMB: Int = 5 * 1024 // 10 GB
 
   def createInputPartitions
+  (
+    cosmosPartitioningConfig: CosmosPartitioningConfig,
+    container: CosmosAsyncContainer,
+    partitionMetadata: Array[PartitionMetadata],
+    defaultMinimalPartitionCount: Int,
+    defaultMaxPartitionSizeInMB: Int,
+    readLimit: ReadLimit
+  ): Array[CosmosInputPartition] = {
+
+    TransientErrorsRetryPolicy.executeWithRetry(() =>
+      createInputPartitionsImpl(
+        cosmosPartitioningConfig,
+        container,
+        partitionMetadata,
+        defaultMinimalPartitionCount,
+        defaultMaxPartitionSizeInMB,
+        readLimit))
+  }
+
+  private[this] def createInputPartitionsImpl
   (
     cosmosPartitioningConfig: CosmosPartitioningConfig,
     container: CosmosAsyncContainer,
@@ -89,8 +109,20 @@ private object CosmosPartitionPlanner extends BasicLoggingTrait {
       .flatten
   }
 
-  // scalastyle:off method.length
   def createInitialOffset
+  (
+    container: CosmosAsyncContainer,
+    changeFeedConfig: CosmosChangeFeedConfig,
+    streamId: Option[String]
+  ): String = {
+
+    TransientErrorsRetryPolicy.executeWithRetry(() =>
+      createInitialOffsetImpl(container, changeFeedConfig, streamId)
+    )
+  }
+
+  // scalastyle:off method.length
+  private[this] def createInitialOffsetImpl
   (
     container: CosmosAsyncContainer,
     changeFeedConfig: CosmosChangeFeedConfig,
@@ -150,10 +182,40 @@ private object CosmosPartitionPlanner extends BasicLoggingTrait {
   }
   // scalastyle:on method.length
 
+  // scalastyle:off parameter.number
+  def getLatestOffset
+  (
+    userConfig: Map[String, String],
+    startOffset: ChangeFeedOffset,
+    readLimit: ReadLimit,
+    maxStaleness: Duration,
+    clientConfiguration: CosmosClientConfiguration,
+    cosmosClientStateHandle: Broadcast[CosmosClientMetadataCachesSnapshot],
+    containerConfig: CosmosContainerConfig,
+    partitioningConfig: CosmosPartitioningConfig,
+    defaultParallelism: Int,
+    container: CosmosAsyncContainer
+  ): ChangeFeedOffset = {
+
+    TransientErrorsRetryPolicy.executeWithRetry(() =>
+      getLatestOffsetImpl(
+        userConfig,
+        startOffset,
+        readLimit,
+        maxStaleness,
+        clientConfiguration,
+        cosmosClientStateHandle,
+        containerConfig,
+        partitioningConfig,
+        defaultParallelism,
+        container))
+  }
+  // scalastyle:on parameter.number
+
   // scalastyle:off method.length
   // scalastyle:off parameter.number
   // Based on a start offset, calculate which is the next end offset
-  def getLatestOffset
+  private[this] def getLatestOffsetImpl
   (
     userConfig: Map[String, String],
     startOffset: ChangeFeedOffset,
@@ -471,6 +533,19 @@ private object CosmosPartitionPlanner extends BasicLoggingTrait {
   }
 
   def getPartitionMetadata(
+                            userConfig: Map[String, String],
+                            cosmosClientConfig: CosmosClientConfiguration,
+                            cosmosClientStateHandle: Option[Broadcast[CosmosClientMetadataCachesSnapshot]],
+                            cosmosContainerConfig: CosmosContainerConfig,
+                            maxStaleness: Option[Duration] = None
+                          ): Array[PartitionMetadata] = {
+
+    TransientErrorsRetryPolicy.executeWithRetry(() =>
+      getPartitionMetadataImpl(
+        userConfig, cosmosClientConfig, cosmosClientStateHandle, cosmosContainerConfig, maxStaleness))
+  }
+
+  private[this] def getPartitionMetadataImpl(
       userConfig: Map[String, String],
       cosmosClientConfig: CosmosClientConfiguration,
       cosmosClientStateHandle: Option[Broadcast[CosmosClientMetadataCachesSnapshot]],

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

@@ -89,6 +89,14 @@ private object CosmosTableSchemaInferrer
   private[spark] def inferSchema(client: CosmosAsyncClient,
                                  userConfig: Map[String, String],
                                  defaultSchema: StructType): StructType = {
+
+    TransientErrorsRetryPolicy.executeWithRetry(() =>
+      inferSchemaImpl(client, userConfig, defaultSchema))
+  }
+
+  private[this] def inferSchemaImpl(client: CosmosAsyncClient,
+                                 userConfig: Map[String, String],
+                                 defaultSchema: StructType): StructType = {
     val cosmosInferenceConfig = CosmosSchemaInferenceConfig.parseCosmosInferenceConfig(userConfig)
     val cosmosReadConfig = CosmosReadConfig.parseCosmosReadConfig(userConfig)
     if (cosmosInferenceConfig.inferSchemaEnabled) {

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

@@ -137,8 +137,28 @@ private object PartitionMetadataCache extends BasicLoggingTrait {
 
   this.startRefreshTimer()
 
-  //scalastyle:off method.length
   private def readPartitionMetadata
+  (
+    userConfig: Map[String, String],
+    cosmosClientConfiguration: CosmosClientConfiguration,
+    cosmosClientStateHandle: Option[Broadcast[CosmosClientMetadataCachesSnapshot]],
+    cosmosContainerConfig: CosmosContainerConfig,
+    feedRange: NormalizedRange,
+    tolerateNotFound: Boolean
+  ): SMono[Option[PartitionMetadata]] = {
+
+    TransientErrorsRetryPolicy.executeWithRetry(() =>
+      readPartitionMetadataImpl(
+        userConfig,
+        cosmosClientConfiguration,
+        cosmosClientStateHandle,
+        cosmosContainerConfig,
+        feedRange,
+        tolerateNotFound))
+  }
+
+  //scalastyle:off method.length
+  private def readPartitionMetadataImpl
   (
     userConfig: Map[String, String],
     cosmosClientConfiguration: CosmosClientConfiguration,

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

@@ -0,0 +1,64 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.cosmos.spark
+
+import com.azure.cosmos.CosmosException
+import com.azure.cosmos.spark.diagnostics.BasicLoggingTrait
+
+import java.util.concurrent.atomic.AtomicLong
+import scala.util.Random
+import scala.util.control.Breaks
+
+private[spark] object TransientErrorsRetryPolicy extends BasicLoggingTrait {
+  private val rnd = Random
+
+  def executeWithRetry[T]
+  (
+    func: () => T,
+    initialMaxRetryIntervalInMs: Int = CosmosConstants.initialMaxRetryIntervalForTransientFailuresInMs,
+    maxRetryIntervalInMs: Int = CosmosConstants.maxRetryIntervalForTransientFailuresInMs,
+    maxRetryCount: Int = Int.MaxValue
+  ): T = {
+    val loop = new Breaks()
+    val retryCount = new AtomicLong(0)
+    var returnValue: Option[T] = None
+
+    loop.breakable {
+      var currentMaxRetryIntervalInMs = Math.min(initialMaxRetryIntervalInMs, maxRetryIntervalInMs)
+      while (true) {
+        val retryIntervalInMs = rnd.nextInt(currentMaxRetryIntervalInMs)
+
+        try {
+          returnValue = Some(func())
+          loop.break
+        }
+        catch {
+          case cosmosException: CosmosException =>
+            if (Exceptions.canBeTransientFailure(cosmosException)) {
+              val retryCountSnapshot = retryCount.incrementAndGet()
+              if (retryCountSnapshot > maxRetryCount) {
+                logError(
+                  s"Too many transient failure retry attempts ($retryCountSnapshot) in " +
+                    s"TransientIORetryPolicy.executeWithRetry",
+                  cosmosException)
+                throw cosmosException
+              } else {
+                logWarning(
+                  s"Transient failure handled in TransientIORetryPolicy.executeWithRetry -" +
+                    s" will be retried (attempt#$retryCountSnapshot) in ${retryIntervalInMs}ms",
+                  cosmosException)
+              }
+            } else {
+              throw cosmosException
+            }
+          case other: Throwable => throw other
+        }
+
+        Thread.sleep(retryIntervalInMs)
+        currentMaxRetryIntervalInMs = Math.min(2 * currentMaxRetryIntervalInMs, maxRetryIntervalInMs)
+      }
+    }
+
+    returnValue.get
+  }
+}