Skip to content

Commit 6539752

Browse files
LuciferYangLuciferYang
committed
[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 apache#56700 from LuciferYang/sdp-resolution-busy-wait. Authored-by: YangJie <yangjie01@baidu.com> Signed-off-by: yangjie01 <yangjie01@baidu.com> (cherry picked from commit 0747e28)
1 parent fe16861 commit 6539752

2 files changed

Lines changed: 194 additions & 127 deletions

File tree

‎sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/DataflowGraphTransformer.scala‎

Lines changed: 136 additions & 127 deletions
Original file line numberDiff line numberDiff line change
@@ -22,10 +22,10 @@ import java.util.concurrent.{
2222
ConcurrentLinkedDeque,
2323
ConcurrentLinkedQueue,
2424
ExecutionException,
25+
ExecutorCompletionService,
2526
Future
2627
}
2728

28-
import scala.collection.mutable.ArrayBuffer
2929
import scala.jdk.CollectionConverters._
3030
import scala.util.control.NoStackTrace
3131

@@ -134,159 +134,168 @@ class DataflowGraphTransformer(graph: DataflowGraph) extends AutoCloseable {
134134
val failedFlowsQueue = new ConcurrentLinkedQueue[ResolutionFailedFlow]()
135135
val failedDependentFlows = new ConcurrentHashMap[TableIdentifier, Seq[ResolutionFailedFlow]]()
136136

137-
var futures = ArrayBuffer[Future[Unit]]()
137+
val completionService = new ExecutorCompletionService[Unit](executor)
138+
var outstanding = 0
138139
val toBeResolvedFlows = new ConcurrentLinkedDeque[Flow]()
139140
toBeResolvedFlows.addAll(flows.asJava)
140141

141-
while (futures.nonEmpty || toBeResolvedFlows.peekFirst() != null) {
142-
val (done, notDone) = futures.partition(_.isDone)
143-
// Explicitly call future.get() to propagate exceptions one by one if any
142+
// Waits on a finished resolution task and propagates its exception, if any.
143+
def reap(finished: Future[Unit]): Unit = {
144144
try {
145-
done.foreach(_.get())
145+
finished.get()
146146
} catch {
147147
case exn: ExecutionException =>
148148
// Computation threw the exception that is the cause of exn
149149
throw exn.getCause
150150
}
151-
futures = notDone
152-
val flowOpt = {
153-
// We only schedule [[batchSize]] number of flows in parallel.
154-
if (futures.size < batchSize) {
155-
Option(toBeResolvedFlows.pollFirst())
156-
} else {
157-
None
158-
}
151+
outstanding -= 1
152+
}
153+
154+
while (outstanding > 0 || toBeResolvedFlows.peekFirst() != null) {
155+
// Reap every resolution task that has already finished, without blocking.
156+
var finished = completionService.poll()
157+
while (finished != null) {
158+
reap(finished)
159+
finished = completionService.poll()
159160
}
160-
flowOpt.foreach { flow =>
161-
futures.append(
162-
executor.submit(
163-
() =>
161+
// We only schedule [[batchSize]] number of flows in parallel.
162+
if (outstanding < batchSize && toBeResolvedFlows.peekFirst() != null) {
163+
val flow = toBeResolvedFlows.pollFirst()
164+
outstanding += 1
165+
completionService.submit(
166+
() =>
167+
try {
164168
try {
165-
try {
166-
// Note: Flow don't need their inputs passed, so for now we send empty Seq.
167-
val result = transformer(flow, Seq.empty)
168-
require(
169-
result.forall(_.isInstanceOf[ResolvedFlow]),
170-
"transformer must return a Seq[Flow]"
171-
)
169+
// Note: Flow don't need their inputs passed, so for now we send empty Seq.
170+
val result = transformer(flow, Seq.empty)
171+
require(
172+
result.forall(_.isInstanceOf[ResolvedFlow]),
173+
"transformer must return a Seq[Flow]"
174+
)
172175

173-
val transformedFlows = result.map(_.asInstanceOf[ResolvedFlow])
174-
resolvedFlowsMap.put(flow.identifier, transformedFlows)
175-
resolvedFlows.addAll(transformedFlows.asJava)
176-
} catch {
177-
case e: TransformNodeRetryableException =>
178-
val datasetIdentifier = e.datasetIdentifier
179-
failedDependentFlows.compute(
180-
datasetIdentifier,
181-
(_, flows) => {
182-
// Don't add the input flow back but the failed flow object
183-
// back which has relevant failure information.
184-
val failedFlow = e.failedNode
185-
if (flows == null) {
186-
Seq(failedFlow)
187-
} else {
188-
flows :+ failedFlow
189-
}
176+
val transformedFlows = result.map(_.asInstanceOf[ResolvedFlow])
177+
resolvedFlowsMap.put(flow.identifier, transformedFlows)
178+
resolvedFlows.addAll(transformedFlows.asJava)
179+
} catch {
180+
case e: TransformNodeRetryableException =>
181+
val datasetIdentifier = e.datasetIdentifier
182+
failedDependentFlows.compute(
183+
datasetIdentifier,
184+
(_, flows) => {
185+
// Don't add the input flow back but the failed flow object
186+
// back which has relevant failure information.
187+
val failedFlow = e.failedNode
188+
if (flows == null) {
189+
Seq(failedFlow)
190+
} else {
191+
flows :+ failedFlow
190192
}
191-
)
192-
// Between the time the flow started and finished resolving, perhaps the
193-
// dependent dataset was resolved
194-
resolvedFlowDestinationsMap.computeIfPresent(
195-
datasetIdentifier,
196-
(_, resolved) => {
197-
if (resolved) {
198-
// Check if the dataset that the flow is dependent on has been resolved
199-
// and if so, remove all dependent flows from the failedDependentFlows and
200-
// add them to the toBeResolvedFlows queue for retry.
201-
failedDependentFlows.computeIfPresent(
202-
datasetIdentifier,
203-
(_, toRetryFlows) => {
204-
toRetryFlows.foreach(toBeResolvedFlows.addFirst(_))
205-
null
206-
}
207-
)
208-
}
209-
resolved
193+
}
194+
)
195+
// Between the time the flow started and finished resolving, perhaps the
196+
// dependent dataset was resolved
197+
resolvedFlowDestinationsMap.computeIfPresent(
198+
datasetIdentifier,
199+
(_, resolved) => {
200+
if (resolved) {
201+
// Check if the dataset that the flow is dependent on has been resolved
202+
// and if so, remove all dependent flows from the failedDependentFlows and
203+
// add them to the toBeResolvedFlows queue for retry.
204+
failedDependentFlows.computeIfPresent(
205+
datasetIdentifier,
206+
(_, toRetryFlows) => {
207+
toRetryFlows.foreach(toBeResolvedFlows.addFirst(_))
208+
null
209+
}
210+
)
210211
}
212+
resolved
213+
}
214+
)
215+
case other: Throwable => throw other
216+
}
217+
// If all flows to this particular destination are resolved, move to the destination
218+
// node transformer
219+
if (flowsTo(flow.destinationIdentifier).forall({ flowToDestination =>
220+
resolvedFlowsMap.containsKey(flowToDestination.identifier)
221+
})) {
222+
// If multiple flows completed in parallel, ensure we resolve the destination only
223+
// once by electing a leader via computeIfAbsent
224+
var isCurrentThreadLeader = false
225+
resolvedFlowDestinationsMap.computeIfAbsent(flow.destinationIdentifier, _ => {
226+
isCurrentThreadLeader = true
227+
// Set initial value as false as flow destination is not resolved yet.
228+
false
229+
})
230+
if (isCurrentThreadLeader) {
231+
if (tableMap.contains(flow.destinationIdentifier)) {
232+
val transformed =
233+
transformer(
234+
tableMap(flow.destinationIdentifier),
235+
flowsTo(flow.destinationIdentifier)
236+
)
237+
resolvedTables.addAll(
238+
transformed.collect { case t: Table => t }.asJava
211239
)
212-
case other: Throwable => throw other
213-
}
214-
// If all flows to this particular destination are resolved, move to the destination
215-
// node transformer
216-
if (flowsTo(flow.destinationIdentifier).forall({ flowToDestination =>
217-
resolvedFlowsMap.containsKey(flowToDestination.identifier)
218-
})) {
219-
// If multiple flows completed in parallel, ensure we resolve the destination only
220-
// once by electing a leader via computeIfAbsent
221-
var isCurrentThreadLeader = false
222-
resolvedFlowDestinationsMap.computeIfAbsent(flow.destinationIdentifier, _ => {
223-
isCurrentThreadLeader = true
224-
// Set initial value as false as flow destination is not resolved yet.
225-
false
226-
})
227-
if (isCurrentThreadLeader) {
228-
if (tableMap.contains(flow.destinationIdentifier)) {
240+
resolvedFlows.addAll(
241+
transformed.collect { case f: ResolvedFlow => f }.asJava
242+
)
243+
} else if (viewMap.contains(flow.destinationIdentifier)) {
244+
resolvedViews.addAll {
229245
val transformed =
230246
transformer(
231-
tableMap(flow.destinationIdentifier),
247+
viewMap(flow.destinationIdentifier),
232248
flowsTo(flow.destinationIdentifier)
233249
)
234-
resolvedTables.addAll(
235-
transformed.collect { case t: Table => t }.asJava
236-
)
237-
resolvedFlows.addAll(
238-
transformed.collect { case f: ResolvedFlow => f }.asJava
239-
)
240-
} else if (viewMap.contains(flow.destinationIdentifier)) {
241-
resolvedViews.addAll {
242-
val transformed =
243-
transformer(
244-
viewMap(flow.destinationIdentifier),
245-
flowsTo(flow.destinationIdentifier)
246-
)
247-
transformed.map(_.asInstanceOf[View]).asJava
248-
}
249-
} else if (sinkMap.contains(flow.destinationIdentifier)) {
250-
resolvedSinks.addAll {
251-
val transformed =
252-
transformer(
253-
sinkMap(flow.destinationIdentifier), flowsTo(flow.destinationIdentifier)
254-
)
255-
require(
256-
transformed.forall(_.isInstanceOf[Sink]),
257-
"transformer must return a Seq[Sink]"
250+
transformed.map(_.asInstanceOf[View]).asJava
251+
}
252+
} else if (sinkMap.contains(flow.destinationIdentifier)) {
253+
resolvedSinks.addAll {
254+
val transformed =
255+
transformer(
256+
sinkMap(flow.destinationIdentifier), flowsTo(flow.destinationIdentifier)
258257
)
259-
transformed.map(_.asInstanceOf[Sink]).asJava
260-
}
261-
} else {
262-
throw new IllegalArgumentException(
263-
s"Unsupported destination ${flow.destinationIdentifier.unquotedString}" +
264-
s" in flow: ${flow.displayName} at transformDownNodes"
258+
require(
259+
transformed.forall(_.isInstanceOf[Sink]),
260+
"transformer must return a Seq[Sink]"
265261
)
262+
transformed.map(_.asInstanceOf[Sink]).asJava
266263
}
267-
// Set flow destination as resolved now.
268-
resolvedFlowDestinationsMap.computeIfPresent(
269-
flow.destinationIdentifier,
270-
(_, _) => {
271-
// If there are any other node failures dependent on this destination, retry
272-
// them
273-
failedDependentFlows.computeIfPresent(
274-
flow.destinationIdentifier,
275-
(_, toRetryFlows) => {
276-
toRetryFlows.foreach(toBeResolvedFlows.addFirst(_))
277-
null
278-
}
279-
)
280-
true
281-
}
264+
} else {
265+
throw new IllegalArgumentException(
266+
s"Unsupported destination ${flow.destinationIdentifier.unquotedString}" +
267+
s" in flow: ${flow.displayName} at transformDownNodes"
282268
)
283269
}
270+
// Set flow destination as resolved now.
271+
resolvedFlowDestinationsMap.computeIfPresent(
272+
flow.destinationIdentifier,
273+
(_, _) => {
274+
// If there are any other node failures dependent on this destination, retry
275+
// them
276+
failedDependentFlows.computeIfPresent(
277+
flow.destinationIdentifier,
278+
(_, toRetryFlows) => {
279+
toRetryFlows.foreach(toBeResolvedFlows.addFirst(_))
280+
null
281+
}
282+
)
283+
true
284+
}
285+
)
284286
}
285-
} catch {
286-
case ex: TransformNodeFailedException => failedFlowsQueue.add(ex.failedNode)
287287
}
288-
)
288+
} catch {
289+
case ex: TransformNodeFailedException => failedFlowsQueue.add(ex.failedNode)
290+
}
289291
)
292+
} else if (outstanding > 0) {
293+
// Nothing could be scheduled (slots full, or the queue is drained) but tasks are still
294+
// running: block until the next finishes instead of busy-spinning on Future.isDone. The
295+
// outstanding > 0 guard is required, not redundant: the poll() drain above can take
296+
// outstanding to 0 with an empty queue, and then there is nothing to wait for - the loop
297+
// should just exit rather than block forever in take().
298+
reap(completionService.take())
290299
}
291300
}
292301

‎sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/ConnectValidPipelineSuite.scala‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -512,6 +512,64 @@ class ConnectValidPipelineSuite extends PipelineTest with SharedSparkSession {
512512
assert(g.flow(TableIdentifier("sink_flow")).isInstanceOf[StreamingFlow])
513513
}
514514

515+
test("resolution terminates and resolves all flows when flow count exceeds parallelism") {
516+
val session = spark
517+
import session.implicits._
518+
519+
// DataflowGraphTransformer caps in-flight resolutions at `parallelism` (10). With more
520+
// independent flows than that, the slots fill and the scheduler blocks on a finished task
521+
// (the `take()` branch) once `parallelism` tasks are outstanding - the path the busy-wait
522+
// rewrite changes. Small graphs in the other suites never reach this regime. This asserts the
523+
// outcome only (all flows resolve and the call returns), so it is deterministic and has no
524+
// timing dependence; a regression that deadlocked would hang here until the suite times out.
525+
val numFlows = 25
526+
class P extends TestGraphRegistrationContext(spark) {
527+
(0 until numFlows).foreach { i =>
528+
registerPersistedView(s"v$i", query = dfFlowFunc(Seq(i).toDF("x")))
529+
}
530+
}
531+
val p = new P().resolveToDataflowGraph()
532+
533+
assert(p.resolved, "all flows should resolve when their count exceeds parallelism")
534+
(0 until numFlows).foreach { i =>
535+
assert(
536+
p.resolvedFlow.contains(fullyQualifiedIdentifier(s"v$i")),
537+
s"flow v$i was not resolved")
538+
}
539+
}
540+
541+
test("resolution re-queues retryable flows under load when consumers exceed parallelism") {
542+
val session = spark
543+
import session.implicits._
544+
545+
// A wide fan-out: many consumers reading from one source view, with the consumer count above
546+
// `parallelism` (10), so slots fill and the loop blocks on take(). The consumers are
547+
// registered (and therefore scheduled) before `src`, so the first batch resolves consumers
548+
// whose `src` input is not yet available: each throws TransformNodeRetryableException and is
549+
// parked as a dependent of `src`; once `src` resolves they are re-queued onto the deque and
550+
// retried, re-driving the loop until every consumer resolves. This deterministically exercises
551+
// the retryable re-queue path together with the blocking branch. Asserts only that everything
552+
// resolves and the call returns - no timing assertions.
553+
val numConsumers = 20
554+
class P extends TestGraphRegistrationContext(spark) {
555+
(0 until numConsumers).foreach { i =>
556+
registerPersistedView(s"c$i", query = sqlFlowFunc(spark, "SELECT x FROM src"))
557+
}
558+
registerPersistedView("src", query = dfFlowFunc(Seq(1, 2, 3).toDF("x")))
559+
}
560+
val p = new P().resolveToDataflowGraph()
561+
562+
assert(p.resolved, "source and all consumers should resolve under load")
563+
assert(
564+
p.resolvedFlow.contains(fullyQualifiedIdentifier("src")),
565+
"source flow was not resolved")
566+
(0 until numConsumers).foreach { i =>
567+
assert(
568+
p.resolvedFlow.contains(fullyQualifiedIdentifier(s"c$i")),
569+
s"consumer flow c$i was not resolved")
570+
}
571+
}
572+
515573
/** Verifies the [[DataflowGraph]] has the specified [[Flow]] with the specified schema. */
516574
private def verifyFlowSchema(
517575
pipeline: DataflowGraph,

0 commit comments

Comments
 (0)