diff --git a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/DataflowGraphTransformer.scala b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/DataflowGraphTransformer.scala index 2523c0ae5502..f52d5572054d 100644 --- a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/DataflowGraphTransformer.scala +++ b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/DataflowGraphTransformer.scala @@ -22,10 +22,10 @@ import java.util.concurrent.{ ConcurrentLinkedDeque, ConcurrentLinkedQueue, ExecutionException, + ExecutorCompletionService, Future } -import scala.collection.mutable.ArrayBuffer import scala.jdk.CollectionConverters._ import scala.util.control.NoStackTrace @@ -134,159 +134,168 @@ class DataflowGraphTransformer(graph: DataflowGraph) extends AutoCloseable { val failedFlowsQueue = new ConcurrentLinkedQueue[ResolutionFailedFlow]() val failedDependentFlows = new ConcurrentHashMap[TableIdentifier, Seq[ResolutionFailedFlow]]() - var futures = ArrayBuffer[Future[Unit]]() + val completionService = new ExecutorCompletionService[Unit](executor) + var outstanding = 0 val toBeResolvedFlows = new ConcurrentLinkedDeque[Flow]() toBeResolvedFlows.addAll(flows.asJava) - while (futures.nonEmpty || toBeResolvedFlows.peekFirst() != null) { - val (done, notDone) = futures.partition(_.isDone) - // Explicitly call future.get() to propagate exceptions one by one if any + // Waits on a finished resolution task and propagates its exception, if any. + def reap(finished: Future[Unit]): Unit = { try { - done.foreach(_.get()) + finished.get() } catch { case exn: ExecutionException => // Computation threw the exception that is the cause of exn throw exn.getCause } - futures = notDone - val flowOpt = { - // We only schedule [[batchSize]] number of flows in parallel. - if (futures.size < batchSize) { - Option(toBeResolvedFlows.pollFirst()) - } else { - None - } + outstanding -= 1 + } + + while (outstanding > 0 || toBeResolvedFlows.peekFirst() != null) { + // Reap every resolution task that has already finished, without blocking. + var finished = completionService.poll() + while (finished != null) { + reap(finished) + finished = completionService.poll() } - flowOpt.foreach { flow => - futures.append( - executor.submit( - () => + // We only schedule [[batchSize]] number of flows in parallel. + if (outstanding < batchSize && toBeResolvedFlows.peekFirst() != null) { + val flow = toBeResolvedFlows.pollFirst() + outstanding += 1 + completionService.submit( + () => + try { try { - try { - // Note: Flow don't need their inputs passed, so for now we send empty Seq. - val result = transformer(flow, Seq.empty) - require( - result.forall(_.isInstanceOf[ResolvedFlow]), - "transformer must return a Seq[Flow]" - ) + // Note: Flow don't need their inputs passed, so for now we send empty Seq. + val result = transformer(flow, Seq.empty) + require( + result.forall(_.isInstanceOf[ResolvedFlow]), + "transformer must return a Seq[Flow]" + ) - val transformedFlows = result.map(_.asInstanceOf[ResolvedFlow]) - resolvedFlowsMap.put(flow.identifier, transformedFlows) - resolvedFlows.addAll(transformedFlows.asJava) - } catch { - case e: TransformNodeRetryableException => - val datasetIdentifier = e.datasetIdentifier - failedDependentFlows.compute( - datasetIdentifier, - (_, flows) => { - // Don't add the input flow back but the failed flow object - // back which has relevant failure information. - val failedFlow = e.failedNode - if (flows == null) { - Seq(failedFlow) - } else { - flows :+ failedFlow - } + val transformedFlows = result.map(_.asInstanceOf[ResolvedFlow]) + resolvedFlowsMap.put(flow.identifier, transformedFlows) + resolvedFlows.addAll(transformedFlows.asJava) + } catch { + case e: TransformNodeRetryableException => + val datasetIdentifier = e.datasetIdentifier + failedDependentFlows.compute( + datasetIdentifier, + (_, flows) => { + // Don't add the input flow back but the failed flow object + // back which has relevant failure information. + val failedFlow = e.failedNode + if (flows == null) { + Seq(failedFlow) + } else { + flows :+ failedFlow } - ) - // Between the time the flow started and finished resolving, perhaps the - // dependent dataset was resolved - resolvedFlowDestinationsMap.computeIfPresent( - datasetIdentifier, - (_, resolved) => { - if (resolved) { - // Check if the dataset that the flow is dependent on has been resolved - // and if so, remove all dependent flows from the failedDependentFlows and - // add them to the toBeResolvedFlows queue for retry. - failedDependentFlows.computeIfPresent( - datasetIdentifier, - (_, toRetryFlows) => { - toRetryFlows.foreach(toBeResolvedFlows.addFirst(_)) - null - } - ) - } - resolved + } + ) + // Between the time the flow started and finished resolving, perhaps the + // dependent dataset was resolved + resolvedFlowDestinationsMap.computeIfPresent( + datasetIdentifier, + (_, resolved) => { + if (resolved) { + // Check if the dataset that the flow is dependent on has been resolved + // and if so, remove all dependent flows from the failedDependentFlows and + // add them to the toBeResolvedFlows queue for retry. + failedDependentFlows.computeIfPresent( + datasetIdentifier, + (_, toRetryFlows) => { + toRetryFlows.foreach(toBeResolvedFlows.addFirst(_)) + null + } + ) } + resolved + } + ) + case other: Throwable => throw other + } + // If all flows to this particular destination are resolved, move to the destination + // node transformer + if (flowsTo(flow.destinationIdentifier).forall({ flowToDestination => + resolvedFlowsMap.containsKey(flowToDestination.identifier) + })) { + // If multiple flows completed in parallel, ensure we resolve the destination only + // once by electing a leader via computeIfAbsent + var isCurrentThreadLeader = false + resolvedFlowDestinationsMap.computeIfAbsent(flow.destinationIdentifier, _ => { + isCurrentThreadLeader = true + // Set initial value as false as flow destination is not resolved yet. + false + }) + if (isCurrentThreadLeader) { + if (tableMap.contains(flow.destinationIdentifier)) { + val transformed = + transformer( + tableMap(flow.destinationIdentifier), + flowsTo(flow.destinationIdentifier) + ) + resolvedTables.addAll( + transformed.collect { case t: Table => t }.asJava ) - case other: Throwable => throw other - } - // If all flows to this particular destination are resolved, move to the destination - // node transformer - if (flowsTo(flow.destinationIdentifier).forall({ flowToDestination => - resolvedFlowsMap.containsKey(flowToDestination.identifier) - })) { - // If multiple flows completed in parallel, ensure we resolve the destination only - // once by electing a leader via computeIfAbsent - var isCurrentThreadLeader = false - resolvedFlowDestinationsMap.computeIfAbsent(flow.destinationIdentifier, _ => { - isCurrentThreadLeader = true - // Set initial value as false as flow destination is not resolved yet. - false - }) - if (isCurrentThreadLeader) { - if (tableMap.contains(flow.destinationIdentifier)) { + resolvedFlows.addAll( + transformed.collect { case f: ResolvedFlow => f }.asJava + ) + } else if (viewMap.contains(flow.destinationIdentifier)) { + resolvedViews.addAll { val transformed = transformer( - tableMap(flow.destinationIdentifier), + viewMap(flow.destinationIdentifier), flowsTo(flow.destinationIdentifier) ) - resolvedTables.addAll( - transformed.collect { case t: Table => t }.asJava - ) - resolvedFlows.addAll( - transformed.collect { case f: ResolvedFlow => f }.asJava - ) - } else if (viewMap.contains(flow.destinationIdentifier)) { - resolvedViews.addAll { - val transformed = - transformer( - viewMap(flow.destinationIdentifier), - flowsTo(flow.destinationIdentifier) - ) - transformed.map(_.asInstanceOf[View]).asJava - } - } else if (sinkMap.contains(flow.destinationIdentifier)) { - resolvedSinks.addAll { - val transformed = - transformer( - sinkMap(flow.destinationIdentifier), flowsTo(flow.destinationIdentifier) - ) - require( - transformed.forall(_.isInstanceOf[Sink]), - "transformer must return a Seq[Sink]" + transformed.map(_.asInstanceOf[View]).asJava + } + } else if (sinkMap.contains(flow.destinationIdentifier)) { + resolvedSinks.addAll { + val transformed = + transformer( + sinkMap(flow.destinationIdentifier), flowsTo(flow.destinationIdentifier) ) - transformed.map(_.asInstanceOf[Sink]).asJava - } - } else { - throw new IllegalArgumentException( - s"Unsupported destination ${flow.destinationIdentifier.unquotedString}" + - s" in flow: ${flow.displayName} at transformDownNodes" + require( + transformed.forall(_.isInstanceOf[Sink]), + "transformer must return a Seq[Sink]" ) + transformed.map(_.asInstanceOf[Sink]).asJava } - // Set flow destination as resolved now. - resolvedFlowDestinationsMap.computeIfPresent( - flow.destinationIdentifier, - (_, _) => { - // If there are any other node failures dependent on this destination, retry - // them - failedDependentFlows.computeIfPresent( - flow.destinationIdentifier, - (_, toRetryFlows) => { - toRetryFlows.foreach(toBeResolvedFlows.addFirst(_)) - null - } - ) - true - } + } else { + throw new IllegalArgumentException( + s"Unsupported destination ${flow.destinationIdentifier.unquotedString}" + + s" in flow: ${flow.displayName} at transformDownNodes" ) } + // Set flow destination as resolved now. + resolvedFlowDestinationsMap.computeIfPresent( + flow.destinationIdentifier, + (_, _) => { + // If there are any other node failures dependent on this destination, retry + // them + failedDependentFlows.computeIfPresent( + flow.destinationIdentifier, + (_, toRetryFlows) => { + toRetryFlows.foreach(toBeResolvedFlows.addFirst(_)) + null + } + ) + true + } + ) } - } catch { - case ex: TransformNodeFailedException => failedFlowsQueue.add(ex.failedNode) } - ) + } catch { + case ex: TransformNodeFailedException => failedFlowsQueue.add(ex.failedNode) + } ) + } else if (outstanding > 0) { + // Nothing could be scheduled (slots full, or the queue is drained) but tasks are still + // running: block until the next finishes instead of busy-spinning on Future.isDone. The + // outstanding > 0 guard is required, not redundant: the poll() drain above can take + // outstanding to 0 with an empty queue, and then there is nothing to wait for - the loop + // should just exit rather than block forever in take(). + reap(completionService.take()) } } diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/ConnectValidPipelineSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/ConnectValidPipelineSuite.scala index 4844cd6d6db3..2a636a9564a0 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/ConnectValidPipelineSuite.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/ConnectValidPipelineSuite.scala @@ -512,6 +512,64 @@ class ConnectValidPipelineSuite extends PipelineTest with SharedSparkSession { assert(g.flow(TableIdentifier("sink_flow")).isInstanceOf[StreamingFlow]) } + test("resolution terminates and resolves all flows when flow count exceeds parallelism") { + val session = spark + import session.implicits._ + + // DataflowGraphTransformer caps in-flight resolutions at `parallelism` (10). With more + // independent flows than that, the slots fill and the scheduler blocks on a finished task + // (the `take()` branch) once `parallelism` tasks are outstanding - the path the busy-wait + // rewrite changes. Small graphs in the other suites never reach this regime. This asserts the + // outcome only (all flows resolve and the call returns), so it is deterministic and has no + // timing dependence; a regression that deadlocked would hang here until the suite times out. + val numFlows = 25 + class P extends TestGraphRegistrationContext(spark) { + (0 until numFlows).foreach { i => + registerPersistedView(s"v$i", query = dfFlowFunc(Seq(i).toDF("x"))) + } + } + val p = new P().resolveToDataflowGraph() + + assert(p.resolved, "all flows should resolve when their count exceeds parallelism") + (0 until numFlows).foreach { i => + assert( + p.resolvedFlow.contains(fullyQualifiedIdentifier(s"v$i")), + s"flow v$i was not resolved") + } + } + + test("resolution re-queues retryable flows under load when consumers exceed parallelism") { + val session = spark + import session.implicits._ + + // A wide fan-out: many consumers reading from one source view, with the consumer count above + // `parallelism` (10), so slots fill and the loop blocks on take(). The consumers are + // registered (and therefore scheduled) before `src`, so the first batch resolves consumers + // whose `src` input is not yet available: each throws TransformNodeRetryableException and is + // parked as a dependent of `src`; once `src` resolves they are re-queued onto the deque and + // retried, re-driving the loop until every consumer resolves. This deterministically exercises + // the retryable re-queue path together with the blocking branch. Asserts only that everything + // resolves and the call returns - no timing assertions. + val numConsumers = 20 + class P extends TestGraphRegistrationContext(spark) { + (0 until numConsumers).foreach { i => + registerPersistedView(s"c$i", query = sqlFlowFunc(spark, "SELECT x FROM src")) + } + registerPersistedView("src", query = dfFlowFunc(Seq(1, 2, 3).toDF("x"))) + } + val p = new P().resolveToDataflowGraph() + + assert(p.resolved, "source and all consumers should resolve under load") + assert( + p.resolvedFlow.contains(fullyQualifiedIdentifier("src")), + "source flow was not resolved") + (0 until numConsumers).foreach { i => + assert( + p.resolvedFlow.contains(fullyQualifiedIdentifier(s"c$i")), + s"consumer flow c$i was not resolved") + } + } + /** Verifies the [[DataflowGraph]] has the specified [[Flow]] with the specified schema. */ private def verifyFlowSchema( pipeline: DataflowGraph,