Fix Spark change feed hang after partition splits

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

@@ -7,6 +7,7 @@
 #### Breaking Changes
 
 #### Bugs Fixed
+* Fixed bounded change feed reads hanging after partition splits when stale feed-range metadata collapsed child continuation LSNs into a single `endLsn`. - See [PR 49883](https://github.com/Azure/azure-sdk-for-java/pull/49883)
 
 #### Other Changes
 

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

@@ -7,6 +7,7 @@
 #### Breaking Changes
 
 #### Bugs Fixed
+* Fixed bounded change feed reads hanging after partition splits when stale feed-range metadata collapsed child continuation LSNs into a single `endLsn`. - See [PR 49883](https://github.com/Azure/azure-sdk-for-java/pull/49883)
 
 #### Other Changes
 

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

@@ -7,6 +7,7 @@
 #### Breaking Changes
 
 #### Bugs Fixed
+* Fixed bounded change feed reads hanging after partition splits when stale feed-range metadata collapsed child continuation LSNs into a single `endLsn`. - See [PR 49883](https://github.com/Azure/azure-sdk-for-java/pull/49883)
 
 #### Other Changes
 

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

@@ -7,6 +7,7 @@
 #### Breaking Changes
 
 #### Bugs Fixed
+* Fixed bounded change feed reads hanging after partition splits when stale feed-range metadata collapsed child continuation LSNs into a single `endLsn`. - See [PR 49883](https://github.com/Azure/azure-sdk-for-java/pull/49883)
 
 #### Other Changes
 

sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/implementation/SparkBridgeImplementationInternal.scala

@@ -139,8 +139,14 @@ private[cosmos] object SparkBridgeImplementationInternal extends BasicLoggingTra
 
   def extractContinuationTokensFromChangeFeedStateJson(stateJsonBase64: String): Array[(NormalizedRange, Long)] = {
     assert(!Strings.isNullOrWhiteSpace(stateJsonBase64), s"Argument 'stateJsonBase64' must not be null or empty.")
-    val state = ChangeFeedState.fromString(stateJsonBase64)
-    state
+    extractContinuationTokensFromChangeFeedState(ChangeFeedState.fromString(stateJsonBase64))
+  }
+
+  def extractContinuationTokensFromChangeFeedState(
+    changeFeedState: ChangeFeedState
+  ): Array[(NormalizedRange, Long)] = {
+    assert(changeFeedState != null, "Argument 'changeFeedState' must not be null.")
+    changeFeedState
       .extractContinuationTokens() // already sorted
       .asScala
       .map(t => Tuple2(rangeToNormalizedRange(t.getRange), toLsn(t.getToken)))

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

@@ -414,17 +414,25 @@ private object CosmosPartitionPlanner extends BasicLoggingTrait {
       logDebug(endOffsetDebug.toString)
     }
 
+    createLatestOffset(startOffset, inputPartitions)
+  }
+  // scalastyle:on method.length
+  // scalastyle:on parameter.number
+
+  private[spark] def createLatestOffset
+  (
+    startOffset: ChangeFeedOffset,
+    inputPartitions: Array[CosmosInputPartition]
+  ): ChangeFeedOffset = {
     val changeFeedStateJson = SparkBridgeImplementationInternal
       .createChangeFeedStateJson(
         startOffset.changeFeedState,
         inputPartitions.map(m => (m.feedRange, m.endLsn.get)))
 
     ChangeFeedOffset(changeFeedStateJson, Some(inputPartitions))
   }
-  // scalastyle:on method.length
-  // scalastyle:on parameter.number
 
-  private[this] def getOrderedPartitionMetadataWithStartLsn
+  private[spark] def getOrderedPartitionMetadataWithStartLsn
   (
     stateJson: String,
     latestPartitionMetadata: Array[PartitionMetadata]
@@ -834,12 +842,41 @@ private object CosmosPartitionPlanner extends BasicLoggingTrait {
                   0,
                   None,
                   new AtomicLong(0),
-                  new AtomicLong(0)
+                  new AtomicLong(0),
+                  None
                 ))
               })
           .collectSeq()
+          .map(metadata =>
+            filterPartitionMetadataByFeedRanges(
+              expandPartitionMetadataByLatestLsn(metadata, isChangeFeed),
+              partitionConfig.feedRangeFiler))
       })
       .block()
       .toArray
   }
+
+  private[spark] def expandPartitionMetadataByLatestLsn(
+    metadata: Seq[PartitionMetadata],
+    isChangeFeed: Boolean
+  ): Seq[PartitionMetadata] = {
+    if (isChangeFeed) {
+      metadata.flatMap(_.splitByLatestLsn())
+    } else {
+      metadata
+    }
+  }
+
+  private[spark] def filterPartitionMetadataByFeedRanges(
+    metadata: Seq[PartitionMetadata],
+    feedRangeFilter: Option[Array[NormalizedRange]]
+  ): Seq[PartitionMetadata] = {
+    feedRangeFilter match {
+      case Some(epkRangesInScope) =>
+        metadata.filter(partitionMetadata =>
+          epkRangesInScope.exists(epk =>
+            SparkBridgeImplementationInternal.doRangesOverlap(epk, partitionMetadata.feedRange)))
+      case None => metadata
+    }
+  }
 }

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

@@ -20,6 +20,36 @@ private object PartitionMetadata {
     s"$databaseId/$containerId/${feedRange.min}-${feedRange.max}"
 
   // scalastyle:off parameter.number
+  def apply(userConfig: Map[String, String],
+            cosmosClientConfig: CosmosClientConfiguration,
+            cosmosClientStateHandles: Option[Broadcast[CosmosClientMetadataCachesSnapshots]],
+            cosmosContainerConfig: CosmosContainerConfig,
+            feedRange: NormalizedRange,
+            documentCount: Long,
+            totalDocumentSizeInKB: Long,
+            firstLsn: Option[Long],
+            latestLsn: Long,
+            startLsn: Long,
+            endLsn: Option[Long],
+            lastRetrieved: AtomicLong,
+            lastUpdated: AtomicLong): PartitionMetadata = {
+    new PartitionMetadata(
+      userConfig,
+      cosmosClientConfig,
+      cosmosClientStateHandles,
+      cosmosContainerConfig,
+      feedRange,
+      documentCount,
+      totalDocumentSizeInKB,
+      firstLsn,
+      latestLsn,
+      startLsn,
+      endLsn,
+      lastRetrieved,
+      lastUpdated,
+      None)
+  }
+
   def apply(userConfig: Map[String, String],
             cosmosClientConfig: CosmosClientConfiguration,
             cosmosClientStateHandles: Option[Broadcast[CosmosClientMetadataCachesSnapshots]],
@@ -51,7 +81,8 @@ private object PartitionMetadata {
       startLsn,
       endLsn,
       new AtomicLong(nowEpochMs),
-      new AtomicLong(nowEpochMs)
+      new AtomicLong(nowEpochMs),
+      Some(fromNowContinuationToken)
     )
   }
 }
@@ -70,7 +101,8 @@ private[cosmos] case class PartitionMetadata
   startLsn: Long,
   endLsn: Option[Long],
   lastRetrieved: AtomicLong,
-  lastUpdated: AtomicLong
+  lastUpdated: AtomicLong,
+  fromNowContinuationState: Option[String]
 )  extends BasicLoggingTrait {
 
   requireNotNull(feedRange, "feedRange")
@@ -94,7 +126,8 @@ private[cosmos] case class PartitionMetadata
       startLsn,
       this.endLsn,
       new AtomicLong(this.lastRetrieved.get),
-      new AtomicLong(this.lastUpdated.get)
+      new AtomicLong(this.lastUpdated.get),
+      this.getFromNowContinuationStateForRange(subRange)
     )
   }
 
@@ -112,10 +145,51 @@ private[cosmos] case class PartitionMetadata
       startLsn,
       Some(explicitEndLsn),
       new AtomicLong(this.lastRetrieved.get),
-      new AtomicLong(this.lastUpdated.get)
+      new AtomicLong(this.lastUpdated.get),
+      this.fromNowContinuationState
     )
   }
 
+  def splitByLatestLsn(): Seq[PartitionMetadata] = {
+    this.fromNowContinuationState match {
+      case None => Seq(this)
+      case Some(stateJson) =>
+        val parsedState = SparkBridgeImplementationInternal.parseChangeFeedState(stateJson)
+        val latestLsnsByRange =
+          SparkBridgeImplementationInternal.extractContinuationTokensFromChangeFeedState(parsedState)
+
+        if (latestLsnsByRange.length <= 1) {
+          Seq(this)
+        } else {
+          logInfo(
+            s"FromNow continuation for range '${this.feedRange}' resolved to " +
+              s"${latestLsnsByRange.length} effective ranges. Planning each range independently.")
+
+          val effectiveRangesAndLsns = latestLsnsByRange
+            .flatMap(rangeAndLsn =>
+              this.intersect(this.feedRange, rangeAndLsn._1)
+                .map(effectiveRange => effectiveRange -> rangeAndLsn._2))
+            .toSeq
+          val effectiveStates = SparkBridgeImplementationInternal.extractChangeFeedStateForRanges(
+            parsedState,
+            effectiveRangesAndLsns.map(_._1))
+
+          assert(
+            effectiveStates.length == effectiveRangesAndLsns.length,
+            "Expected one continuation state for every effective range.")
+
+          effectiveRangesAndLsns
+            .zip(effectiveStates)
+            .map { case ((effectiveRange, effectiveLatestLsn), effectiveState) =>
+              this.withFeedRangeAndLatestLsn(
+                effectiveRange,
+                effectiveLatestLsn,
+                effectiveState)
+            }
+        }
+    }
+  }
+
   def getWeightedLsnGap(
                         isChangeFeed: Boolean,
                         partitionMetricsMap: Option[ConcurrentHashMap[NormalizedRange, ChangeFeedMetricsTracker]] = None): Long = {
@@ -180,4 +254,50 @@ private[cosmos] case class PartitionMetadata
       }
     }
   }
+
+  private def withFeedRangeAndLatestLsn(
+    effectiveRange: NormalizedRange,
+    effectiveLatestLsn: Long,
+    effectiveFromNowContinuationState: String
+  ): PartitionMetadata = {
+    new PartitionMetadata(
+      this.userConfig,
+      this.cosmosClientConfig,
+      this.cosmosClientStateHandles,
+      this.cosmosContainerConfig,
+      effectiveRange,
+      0,
+      0,
+      None,
+      effectiveLatestLsn,
+      this.startLsn,
+      this.endLsn,
+      new AtomicLong(this.lastRetrieved.get),
+      new AtomicLong(this.lastUpdated.get),
+      Some(effectiveFromNowContinuationState)
+    )
+  }
+
+  private def getFromNowContinuationStateForRange(
+    effectiveRange: NormalizedRange
+  ): Option[String] = {
+    this.fromNowContinuationState.map(state => {
+      if (effectiveRange == this.feedRange) {
+        state
+      } else {
+        SparkBridgeImplementationInternal.extractChangeFeedStateForRange(
+          SparkBridgeImplementationInternal.parseChangeFeedState(state),
+          effectiveRange)
+      }
+    })
+  }
+
+  private def intersect(
+    left: NormalizedRange,
+    right: NormalizedRange
+  ): Option[NormalizedRange] = {
+    val min = if (left.min.compareTo(right.min) >= 0) left.min else right.min
+    val max = if (left.max.compareTo(right.max) <= 0) left.max else right.max
+    if (min.compareTo(max) < 0) Some(NormalizedRange(min, max)) else None
+  }
 }