Defer broadcast join probe leaf splits

Avoids starting join probe leaf splits for tasks that are awaiting the
build side completion. Otherwise, build side leaf split concurrency will
be starved by probe side splits that start and then immediately block.

Author: pettyjamesm

core/trino-main/src/main/java/io/trino/execution/SqlTaskExecution.java

@@ -105,6 +105,8 @@ public class SqlTaskExecution
     private final List<PlanNodeId> sourceStartOrder;
     @GuardedBy("this")
     private int schedulingPlanNodeOrdinal;
+    @GuardedBy("this")
+    private ListenableFuture<Void> pipelineDependenciesSatisfied = immediateVoidFuture();
 
     @GuardedBy("this")
     private final Map<PlanNodeId, PendingSplitsForPlanNode> pendingSplitsByPlanNode;
@@ -309,7 +311,8 @@ private synchronized Set<PlanNodeId> updateSplitAssignments(List<SplitAssignment
         // update task with new sources
         for (SplitAssignment splitAssignment : unacknowledgedSplitAssignment) {
             if (driverRunnerFactoriesWithSplitLifeCycle.containsKey(splitAssignment.getPlanNodeId())) {
-                schedulePartitionedSource(splitAssignment);
+                mergeIntoPendingSplits(splitAssignment.getPlanNodeId(), splitAssignment.getSplits(), splitAssignment.isNoMoreSplits());
+                schedulePartitionedSourcePendingSplits();
             }
             else {
                 // tell existing drivers about the new splits
@@ -337,15 +340,34 @@ private void mergeIntoPendingSplits(PlanNodeId planNodeId, Set<ScheduledSplit> s
         }
     }
 
-    private synchronized void schedulePartitionedSource(SplitAssignment splitAssignmentUpdate)
+    private synchronized void scheduleSourcePartitionedSplitsAfterPipelineUnblocked()
     {
-        mergeIntoPendingSplits(splitAssignmentUpdate.getPlanNodeId(), splitAssignmentUpdate.getSplits(), splitAssignmentUpdate.isNoMoreSplits());
+        try (SetThreadName _ = new SetThreadName("Task-" + taskId)) {
+            // Enqueue pending splits as split runners after unblocking
+            schedulePartitionedSourcePendingSplits();
+            // Re-check for task completion since we may have just set no more splits
+            checkTaskCompletion();
+        }
+    }
 
+    private synchronized void schedulePartitionedSourcePendingSplits()
+    {
         while (schedulingPlanNodeOrdinal < sourceStartOrder.size()) {
             PlanNodeId schedulingPlanNode = sourceStartOrder.get(schedulingPlanNodeOrdinal);
 
             DriverSplitRunnerFactory partitionedDriverRunnerFactory = driverRunnerFactoriesWithSplitLifeCycle.get(schedulingPlanNode);
 
+            // Avoid creating split runners for pipelines that are awaiting another pipeline completing (e.g. probe side of a join waiting
+            // on the broadcast completion). Otherwise, build side pipelines will have reduced concurrency available.
+            ListenableFuture<Void> pipelineDependenciesSatisfied = partitionedDriverRunnerFactory.getPipelineDependenciesSatisfied();
+            if (!pipelineDependenciesSatisfied.isDone()) {
+                // Only register a single re-schedule listener if we're blocked on pipeline dependencies
+                if (this.pipelineDependenciesSatisfied.isDone()) {
+                    this.pipelineDependenciesSatisfied = pipelineDependenciesSatisfied;
+                    pipelineDependenciesSatisfied.addListener(this::scheduleSourcePartitionedSplitsAfterPipelineUnblocked, notificationExecutor);
+                }
+                break;
+            }
             PendingSplitsForPlanNode pendingSplits = pendingSplitsByPlanNode.get(schedulingPlanNode);
 
             // Enqueue driver runners with split lifecycle for this plan node and driver life cycle combination.
@@ -600,6 +622,11 @@ private DriverSplitRunnerFactory(DriverFactory driverFactory, Tracer tracer, boo
                     .startSpan();
         }
 
+        public ListenableFuture<Void> getPipelineDependenciesSatisfied()
+        {
+            return driverFactory.getPipelineDependenciesSatisfied();
+        }
+
         public DriverSplitRunner createPartitionedDriverRunner(ScheduledSplit partitionedSplit)
         {
             return createDriverRunner(partitionedSplit, partitionedSplit.getSplit().getSplitWeight().getRawValue());

core/trino-main/src/main/java/io/trino/operator/DriverFactory.java

@@ -14,6 +14,8 @@
 package io.trino.operator;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
 import com.google.errorprone.annotations.concurrent.GuardedBy;
 import io.trino.sql.planner.plan.PlanNodeId;
 import jakarta.annotation.Nullable;
@@ -26,6 +28,7 @@
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
 import static java.util.Objects.requireNonNull;
 
 public class DriverFactory
@@ -35,6 +38,7 @@ public class DriverFactory
     private final boolean outputDriver;
     private final Optional<PlanNodeId> sourceId;
     private final OptionalInt driverInstances;
+    private final ListenableFuture<Void> pipelineDependenciesSatisfied;
 
     // must synchronize between createDriver() and noMoreDrivers(), but isNoMoreDrivers() is safe without synchronizing
     @GuardedBy("this")
@@ -57,6 +61,11 @@ public DriverFactory(int pipelineId, boolean inputDriver, boolean outputDriver,
                 .collect(toImmutableList());
         checkArgument(sourceIds.size() <= 1, "Expected at most one source operator in driver factory, but found %s", sourceIds);
         this.sourceId = sourceIds.isEmpty() ? Optional.empty() : Optional.of(sourceIds.get(0));
+        List<ListenableFuture<Void>> pipelineDependencies = operatorFactories.stream()
+                .map(OperatorFactory::pipelineDependenciesSatisfied)
+                .filter(future -> !future.isDone())
+                .collect(toImmutableList());
+        this.pipelineDependenciesSatisfied = pipelineDependencies.isEmpty() ? Futures.immediateVoidFuture() : Futures.whenAllComplete(pipelineDependencies).call(() -> null, directExecutor());
     }
 
     public int getPipelineId()
@@ -74,6 +83,11 @@ public boolean isOutputDriver()
         return outputDriver;
     }
 
+    public ListenableFuture<Void> getPipelineDependenciesSatisfied()
+    {
+        return pipelineDependenciesSatisfied;
+    }
+
     /**
      * return the sourceId of this DriverFactory.
      * A DriverFactory doesn't always have source node.

core/trino-main/src/main/java/io/trino/operator/OperatorFactory.java

@@ -13,6 +13,10 @@
  */
 package io.trino.operator;
 
+import com.google.common.util.concurrent.ListenableFuture;
+
+import static com.google.common.util.concurrent.Futures.immediateVoidFuture;
+
 public interface OperatorFactory
 {
     Operator createOperator(DriverContext driverContext);
@@ -27,4 +31,14 @@ public interface OperatorFactory
     void noMoreOperators();
 
     OperatorFactory duplicate();
+
+    /**
+     * Returns a future indicating that any dependencies operators have on other pipelines has been satisfied and that leaf splits
+     * should be allowed to start for this operator. This is used to prevent join probe splits from starting before the build side
+     * of a join is ready when the two are in the same stage (i.e.: broadcast join on top of a table scan).
+     */
+    default ListenableFuture<Void> pipelineDependenciesSatisfied()
+    {
+        return immediateVoidFuture();
+    }
 }

core/trino-main/src/main/java/io/trino/operator/WorkProcessorOperatorAdapter.java

@@ -14,6 +14,7 @@
 package io.trino.operator;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 import io.trino.memory.context.MemoryTrackingContext;
 import io.trino.operator.join.JoinOperatorFactory;
@@ -86,6 +87,21 @@ public Optional<OperatorFactory> createOuterOperatorFactory()
             return lookupJoin.createOuterOperatorFactory();
         }
 
+        @Override
+        public ListenableFuture<Void> buildPipelineReady()
+        {
+            if (!(operatorFactory instanceof JoinOperatorFactory lookupJoin)) {
+                return Futures.immediateVoidFuture();
+            }
+            return lookupJoin.buildPipelineReady();
+        }
+
+        @Override
+        public ListenableFuture<Void> pipelineDependenciesSatisfied()
+        {
+            return buildPipelineReady();
+        }
+
         @VisibleForTesting
         public WorkProcessorOperatorFactory getWorkProcessorOperatorFactory()
         {

core/trino-main/src/main/java/io/trino/operator/join/JoinBridgeManager.java

@@ -43,6 +43,7 @@ public static JoinBridgeManager<PartitionedLookupSourceFactory> lookupAllAtOnce(
     private final List<Type> buildOutputTypes;
     private final boolean buildOuter;
     private final T joinBridge;
+    private final ListenableFuture<Void> whenBuildFinishes;
 
     private final AtomicBoolean initialized = new AtomicBoolean();
     private JoinLifecycle joinLifecycle;
@@ -57,6 +58,7 @@ public JoinBridgeManager(
         this.buildOuter = buildOuter;
         this.joinBridge = requireNonNull(joinBridge, "joinBridge is null");
         this.buildOutputTypes = requireNonNull(buildOutputTypes, "buildOutputTypes is null");
+        this.whenBuildFinishes = requireNonNull(joinBridge.whenBuildFinishes(), "whenBuildFinishes is null");
     }
 
     private void initializeIfNecessary()
@@ -67,7 +69,7 @@ private void initializeIfNecessary()
                     return;
                 }
                 int finalProbeFactoryCount = probeFactoryCount.get();
-                joinLifecycle = new JoinLifecycle(joinBridge, finalProbeFactoryCount, buildOuter ? 1 : 0);
+                joinLifecycle = new JoinLifecycle(whenBuildFinishes, joinBridge, finalProbeFactoryCount, buildOuter ? 1 : 0);
                 initialized.set(true);
             }
         }
@@ -83,6 +85,11 @@ public void incrementProbeFactoryCount()
         probeFactoryCount.increment();
     }
 
+    public ListenableFuture<Void> getBuildFinishedFuture()
+    {
+        return whenBuildFinishes;
+    }
+
     public T getJoinBridge()
     {
         initializeIfNecessary();
@@ -139,7 +146,7 @@ private static class JoinLifecycle
         private final ListenableFuture<Void> whenBuildAndProbeFinishes;
         private final ListenableFuture<Void> whenAllFinishes;
 
-        public JoinLifecycle(JoinBridge joinBridge, int probeFactoryCount, int outerFactoryCount)
+        private JoinLifecycle(ListenableFuture<Void> whenBuildFinishes, JoinBridge joinBridge, int probeFactoryCount, int outerFactoryCount)
         {
             // When all probe and lookup-outer operators finish, destroy the join bridge (freeing the memory)
             // * Each LookupOuterOperatorFactory count as 1
@@ -152,7 +159,7 @@ public JoinLifecycle(JoinBridge joinBridge, int probeFactoryCount, int outerFact
             // * Each probe operator count as 1
             probeReferenceCount = new ReferenceCount(probeFactoryCount);
 
-            whenBuildAndProbeFinishes = Futures.whenAllSucceed(joinBridge.whenBuildFinishes(), probeReferenceCount.getFreeFuture()).call(() -> null, directExecutor());
+            whenBuildAndProbeFinishes = Futures.whenAllSucceed(whenBuildFinishes, probeReferenceCount.getFreeFuture()).call(() -> null, directExecutor());
             whenAllFinishes = Futures.whenAllSucceed(whenBuildAndProbeFinishes, outerReferenceCount.getFreeFuture()).call(() -> null, directExecutor());
             whenAllFinishes.addListener(joinBridge::destroy, directExecutor());
         }

core/trino-main/src/main/java/io/trino/operator/join/JoinOperatorFactory.java

@@ -13,11 +13,18 @@
  */
 package io.trino.operator.join;
 
+import com.google.common.util.concurrent.ListenableFuture;
 import io.trino.operator.OperatorFactory;
 
 import java.util.Optional;
 
 public interface JoinOperatorFactory
 {
     Optional<OperatorFactory> createOuterOperatorFactory();
+
+    /**
+     * Future that indicates when the build side of the join has been completed and probe processing
+     * can begin. Used by {@link OperatorFactory#pipelineDependenciesSatisfied()}.
+     */
+    ListenableFuture<Void> buildPipelineReady();
 }

core/trino-main/src/main/java/io/trino/operator/join/LookupJoinOperatorFactory.java

@@ -15,6 +15,7 @@
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.primitives.Ints;
+import com.google.common.util.concurrent.ListenableFuture;
 import io.trino.operator.HashGenerator;
 import io.trino.operator.JoinOperatorType;
 import io.trino.operator.OperatorFactory;
@@ -164,6 +165,12 @@ public String getOperatorType()
         return LookupJoinOperator.class.getSimpleName();
     }
 
+    @Override
+    public ListenableFuture<Void> buildPipelineReady()
+    {
+        return joinBridgeManager.getBuildFinishedFuture();
+    }
+
     @Override
     public WorkProcessorOperator create(ProcessorContext processorContext, WorkProcessor<Page> sourcePages)
     {

core/trino-main/src/main/java/io/trino/operator/join/unspilled/LookupJoinOperatorFactory.java

@@ -14,6 +14,7 @@
 package io.trino.operator.join.unspilled;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.util.concurrent.ListenableFuture;
 import io.trino.operator.JoinOperatorType;
 import io.trino.operator.OperatorFactory;
 import io.trino.operator.ProcessorContext;
@@ -133,6 +134,12 @@ public String getOperatorType()
         return LookupJoinOperator.class.getSimpleName();
     }
 
+    @Override
+    public ListenableFuture<Void> buildPipelineReady()
+    {
+        return joinBridgeManager.getBuildFinishedFuture();
+    }
+
     @Override
     public WorkProcessorOperator create(ProcessorContext processorContext, WorkProcessor<Page> sourcePages)
     {