xiangfu0 commented on code in PR #19944: URL: https://github.com/apache/pinot/pull/19944#discussion_r4343867199
########## pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/SpoolDeadlockGuard.java: ########## @@ -0,0 +1,434 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.planner.logical; + +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.BitSet; +import java.util.Collections; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import javax.annotation.Nullable; +import org.apache.pinot.calcite.rel.logical.PinotRelExchangeType; +import org.apache.pinot.query.planner.plannode.AggregateNode; +import org.apache.pinot.query.planner.plannode.JoinNode; +import org.apache.pinot.query.planner.plannode.MailboxReceiveNode; +import org.apache.pinot.query.planner.plannode.MailboxSendNode; +import org.apache.pinot.query.planner.plannode.PlanNode; +import org.apache.pinot.query.planner.plannode.SetOpNode; +import org.apache.pinot.query.planner.plannode.SortNode; +import org.apache.pinot.query.planner.plannode.WindowNode; + + +/// Detects plans where a spool can block the query until it times out. +/// +/// A spool is a stage that sends its output to more than one consumer stage. It sends each block to all its +/// consumers, so it moves at the pace of the slowest consumer: when the mailbox of one consumer is full, the spool +/// stops for everybody. Some operators read one input only after another input finishes: +/// +/// - A join reads its left (probe) input only after its right (build) input finishes. +/// - `INTERSECT` and `EXCEPT` read their left input only after their right input finishes. +/// - `UNION` reads its inputs in order: input `i` only after inputs `0` to `i - 1` finish. +/// - The main operator chain of a stage reads nothing until the pipeline breaker of that stage finishes, and the +/// pipeline breaker reads its own receives in round-robin with blocking reads. +/// +/// Such an operator *defers* the receives in its later input behind the receives in its earlier input. While a +/// deferred receive waits, nobody reads it: +/// +/// 1. Its mailbox fills up and its sender stops. +/// 2. The sender stops reading its inputs, and the stop moves up through every stage on a *streaming* path. +/// 3. Each spool on that path stops for all its consumers. +/// +/// A waiting receive therefore *blocks* every deferral whose earlier input needs one of the stopped stages: that +/// earlier input cannot finish, so the receives that this other deferral defers wait too. The query blocks when the +/// deferrals block each other in a cycle. The simplest cycle has one deferral: a stage that the deferred receive stops +/// is necessary for the earlier input of the same operator. Longer cycles span several operators, for example two +/// joins where each one reads, as its build input, the spool that feeds the probe input of the other. +/// +/// An operator that reads its whole input before it emits its first row ends a streaming path: an aggregate, a +/// window, a sort with collations, the build input of a join, and the right input of `INTERSECT` and `EXCEPT`. The +/// stop does not go through it, because it has already read everything when the stop comes. +/// +/// The check is conservative: it does not use the size of the data, so it also refuses plans whose data fits in the +/// mailbox buffers. It also does not trust an early end of the input (for example a `LIMIT`) to free a spool. +/// +/// This class is stateless and thread-safe. +public final class SpoolDeadlockGuard { + /// Marks the absence of a shared stage in the matrix of [Analysis#findRisk]. Stage ids are never negative. + private static final int NO_STAGE = -1; + + private SpoolDeadlockGuard() { + } + + /// Returns a cycle of deferrals that can block the plan that starts at `root`, or null if there is none. + /// + /// `root` is the root of the plan after fragmentation: the root of stage 1, or any [MailboxSendNode]. + @Nullable + public static DeadlockRisk findRisk(PlanNode root) { + return new Analysis().findRisk(root); + } + + /// A cycle of deferrals that can block the query. + /// + /// Each element of `cycle` is a wait that cannot end until the next wait (the first one after the last) ends. + public record DeadlockRisk(List<Wait> cycle) { + public DeadlockRisk { + cycle = List.copyOf(cycle); + } + + @Override + public String toString() { + if (cycle.size() == 1) { + Wait wait = cycle.get(0); + return wait + ", and both depend on stage " + wait.sharedStageId(); + } + StringBuilder builder = new StringBuilder("Circular wait: "); + for (int i = 0; i < cycle.size(); i++) { + Wait wait = cycle.get(i); + if (i > 0) { + builder.append("; "); + } + builder.append(wait).append(", and stage ").append(wait.blockingSenderStageId()) + .append(" depends on stage ").append(wait.sharedStageId()).append(", which the next wait stops"); + } + return builder.toString(); + } + } + + /// In stage `stageId`, the `operator` reads the receive from `waitingSenderStageId` only after the receive from + /// `blockingSenderStageId` finishes. The blocking receive needs `sharedStageId`, a stage that the next wait of the + /// cycle stops. + public record Wait(int stageId, String operator, int waitingSenderStageId, int blockingSenderStageId, + int sharedStageId) { + @Override + public String toString() { + return "In stage " + stageId + ", " + operator + " reads the input from stage " + waitingSenderStageId + + " only after the input from stage " + blockingSenderStageId + " finishes"; + } + } + + /// A pair of receives in one stage where `waiting` is read only after `blocking` finishes. + private record Deferral(MailboxReceiveNode waiting, MailboxReceiveNode blocking, String operator) { + } + + /// The receives of one subtree inside a stage. The lists are mutable: [Analysis#collect] fills them while it walks + /// the subtree. + /// + /// @param all every receive in the subtree + /// @param streaming the receives that reach the root of the subtree through a streaming path + private record Receives(List<MailboxReceiveNode> all, List<MailboxReceiveNode> streaming) { + Receives() { + this(new ArrayList<>(), new ArrayList<>()); + } + + void addAll(Receives other) { + all.addAll(other.all); + streaming.addAll(other.streaming); + } + } + + /// What the analysis knows about one stage. + private record StageInfo(Receives receives, List<Deferral> deferrals) { + } + + /// One run of the check. Not thread-safe; each call to [SpoolDeadlockGuard#findRisk] creates its own instance. + private static final class Analysis { + private final Map<PlanNode, StageInfo> _stageInfos = new IdentityHashMap<>(); + /// For each sender, the ids of the stages that stop when nobody reads the output of that sender. + private final Map<MailboxSendNode, BitSet> _stopped = new IdentityHashMap<>(); + /// For each sender, the ids of the stages that must finish before that sender finishes. + private final Map<MailboxSendNode, BitSet> _needed = new IdentityHashMap<>(); + + @Nullable + DeadlockRisk findRisk(PlanNode root) { + List<Deferral> deferrals = collectDeferrals(root); + int numDeferrals = deferrals.size(); + // shared[i][j] is the id of a stage that the waiting receive of deferral i stops and that the blocking receive + // of deferral j needs, or NO_STAGE. When it is not NO_STAGE, deferral i blocks deferral j. + int[][] shared = new int[numDeferrals][numDeferrals]; + for (int i = 0; i < numDeferrals; i++) { + BitSet stopped = getStopped(deferrals.get(i).waiting().getSender()); + for (int j = 0; j < numDeferrals; j++) { + shared[i][j] = firstShared(stopped, getNeeded(deferrals.get(j).blocking().getSender())); + } + } + for (int i = 0; i < numDeferrals; i++) { + if (shared[i][i] != NO_STAGE) { + return new DeadlockRisk(List.of(toWait(deferrals.get(i), shared[i][i]))); + } + } + List<Integer> cycle = findCycle(shared); + if (cycle == null) { + return null; + } + // In the cycle, deferral cycle[k] blocks deferral cycle[k + 1]. A wait describes what keeps it waiting: its + // blocking receive needs a stage that the next wait stops, so the waits are listed in the reverse order. + List<Wait> waits = new ArrayList<>(cycle.size()); + for (int k = cycle.size() - 1; k >= 0; k--) { + int blocked = cycle.get(k); + int blocker = cycle.get(k == 0 ? cycle.size() - 1 : k - 1); + waits.add(toWait(deferrals.get(blocked), shared[blocker][blocked])); + } + return new DeadlockRisk(waits); + } + + private static Wait toWait(Deferral deferral, int sharedStageId) { + // Both receives of a deferral belong to the stage of the operator that defers one behind the other. + return new Wait(deferral.waiting().getStageId(), deferral.operator(), deferral.waiting().getSenderStageId(), + deferral.blocking().getSenderStageId(), sharedStageId); + } + + /// Returns the deferrals of every stage reachable from `root`. Two deferrals with the same pair of senders behave + /// the same, so only the first one is kept. + /// + /// Returns an empty list when no stage sends to more than one stage. Without a spool, the stages form a tree: the + /// subtrees of two inputs never share a stage, so no deferral can block itself or take part in a cycle. + private List<Deferral> collectDeferrals(PlanNode root) { + List<Deferral> result = new ArrayList<>(); + boolean hasSpool = false; + Map<MailboxSendNode, Set<MailboxSendNode>> seenPairs = new IdentityHashMap<>(); + ArrayDeque<PlanNode> pending = new ArrayDeque<>(); + Set<PlanNode> visited = Collections.newSetFromMap(new IdentityHashMap<>()); + pending.add(root); + visited.add(root); + while (!pending.isEmpty()) { + PlanNode stageRoot = pending.poll(); + hasSpool |= stageRoot instanceof MailboxSendNode && ((MailboxSendNode) stageRoot).isMultiSend(); + StageInfo info = getStageInfo(stageRoot); + for (Deferral deferral : info.deferrals()) { + Set<MailboxSendNode> blockingSenders = seenPairs.computeIfAbsent(deferral.waiting().getSender(), + k -> Collections.newSetFromMap(new IdentityHashMap<>())); + if (blockingSenders.add(deferral.blocking().getSender())) { + result.add(deferral); + } + } + for (MailboxReceiveNode receive : info.receives().all()) { + MailboxSendNode sender = receive.getSender(); + if (visited.add(sender)) { + pending.add(sender); + } + } + } + return hasSpool ? result : List.of(); + } + + /// Returns the nodes of a cycle in the graph where `i -> j` exists when `edges[i][j]` is not [#NO_STAGE], in the + /// order of the edges, or null if the graph has no cycle. + @Nullable + private static List<Integer> findCycle(int[][] edges) { + int numNodes = edges.length; + // 0 = not visited, 1 = on the current path, 2 = done + int[] state = new int[numNodes]; + int[] parent = new int[numNodes]; + for (int start = 0; start < numNodes; start++) { + if (state[start] != 0) { + continue; + } + // Iterative depth-first search: next[n] is the next successor of n to try. + int[] next = new int[numNodes]; + ArrayDeque<Integer> path = new ArrayDeque<>(); + path.push(start); + state[start] = 1; + parent[start] = -1; + while (!path.isEmpty()) { + int node = path.peek(); + if (next[node] == numNodes) { + state[node] = 2; + path.pop(); + continue; + } + int successor = next[node]++; + if (edges[node][successor] == NO_STAGE) { + continue; + } + if (state[successor] == 1) { + List<Integer> cycle = new ArrayList<>(); + for (int n = node; n != successor; n = parent[n]) { + cycle.add(n); + } + cycle.add(successor); + Collections.reverse(cycle); + return cycle; + } + if (state[successor] == 0) { + state[successor] = 1; + parent[successor] = node; + path.push(successor); + } + } + } + return null; + } + + /// Returns the lowest stage id present in both sets, or [#NO_STAGE]. + private static int firstShared(BitSet stopped, BitSet needed) { + if (!stopped.intersects(needed)) { + return NO_STAGE; + } + BitSet shared = (BitSet) stopped.clone(); + shared.and(needed); + return shared.nextSetBit(0); + } + + /// The sets are indexed by stage id. Stage ids are unique among the stages reachable from the root, and they are + /// small consecutive integers, so a [BitSet] is compact and makes each intersection a word-level operation. + private BitSet getStopped(MailboxSendNode sender) { + BitSet result = _stopped.get(sender); + if (result == null) { + result = new BitSet(); + result.set(sender.getStageId()); + for (MailboxReceiveNode receive : getStageInfo(sender).receives().streaming()) { + result.or(getStopped(receive.getSender())); + } + _stopped.put(sender, result); + } + return result; + } + + private BitSet getNeeded(MailboxSendNode sender) { + BitSet result = _needed.get(sender); + if (result == null) { + result = new BitSet(); + result.set(sender.getStageId()); + for (MailboxReceiveNode receive : getStageInfo(sender).receives().all()) { + result.or(getNeeded(receive.getSender())); + } + _needed.put(sender, result); + } + return result; + } + + private StageInfo getStageInfo(PlanNode stageRoot) { + StageInfo info = _stageInfos.get(stageRoot); + if (info == null) { + List<Deferral> deferrals = new ArrayList<>(); + Receives receives = new Receives(); + PlanNode treeRoot = stageRoot instanceof MailboxSendNode ? stageRoot.getInputs().get(0) : stageRoot; + receives.addAll(collect(treeRoot, deferrals)); + addPipelineBreakerDeferrals(receives, deferrals); + info = new StageInfo(receives, deferrals); + _stageInfos.put(stageRoot, info); + } + return info; + } + + /// The main operator chain of a stage reads nothing until the pipeline breaker reads all its receives. + /// + /// The pipeline breaker itself reads its receives in round-robin with blocking reads, so a receive that has no data + /// can keep the others unread. Each pair of pipeline-breaker receives is therefore treated as deferred both ways. + private static void addPipelineBreakerDeferrals(Receives receives, List<Deferral> deferrals) { + List<MailboxReceiveNode> breakers = new ArrayList<>(); + for (MailboxReceiveNode receive : receives.all()) { + if (receive.getExchangeType() == PinotRelExchangeType.PIPELINE_BREAKER) { + breakers.add(receive); + } + } + if (breakers.isEmpty()) { + return; + } + for (MailboxReceiveNode receive : receives.all()) { + for (MailboxReceiveNode breaker : breakers) { + if (receive != breaker) { + deferrals.add(new Deferral(receive, breaker, "the pipeline breaker")); + } + } + } + receives.streaming().removeAll(breakers); Review Comment: MINOR (robustness): this `removeAll` is the one place in `Analysis` that relies on `MailboxReceiveNode.equals` (value equality) instead of identity — everything else here deliberately uses `IdentityHashMap`/identity sets. It's safe today only because `equals` includes `_exchangeType`, so a non-breaker can never compare equal to a breaker. It'd be more robust (and consistent with the rest of the class) to remove by identity, e.g. `receives.streaming().removeIf(r -> breakers.stream().anyMatch(b -> b == r))` or keep breakers in an identity set. ########## pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/EquivalentStagesReplacer.java: ########## @@ -55,33 +77,58 @@ public interface OnSubstitution { } private static class Replacer extends PlanNodeVisitor.DepthFirstVisitor<Void, GroupedStages> { + private final PlanNode _root; private final OnSubstitution _listener; + private final boolean _guarded; + private final List<SpoolSubstitution> _substitutions = new ArrayList<>(); + /// A stage can be visited more than once (for example the leader of a group), so a receive is decided only once. + private final Set<MailboxReceiveNode> _visitedReceives = Collections.newSetFromMap(new IdentityHashMap<>()); - public Replacer(OnSubstitution listener) { + public Replacer(PlanNode root, OnSubstitution listener, boolean guarded) { + _root = root; _listener = listener; + _guarded = guarded; } @Override public Void visitMailboxReceive(MailboxReceiveNode node, GroupedStages equivalenceGroups) { + if (!_visitedReceives.add(node)) { + return null; + } MailboxSendNode sender = node.getSender(); MailboxSendNode leader = equivalenceGroups.getGroup(sender).first(); - if (canSubstitute(sender, leader)) { - // we don't want to visit the children of the node given it is going to be pruned - node.setSender(leader); - leader.addReceiver(node); - _listener.onSubstitution(node.getStageId(), sender.getStageId(), leader.getStageId()); - } else { + if (sender == leader) { visitMailboxSend(leader, equivalenceGroups); Review Comment: MINOR (perf): under `auto`, the guard re-runs `findRisk` over the whole plan after every accepted substitution, so this is roughly O(substitutions x plan-size), and `findRisk` itself builds an O(deferrals^2) matrix where deferrals can grow with the cartesian product of receives per join/union. For plans with many equivalent stages that could add noticeable planning latency. It's gated behind spools being present and only on `auto`, so probably fine in practice — but worth confirming on a worst-case plan (many shared stages feeding several joins) before this is the default. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
