!788 [I3CRBJ] Avoid closing split source if it will be used after resume

Merge pull request !788 from sandygao/I3CRBJ-avoid-closing-split-source

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,

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);

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;
+    }
 }

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();