diff --git a/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md index 9bbb33441cfd..2473af79929b 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md +++ b/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 diff --git a/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md index 8866ab997a2f..3cb4c50f2d15 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md +++ b/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 diff --git a/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md index b873c993788e..cd52e82f2e55 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md +++ b/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 diff --git a/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md index 5cd9f73f9cfb..4fc942f05da1 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md +++ b/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 diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/implementation/SparkBridgeImplementationInternal.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/implementation/SparkBridgeImplementationInternal.scala index f0cb2d8004da..ebba8f9e221d 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/implementation/SparkBridgeImplementationInternal.scala +++ b/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))) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionPlanner.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionPlanner.scala index 1c9751649766..9b7ffbf1cbd9 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionPlanner.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionPlanner.scala @@ -414,6 +414,16 @@ 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, @@ -421,10 +431,8 @@ private object CosmosPartitionPlanner extends BasicLoggingTrait { 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,48 @@ 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.foreach(partitionMetadata => { + if (partitionMetadata.fromNowContinuationState.isEmpty && partitionMetadata.latestLsn > 0) { + logWarning( + s"Change-feed metadata for range '${partitionMetadata.feedRange}' is missing the FromNow continuation " + + "state; composite split expansion cannot be performed.") + } + }) + 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 + } + } } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/PartitionMetadata.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/PartitionMetadata.scala index f290f00e129e..aa0e9d9ef506 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/PartitionMetadata.scala +++ b/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,55 @@ 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 + if (effectiveRangesAndLsns.length != latestLsnsByRange.length) { + throw new IllegalStateException( + s"FromNow continuation for range '${this.feedRange}' contained non-overlapping effective ranges.") + } + val effectiveStates = SparkBridgeImplementationInternal.extractChangeFeedStateForRanges( + parsedState, + effectiveRangesAndLsns.map(_._1)) + + require( + 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 +258,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 + } } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionPlannerSpec.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionPlannerSpec.scala index f745bae6b56d..6bcfb6a2337a 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionPlannerSpec.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionPlannerSpec.scala @@ -4,16 +4,50 @@ package com.azure.cosmos.spark import com.azure.core.management.AzureEnvironment import com.azure.cosmos.changeFeedMetrics.ChangeFeedMetricsTracker +import com.azure.cosmos.implementation.SparkBridgeImplementationInternal import com.azure.cosmos.{ReadConsistencyStrategy, spark} +import com.fasterxml.jackson.databind.ObjectMapper import org.apache.spark.sql.connector.read.streaming.ReadLimit +import java.nio.charset.StandardCharsets import java.time.Instant import java.util.UUID import java.util.concurrent.ConcurrentHashMap import java.util.concurrent.atomic.AtomicLong +import java.util.Base64 + +// scalastyle:off underscore.import +import scala.collection.JavaConverters._ +// scalastyle:on underscore.import class CosmosPartitionPlannerSpec extends UnitSpec { private[this] val rnd = scala.util.Random + private[this] val plannerClientConfig = CosmosClientConfiguration( + UUID.randomUUID().toString, + UUID.randomUUID().toString, + CosmosMasterKeyAuthConfig(UUID.randomUUID().toString), + None, + UUID.randomUUID().toString, + useGatewayMode = false, + enforceNativeTransport = false, + proactiveConnectionInitialization = None, + proactiveConnectionInitializationDurationInSeconds = 120, + httpConnectionPoolSize = 1000, + readConsistencyStrategy = ReadConsistencyStrategy.EVENTUAL, + disableTcpConnectionEndpointRediscovery = false, + preferredRegionsList = Option.empty, + subscriptionId = None, + tenantId = None, + resourceGroupName = None, + azureEnvironmentEndpoints = AzureEnvironment.AZURE.getEndpoints, + sparkEnvironmentInfo = "", + clientBuilderInterceptors = None, + clientInterceptors = None, + sampledDiagnosticsLoggerConfig = None, + azureMonitorConfig = None, + additionalHeaders = None) + private[this] val plannerContainerConfig = + CosmosContainerConfig(UUID.randomUUID().toString, UUID.randomUUID().toString, None) it should "calculateEndLsn without readLimit" in { @@ -625,4 +659,212 @@ class CosmosPartitionPlannerSpec extends UnitSpec { // Should stay at start LSN since no progress can be made calculate(0).endLsn.get shouldEqual 2050 } + + it should "expand composite latest metadata only for change feed planning" in { + val ranges = Seq( + NormalizedRange("", "AA"), + NormalizedRange("AA", "BB"), + NormalizedRange("BB", "FF")) + val state = createChangeFeedState(Seq( + ranges(0) -> 180L, + ranges(1) -> 160L, + ranges(2) -> 130L)) + val metadata = createMetadata( + NormalizedRange("", "FF"), + latestLsn = 180L, + startLsn = 100L, + fromNowContinuationState = Some(state), + documentCount = 42L, + totalDocumentSizeInKB = 84L) + + val itemQueryMetadata = CosmosPartitionPlanner.expandPartitionMetadataByLatestLsn( + Seq(metadata), + isChangeFeed = false) + itemQueryMetadata should contain only (metadata) + itemQueryMetadata.head should be theSameInstanceAs metadata + itemQueryMetadata.head.documentCount shouldEqual 42L + itemQueryMetadata.head.totalDocumentSizeInKB shouldEqual 84L + + val changeFeedMetadata = CosmosPartitionPlanner.expandPartitionMetadataByLatestLsn( + Seq(metadata), + isChangeFeed = true) + changeFeedMetadata.map(m => m.feedRange -> m.latestLsn) should contain theSameElementsInOrderAs + Seq(ranges(0) -> 180L, ranges(1) -> 160L, ranges(2) -> 130L) + changeFeedMetadata.map(_.documentCount).sum shouldEqual 0L + changeFeedMetadata.map(_.totalDocumentSizeInKB).sum shouldEqual 0L + changeFeedMetadata.foreach(_.firstLsn shouldBe empty) + } + + it should "exclude out-of-scope split children before allocating ReadMaxRows" in { + val includedRange = NormalizedRange("", "AA") + val excludedRange = NormalizedRange("AA", "FF") + val state = createChangeFeedState(Seq( + includedRange -> 180L, + excludedRange -> 1000L)) + val metadata = createMetadata( + NormalizedRange("", "FF"), + latestLsn = 180L, + startLsn = 100L, + fromNowContinuationState = Some(state)) + val expandedMetadata = CosmosPartitionPlanner.expandPartitionMetadataByLatestLsn( + Seq(metadata), + isChangeFeed = true) + + val filteredMetadata = CosmosPartitionPlanner.filterPartitionMetadataByFeedRanges( + expandedMetadata, + Some(Array(includedRange))) + filteredMetadata.map(_.feedRange) should contain only includedRange + + val planned = CosmosPartitionPlanner.calculateEndLsn( + filteredMetadata.toArray, + ReadLimit.maxRows(10), + isChangeFeed = true) + planned.map(m => m.feedRange -> m.endLsn.get) should contain only (includedRange -> 110L) + } + + it should "reconcile split metadata and preserve v1 offsets with ReadAllAvailable" in { + val parentRange = NormalizedRange("", "FF") + val childRanges = Seq(NormalizedRange("", "AA"), NormalizedRange("AA", "FF")) + val startState = createChangeFeedState(Seq(parentRange -> 100L)) + val latestState = createChangeFeedState(Seq(childRanges(0) -> 180L, childRanges(1) -> 130L)) + val latestParentMetadata = createMetadata( + parentRange, + latestLsn = 180L, + startLsn = 0L, + fromNowContinuationState = Some(latestState)) + val latestChildMetadata = CosmosPartitionPlanner.expandPartitionMetadataByLatestLsn( + Seq(latestParentMetadata), + isChangeFeed = true) + + val reconciled = CosmosPartitionPlanner.getOrderedPartitionMetadataWithStartLsn( + startState, + latestChildMetadata.toArray) + reconciled.map(m => (m.feedRange, m.startLsn, m.latestLsn)) should contain theSameElementsInOrderAs + Seq((childRanges(0), 100L, 180L), (childRanges(1), 100L, 130L)) + + val planned = CosmosPartitionPlanner.calculateEndLsn( + reconciled, + ReadLimit.allAvailable(), + isChangeFeed = true) + planned.map(m => m.feedRange -> m.endLsn.get) should contain theSameElementsInOrderAs + Seq(childRanges(0) -> 180L, childRanges(1) -> 130L) + + val inputPartitions = planned.map(m => CosmosInputPartition(m.feedRange, m.endLsn)) + val latestOffset = CosmosPartitionPlanner.createLatestOffset( + ChangeFeedOffset(startState, None), + inputPartitions) + val persistedOffset = new ObjectMapper().readTree(latestOffset.json()) + persistedOffset.fieldNames().asScala.toSet shouldEqual Set("id", "state", "partitions") + persistedOffset.get("id").asText() shouldEqual ChangeFeedOffset.V1Identifier + SparkBridgeImplementationInternal + .extractContinuationTokensFromChangeFeedStateJson(latestOffset.changeFeedState) should + contain theSameElementsInOrderAs Seq(childRanges(0) -> 180L, childRanges(1) -> 130L) + + val roundTrippedOffset = ChangeFeedOffset.fromJson(latestOffset.json()) + roundTrippedOffset.inputPartitions.get.map(p => p.feedRange -> p.endLsn.get) should + contain theSameElementsInOrderAs Seq(childRanges(0) -> 180L, childRanges(1) -> 130L) + } + + it should "cap divergent child ranges independently with ReadMaxRows" in { + val parentRange = NormalizedRange("", "FF") + val childRanges = Seq(NormalizedRange("", "AA"), NormalizedRange("AA", "FF")) + val startState = createChangeFeedState(Seq(parentRange -> 100L)) + val latestState = createChangeFeedState(Seq(childRanges(0) -> 180L, childRanges(1) -> 130L)) + val latestParentMetadata = createMetadata( + parentRange, + latestLsn = 180L, + startLsn = 0L, + fromNowContinuationState = Some(latestState)) + val latestChildMetadata = CosmosPartitionPlanner.expandPartitionMetadataByLatestLsn( + Seq(latestParentMetadata), + isChangeFeed = true) + val reconciled = CosmosPartitionPlanner.getOrderedPartitionMetadataWithStartLsn( + startState, + latestChildMetadata.toArray) + + val planned = CosmosPartitionPlanner.calculateEndLsn( + reconciled, + ReadLimit.maxRows(20), + isChangeFeed = true) + + planned.map(m => m.feedRange -> m.endLsn.get) should contain theSameElementsInOrderAs + Seq(childRanges(0) -> 114L, childRanges(1) -> 105L) + planned.foreach(metadata => { + metadata.endLsn.get should be > metadata.startLsn + metadata.endLsn.get should be <= metadata.latestLsn + }) + } + + it should "reconcile child checkpoints with merged latest metadata" in { + val parentRange = NormalizedRange("", "FF") + val childRanges = Seq(NormalizedRange("", "AA"), NormalizedRange("AA", "FF")) + val startState = createChangeFeedState(Seq(childRanges(0) -> 150L, childRanges(1) -> 100L)) + val mergedMetadata = createMetadata(parentRange, latestLsn = 200L, startLsn = 0L) + + val reconciled = CosmosPartitionPlanner.getOrderedPartitionMetadataWithStartLsn( + startState, + Array(mergedMetadata)) + reconciled.map(m => (m.feedRange, m.startLsn, m.latestLsn)) should contain theSameElementsInOrderAs + Seq((childRanges(0), 150L, 200L), (childRanges(1), 100L, 200L)) + + val planned = CosmosPartitionPlanner.calculateEndLsn( + reconciled, + ReadLimit.allAvailable(), + isChangeFeed = true) + planned.map(m => m.feedRange -> m.endLsn.get) should contain theSameElementsInOrderAs + Seq(childRanges(0) -> 200L, childRanges(1) -> 200L) + } + + private[this] def createMetadata( + feedRange: NormalizedRange, + latestLsn: Long, + startLsn: Long, + fromNowContinuationState: Option[String] = None, + documentCount: Long = 0L, + totalDocumentSizeInKB: Long = 0L + ): PartitionMetadata = { + val nowEpochMs = Instant.now.toEpochMilli + new PartitionMetadata( + Map.empty, + plannerClientConfig, + None, + plannerContainerConfig, + feedRange, + documentCount, + totalDocumentSizeInKB, + Some(0L), + latestLsn, + startLsn, + None, + new AtomicLong(nowEpochMs), + new AtomicLong(nowEpochMs), + fromNowContinuationState) + } + + private[this] def createChangeFeedState( + latestLsns: Seq[(NormalizedRange, Long)] + ): String = { + val collectionRid = UUID.randomUUID().toString + val continuationTokens = latestLsns + .map(rangeAndLsn => + s"""{"token":"\\"${rangeAndLsn._2}\\"","range":""" + + s"""{"min":"${rangeAndLsn._1.min}","max":"${rangeAndLsn._1.max}"}}""") + .mkString(",") + val json = String.format( + "{\"V\":1," + + "\"Rid\":\"%s\"," + + "\"Mode\":\"INCREMENTAL\"," + + "\"StartFrom\":{\"Type\":\"BEGINNING\"}," + + "\"Continuation\":%s}", + collectionRid, + String.format( + "{\"V\":1," + + "\"Rid\":\"%s\"," + + "\"Continuation\":[%s]," + + "\"Range\":{\"min\":\"\",\"max\":\"FF\"}}", + collectionRid, + continuationTokens)) + + Base64.getUrlEncoder.encodeToString(json.getBytes(StandardCharsets.UTF_8)) + } } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/PartitionMetadataSpec.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/PartitionMetadataSpec.scala index c17e7b02dad8..515f2109ca12 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/PartitionMetadataSpec.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/PartitionMetadataSpec.scala @@ -4,7 +4,9 @@ package com.azure.cosmos.spark import com.azure.core.management.AzureEnvironment import com.azure.cosmos.changeFeedMetrics.ChangeFeedMetricsTracker +import com.azure.cosmos.implementation.SparkBridgeImplementationInternal import com.azure.cosmos.{ReadConsistencyStrategy, spark} +import org.apache.spark.sql.connector.read.streaming.ReadLimit import java.nio.charset.StandardCharsets import java.time.Instant @@ -147,6 +149,101 @@ class PartitionMetadataSpec extends UnitSpec { viaApply.lastUpdated.get shouldEqual viaApply.lastRetrieved.get } + it should "preserve and plan every range from a composite FromNow continuation" in { + val lowerRange = NormalizedRange("", "AA") + val middleRange = NormalizedRange("AA", "BB") + val upperRange = NormalizedRange("BB", "FF") + val fromNowState = createChangeFeedState( + Seq(lowerRange -> 179L, middleRange -> 159L, upperRange -> 129L), + "INCREMENTAL") + val metadata = PartitionMetadata( + Map[String, String](), + clientCfg, + None, + contCfg, + NormalizedRange("", "FF"), + documentCount = 120, + totalDocumentSizeInKB = 120, + firstLsn = Some(1L), + fromNowContinuationToken = fromNowState, + startLsn = 2L) + + metadata.latestLsn shouldEqual 179L + + val splitMetadata = metadata.splitByLatestLsn() + splitMetadata.map(m => m.feedRange -> m.latestLsn) should contain theSameElementsInOrderAs + Seq(lowerRange -> 179L, middleRange -> 159L, upperRange -> 129L) + splitMetadata.foreach(m => { + m.documentCount shouldEqual 0 + m.totalDocumentSizeInKB shouldEqual 0 + m.firstLsn shouldBe empty + }) + + splitMetadata.foreach(m => { + m.fromNowContinuationState shouldBe defined + SparkBridgeImplementationInternal + .extractContinuationTokensFromChangeFeedStateJson(m.fromNowContinuationState.get) + .map(t => t._1 -> t._2) should contain only (m.feedRange -> m.latestLsn) + }) + + val planned = CosmosPartitionPlanner.calculateEndLsn( + splitMetadata.toArray, + ReadLimit.allAvailable(), + isChangeFeed = true) + + planned.map(m => m.feedRange -> m.endLsn.get) should contain theSameElementsInOrderAs + Seq(lowerRange -> 179L, middleRange -> 159L, upperRange -> 129L) + } + + it should "reuse FromNow state for the same range and project it for a subrange" in { + val lowerRange = NormalizedRange("", "AA") + val upperRange = NormalizedRange("AA", "FF") + val fromNowState = createChangeFeedState( + Seq(lowerRange -> 179L, upperRange -> 129L), + "INCREMENTAL") + val metadata = PartitionMetadata( + Map[String, String](), + clientCfg, + None, + contCfg, + NormalizedRange("", "FF"), + documentCount = 120, + totalDocumentSizeInKB = 120, + firstLsn = Some(1L), + fromNowContinuationToken = fromNowState, + startLsn = 2L) + + val sameRangeClone = metadata.cloneForSubRange(metadata.feedRange, 3L) + sameRangeClone.fromNowContinuationState shouldEqual Some(fromNowState) + + val subRangeClone = metadata.cloneForSubRange(lowerRange, 3L) + SparkBridgeImplementationInternal + .extractContinuationTokensFromChangeFeedStateJson(subRangeClone.fromNowContinuationState.get) should + contain only (lowerRange -> 179L) + } + + it should "reject composite FromNow ranges that do not intersect the metadata range" in { + val fromNowState = createChangeFeedState( + Seq( + NormalizedRange("", "55") -> 179L, + NormalizedRange("BB", "FF") -> 129L), + "INCREMENTAL") + val metadata = PartitionMetadata( + Map[String, String](), + clientCfg, + None, + contCfg, + NormalizedRange("", "AA"), + documentCount = 120, + totalDocumentSizeInKB = 120, + firstLsn = Some(1L), + fromNowContinuationToken = fromNowState, + startLsn = 2L) + + val error = intercept[IllegalStateException](metadata.splitByLatestLsn()) + error.getMessage should include ("contained non-overlapping effective ranges") + } + it should "create instance with valid parameters via apply in full fidelity mode" in { val clientConfig = spark.CosmosClientConfiguration( @@ -1109,8 +1206,20 @@ class PartitionMetadataSpec extends UnitSpec { //scalastyle:on null //scalastyle:on multiple.string.literals - private[this] def createChangeFeedState(latestLsn: Long, mode: String) = { + private[this] def createChangeFeedState(latestLsn: Long, mode: String): String = { + createChangeFeedState(Seq(NormalizedRange("", "FF") -> latestLsn), mode) + } + + private[this] def createChangeFeedState( + latestLsns: Seq[(NormalizedRange, Long)], + mode: String + ): String = { val collectionRid = UUID.randomUUID().toString + val continuationTokens = latestLsns + .map(rangeAndLsn => + s"""{"token":"\\"${rangeAndLsn._2}\\"","range":""" + + s"""{"min":"${rangeAndLsn._1.min}","max":"${rangeAndLsn._1.max}"}}""") + .mkString(",") val json = String.format( "{\"V\":1," + @@ -1121,14 +1230,12 @@ class PartitionMetadataSpec extends UnitSpec { collectionRid, mode, String.format( - "{\"V\":1," + - "\"Rid\":\"%s\"," + - "\"Continuation\":[" + - "{\"token\":\"\\\"%s\\\"\",\"range\":{\"min\":\"\",\"max\":\"FF\"}}" + - "]," + - "\"PKRangeId\":\"0\"}", - collectionRid, - String.valueOf(latestLsn))) + "{\"V\":1," + + "\"Rid\":\"%s\"," + + "\"Continuation\":[%s]," + + "\"Range\":{\"min\":\"\",\"max\":\"FF\"}}", + collectionRid, + continuationTokens)) Base64.getUrlEncoder.encodeToString(json.getBytes(StandardCharsets.UTF_8)) } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/SparkE2EChangeFeedSplitITest.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/SparkE2EChangeFeedSplitITest.scala index 0846b2cf6165..8afbd88dbd0f 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/SparkE2EChangeFeedSplitITest.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/SparkE2EChangeFeedSplitITest.scala @@ -3,13 +3,17 @@ package com.azure.cosmos.spark import com.azure.cosmos.CosmosClientBuilder -import com.azure.cosmos.implementation.{TestConfigurations, Utils} -import com.azure.cosmos.models.ThroughputProperties +import com.azure.cosmos.implementation.{SparkBridgeImplementationInternal, TestConfigurations, Utils} +import com.azure.cosmos.models.{PartitionKey, PartitionKeyDefinition, ThroughputProperties} import com.azure.cosmos.spark.diagnostics.BasicLoggingTrait import org.apache.hadoop.fs.Path +import org.apache.spark.sql.{DataFrame, Row} import java.nio.file.Paths import java.util.UUID +import java.util.concurrent.{Callable, ExecutionException, Executors, ThreadFactory, TimeUnit, TimeoutException} +import scala.collection.mutable +import scala.jdk.CollectionConverters._ class SparkE2EChangeFeedSplitITest extends IntegrationSpec @@ -20,6 +24,9 @@ class SparkE2EChangeFeedSplitITest //scalastyle:off multiple.string.literals //scalastyle:off magic.number + private val MaxSplitWaitAttempts = 300 + private val MaxPartitionKeySearchAttempts = 100000 + private val TestTimeoutInMinutes = 5L "spark change feed query (incremental)" should "honor checkpoint location and read limit after partition split" in { val cosmosEndpoint = TestConfigurations.HOST @@ -29,10 +36,14 @@ class SparkE2EChangeFeedSplitITest logWarning("Skipping this test on emulator, because emulator doesn't allow splitting partitions") } else { val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainer) + val initialFeedRanges = container.getFeedRanges.block().asScala + .map(feedRange => SparkBridgeImplementationInternal.toNormalizedRange(feedRange)) + .sortBy(_.min) + .toSeq for (sequenceNumber <- 1 to 50) { val objectNode = Utils.getSimpleObjectMapper.createObjectNode() - objectNode.put("name", "Shrodigner's cat") + objectNode.put("name", "Schrodinger's cat") objectNode.put("type", "cat") objectNode.put("age", 20) objectNode.put("sequenceNumber", sequenceNumber) @@ -55,14 +66,16 @@ class SparkE2EChangeFeedSplitITest ) val df1 = spark.read.format("cosmos.oltp.changeFeed").options(cfg).load() - val rowsArray1 = df1.collect() + val rowsArray1 = collectWithTimeout(df1, "initial bounded change feed batch") // technically possible that even with 50 documents randomly distributed across 3 partitions some // has no documents // rowsArray should have size df.rdd.getNumPartitions rowsArray1.length > 0 shouldEqual true rowsArray1.length <= df1.rdd.getNumPartitions shouldEqual true - val initialCount = rowsArray1.length + val firstBatchSequenceNumbers = getSequenceNumbers(rowsArray1) + firstBatchSequenceNumbers.distinct.size shouldEqual firstBatchSequenceNumbers.size + df1.rdd.getNumPartitions shouldEqual initialFeedRanges.size df1.schema.equals( ChangeFeedTable.defaultIncrementalChangeFeedSchemaForInferenceDisabled) shouldEqual true @@ -87,48 +100,232 @@ class SparkE2EChangeFeedSplitITest .key(cosmosMasterKey) .buildClient() - val initialThroughput = separateClient - .getDatabase(cosmosDatabase) - .getContainer(cosmosContainer) - .readThroughput() - .getProperties - .getManualThroughput + try { + val splitContainer = separateClient + .getDatabase(cosmosDatabase) + .getContainer(cosmosContainer) + + val initialThroughput = splitContainer + .readThroughput() + .getProperties + .getManualThroughput + val newThroughputToForceSplits = (Math.ceil(initialThroughput.toDouble / 6000) * 2 * 10000).toInt + val response = splitContainer + .replaceThroughput(ThroughputProperties.createManualThroughput(newThroughputToForceSplits)) + + response.getStatusCode shouldEqual 200 + + val (splitParent, childRanges) = waitForSplit(splitContainer, initialFeedRanges) + logInfo(s"Cached parent '$splitParent' split into '${childRanges.mkString(",")}'.") + val partitionKeyDefinition = splitContainer.read().getProperties.getPartitionKeyDefinition + + for (sequenceNumber <- 51 to 100) { + val id = getPartitionKeyValueInRange(childRanges.head, partitionKeyDefinition) + val objectNode = Utils.getSimpleObjectMapper.createObjectNode() + objectNode.put("name", "Schrodinger's cat") + objectNode.put("type", "cat") + objectNode.put("age", 20) + objectNode.put("sequenceNumber", sequenceNumber) + objectNode.put("id", id) + splitContainer.createItem(objectNode) + } + + val cfgWithoutItemCountPerTriggerHint = cfg.filter( + keyValuePair => !keyValuePair._1.equals("spark.cosmos.changeFeed.itemCountPerTriggerHint")) + val expectedSecondBatchSequenceNumbers = (1 to 100).toSet -- firstBatchSequenceNumbers.toSet + val secondBatchSequenceNumbers = mutable.ArrayBuffer.empty[Int] + val postSplitDeadline = System.nanoTime() + TimeUnit.MINUTES.toNanos(TestTimeoutInMinutes) + var postSplitBatchNumber = 0 + + while (secondBatchSequenceNumbers.toSet != expectedSecondBatchSequenceNumbers && + System.nanoTime() < postSplitDeadline) { + + postSplitBatchNumber += 1 + val remainingNanos = Math.max(1L, postSplitDeadline - System.nanoTime()) + val dataFrame = spark.read + .format("cosmos.oltp.changeFeed") + .options(cfgWithoutItemCountPerTriggerHint) + .load() + val rows = collectWithTimeout( + dataFrame, + s"post-split change feed batch '$postSplitBatchNumber'", + remainingNanos) + val batchSequenceNumbers = getSequenceNumbers(rows) + batchSequenceNumbers.distinct.size shouldEqual batchSequenceNumbers.size + secondBatchSequenceNumbers.toSet.intersect(batchSequenceNumbers.toSet) shouldBe empty + batchSequenceNumbers.toSet.subsetOf(expectedSecondBatchSequenceNumbers) shouldEqual true + secondBatchSequenceNumbers ++= batchSequenceNumbers + + if (secondBatchSequenceNumbers.toSet != expectedSecondBatchSequenceNumbers) { + promoteLatestOffset(hdfs, latestOffsetFileLocation, startOffsetFileLocation) + } + } - val newThroughputToForceSplits = (Math.ceil(initialThroughput.toDouble / 6000) * 2 * 10000).toInt + secondBatchSequenceNumbers.toSet shouldEqual expectedSecondBatchSequenceNumbers + } finally { + separateClient.close() + } + } + } - val response = separateClient - .getDatabase(cosmosDatabase) - .getContainer(cosmosContainer) - .replaceThroughput(ThroughputProperties.createManualThroughput(newThroughputToForceSplits)) + private def waitForSplit( + splitContainer: com.azure.cosmos.CosmosContainer, + initialFeedRanges: Seq[NormalizedRange] + ): (NormalizedRange, Seq[NormalizedRange]) = { + val deadline = System.nanoTime() + TimeUnit.MINUTES.toNanos(TestTimeoutInMinutes) + var attempts = 0 + var currentRanges = Seq.empty[NormalizedRange] + var splitResult = Option.empty[(NormalizedRange, Seq[NormalizedRange])] - response.getStatusCode shouldEqual 200 + while (splitResult.isEmpty && + attempts < MaxSplitWaitAttempts && + System.nanoTime() < deadline) { - val initialPartitionCount = df1.rdd.getNumPartitions - var currentPartitionCount = separateClient - .getDatabase(cosmosDatabase) - .getContainer(cosmosContainer) - .getFeedRanges - .size() + attempts += 1 + currentRanges = splitContainer.getFeedRanges.asScala + .map(SparkBridgeImplementationInternal.toNormalizedRange) + .sortBy(_.min) + .toSeq + splitResult = initialFeedRanges + .map(parent => parent -> currentRanges.filter(child => isStrictSubRange(parent, child))) + .find(_._2.size >= 2) - while (currentPartitionCount < initialPartitionCount * 2) { - logInfo(s"Offer replace still pending current Partition count '$currentPartitionCount' - " + - s"target Partition count '${initialPartitionCount * 2}'- waiting for 1 second...") + if (splitResult.isEmpty && attempts < MaxSplitWaitAttempts && System.nanoTime() < deadline) { + logInfo("Offer replace still pending a split of one cached parent range - waiting for 1 second...") Thread.sleep(1000) + } + } - currentPartitionCount = separateClient - .getDatabase(cosmosDatabase) - .getContainer(cosmosContainer) - .getFeedRanges - .size() + splitResult.getOrElse( + fail(s"No cached parent range split after '$attempts' attempts. " + + s"Initial ranges: '${initialFeedRanges.mkString(",")}', current ranges: '${currentRanges.mkString(",")}'.")) + } + + private def collectWithTimeout( + dataFrame: DataFrame, + operation: String, + timeoutNanos: Long = TimeUnit.MINUTES.toNanos(TestTimeoutInMinutes) + ): Array[Row] = { + val jobGroupId = s"spark-change-feed-split-${UUID.randomUUID()}" + val collectExecutor = Executors.newSingleThreadExecutor(new ThreadFactory { + override def newThread(runnable: Runnable): Thread = { + val thread = new Thread(runnable, s"$jobGroupId-worker") + thread.setDaemon(true) + thread + } + }) + val collectFuture = collectExecutor.submit(new Callable[Array[Row]] { + override def call(): Array[Row] = { + spark.sparkContext.setJobGroup(jobGroupId, operation, interruptOnCancel = true) + try { + dataFrame.collect() + } finally { + spark.sparkContext.clearJobGroup() + } } + }) + var restoreInterrupt = false - val cfgWithoutItemCountPerTriggerHint = cfg.filter(keyValuePair => !keyValuePair._1.equals("spark.cosmos.changeFeed.itemCountPerTriggerHint")) - val df2 = spark.read.format("cosmos.oltp.changeFeed").options(cfgWithoutItemCountPerTriggerHint).load() - val rowsArray2 = df2.collect() - rowsArray2 should have size 50 - initialCount + try { + collectFuture.get(timeoutNanos, TimeUnit.NANOSECONDS) + } catch { + case _: TimeoutException => + spark.sparkContext.cancelJobGroup(jobGroupId) + collectFuture.cancel(true) + fail(s"Timed out while running '$operation'.") + case error: InterruptedException => + spark.sparkContext.cancelJobGroup(jobGroupId) + collectFuture.cancel(true) + restoreInterrupt = true + throw error + case error: ExecutionException => + error.getCause match { + case runtimeException: RuntimeException => throw runtimeException + case fatalError: Error => throw fatalError + case cause => throw new RuntimeException(cause) + } + } finally { + collectExecutor.shutdownNow() + try { + if (!collectExecutor.awaitTermination(30, TimeUnit.SECONDS)) { + logWarning(s"Collection worker for '$operation' did not terminate within 30 seconds.") + } + } catch { + case _: InterruptedException => + restoreInterrupt = true + } + if (restoreInterrupt) { + Thread.currentThread().interrupt() + } + } + } + + private def promoteLatestOffset( + hdfs: org.apache.hadoop.fs.FileSystem, + latestOffsetFileLocation: String, + startOffsetFileLocation: String + ): Unit = { + val latestOffsetPath = new Path(latestOffsetFileLocation) + val startOffsetPath = new Path(startOffsetFileLocation) + hdfs.exists(latestOffsetPath) shouldEqual true + if (hdfs.exists(startOffsetPath)) { + hdfs.delete(startOffsetPath, false) shouldEqual true } + hdfs.copyToLocalFile(true, latestOffsetPath, startOffsetPath) + hdfs.exists(latestOffsetPath) shouldEqual false + hdfs.exists(startOffsetPath) shouldEqual true + } + + private def getSequenceNumbers(rows: Array[Row]): Seq[Int] = { + val objectMapper = Utils.getSimpleObjectMapper + rows.map(row => + objectMapper + .readTree(row.getAs[String](CosmosTableSchemaInferrer.RawJsonBodyAttributeName)) + .get("sequenceNumber") + .asInt()) + } + + private def isStrictSubRange(parent: NormalizedRange, child: NormalizedRange): Boolean = { + child != parent && + child.min.compareTo(parent.min) >= 0 && + child.max.compareTo(parent.max) <= 0 + } + + private def getPartitionKeyValueInRange( + targetRange: NormalizedRange, + partitionKeyDefinition: PartitionKeyDefinition + ): String = { + var value = UUID.randomUUID().toString + var effectiveRange = SparkBridgeImplementationInternal.partitionKeyToNormalizedRange( + new PartitionKey(value), + partitionKeyDefinition) + var attempts = 1 + + while (!isPartitionKeyInRange(effectiveRange, targetRange) && + attempts < MaxPartitionKeySearchAttempts) { + + attempts += 1 + value = UUID.randomUUID().toString + effectiveRange = SparkBridgeImplementationInternal.partitionKeyToNormalizedRange( + new PartitionKey(value), + partitionKeyDefinition) + } + + if (isPartitionKeyInRange(effectiveRange, targetRange)) { + value + } else { + fail(s"Unable to find a partition key in range '$targetRange' after '$attempts' attempts.") + } + } + + private def isPartitionKeyInRange( + effectiveRange: NormalizedRange, + targetRange: NormalizedRange + ): Boolean = { + effectiveRange.min.compareTo(targetRange.min) >= 0 && + effectiveRange.min.compareTo(targetRange.max) < 0 } //scalastyle:on magic.number //scalastyle:on multiple.string.literals -} \ No newline at end of file +} diff --git a/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md index c2dc484ba917..87802925d456 100644 --- a/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_4-0_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 diff --git a/sdk/cosmos/azure-cosmos-spark_4-1_2-13/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_4-1_2-13/CHANGELOG.md index ae8e4bb6bd25..5f81671810eb 100644 --- a/sdk/cosmos/azure-cosmos-spark_4-1_2-13/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_4-1_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