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,9 @@
  */
 package io.trino.operator;
 
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+
 public interface OperatorFactory
 {
     Operator createOperator(DriverContext driverContext);
@@ -27,4 +30,9 @@ public interface OperatorFactory
     void noMoreOperators();
 
     OperatorFactory duplicate();
+
+    default ListenableFuture<Void> pipelineDependenciesSatisfied()
+    {
+        return Futures.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

@@ -17,6 +17,7 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.SettableFuture;
 import io.trino.operator.ReferenceCount;
 import io.trino.spi.type.Type;
 
@@ -43,6 +44,7 @@ public static JoinBridgeManager<PartitionedLookupSourceFactory> lookupAllAtOnce(
     private final List<Type> buildOutputTypes;
     private final boolean buildOuter;
     private final T joinBridge;
+    private final SettableFuture<Void> whenBuildFinishes = SettableFuture.create();
 
     private final AtomicBoolean initialized = new AtomicBoolean();
     private JoinLifecycle joinLifecycle;
@@ -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,8 +146,9 @@ private static class JoinLifecycle
         private final ListenableFuture<Void> whenBuildAndProbeFinishes;
         private final ListenableFuture<Void> whenAllFinishes;
 
-        public JoinLifecycle(JoinBridge joinBridge, int probeFactoryCount, int outerFactoryCount)
+        private JoinLifecycle(SettableFuture<Void> notifyOnBuildFinished, JoinBridge joinBridge, int probeFactoryCount, int outerFactoryCount)
         {
+            requireNonNull(notifyOnBuildFinished, "notifyOnBuildFinished is null");
             // When all probe and lookup-outer operators finish, destroy the join bridge (freeing the memory)
             // * Each LookupOuterOperatorFactory count as 1
             //   * There is at most 1 LookupOuterOperatorFactory
@@ -152,7 +160,9 @@ 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());
+            ListenableFuture<Void> whenBuildFinishes = joinBridge.whenBuildFinishes();
+            notifyOnBuildFinished.setFuture(whenBuildFinishes);
+            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,14 @@
  */
 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();
+
+    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)
     {