Add flag if task is speculative to TaskStatus

losipiuk · losipiuk · commit a0e649e48f9d · 2023-05-11T14:21:33.000+02:00
diff --git a/core/trino-main/src/main/java/io/trino/execution/MemoryTrackingRemoteTaskFactory.java b/core/trino-main/src/main/java/io/trino/execution/MemoryTrackingRemoteTaskFactory.java
@@ -49,6 +49,7 @@ public RemoteTask createRemoteTask(
             Span stageSpan,
             TaskId taskId,
             InternalNode node,
+            boolean speculative,
             PlanFragment fragment,
             Multimap<PlanNodeId, Split> initialSplits,
             OutputBuffers outputBuffers,
@@ -62,6 +63,7 @@ public RemoteTask createRemoteTask(
                 stageSpan,
                 taskId,
                 node,
+                speculative,
                 fragment,
                 initialSplits,
                 outputBuffers,
diff --git a/core/trino-main/src/main/java/io/trino/execution/RemoteTask.java b/core/trino-main/src/main/java/io/trino/execution/RemoteTask.java
@@ -39,6 +39,8 @@ public interface RemoteTask
 
     void setOutputBuffers(OutputBuffers outputBuffers);
 
+    void setSpeculative(boolean speculative);
+
     /**
      * Listener is always notified asynchronously using a dedicated notification thread pool so, care should
      * be taken to avoid leaking {@code this} when adding a listener in a constructor. Additionally, it is
diff --git a/core/trino-main/src/main/java/io/trino/execution/RemoteTaskFactory.java b/core/trino-main/src/main/java/io/trino/execution/RemoteTaskFactory.java
@@ -35,6 +35,7 @@ RemoteTask createRemoteTask(
             Span stageSpan,
             TaskId taskId,
             InternalNode node,
+            boolean speculative,
             PlanFragment fragment,
             Multimap<PlanNodeId, Split> initialSplits,
             OutputBuffers outputBuffers,
diff --git a/core/trino-main/src/main/java/io/trino/execution/SqlStage.java b/core/trino-main/src/main/java/io/trino/execution/SqlStage.java
@@ -246,7 +246,8 @@ public synchronized Optional<RemoteTask> createTask(
             OutputBuffers outputBuffers,
             Multimap<PlanNodeId, Split> splits,
             Set<PlanNodeId> noMoreSplits,
-            Optional<DataSize> estimatedMemory)
+            Optional<DataSize> estimatedMemory,
+            boolean speculative)
     {
         if (stateMachine.getState().isDone()) {
             return Optional.empty();
@@ -261,6 +262,7 @@ public synchronized Optional<RemoteTask> createTask(
                 stateMachine.getStageSpan(),
                 taskId,
                 node,
+                speculative,
                 stateMachine.getFragment().withBucketToPartition(bucketToPartition),
                 splits,
                 outputBuffers,
diff --git a/core/trino-main/src/main/java/io/trino/execution/SqlTask.java b/core/trino-main/src/main/java/io/trino/execution/SqlTask.java
@@ -88,6 +88,7 @@ public class SqlTask
     private final String taskInstanceId;
     private final URI location;
     private final String nodeId;
+    private final AtomicBoolean speculative = new AtomicBoolean(false);
     private final TaskStateMachine taskStateMachine;
     private final OutputBuffer outputBuffer;
     private final QueryContext queryContext;
@@ -376,6 +377,7 @@ else if (taskHolder.getTaskExecution() != null) {
                 state,
                 location,
                 nodeId,
+                speculative.get(),
                 failures,
                 queuedPartitionedDrivers,
                 runningPartitionedDrivers,
@@ -468,7 +470,8 @@ public TaskInfo updateTask(
             Optional<PlanFragment> fragment,
             List<SplitAssignment> splitAssignments,
             OutputBuffers outputBuffers,
-            Map<DynamicFilterId, Domain> dynamicFilterDomains)
+            Map<DynamicFilterId, Domain> dynamicFilterDomains,
+            boolean speculative)
     {
         try {
             // trace token must be set first to make sure failure injection for getTaskResults requests works as expected
@@ -495,6 +498,9 @@ public TaskInfo updateTask(
                 taskExecution.addSplitAssignments(splitAssignments);
                 taskExecution.getTaskContext().addDynamicFilter(dynamicFilterDomains);
             }
+
+            // update speculative flag
+            this.speculative.set(speculative);
         }
         catch (Error e) {
             failed(e);
diff --git a/core/trino-main/src/main/java/io/trino/execution/SqlTaskManager.java b/core/trino-main/src/main/java/io/trino/execution/SqlTaskManager.java
@@ -470,10 +470,11 @@ public TaskInfo updateTask(
             Optional<PlanFragment> fragment,
             List<SplitAssignment> splitAssignments,
             OutputBuffers outputBuffers,
-            Map<DynamicFilterId, Domain> dynamicFilterDomains)
+            Map<DynamicFilterId, Domain> dynamicFilterDomains,
+            boolean speculative)
     {
         try {
-            return versionEmbedder.embedVersion(() -> doUpdateTask(session, taskId, stageSpan, fragment, splitAssignments, outputBuffers, dynamicFilterDomains)).call();
+            return versionEmbedder.embedVersion(() -> doUpdateTask(session, taskId, stageSpan, fragment, splitAssignments, outputBuffers, dynamicFilterDomains, speculative)).call();
         }
         catch (Exception e) {
             throwIfUnchecked(e);
@@ -489,7 +490,8 @@ private TaskInfo doUpdateTask(
             Optional<PlanFragment> fragment,
             List<SplitAssignment> splitAssignments,
             OutputBuffers outputBuffers,
-            Map<DynamicFilterId, Domain> dynamicFilterDomains)
+            Map<DynamicFilterId, Domain> dynamicFilterDomains,
+            boolean speculative)
     {
         requireNonNull(session, "session is null");
         requireNonNull(taskId, "taskId is null");
@@ -528,7 +530,7 @@ private TaskInfo doUpdateTask(
                 });
 
         sqlTask.recordHeartbeat();
-        return sqlTask.updateTask(session, stageSpan, fragment, splitAssignments, outputBuffers, dynamicFilterDomains);
+        return sqlTask.updateTask(session, stageSpan, fragment, splitAssignments, outputBuffers, dynamicFilterDomains, speculative);
     }
 
     /**
diff --git a/core/trino-main/src/main/java/io/trino/execution/TaskInfo.java b/core/trino-main/src/main/java/io/trino/execution/TaskInfo.java
@@ -130,10 +130,10 @@ public String toString()
                 .toString();
     }
 
-    public static TaskInfo createInitialTask(TaskId taskId, URI location, String nodeId, Optional<List<PipelinedBufferInfo>> pipelinedBufferStates, TaskStats taskStats)
+    public static TaskInfo createInitialTask(TaskId taskId, URI location, String nodeId, boolean speculative, Optional<List<PipelinedBufferInfo>> pipelinedBufferStates, TaskStats taskStats)
     {
         return new TaskInfo(
-                initialTaskStatus(taskId, location, nodeId),
+                initialTaskStatus(taskId, location, nodeId, speculative),
                 DateTime.now(),
                 new OutputBufferInfo(
                         "UNINITIALIZED",
diff --git a/core/trino-main/src/main/java/io/trino/execution/TaskStatus.java b/core/trino-main/src/main/java/io/trino/execution/TaskStatus.java
@@ -52,6 +52,7 @@ public class TaskStatus
     private final TaskState state;
     private final URI self;
     private final String nodeId;
+    private final boolean speculative;
 
     private final int queuedPartitionedDrivers;
     private final long queuedPartitionedSplitsWeight;
@@ -80,6 +81,7 @@ public TaskStatus(
             @JsonProperty("state") TaskState state,
             @JsonProperty("self") URI self,
             @JsonProperty("nodeId") String nodeId,
+            @JsonProperty("speculative") boolean speculative,
             @JsonProperty("failures") List<ExecutionFailureInfo> failures,
             @JsonProperty("queuedPartitionedDrivers") int queuedPartitionedDrivers,
             @JsonProperty("runningPartitionedDrivers") int runningPartitionedDrivers,
@@ -104,6 +106,7 @@ public TaskStatus(
         this.state = requireNonNull(state, "state is null");
         this.self = requireNonNull(self, "self is null");
         this.nodeId = requireNonNull(nodeId, "nodeId is null");
+        this.speculative = speculative;
 
         checkArgument(queuedPartitionedDrivers >= 0, "queuedPartitionedDrivers must be positive");
         this.queuedPartitionedDrivers = queuedPartitionedDrivers;
@@ -169,6 +172,12 @@ public String getNodeId()
         return nodeId;
     }
 
+    @JsonProperty
+    public boolean isSpeculative()
+    {
+        return speculative;
+    }
+
     @JsonProperty
     public List<ExecutionFailureInfo> getFailures()
     {
@@ -268,7 +277,7 @@ public String toString()
                 .toString();
     }
 
-    public static TaskStatus initialTaskStatus(TaskId taskId, URI location, String nodeId)
+    public static TaskStatus initialTaskStatus(TaskId taskId, URI location, String nodeId, boolean speculative)
     {
         return new TaskStatus(
                 taskId,
@@ -277,6 +286,7 @@ public static TaskStatus initialTaskStatus(TaskId taskId, URI location, String n
                 PLANNED,
                 location,
                 nodeId,
+                speculative,
                 ImmutableList.of(),
                 0,
                 0,
@@ -303,6 +313,7 @@ public static TaskStatus failWith(TaskStatus taskStatus, TaskState state, List<E
                 state,
                 taskStatus.getSelf(),
                 taskStatus.getNodeId(),
+                false,
                 exceptions,
                 taskStatus.getQueuedPartitionedDrivers(),
                 taskStatus.getRunningPartitionedDrivers(),
diff --git a/core/trino-main/src/main/java/io/trino/execution/scheduler/EventDrivenFaultTolerantQueryScheduler.java b/core/trino-main/src/main/java/io/trino/execution/scheduler/EventDrivenFaultTolerantQueryScheduler.java
@@ -1403,7 +1403,8 @@ public Optional<RemoteTask> schedule(int partitionId, ExchangeSinkInstanceHandle
                     outputBuffers,
                     splits,
                     noMoreSplits,
-                    Optional.of(partition.getMemoryRequirements().getRequiredMemory()));
+                    Optional.of(partition.getMemoryRequirements().getRequiredMemory()),
+                    false);  // TODO pass correct value here
             task.ifPresent(remoteTask -> {
                 partition.addTask(remoteTask, outputBuffers);
                 runningPartitions.add(partitionId);
diff --git a/core/trino-main/src/main/java/io/trino/execution/scheduler/PipelinedStageExecution.java b/core/trino-main/src/main/java/io/trino/execution/scheduler/PipelinedStageExecution.java
@@ -299,7 +299,8 @@ public synchronized Optional<RemoteTask> scheduleTask(
                 outputBuffers,
                 initialSplits,
                 ImmutableSet.of(),
-                Optional.empty());
+                Optional.empty(),
+                false);
 
         if (optionalTask.isEmpty()) {
             return Optional.empty();
diff --git a/core/trino-main/src/main/java/io/trino/server/HttpRemoteTaskFactory.java b/core/trino-main/src/main/java/io/trino/server/HttpRemoteTaskFactory.java
@@ -141,6 +141,7 @@ public RemoteTask createRemoteTask(
             Span stageSpan,
             TaskId taskId,
             InternalNode node,
+            boolean speculative,
             PlanFragment fragment,
             Multimap<PlanNodeId, Split> initialSplits,
             OutputBuffers outputBuffers,
@@ -154,6 +155,7 @@ public RemoteTask createRemoteTask(
                 stageSpan,
                 taskId,
                 node.getNodeIdentifier(),
+                speculative,
                 locationFactory.createTaskLocation(node, taskId),
                 fragment,
                 initialSplits,
diff --git a/core/trino-main/src/main/java/io/trino/server/TaskResource.java b/core/trino-main/src/main/java/io/trino/server/TaskResource.java
@@ -158,7 +158,8 @@ public void createOrUpdateTask(
                 taskUpdateRequest.getFragment(),
                 taskUpdateRequest.getSplitAssignments(),
                 taskUpdateRequest.getOutputIds(),
-                taskUpdateRequest.getDynamicFilterDomains());
+                taskUpdateRequest.getDynamicFilterDomains(),
+                taskUpdateRequest.isSpeculative());
 
         if (shouldSummarize(uriInfo)) {
             taskInfo = taskInfo.summarize();
diff --git a/core/trino-main/src/main/java/io/trino/server/TaskUpdateRequest.java b/core/trino-main/src/main/java/io/trino/server/TaskUpdateRequest.java
@@ -43,6 +43,7 @@ public class TaskUpdateRequest
     private final OutputBuffers outputIds;
     private final Map<DynamicFilterId, Domain> dynamicFilterDomains;
     private final Optional<Slice> exchangeEncryptionKey;
+    private final boolean speculative;
 
     @JsonCreator
     public TaskUpdateRequest(
@@ -53,7 +54,8 @@ public TaskUpdateRequest(
             @JsonProperty("splitAssignments") List<SplitAssignment> splitAssignments,
             @JsonProperty("outputIds") OutputBuffers outputIds,
             @JsonProperty("dynamicFilterDomains") Map<DynamicFilterId, Domain> dynamicFilterDomains,
-            @JsonProperty("exchangeEncryptionKey") Optional<Slice> exchangeEncryptionKey)
+            @JsonProperty("exchangeEncryptionKey") Optional<Slice> exchangeEncryptionKey,
+            @JsonProperty("speculative") boolean speculative)
     {
         requireNonNull(session, "session is null");
         requireNonNull(extraCredentials, "extraCredentials is null");
@@ -72,6 +74,7 @@ public TaskUpdateRequest(
         this.outputIds = outputIds;
         this.dynamicFilterDomains = dynamicFilterDomains;
         this.exchangeEncryptionKey = exchangeEncryptionKey;
+        this.speculative = speculative;
     }
 
     @JsonProperty
@@ -122,6 +125,12 @@ public Optional<Slice> getExchangeEncryptionKey()
         return exchangeEncryptionKey;
     }
 
+    @JsonProperty
+    public boolean isSpeculative()
+    {
+        return speculative;
+    }
+
     @Override
     public String toString()
     {
diff --git a/core/trino-main/src/main/java/io/trino/server/remotetask/HttpRemoteTask.java b/core/trino-main/src/main/java/io/trino/server/remotetask/HttpRemoteTask.java
@@ -135,6 +135,7 @@ public final class HttpRemoteTask
     private final Session session;
     private final Span stageSpan;
     private final String nodeId;
+    private final AtomicBoolean speculative;
     private final PlanFragment planFragment;
 
     private final AtomicLong nextSplitId = new AtomicLong();
@@ -207,6 +208,7 @@ public HttpRemoteTask(
             Span stageSpan,
             TaskId taskId,
             String nodeId,
+            boolean speculative,
             URI location,
             PlanFragment planFragment,
             Multimap<PlanNodeId, Split> initialSplits,
@@ -254,6 +256,7 @@ public HttpRemoteTask(
             this.session = session;
             this.stageSpan = stageSpan;
             this.nodeId = nodeId;
+            this.speculative = new AtomicBoolean(speculative);
             this.planFragment = planFragment;
             this.outputBuffers.set(outputBuffers);
             this.httpClient = httpClient;
@@ -311,7 +314,7 @@ public HttpRemoteTask(
                                 .collect(toImmutableList()));
             }
 
-            TaskInfo initialTask = createInitialTask(taskId, location, nodeId, pipelinedBufferStates, new TaskStats(DateTime.now(), null));
+            TaskInfo initialTask = createInitialTask(taskId, location, nodeId, this.speculative.get(), pipelinedBufferStates, new TaskStats(DateTime.now(), null));
 
             this.dynamicFiltersFetcher = new DynamicFiltersFetcher(
                     this::fatalUnacknowledgedFailure,
@@ -493,6 +496,16 @@ public void setOutputBuffers(OutputBuffers newOutputBuffers)
         }
     }
 
+    @Override
+    public void setSpeculative(boolean speculative)
+    {
+        checkArgument(!speculative, "we can only move task from speculative to non-speculative");
+        if (this.speculative.compareAndSet(true, speculative)) {
+            // versioning should be not needed here as we can only migrate task from speculative to non-speculative; so out of order requests do not matter
+            triggerUpdate();
+        }
+    }
+
     @Override
     public PartitionedSplitsInfo getPartitionedSplitsInfo()
     {
@@ -734,7 +747,8 @@ private void sendUpdate()
                 splitAssignments,
                 outputBuffers.get(),
                 dynamicFilterDomains.getDynamicFilterDomains(),
-                session.getExchangeEncryptionKey());
+                session.getExchangeEncryptionKey(),
+                speculative.get());
         byte[] taskUpdateRequestJson = taskUpdateRequestCodec.toJsonBytes(updateRequest);
 
         // try to adjust batch size to meet expected request size
diff --git a/core/trino-main/src/test/java/io/trino/execution/MockRemoteTaskFactory.java b/core/trino-main/src/test/java/io/trino/execution/MockRemoteTaskFactory.java
@@ -133,6 +133,7 @@ public MockRemoteTask createTableScanTask(TaskId taskId, InternalNode newNode, L
                 Span.getInvalid(),
                 taskId,
                 newNode,
+                false,
                 testFragment,
                 initialSplits.build(),
                 PipelinedOutputBuffers.createInitial(BROADCAST),
@@ -148,6 +149,7 @@ public MockRemoteTask createRemoteTask(
             Span stageSpan,
             TaskId taskId,
             InternalNode node,
+            boolean speculative,
             PlanFragment fragment,
             Multimap<PlanNodeId, Split> initialSplits,
             OutputBuffers outputBuffers,
@@ -277,6 +279,7 @@ public synchronized TaskStatus getTaskStatus()
                     state,
                     location,
                     nodeId,
+                    false,
                     failures,
                     queuedSplitsInfo.getCount(),
                     combinedSplitsInfo.getCount() - queuedSplitsInfo.getCount(),
@@ -399,6 +402,12 @@ public void setOutputBuffers(OutputBuffers outputBuffers)
             outputBuffer.setOutputBuffers(outputBuffers);
         }
 
+        @Override
+        public void setSpeculative(boolean speculative)
+        {
+            // ignore
+        }
+
         @Override
         public void addStateChangeListener(StateChangeListener<TaskStatus> stateChangeListener)
         {
diff --git a/core/trino-main/src/test/java/io/trino/execution/TaskTestUtils.java b/core/trino-main/src/test/java/io/trino/execution/TaskTestUtils.java
@@ -183,7 +183,7 @@ public static LocalExecutionPlanner createTestingPlanner()
 
     public static TaskInfo updateTask(SqlTask sqlTask, List<SplitAssignment> splitAssignments, OutputBuffers outputBuffers)
     {
-        return sqlTask.updateTask(TEST_SESSION, Span.getInvalid(), Optional.of(PLAN_FRAGMENT), splitAssignments, outputBuffers, ImmutableMap.of());
+        return sqlTask.updateTask(TEST_SESSION, Span.getInvalid(), Optional.of(PLAN_FRAGMENT), splitAssignments, outputBuffers, ImmutableMap.of(), false);
     }
 
     public static SplitMonitor createTestSplitMonitor()
diff --git a/core/trino-main/src/test/java/io/trino/execution/TestSqlStage.java b/core/trino-main/src/test/java/io/trino/execution/TestSqlStage.java
@@ -150,7 +150,8 @@ private void testFinalStageInfoInternal()
                             PipelinedOutputBuffers.createInitial(ARBITRARY),
                             initialSplits,
                             ImmutableSet.of(),
-                            Optional.empty());
+                            Optional.empty(),
+                            false);
                     if (created.isPresent()) {
                         if (created.get() instanceof MockRemoteTaskFactory.MockRemoteTask mockTask) {
                             mockTask.start();
diff --git a/core/trino-main/src/test/java/io/trino/execution/TestSqlTask.java b/core/trino-main/src/test/java/io/trino/execution/TestSqlTask.java
diff --git a/core/trino-main/src/test/java/io/trino/execution/TestSqlTaskManager.java b/core/trino-main/src/test/java/io/trino/execution/TestSqlTaskManager.java
diff --git a/core/trino-main/src/test/java/io/trino/execution/TestingRemoteTaskFactory.java b/core/trino-main/src/test/java/io/trino/execution/TestingRemoteTaskFactory.java
diff --git a/core/trino-main/src/test/java/io/trino/execution/scheduler/TestScaledWriterScheduler.java b/core/trino-main/src/test/java/io/trino/execution/scheduler/TestScaledWriterScheduler.java
diff --git a/core/trino-main/src/test/java/io/trino/memory/TestLeastWastedEffortTaskLowMemoryKiller.java b/core/trino-main/src/test/java/io/trino/memory/TestLeastWastedEffortTaskLowMemoryKiller.java
diff --git a/core/trino-main/src/test/java/io/trino/server/remotetask/TestHttpRemoteTask.java b/core/trino-main/src/test/java/io/trino/server/remotetask/TestHttpRemoteTask.java

Original file line number	Diff line number	Diff line change
`@@ -130,10 +130,10 @@ public String toString()`
`130`	`130`	`.toString();`
`131`	`131`	`}`
`132`	`132`
`133`		`- public static TaskInfo createInitialTask(TaskId taskId, URI location, String nodeId, Optional<List<PipelinedBufferInfo>> pipelinedBufferStates, TaskStats taskStats)`
	`133`	`+ public static TaskInfo createInitialTask(TaskId taskId, URI location, String nodeId, boolean speculative, Optional<List<PipelinedBufferInfo>> pipelinedBufferStates, TaskStats taskStats)`
`134`	`134`	`{`
`135`	`135`	`return new TaskInfo(`
`136`		`- initialTaskStatus(taskId, location, nodeId),`
	`136`	`+ initialTaskStatus(taskId, location, nodeId, speculative),`
`137`	`137`	`DateTime.now(),`
`138`	`138`	`new OutputBufferInfo(`
`139`	`139`	`"UNINITIALIZED",`