Repository: prestodb/presto
Issue: Detect quick worker restarts and fail queries immediately

When a worker restart, the engine may not notice the change and wait for work that will never complete.  An easy fix for this is to add the node instance id to task info and then the `HttpRemoteTask` would see the change and fail the task.

This issue is identified by having a task stuck in the RUNNING state with 0 splits assigned.  If there are multiple queries, they will all be stuck on the same node or nodes that rebooted.


Changed files:
presto-main/src/main/java/com/facebook/presto/execution/SqlTask.java
presto-main/src/main/java/com/facebook/presto/util/Failures.java
presto-main/src/test/java/com/facebook/presto/execution/MockRemoteTaskFactory.java
presto-main/src/test/java/com/facebook/presto/server/MockTaskManager.java
presto-main/src/test/java/com/facebook/presto/execution/TestSqlTask.java
presto-main/src/test/java/com/facebook/presto/execution/TestSqlTaskManager.java
presto-main/src/main/java/com/facebook/presto/server/HttpRemoteTask.java
presto-main/src/main/java/com/facebook/presto/execution/TaskInfo.java
presto-main/src/test/java/com/facebook/presto/execution/TestSqlStageExecution.java
presto-spi/src/main/java/com/facebook/presto/spi/StandardErrorCode.java
presto-main/src/main/java/com/facebook/presto/execution/SqlTaskManager.java

diff --git a/presto-main/src/main/java/com/facebook/presto/execution/SqlTask.java b/presto-main/src/main/java/com/facebook/presto/execution/SqlTask.java
index 1ee5413c8d..12d7dfb1ef 100644
--- a/presto-main/src/main/java/com/facebook/presto/execution/SqlTask.java
+++ b/presto-main/src/main/java/com/facebook/presto/execution/SqlTask.java
@@ -35,6 +35,7 @@ import javax.annotation.Nullable;
 
 import java.net.URI;
 import java.util.List;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.atomic.AtomicLong;
@@ -49,6 +50,7 @@ public class SqlTask
     private static final Logger log = Logger.get(SqlTask.class);
 
     private final TaskId taskId;
+    private final String nodeInstanceId;
     private final URI location;
     private final TaskStateMachine taskStateMachine;
     private final SharedBuffer sharedBuffer;
@@ -62,6 +64,7 @@ public class SqlTask
 
     public SqlTask(
             TaskId taskId,
+            String nodeInstanceId,
             URI location,
             SqlTaskExecutionFactory sqlTaskExecutionFactory,
             ExecutorService taskNotificationExecutor,
@@ -69,6 +72,7 @@ public class SqlTask
             DataSize maxBufferSize)
     {
         this.taskId = checkNotNull(taskId, "taskId is null");
+        this.nodeInstanceId = checkNotNull(nodeInstanceId, "nodeInstanceId is null");
         this.location = checkNotNull(location, "location is null");
         this.sqlTaskExecutionFactory = checkNotNull(sqlTaskExecutionFactory, "sqlTaskExecutionFactory is null");
         checkNotNull(taskNotificationExecutor, "taskNotificationExecutor is null");
@@ -174,6 +178,7 @@ public class SqlTask
 
         return new TaskInfo(
                 taskStateMachine.getTaskId(),
+                Optional.of(nodeInstanceId),
                 versionNumber,
                 state,
                 location,
diff --git a/presto-main/src/main/java/com/facebook/presto/execution/SqlTaskManager.java b/presto-main/src/main/java/com/facebook/presto/execution/SqlTaskManager.java
index 480d0a9f8e..1727c91682 100644
--- a/presto-main/src/main/java/com/facebook/presto/execution/SqlTaskManager.java
+++ b/presto-main/src/main/java/com/facebook/presto/execution/SqlTaskManager.java
@@ -27,6 +27,7 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.ListenableFuture;
 import io.airlift.concurrent.ThreadPoolExecutorMBean;
 import io.airlift.log.Logger;
+import io.airlift.node.NodeInfo;
 import io.airlift.units.DataSize;
 import io.airlift.units.Duration;
 import org.joda.time.DateTime;
@@ -78,8 +79,10 @@ public class SqlTaskManager
             final LocationFactory locationFactory,
             TaskExecutor taskExecutor,
             QueryMonitor queryMonitor,
+            NodeInfo nodeInfo,
             TaskManagerConfig config)
     {
+        checkNotNull(nodeInfo, "nodeInfo is null");
         checkNotNull(config, "config is null");
         infoCacheTime = config.getInfoMaxAge();
         clientTimeout = config.getClientTimeout();
@@ -102,6 +105,7 @@ public class SqlTaskManager
             {
                 return new SqlTask(
                         taskId,
+                        nodeInfo.getInstanceId(),
                         locationFactory.createLocalTaskLocation(taskId),
                         sqlTaskExecutionFactory,
                         taskNotificationExecutor,
diff --git a/presto-main/src/main/java/com/facebook/presto/execution/TaskInfo.java b/presto-main/src/main/java/com/facebook/presto/execution/TaskInfo.java
index d68b08c45e..f20dad9dd4 100644
--- a/presto-main/src/main/java/com/facebook/presto/execution/TaskInfo.java
+++ b/presto-main/src/main/java/com/facebook/presto/execution/TaskInfo.java
@@ -24,6 +24,7 @@ import javax.annotation.concurrent.Immutable;
 
 import java.net.URI;
 import java.util.List;
+import java.util.Optional;
 import java.util.Set;
 
 import static com.google.common.base.MoreObjects.toStringHelper;
@@ -50,6 +51,7 @@ public class TaskInfo
     public static final long MAX_VERSION = Long.MAX_VALUE;
 
     private final TaskId taskId;
+    private final Optional<String> nodeInstanceId;
     private final long version;
     private final TaskState state;
     private final URI self;
@@ -61,6 +63,7 @@ public class TaskInfo
 
     @JsonCreator
     public TaskInfo(@JsonProperty("taskId") TaskId taskId,
+            @JsonProperty("nodeInstanceId") Optional<String> nodeInstanceId,
             @JsonProperty("version") long version,
             @JsonProperty("state") TaskState state,
             @JsonProperty("self") URI self,
@@ -71,6 +74,7 @@ public class TaskInfo
             @JsonProperty("failures") List<ExecutionFailureInfo> failures)
     {
         this.taskId = checkNotNull(taskId, "taskId is null");
+        this.nodeInstanceId = checkNotNull(nodeInstanceId, "nodeInstanceId is null");
         this.version = checkNotNull(version, "version is null");
         this.state = checkNotNull(state, "state is null");
         this.self = checkNotNull(self, "self is null");
@@ -93,6 +97,12 @@ public class TaskInfo
         return taskId;
     }
 
+    @JsonProperty
+    public Optional<String> getNodeInstanceId()
+    {
+        return nodeInstanceId;
+    }
+
     @JsonProperty
     public long getVersion()
     {
@@ -143,7 +153,7 @@ public class TaskInfo
 
     public TaskInfo summarize()
     {
-        return new TaskInfo(taskId, version, state, self, lastHeartbeat, outputBuffers, noMoreSplits, stats.summarize(), failures);
+        return new TaskInfo(taskId, nodeInstanceId, version, state, self, lastHeartbeat, outputBuffers, noMoreSplits, stats.summarize(), failures);
     }
 
     @Override
diff --git a/presto-main/src/main/java/com/facebook/presto/server/HttpRemoteTask.java b/presto-main/src/main/java/com/facebook/presto/server/HttpRemoteTask.java
index 473996c505..40e59dab43 100644
--- a/presto-main/src/main/java/com/facebook/presto/server/HttpRemoteTask.java
+++ b/presto-main/src/main/java/com/facebook/presto/server/HttpRemoteTask.java
@@ -67,6 +67,7 @@ import java.net.URI;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map.Entry;
+import java.util.Optional;
 import java.util.Queue;
 import java.util.Set;
 import java.util.concurrent.CancellationException;
@@ -82,7 +83,9 @@ import java.util.concurrent.atomic.AtomicReference;
 
 import static com.facebook.presto.spi.StandardErrorCode.REMOTE_TASK_ERROR;
 import static com.facebook.presto.spi.StandardErrorCode.TOO_MANY_REQUESTS_FAILED;
+import static com.facebook.presto.spi.StandardErrorCode.WORKER_RESTARTED;
 import static com.facebook.presto.util.Failures.WORKER_NODE_ERROR;
+import static com.facebook.presto.util.Failures.WORKER_RESTARTED_ERROR;
 import static com.facebook.presto.util.Failures.toFailure;
 import static com.facebook.presto.util.ImmutableCollectors.toImmutableList;
 import static com.google.common.base.MoreObjects.toStringHelper;
@@ -188,6 +191,7 @@ public class HttpRemoteTask
 
             taskInfo = new StateMachine<>("task " + taskId, executor, new TaskInfo(
                     taskId,
+                    Optional.empty(),
                     TaskInfo.MIN_VERSION,
                     TaskState.PLANNED,
                     location,
@@ -322,7 +326,14 @@ public class HttpRemoteTask
         }
 
         // change to new value if old value is not changed and new value has a newer version
+        AtomicBoolean workerRestarted = new AtomicBoolean();
         taskInfo.setIf(newValue, oldValue -> {
+            // did the worker restart
+            if (oldValue.getNodeInstanceId().isPresent() && !oldValue.getNodeInstanceId().equals(newValue.getNodeInstanceId())) {
+                workerRestarted.set(true);
+                return false;
+            }
+
             if (oldValue.getState().isDone()) {
                 // never update if the task has reached a terminal state
                 return false;
@@ -334,6 +345,12 @@ public class HttpRemoteTask
             return true;
         });
 
+        if (workerRestarted.get()) {
+            PrestoException exception = new PrestoException(WORKER_RESTARTED, format("%s (%s)", WORKER_RESTARTED_ERROR, newValue.getSelf()));
+            failTask(exception);
+            abort();
+        }
+
         // remove acknowledged splits, which frees memory
         for (TaskSource source : sources) {
             PlanNodeId planNodeId = source.getPlanNodeId();
@@ -465,6 +482,7 @@ public class HttpRemoteTask
             TaskInfo taskInfo = getTaskInfo();
             URI uri = taskInfo.getSelf();
             updateTaskInfo(new TaskInfo(taskInfo.getTaskId(),
+                    taskInfo.getNodeInstanceId(),
                     TaskInfo.MAX_VERSION,
                     TaskState.ABORTED,
                     uri,
@@ -521,6 +539,7 @@ public class HttpRemoteTask
             log.debug(cause, "Remote task failed: %s", taskInfo.getSelf());
         }
         updateTaskInfo(new TaskInfo(taskInfo.getTaskId(),
+                taskInfo.getNodeInstanceId(),
                 TaskInfo.MAX_VERSION,
                 TaskState.FAILED,
                 taskInfo.getSelf(),
diff --git a/presto-main/src/main/java/com/facebook/presto/util/Failures.java b/presto-main/src/main/java/com/facebook/presto/util/Failures.java
index d84d0d8af8..e9b52ae132 100644
--- a/presto-main/src/main/java/com/facebook/presto/util/Failures.java
+++ b/presto-main/src/main/java/com/facebook/presto/util/Failures.java
@@ -34,11 +34,13 @@ import static java.util.Arrays.asList;
 
 public final class Failures
 {
-    public static final String WORKER_NODE_ERROR = "" +
-            "Encountered too many errors talking to a worker node. " +
-            "The node may have crashed or be under too much load. " +
+    private static final String NODE_CRASHED_ERROR = "The node may have crashed or be under too much load. " +
             "This is probably a transient issue, so please retry your query in a few minutes.";
 
+    public static final String WORKER_NODE_ERROR = "Encountered too many errors talking to a worker node. " + NODE_CRASHED_ERROR;
+
+    public static final String WORKER_RESTARTED_ERROR = "A worker node running your query has restarted. " + NODE_CRASHED_ERROR;
+
     private Failures() {}
 
     public static ExecutionFailureInfo toFailure(Throwable failure)
diff --git a/presto-main/src/test/java/com/facebook/presto/execution/MockRemoteTaskFactory.java b/presto-main/src/test/java/com/facebook/presto/execution/MockRemoteTaskFactory.java
index d31eccc264..fb79a6abc8 100644
--- a/presto-main/src/test/java/com/facebook/presto/execution/MockRemoteTaskFactory.java
+++ b/presto-main/src/test/java/com/facebook/presto/execution/MockRemoteTaskFactory.java
@@ -162,6 +162,7 @@ public class MockRemoteTaskFactory
 
             return new TaskInfo(
                     taskStateMachine.getTaskId(),
+                    Optional.empty(),
                     nextTaskInfoVersion.getAndIncrement(),
                     state,
                     location,
diff --git a/presto-main/src/test/java/com/facebook/presto/execution/TestSqlStageExecution.java b/presto-main/src/test/java/com/facebook/presto/execution/TestSqlStageExecution.java
index dc34df9d98..b0d045974c 100644
--- a/presto-main/src/test/java/com/facebook/presto/execution/TestSqlStageExecution.java
+++ b/presto-main/src/test/java/com/facebook/presto/execution/TestSqlStageExecution.java
@@ -442,6 +442,7 @@ public class TestSqlStageExecution
 
                 return new TaskInfo(
                         taskStateMachine.getTaskId(),
+                        Optional.empty(),
                         nextTaskInfoVersion.getAndIncrement(),
                         state,
                         location,
diff --git a/presto-main/src/test/java/com/facebook/presto/execution/TestSqlTask.java b/presto-main/src/test/java/com/facebook/presto/execution/TestSqlTask.java
index 04d77f8b58..b363827a20 100644
--- a/presto-main/src/test/java/com/facebook/presto/execution/TestSqlTask.java
+++ b/presto-main/src/test/java/com/facebook/presto/execution/TestSqlTask.java
@@ -278,6 +278,7 @@ public class TestSqlTask
 
         return new SqlTask(
                 taskId,
+                "test",
                 location,
                 sqlTaskExecutionFactory,
                 taskNotificationExecutor,
diff --git a/presto-main/src/test/java/com/facebook/presto/execution/TestSqlTaskManager.java b/presto-main/src/test/java/com/facebook/presto/execution/TestSqlTaskManager.java
index 701501ca26..22b7f83597 100644
--- a/presto-main/src/test/java/com/facebook/presto/execution/TestSqlTaskManager.java
+++ b/presto-main/src/test/java/com/facebook/presto/execution/TestSqlTaskManager.java
@@ -245,6 +245,7 @@ public class TestSqlTaskManager
                 new MockLocationFactory(),
                 taskExecutor,
                 new QueryMonitor(new ObjectMapperProvider().get(), new NullEventClient(), new NodeInfo("test"), new NodeVersion("testVersion")),
+                new NodeInfo("test"),
                 config);
     }
 
diff --git a/presto-main/src/test/java/com/facebook/presto/server/MockTaskManager.java b/presto-main/src/test/java/com/facebook/presto/server/MockTaskManager.java
index b3870629ad..113221fb3b 100644
--- a/presto-main/src/test/java/com/facebook/presto/server/MockTaskManager.java
+++ b/presto-main/src/test/java/com/facebook/presto/server/MockTaskManager.java
@@ -44,6 +44,7 @@ import java.net.URI;
 import java.util.Collections;
 import java.util.List;
 import java.util.NoSuchElementException;
+import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ScheduledExecutorService;
@@ -271,6 +272,7 @@ public class MockTaskManager
 
             return new TaskInfo(
                     taskStateMachine.getTaskId(),
+                    Optional.empty(),
                     nextTaskInfoVersion.getAndIncrement(),
                     state,
                     location,
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/StandardErrorCode.java b/presto-spi/src/main/java/com/facebook/presto/spi/StandardErrorCode.java
index 9ed35f16c4..e54d5e3e8a 100644
--- a/presto-spi/src/main/java/com/facebook/presto/spi/StandardErrorCode.java
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/StandardErrorCode.java
@@ -42,6 +42,7 @@ public enum StandardErrorCode
     NO_NODES_AVAILABLE(0x0001_0005),
     REMOTE_TASK_ERROR(0x0001_0006),
     COMPILER_ERROR(0x0001_0007),
+    WORKER_RESTARTED(0x0001_0008),
 
     INSUFFICIENT_RESOURCES(0x0002_0000),
     EXCEEDED_MEMORY_LIMIT(0x0002_0001),