From 6539752a2d3123239ed1775e623f22d26d4e90ea Mon Sep 17 00:00:00 2001 From: YangJie Date: Mon, 27 Jul 2026 23:31:18 +0800 Subject: [PATCH] [SPARK-57638][SQL] Avoid busy-waiting in Declarative Pipelines flow resolution `DataflowGraphTransformer.transformDownNodes` resolves flows on a bounded thread pool and drives them from a `while` loop that, each pass, partitioned the in-flight futures with the non-blocking `future.isDone`, reaped the completed ones, and scheduled a new flow if a slot was free. When all slots were in flight (or the queue was drained and only the last futures remained) and none had completed, the pass reaped nothing and scheduled nothing, then looped again immediately - busy-spinning on `isDone` and pinning a core for the duration of resolution. This drives the loop with an `ExecutorCompletionService` instead: completed tasks are drained with the non-blocking `poll()`, and when nothing can be scheduled but tasks are still running, the loop blocks on `take()` until the next one finishes rather than spinning. Behavior is otherwise unchanged - the same flows are scheduled in the same order, exceptions are still propagated via `Future.get()`, and an `outstanding` counter replaces the `ArrayBuffer[Future]` for slot bookkeeping. Resolving a graph with more flows than the parallelism (10) kept one CPU core busy at 100% doing no useful work for the whole resolution, which is wasteful and shows up as unexplained driver CPU. No. Two new cases in `ConnectValidPipelineSuite` cover the regime this PR changes - more flows than `parallelism` (10), so the slots fill and the loop reaches the blocking `take()` branch that replaces the busy-wait. The small graphs in the existing suites never get there. - `resolution terminates and resolves all flows when flow count exceeds parallelism` - 25 independent flows. - `resolution re-queues retryable flows under load when consumers exceed parallelism` - 20 consumers registered before their source `src`, so the first batch throws `TransformNodeRetryableException`, parks as dependents of `src`, and is re-queued once `src` resolves; this exercises the retryable re-queue path together with the blocking branch. Both assert only the outcome (every flow resolves and the call returns), so they are deterministic and have no timing dependence - a regression that deadlocked would hang until the suite times out. Asserting the absence of a busy-wait directly is not included, since that requires CPU-time or timing measurements that are flaky in CI. Existing graph-resolution suites (`ConnectValidPipelineSuite`, `ConnectInvalidPipelineSuite`, `SqlPipelineSuite`, `TriggeredGraphExecutionSuite`, `MaterializeTablesSuite`) still pass; the change only affects how the loop waits, not what it resolves. Generated-by: Claude Code (Claude Opus 4.8) Closes #56700 from LuciferYang/sdp-resolution-busy-wait. Authored-by: YangJie Signed-off-by: yangjie01 (cherry picked from commit 0747e287579272958600e3db3ccf03044378576f) --- .../graph/DataflowGraphTransformer.scala | 263 +++++++++--------- .../graph/ConnectValidPipelineSuite.scala | 58 ++++ 2 files changed, 194 insertions(+), 127 deletions(-) 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 2523c0ae5502a..f52d5572054d2 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 4844cd6d6db36..2a636a9564a04 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,