From 809864577d0fe145724bf3d3685419566be410a6 Mon Sep 17 00:00:00 2001
From: Sandy Gao <shudigao@hotmail.com>
Date: Fri, 26 Mar 2021 12:11:48 -0400
Subject: [PATCH] Don't close scheduler and split sources before rescheduling

- Otherwise split sources become unusable after rescheduling
- Detect issue in snapshot configuration
- Update log when resume fails to be more accurate
---
 .../java/io/prestosql/execution/SqlQueryExecution.java | 10 ++++++++++
 .../execution/scheduler/SqlQueryScheduler.java         |  8 ++++++--
 .../io/prestosql/snapshot/QuerySnapshotManager.java    |  7 ++++++-
 .../main/java/io/prestosql/snapshot/SnapshotUtils.java | 10 ++++++++++
 4 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/presto-main/src/main/java/io/prestosql/execution/SqlQueryExecution.java b/presto-main/src/main/java/io/prestosql/execution/SqlQueryExecution.java
index 180da0262..eaa42cb17 100644
--- a/presto-main/src/main/java/io/prestosql/execution/SqlQueryExecution.java
+++ b/presto-main/src/main/java/io/prestosql/execution/SqlQueryExecution.java
@@ -736,6 +736,16 @@ public class SqlQueryExecution
             reasons.add("Requires more than 1 worker nodes");
         }
 
+        if (!snapshotManager.getSnapshotUtils().hasStoreClient()) {
+            String snapshotProfile = snapshotManager.getSnapshotUtils().getSnapshotProfile();
+            if (snapshotProfile == null) {
+                reasons.add("Property hetu.experimental.snapshot.profile is not specified");
+            }
+            else {
+                reasons.add("Specified value '" + snapshotProfile + "' for property hetu.experimental.snapshot.profile is not valid");
+            }
+        }
+
         if (!reasons.isEmpty()) {
             // Disable snapshot support in the session. If this value has been used before this point,
             // then we may need to remedy those places to disable snapshot as well. Fortunately,
diff --git a/presto-main/src/main/java/io/prestosql/execution/scheduler/SqlQueryScheduler.java b/presto-main/src/main/java/io/prestosql/execution/scheduler/SqlQueryScheduler.java
index f6e767c4e..9492d49c6 100644
--- a/presto-main/src/main/java/io/prestosql/execution/scheduler/SqlQueryScheduler.java
+++ b/presto-main/src/main/java/io/prestosql/execution/scheduler/SqlQueryScheduler.java
@@ -288,10 +288,10 @@ public class SqlQueryScheduler
     // either as a result of a task failure, or a failed attempt to resume the query
     public void cancelToResume()
     {
+        queryStateMachine.transitionToRescheduling();
         for (SqlStageExecution stageExecution : stages.values()) {
             stageExecution.cancelToResume();
         }
-        queryStateMachine.transitionToRescheduling();
     }
 
     // this is a separate method to ensure that the `this` reference is not leaked during construction
@@ -740,7 +740,11 @@ public class SqlQueryScheduler
             RuntimeException closeError = new RuntimeException();
             for (StageScheduler scheduler : stageSchedulers.values()) {
                 try {
-                    scheduler.close();
+                    // Snapshot: when trying to reschedule, then don't close the scheduler (and more importantly, split sources in it)
+                    QueryState state = queryStateMachine.getQueryState();
+                    if (state != QueryState.RESCHEDULING && state != QueryState.RESUMING) {
+                        scheduler.close();
+                    }
                 }
                 catch (Throwable t) {
                     queryStateMachine.transitionToFailed(t);
diff --git a/presto-main/src/main/java/io/prestosql/snapshot/QuerySnapshotManager.java b/presto-main/src/main/java/io/prestosql/snapshot/QuerySnapshotManager.java
index ee8383a72..6d0507ec1 100644
--- a/presto-main/src/main/java/io/prestosql/snapshot/QuerySnapshotManager.java
+++ b/presto-main/src/main/java/io/prestosql/snapshot/QuerySnapshotManager.java
@@ -217,9 +217,9 @@ public class QuerySnapshotManager
         {
             public void run()
             {
-                LOG.warn("Snapshot restore timed out, failed to restore snapshot for %s, snapshot %d", queryId.getId(), lastTriedId.getAsLong());
                 synchronized (this) {
                     if (retryTimer.isPresent()) {
+                        LOG.warn("Snapshot restore timed out, failed to restore snapshot for %s, snapshot %s", queryId.getId(), lastTriedId.toString());
                         retryTimer = Optional.empty();
                     }
                     else {
@@ -428,4 +428,9 @@ public class QuerySnapshotManager
     {
         return restoreResult;
     }
+
+    public SnapshotUtils getSnapshotUtils()
+    {
+        return snapshotUtils;
+    }
 }
diff --git a/presto-main/src/main/java/io/prestosql/snapshot/SnapshotUtils.java b/presto-main/src/main/java/io/prestosql/snapshot/SnapshotUtils.java
index 9d974afc4..686f22d49 100644
--- a/presto-main/src/main/java/io/prestosql/snapshot/SnapshotUtils.java
+++ b/presto-main/src/main/java/io/prestosql/snapshot/SnapshotUtils.java
@@ -92,6 +92,16 @@ public class SnapshotUtils
         return isCoordinator;
     }
 
+    public boolean hasStoreClient()
+    {
+        return snapshotStoreClient != null;
+    }
+
+    public String getSnapshotProfile()
+    {
+        return snapshotConfig.getSnapshotProfile();
+    }
+
     public void initialize()
     {
         snapshotStoreClient = buildSnapshotStoreClient();