diff --git a/presto-cli/src/main/java/io/prestosql/cli/StatusPrinter.java b/presto-cli/src/main/java/io/prestosql/cli/StatusPrinter.java index 2a3f0f551..43641b909 100644 --- a/presto-cli/src/main/java/io/prestosql/cli/StatusPrinter.java +++ b/presto-cli/src/main/java/io/prestosql/cli/StatusPrinter.java @@ -19,6 +19,7 @@ import io.airlift.log.Logger; import io.airlift.units.DataSize; import io.airlift.units.Duration; import io.prestosql.client.QueryStatusInfo; +import io.prestosql.client.SnapshotStats; import io.prestosql.client.StageStats; import io.prestosql.client.StatementClient; import io.prestosql.client.StatementStats; @@ -217,6 +218,33 @@ Spilled: 20GB readTime.getValue(SECONDS)); reprintLine(summary); } + + // Snapshot Capture stats All: 100MB/22s/18s, Last: 40MB/10s/7s + SnapshotStats snapshotStats = stats.getSnapshotStats(); + // snapshotStats should be null in case snapshot feature is disabled + if (snapshotStats != null) { + Duration allCaptureCPUTime = millis(snapshotStats.getTotalCaptureCpuTime()); + Duration allCaptureWallTime = millis(snapshotStats.getTotalCaptureWallTime()); + Duration lastCaptureCPUTime = millis(snapshotStats.getLastCaptureCpuTime()); + Duration lastCaptureWallTime = millis(snapshotStats.getLastCaptureWallTime()); + String allSnapshotsSize = FormatUtils.formatDataSize(bytes(snapshotStats.getAllCaptureSize()), true); + String lastSnapshotSize = FormatUtils.formatDataSize(bytes(snapshotStats.getLastCaptureSize()), true); + String captureSummary = String.format("Snapshot Capture: All: %s/%.1fs/%.1fs, Last: %s/%.1fs/%.1fs", + allSnapshotsSize, allCaptureCPUTime.getValue(SECONDS), allCaptureWallTime.getValue(SECONDS), + lastSnapshotSize, lastCaptureCPUTime.getValue(SECONDS), lastCaptureWallTime.getValue(SECONDS)); + reprintLine(captureSummary); + + // Snapshot restore stats: 1/100MB/22s/18s + long restoreCount = snapshotStats.getSuccessRestoreCount(); + if (restoreCount > 0) { + Duration allRestoreCPUTime = millis(snapshotStats.getTotalRestoreCpuTime()); + Duration allRestoreWallTime = millis(snapshotStats.getTotalRestoreWallTime()); + String allRestoreSize = FormatUtils.formatDataSize(bytes(snapshotStats.getTotalRestoreSize()), true); + String restoreSummary = String.format(Locale.ROOT, "Snapshot Restore: %d/%s/%.1fs/%.1fs", restoreCount, + allRestoreSize, allRestoreCPUTime.getValue(SECONDS), allRestoreWallTime.getValue(SECONDS)); + reprintLine(restoreSummary); + } + } } // 0:32 [2.12GB, 15M rows] [67MB/s, 463K rows/s] diff --git a/presto-client/src/main/java/io/prestosql/client/SnapshotStats.java b/presto-client/src/main/java/io/prestosql/client/SnapshotStats.java new file mode 100644 index 000000000..f4585c1eb --- /dev/null +++ b/presto-client/src/main/java/io/prestosql/client/SnapshotStats.java @@ -0,0 +1,283 @@ +/* + * Copyright (C) 2018-2022. Huawei Technologies Co., Ltd. All rights reserved. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.prestosql.client; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; + +import javax.annotation.concurrent.Immutable; + +import static com.google.common.base.MoreObjects.toStringHelper; + +@Immutable +public class SnapshotStats +{ + // id of last successful snapshot + private final long lastCaptureSnapshotId; + // Total CPU time (ms) taken for capturing all snapshots of query + private final long totalCaptureCpuTime; + // Total Wall time (ms) taken for capturing all snapshots of query + private final long totalCaptureWallTime; + // Cpu time (ms) taken for capturing last successful snapshot of query + private final long lastCaptureCpuTime; + // Wall time (ms) taken for capturing last successful snapshot of query + private final long lastCaptureWallTime; + // Size of all snapshots (bytes) + private final long allCaptureSize; + // Size (bytes) of last successful snapshot + private final long lastCaptureSize; + // Snapshot id of last restore + private final long lastRestoreSnapshotId; + // Number of successful restores in current query + private final long successRestoreCount; + // Total Wall time (ms) for all restores happened during query + private final long totalRestoreWallTime; + // Total restored size during restore + private final long totalRestoreSize; + // Total Cpu time (ms) for loading state during restore + private final long totalRestoreCpuTime; + + @JsonCreator + public SnapshotStats( + @JsonProperty("lastCaptureSnapshotId") long lastCaptureSnapshotId, + @JsonProperty("totalCaptureCpuTime") long totalCaptureCpuTime, + @JsonProperty("totalCaptureWallTime") long totalCaptureWallTime, + @JsonProperty("lastCaptureCpuTime") long lastCaptureCpuTime, + @JsonProperty("lastCaptureWallTime") long lastCaptureWallTime, + @JsonProperty("allCaptureSize") long allCaptureSize, + @JsonProperty("lastCaptureSize") long lastCaptureSize, + @JsonProperty("lastRestoreSnapshotId") long lastRestoreSnapshotId, + @JsonProperty("successRestoreCount") long successRestoreCount, + @JsonProperty("totalRestoreWallTime") long totalRestoreWallTime, + @JsonProperty("totalRestoreSize") long totalRestoreSize, + @JsonProperty("totalRestoreCpuTime") long totalRestoreCpuTime) + { + this.lastCaptureSnapshotId = lastCaptureSnapshotId; + this.totalCaptureCpuTime = totalCaptureCpuTime; + this.totalCaptureWallTime = totalCaptureWallTime; + this.lastCaptureCpuTime = lastCaptureCpuTime; + this.lastCaptureWallTime = lastCaptureWallTime; + this.allCaptureSize = allCaptureSize; + this.lastCaptureSize = lastCaptureSize; + this.lastRestoreSnapshotId = lastRestoreSnapshotId; + this.successRestoreCount = successRestoreCount; + this.totalRestoreWallTime = totalRestoreWallTime; + this.totalRestoreSize = totalRestoreSize; + this.totalRestoreCpuTime = totalRestoreCpuTime; + } + + @JsonProperty + public long getLastCaptureSnapshotId() + { + return lastCaptureSnapshotId; + } + + @JsonProperty + public long getTotalCaptureCpuTime() + { + return totalCaptureCpuTime; + } + + @JsonProperty + public long getTotalCaptureWallTime() + { + return totalCaptureWallTime; + } + + @JsonProperty + public long getLastCaptureCpuTime() + { + return lastCaptureCpuTime; + } + + @JsonProperty + public long getLastCaptureWallTime() + { + return lastCaptureWallTime; + } + + @JsonProperty + public long getAllCaptureSize() + { + return allCaptureSize; + } + + @JsonProperty + public long getLastCaptureSize() + { + return lastCaptureSize; + } + + @JsonProperty + public long getLastRestoreSnapshotId() + { + return lastRestoreSnapshotId; + } + + @JsonProperty + public long getSuccessRestoreCount() + { + return successRestoreCount; + } + + @JsonProperty + public long getTotalRestoreWallTime() + { + return totalRestoreWallTime; + } + + @JsonProperty + public long getTotalRestoreSize() + { + return totalRestoreSize; + } + + @JsonProperty + public long getTotalRestoreCpuTime() + { + return totalRestoreCpuTime; + } + + @Override + public String toString() + { + return toStringHelper(this) + .add("lastCaptureSnapshotId", lastCaptureSnapshotId) + .add("totalCaptureCpuTime", totalCaptureCpuTime) + .add("totalCaptureWallTime", totalCaptureWallTime) + .add("lastCaptureCpuTime", lastCaptureCpuTime) + .add("lastCaptureWallTime", lastCaptureWallTime) + .add("allCaptureSize", allCaptureSize) + .add("lastCaptureSize", lastCaptureSize) + .add("lastRestoreSnapshotId", lastRestoreSnapshotId) + .add("successRestoreCount", successRestoreCount) + .add("totalRestoreWallTime", totalRestoreWallTime) + .add("totalRestoreSize", totalRestoreSize) + .add("totalRestoreCpuTime", totalRestoreCpuTime) + .toString(); + } + + public static SnapshotStats.Builder builder() + { + return new SnapshotStats.Builder(); + } + + public static class Builder + { + private long lastCaptureSnapshotId; + private long totalCaptureCpuTime; + private long totalCaptureWallTime; + private long lastCaptureCpuTime; + private long lastCaptureWallTime; + private long allCaptureSize; + private long lastCaptureSize; + private long lastRestoreSnapshotId; + private long successRestoreCount; + private long totalRestoreWallTime; + private long totalRestoreSize; + private long totalRestoreCpuTime; + + private Builder() {} + + public Builder setLastCaptureSnapshotId(long lastCaptureSnapshotId) + { + this.lastCaptureSnapshotId = lastCaptureSnapshotId; + return this; + } + + public Builder setTotalCpuTimeMillis(long totalCaptureCpuTime) + { + this.totalCaptureCpuTime = totalCaptureCpuTime; + return this; + } + + public Builder setTotalWallTimeMillis(long totalCaptureWallTime) + { + this.totalCaptureWallTime = totalCaptureWallTime; + return this; + } + + public Builder setLastSnapshotCpuTimeMillis(long lastCaptureCpuTime) + { + this.lastCaptureCpuTime = lastCaptureCpuTime; + return this; + } + + public Builder setLastSnapshotWallTimeMillis(long lastCaptureWallTime) + { + this.lastCaptureWallTime = lastCaptureWallTime; + return this; + } + + public Builder setAllSnapshotsSizeBytes(long allCaptureSize) + { + this.allCaptureSize = allCaptureSize; + return this; + } + + public Builder setLastSnapshotSizeBytes(long lastCaptureSize) + { + this.lastCaptureSize = lastCaptureSize; + return this; + } + + public Builder setLastRestoreSnapshotId(long lastRestoreSnapshotId) + { + this.lastRestoreSnapshotId = lastRestoreSnapshotId; + return this; + } + + public Builder setSuccessRestoreCount(long successRestoreCount) + { + this.successRestoreCount = successRestoreCount; + return this; + } + + public Builder setTotalRestoreWallTime(long totalRestoreWallTime) + { + this.totalRestoreWallTime = totalRestoreWallTime; + return this; + } + + public Builder setTotalRestoreSize(long totalRestoreSize) + { + this.totalRestoreSize = totalRestoreSize; + return this; + } + + public Builder setTotalRestoreCpuTime(long totalRestoreCpuTime) + { + this.totalRestoreCpuTime = totalRestoreCpuTime; + return this; + } + + public SnapshotStats build() + { + return new SnapshotStats( + lastCaptureSnapshotId, + totalCaptureCpuTime, + totalCaptureWallTime, + lastCaptureCpuTime, + lastCaptureWallTime, + allCaptureSize, + lastCaptureSize, + lastRestoreSnapshotId, + successRestoreCount, + totalRestoreWallTime, + totalRestoreSize, + totalRestoreCpuTime); + } + } +} diff --git a/presto-client/src/main/java/io/prestosql/client/StatementStats.java b/presto-client/src/main/java/io/prestosql/client/StatementStats.java index 730602eb7..893429faa 100644 --- a/presto-client/src/main/java/io/prestosql/client/StatementStats.java +++ b/presto-client/src/main/java/io/prestosql/client/StatementStats.java @@ -48,6 +48,7 @@ public class StatementStats private long elapsedSpillWriteTimeMillis; private int spilledNodes; private final StageStats rootStage; + private final SnapshotStats snapshotStats; @JsonCreator public StatementStats( @@ -70,7 +71,8 @@ public class StatementStats @JsonProperty("elapsedSpillReadTimeMillis") long elapsedSpillReadTimeMillis, @JsonProperty("elapsedSpillWriteTimeMillis") long elapsedSpillWriteTimeMillis, @JsonProperty("spilledNodes") int spilledNodes, - @JsonProperty("rootStage") StageStats rootStage) + @JsonProperty("rootStage") StageStats rootStage, + @JsonProperty("snapshotStats") SnapshotStats snapshotStats) { this.state = requireNonNull(state, "state is null"); this.queued = queued; @@ -92,6 +94,7 @@ public class StatementStats this.elapsedSpillWriteTimeMillis = elapsedSpillWriteTimeMillis; this.spilledNodes = spilledNodes; this.rootStage = rootStage; + this.snapshotStats = snapshotStats; } @JsonProperty @@ -224,6 +227,12 @@ public class StatementStats return spilledNodes; } + @JsonProperty + public SnapshotStats getSnapshotStats() + { + return snapshotStats; + } + @Override public String toString() { @@ -247,6 +256,7 @@ public class StatementStats .add("elapsedSpillReadTime", elapsedSpillReadTimeMillis) .add("elapsedSpillWriteTime", elapsedSpillWriteTimeMillis) .add("rootStage", rootStage) + .add("snapshotStats", snapshotStats) .toString(); } @@ -277,6 +287,7 @@ public class StatementStats private long spillReadTimeMillis; private long spillWriteTimeMillis; private int spilledNodes; + private SnapshotStats snapshotStats; private Builder() {} @@ -400,6 +411,12 @@ public class StatementStats return this; } + public Builder setSnapshotStats(SnapshotStats snapshotStats) + { + this.snapshotStats = snapshotStats; + return this; + } + public StatementStats build() { return new StatementStats( @@ -422,7 +439,8 @@ public class StatementStats spillReadTimeMillis, spillWriteTimeMillis, spilledNodes, - rootStage); + rootStage, + snapshotStats); } } } diff --git a/presto-jdbc/src/test/java/io/prestosql/jdbc/TestProgressMonitor.java b/presto-jdbc/src/test/java/io/prestosql/jdbc/TestProgressMonitor.java index c07063a40..2366da02f 100644 --- a/presto-jdbc/src/test/java/io/prestosql/jdbc/TestProgressMonitor.java +++ b/presto-jdbc/src/test/java/io/prestosql/jdbc/TestProgressMonitor.java @@ -90,7 +90,7 @@ public class TestProgressMonitor nextUriId == null ? null : server.url(format("/v1/statement/%s/%s", queryId, nextUriId)).uri(), responseColumns, data, - new StatementStats(state, state.equals("QUEUED"), true, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, null), + new StatementStats(state, state.equals("QUEUED"), true, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, null, null), null, ImmutableList.of(), null, diff --git a/presto-main/src/main/java/io/prestosql/execution/QueryExecution.java b/presto-main/src/main/java/io/prestosql/execution/QueryExecution.java index 370a19fe4..694251fd2 100644 --- a/presto-main/src/main/java/io/prestosql/execution/QueryExecution.java +++ b/presto-main/src/main/java/io/prestosql/execution/QueryExecution.java @@ -25,6 +25,7 @@ import io.prestosql.execution.warnings.WarningCollector; import io.prestosql.memory.VersionedMemoryPoolId; import io.prestosql.operator.TaskLocation; import io.prestosql.server.BasicQueryInfo; +import io.prestosql.snapshot.QuerySnapshotManager; import io.prestosql.spi.type.Type; import io.prestosql.sql.planner.Plan; @@ -55,6 +56,11 @@ public interface QueryExecution Duration getTotalCpuTime(); + default QuerySnapshotManager getQuerySnapshotManager() + { + return null; + } + DataSize getUserMemoryReservation(); DataSize getTotalMemoryReservation(); diff --git a/presto-main/src/main/java/io/prestosql/execution/QueryManager.java b/presto-main/src/main/java/io/prestosql/execution/QueryManager.java index 99e46cfe9..4ba2c1878 100644 --- a/presto-main/src/main/java/io/prestosql/execution/QueryManager.java +++ b/presto-main/src/main/java/io/prestosql/execution/QueryManager.java @@ -17,6 +17,7 @@ import com.google.common.util.concurrent.ListenableFuture; import io.prestosql.Session; import io.prestosql.execution.StateMachine.StateChangeListener; import io.prestosql.server.BasicQueryInfo; +import io.prestosql.snapshot.QuerySnapshotManager; import io.prestosql.spi.QueryId; import java.util.List; @@ -113,4 +114,9 @@ public interface QueryManager QueryManagerStats getStats(); default void checkForQueryPruning(QueryId queryId, QueryInfo queryInfo) {} + + default QuerySnapshotManager getQuerySnapshotManager(QueryId queryId) + { + return null; + } } 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 a7945a710..145dc7e9b 100644 --- a/presto-main/src/main/java/io/prestosql/execution/SqlQueryExecution.java +++ b/presto-main/src/main/java/io/prestosql/execution/SqlQueryExecution.java @@ -300,6 +300,12 @@ public class SqlQueryExecution stateMachine.setMemoryPool(poolId); } + @Override + public QuerySnapshotManager getQuerySnapshotManager() + { + return snapshotManager; + } + @Override public DataSize getUserMemoryReservation() { @@ -547,7 +553,7 @@ public class SqlQueryExecution } catch (PrestoException e) { if (e.getErrorCode() == NO_NODES_AVAILABLE.toErrorCode()) { - // Not enough worker to resume all tasks. Retrying from any saves snapshot likely wont' work either. + // Not enough worker to resume all tasks. Retrying from any saved snapshot likely wont' work either. // Clear ongoing and existing snapshots and restart. snapshotManager.invalidateAllSnapshots(); scheduler = createResumeScheduler(plan, rootOutputBuffers); diff --git a/presto-main/src/main/java/io/prestosql/execution/SqlQueryManager.java b/presto-main/src/main/java/io/prestosql/execution/SqlQueryManager.java index d5ea8dc69..31720e35d 100644 --- a/presto-main/src/main/java/io/prestosql/execution/SqlQueryManager.java +++ b/presto-main/src/main/java/io/prestosql/execution/SqlQueryManager.java @@ -30,6 +30,7 @@ import io.prestosql.memory.ClusterMemoryManager; import io.prestosql.metadata.SessionPropertyManager; import io.prestosql.queryeditorui.QueryEditorUIModule; import io.prestosql.server.BasicQueryInfo; +import io.prestosql.snapshot.QuerySnapshotManager; import io.prestosql.spi.ErrorType; import io.prestosql.spi.PrestoException; import io.prestosql.spi.QueryId; @@ -475,6 +476,12 @@ public class SqlQueryManager } } + @Override + public QuerySnapshotManager getQuerySnapshotManager(QueryId queryId) + { + return queryTracker.getQuery(queryId).getQuerySnapshotManager(); + } + private boolean isIndexCreationQuery(QueryInfo queryInfo) { return queryInfo.getQuery().toUpperCase(Locale.ROOT).startsWith("CREATE INDEX"); diff --git a/presto-main/src/main/java/io/prestosql/execution/SqlTask.java b/presto-main/src/main/java/io/prestosql/execution/SqlTask.java index f8facf248..7e9b15dfc 100644 --- a/presto-main/src/main/java/io/prestosql/execution/SqlTask.java +++ b/presto-main/src/main/java/io/prestosql/execution/SqlTask.java @@ -40,7 +40,7 @@ import io.prestosql.operator.PipelineStatus; import io.prestosql.operator.TaskContext; import io.prestosql.operator.TaskStats; import io.prestosql.snapshot.RestoreResult; -import io.prestosql.snapshot.SnapshotResult; +import io.prestosql.snapshot.SnapshotInfo; import io.prestosql.snapshot.TaskSnapshotManager; import io.prestosql.spi.plan.PlanNodeId; import io.prestosql.sql.planner.PlanFragment; @@ -264,7 +264,7 @@ public class SqlTask Set completedDriverGroups = ImmutableSet.of(); long fullGcCount = 0; Duration fullGcTime = new Duration(0, MILLISECONDS); - Map snapshotCaptureResult = ImmutableMap.of(); + Map snapshotCaptureResult = ImmutableMap.of(); Optional snapshotRestoreResult = Optional.empty(); TaskInfo finalTaskInfo = taskHolder.getFinalTaskInfo(); if (finalTaskInfo != null) { diff --git a/presto-main/src/main/java/io/prestosql/execution/TaskStatus.java b/presto-main/src/main/java/io/prestosql/execution/TaskStatus.java index 4e67fa2f6..197bef368 100644 --- a/presto-main/src/main/java/io/prestosql/execution/TaskStatus.java +++ b/presto-main/src/main/java/io/prestosql/execution/TaskStatus.java @@ -21,7 +21,7 @@ import com.google.common.collect.ImmutableSet; import io.airlift.units.DataSize; import io.airlift.units.Duration; import io.prestosql.snapshot.RestoreResult; -import io.prestosql.snapshot.SnapshotResult; +import io.prestosql.snapshot.SnapshotInfo; import java.net.URI; import java.util.List; @@ -80,7 +80,7 @@ public class TaskStatus // snapshotCaptureResult and snapshotRestoreResult are used to store result of snapshot capture and restore. They are empty when the following conditions happened: // (1) Snapshot is not enabled // (2) Snapshot is enabled but no data for capture/restore result - private final Map snapshotCaptureResult; + private final Map snapshotCaptureResult; private final Optional snapshotRestoreResult; @JsonCreator @@ -102,7 +102,7 @@ public class TaskStatus @JsonProperty("revocableMemoryReservation") DataSize revocableMemoryReservation, @JsonProperty("fullGcCount") long fullGcCount, @JsonProperty("fullGcTime") Duration fullGcTime, - @JsonProperty("snapshotCaptureResult") Map snapshotCaptureResult, + @JsonProperty("snapshotCaptureResult") Map snapshotCaptureResult, @JsonProperty("snapshotRestoreResult") Optional snapshotRestoreResult) { this.taskId = requireNonNull(taskId, "taskId is null"); @@ -241,7 +241,7 @@ public class TaskStatus } @JsonProperty - public Map getSnapshotCaptureResult() + public Map getSnapshotCaptureResult() { return snapshotCaptureResult; } 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 f6b179fa5..2f352a8a9 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 @@ -311,6 +311,7 @@ public class SqlQueryScheduler public synchronized void cancelToResume() { if (!resumed) { + snapshotManager.setRestoreStartTime(System.currentTimeMillis()); // Resume at most once for each scheduler resumed = true; new Thread(() -> { diff --git a/presto-main/src/main/java/io/prestosql/server/protocol/PagePublisherQueryManager.java b/presto-main/src/main/java/io/prestosql/server/protocol/PagePublisherQueryManager.java index e4c97e73b..f17dc2ed8 100644 --- a/presto-main/src/main/java/io/prestosql/server/protocol/PagePublisherQueryManager.java +++ b/presto-main/src/main/java/io/prestosql/server/protocol/PagePublisherQueryManager.java @@ -62,7 +62,7 @@ public class PagePublisherQueryManager private final Set queries = Sets.newConcurrentHashSet(); private final Map queryRunners = new ConcurrentHashMap<>(); private static final DataCenterQueryResults FINISHED_RESULTS_DONOT_USE_HEADER = new DataCenterQueryResults("", URI.create(""), null, null, null, null, - new StatementStats("FINISHED", false, false, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, null), null, + new StatementStats("FINISHED", false, false, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, null, null), null, Collections.emptyList(), null, false); private final DispatchManager dispatchManager; diff --git a/presto-main/src/main/java/io/prestosql/server/protocol/PagePublisherQueryRunner.java b/presto-main/src/main/java/io/prestosql/server/protocol/PagePublisherQueryRunner.java index 487ef7f84..5e3cc2abe 100644 --- a/presto-main/src/main/java/io/prestosql/server/protocol/PagePublisherQueryRunner.java +++ b/presto-main/src/main/java/io/prestosql/server/protocol/PagePublisherQueryRunner.java @@ -69,13 +69,13 @@ public class PagePublisherQueryRunner private static final Logger LOGGER = Logger.get(PagePublisherQueryRunner.class); private static final DataCenterQueryResults RUNNING_RESULTS = new DataCenterQueryResults("", URI.create(""), null, URI.create(""), null, null, - new StatementStats("RUNNING", false, false, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, null), null, + new StatementStats("RUNNING", false, false, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, null, null), null, Collections.emptyList(), null, true); private static final DataCenterQueryResults FINISHED_RESULTS = new DataCenterQueryResults("", URI.create(""), null, null, null, null, - new StatementStats("FINISHED", false, false, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, null), null, + new StatementStats("FINISHED", false, false, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, null, null), null, Collections.emptyList(), null, true); private static final DataCenterQueryResults FAILED_RESULTS = new DataCenterQueryResults("", URI.create(""), null, null, null, null, - new StatementStats("FAILED", false, false, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, null), null, + new StatementStats("FAILED", false, false, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, null, null), null, Collections.emptyList(), null, true); private static final Ordering> WAIT_ORDERING = Ordering.natural().nullsLast(); private static final Duration MAX_WAIT_TIME = new Duration(1, SECONDS); diff --git a/presto-main/src/main/java/io/prestosql/server/protocol/Query.java b/presto-main/src/main/java/io/prestosql/server/protocol/Query.java index ecd870de6..0b8043dd2 100644 --- a/presto-main/src/main/java/io/prestosql/server/protocol/Query.java +++ b/presto-main/src/main/java/io/prestosql/server/protocol/Query.java @@ -36,6 +36,7 @@ import io.prestosql.client.NamedClientTypeSignature; import io.prestosql.client.QueryError; import io.prestosql.client.QueryResults; import io.prestosql.client.RowFieldName; +import io.prestosql.client.SnapshotStats; import io.prestosql.client.StageStats; import io.prestosql.client.StatementStats; import io.prestosql.client.Warning; @@ -50,6 +51,9 @@ import io.prestosql.execution.TaskInfo; import io.prestosql.operator.ExchangeClient; import io.prestosql.operator.PipelineStats; import io.prestosql.operator.TaskLocation; +import io.prestosql.snapshot.QuerySnapshotManager; +import io.prestosql.snapshot.RestoreResult; +import io.prestosql.snapshot.SnapshotInfo; import io.prestosql.spi.ErrorCode; import io.prestosql.spi.Page; import io.prestosql.spi.PageBuilder; @@ -86,6 +90,7 @@ import java.util.OptionalLong; import java.util.Set; import java.util.concurrent.Executor; import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.atomic.AtomicLong; import java.util.stream.Collectors; import static com.google.common.base.MoreObjects.firstNonNull; @@ -96,6 +101,7 @@ import static com.google.common.util.concurrent.Futures.immediateFuture; import static com.google.common.util.concurrent.MoreExecutors.directExecutor; import static io.airlift.concurrent.MoreFutures.addTimeout; import static io.prestosql.SystemSessionProperties.isExchangeCompressionEnabled; +import static io.prestosql.SystemSessionProperties.isSnapshotEnabled; import static io.prestosql.execution.QueryState.FAILED; import static io.prestosql.execution.QueryState.RESCHEDULING; import static io.prestosql.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR; @@ -810,7 +816,7 @@ public class Query throw new IllegalArgumentException("Unsupported kind: " + parameter.getKind()); } - private static StatementStats toStatementStats(QueryInfo queryInfo) + private StatementStats toStatementStats(QueryInfo queryInfo) { QueryStats queryStats = queryInfo.getQueryStats(); //Dont print any more stats for Async Query, @@ -837,9 +843,77 @@ public class Query .setSpilledWriteTimeMillis(queryStats.getSpilledWriteTime().toMillis()) .setSpilledNodes(globalUniqueNodes(outputStage, true).size()) .setRootStage(toStageStats(outputStage)) + .setSnapshotStats(toSnapshotStats(queryInfo.getQueryId())) .build(); } + private SnapshotStats toSnapshotStats(QueryId queryId) + { + if (!isSnapshotEnabled(session) || queryId == null) { + return null; + } + AtomicLong totalCpuTimeMillis = new AtomicLong(0L); + AtomicLong lastSnapshotCpuTimeMillis = new AtomicLong(0L); + AtomicLong allSnapshotsSizeBytes = new AtomicLong(0L); + AtomicLong lastSnapshotSizeBytes = new AtomicLong(0L); + AtomicLong totalWallTimeMillis = new AtomicLong(0L); + AtomicLong lastWallTimeMillis = new AtomicLong(0L); + QuerySnapshotManager querySnapshotManager = queryManager.getQuerySnapshotManager(queryId); + if (querySnapshotManager != null) { + long lastSnapshotId = querySnapshotManager.collectSnapshotCaptureStats(eachSize -> (eachWallTime, eachCpuTime) -> { + allSnapshotsSizeBytes.addAndGet(eachSize); + totalWallTimeMillis.addAndGet(eachWallTime); + totalCpuTimeMillis.addAndGet(eachCpuTime); + }, lastSize -> (lastWallTime, lastCpuTime) -> { + lastSnapshotSizeBytes.set(lastSize); + lastSnapshotCpuTimeMillis.set(lastCpuTime); + lastWallTimeMillis.set(lastWallTime); + }); + if (lastSnapshotId > 0) { + SnapshotStats.Builder builder = SnapshotStats.builder(); + log.debug("SnapshotMetrics: totalWallTimeMillis: [%s]ms, lastWallTimeMillis: [%s]ms", totalWallTimeMillis.toString(), lastWallTimeMillis.toString()); + log.debug("SnapshotMetrics: allSnapshotsSizeBytes: [%d], lastSnapshotSizeBytes: [%d]", allSnapshotsSizeBytes.get(), lastSnapshotSizeBytes.get()); + log.debug("SnapshotMetrics: totalCpuTimeMillis: [%d]ms, lastSnapshotCpuTimeMillis: [%d]ms", totalCpuTimeMillis.get(), lastSnapshotCpuTimeMillis.get()); + builder.setLastCaptureSnapshotId(lastSnapshotId) + .setAllSnapshotsSizeBytes(allSnapshotsSizeBytes.get()) + .setLastSnapshotSizeBytes(lastSnapshotSizeBytes.get()) + .setTotalWallTimeMillis(totalWallTimeMillis.get()) + .setLastSnapshotWallTimeMillis(lastWallTimeMillis.get()) + .setTotalCpuTimeMillis(totalCpuTimeMillis.get()) + .setLastSnapshotCpuTimeMillis(lastSnapshotCpuTimeMillis.get()); + + // Restore stats + AtomicLong totalRestoreWallTime = new AtomicLong(0L); + AtomicLong totalRestoreCpuTime = new AtomicLong(0L); + AtomicLong totalRestoreSize = new AtomicLong(0L); + long lastRestoreSnapshotId = 0; + int restoreCount = 0; + List restoreStats = querySnapshotManager.getRestoreStats(); + restoreCount = restoreStats.size(); + log.debug("SnapshotMetrics: restoreCount: [%d]", restoreCount); + // Add restore stats if restore is happened + if (restoreCount > 0) { + lastRestoreSnapshotId = restoreStats.get(restoreCount - 1).getSnapshotId(); + restoreStats.forEach(restoreResult -> { + SnapshotInfo info = restoreResult.getSnapshotInfo(); + totalRestoreWallTime.addAndGet(info.getEndTime() - info.getBeginTime()); + totalRestoreCpuTime.addAndGet(info.getCpuTime()); + totalRestoreSize.addAndGet(info.getSizeBytes()); + }); + log.debug("SnapshotMetrics: totalRestoreWallTime: [%d]ms, totalRestoreCpuTime: [%d]ms", totalRestoreWallTime.get(), totalRestoreCpuTime.get()); + log.debug("SnapshotMetrics: totalRestoreSize: [%d], lastRestoreSnapshotId: [%d]", totalRestoreSize.get(), lastRestoreSnapshotId); + builder.setSuccessRestoreCount(restoreCount) + .setLastRestoreSnapshotId(lastRestoreSnapshotId) + .setTotalRestoreWallTime(totalRestoreWallTime.get()) + .setTotalCpuTimeMillis(totalRestoreCpuTime.get()) + .setTotalRestoreSize(totalRestoreSize.get()); + } + return builder.build(); + } + } + return null; + } + private static StageStats toStageStats(StageInfo stageInfo) { if (stageInfo == null) { diff --git a/presto-main/src/main/java/io/prestosql/server/remotetask/ContinuousTaskStatusFetcher.java b/presto-main/src/main/java/io/prestosql/server/remotetask/ContinuousTaskStatusFetcher.java index f865c97fb..a10cd0cdb 100644 --- a/presto-main/src/main/java/io/prestosql/server/remotetask/ContinuousTaskStatusFetcher.java +++ b/presto-main/src/main/java/io/prestosql/server/remotetask/ContinuousTaskStatusFetcher.java @@ -28,7 +28,7 @@ import io.prestosql.protocol.BaseResponse; import io.prestosql.protocol.Codec; import io.prestosql.snapshot.QuerySnapshotManager; import io.prestosql.snapshot.RestoreResult; -import io.prestosql.snapshot.SnapshotResult; +import io.prestosql.snapshot.SnapshotInfo; import io.prestosql.spi.PrestoException; import javax.annotation.concurrent.GuardedBy; @@ -285,7 +285,7 @@ class ContinuousTaskStatusFetcher stats.statusRoundTripMillis(nanosSince(currentRequestStartNanos).toMillis()); } - private void updateSnapshots(Map captureResult, Optional restoreResult) + private void updateSnapshots(Map captureResult, Optional restoreResult) { snapshotManager.updateQueryCapture(taskId, captureResult); snapshotManager.updateQueryRestore(taskId, restoreResult); diff --git a/presto-main/src/main/java/io/prestosql/snapshot/MultiInputSnapshotState.java b/presto-main/src/main/java/io/prestosql/snapshot/MultiInputSnapshotState.java index 2324a99e5..8e57e1353 100644 --- a/presto-main/src/main/java/io/prestosql/snapshot/MultiInputSnapshotState.java +++ b/presto-main/src/main/java/io/prestosql/snapshot/MultiInputSnapshotState.java @@ -14,6 +14,7 @@ */ package io.prestosql.snapshot; +import com.google.common.base.Stopwatch; import com.google.common.collect.Iterators; import io.airlift.log.Logger; import io.hetu.core.transport.execution.buffer.PagesSerde; @@ -30,6 +31,7 @@ import java.util.Iterator; import java.util.List; import java.util.Optional; import java.util.Set; +import java.util.concurrent.TimeUnit; import java.util.function.Function; import java.util.function.Supplier; @@ -340,11 +342,13 @@ public class MultiInputSnapshotState LOG.error("BUG! State of component %s has never been stored successfully before snapshot %d", restorableId, snapshotId); } else { + Stopwatch timer = Stopwatch.createStarted(); List storedStates = (List) storedState.get(); pendingPages = storedStates.listIterator(); restorable.restore(pendingPages.next(), pagesSerde); LOG.debug("Successfully restored state to snapshot %d for %s", snapshotId, restorableId); - snapshotManager.succeededToRestore(componentId); + timer.stop(); + snapshotManager.succeededToRestore(componentId, timer.elapsed(TimeUnit.MILLISECONDS)); } } catch (Exception e) { @@ -390,7 +394,7 @@ public class MultiInputSnapshotState snapshot = new SnapshotState(marker); pendingMarkers.add(marker); try { - snapshot.states.add(restorable.capture(pagesSerde)); + snapshot.addState(restorable, pagesSerde); } catch (Exception e) { LOG.warn(e, "Failed to capture and store snapshot state"); @@ -418,10 +422,10 @@ public class MultiInputSnapshotState SnapshotStateId componentId = snapshotStateIdGenerator.apply(snapshotId); try { if (restorable.supportsConsolidatedWrites()) { - snapshotManager.storeConsolidatedState(componentId, snapshot.states); + snapshotManager.storeConsolidatedState(componentId, snapshot.states, snapshot.serTime); } else { - snapshotManager.storeState(componentId, snapshot.states); + snapshotManager.storeState(componentId, snapshot.states, snapshot.serTime); } snapshotManager.succeededToCapture(componentId); LOG.debug("Successfully saved state to snapshot %d for %s", snapshotId, restorableId); @@ -456,10 +460,7 @@ public class MultiInputSnapshotState // For all pending snapshots that have not received marker from this channel, // need to capture the input as part of channel state. if (!snapshot.markedChannels.contains(channel)) { - if (channelSnapshot == null) { - channelSnapshot = pagesSerde.serialize(page).capture(pagesSerde); - } - snapshot.states.add(channelSnapshot); + snapshot.addInputState(page, pagesSerde); } } @@ -509,12 +510,33 @@ public class MultiInputSnapshotState // First entry is snapshot of the operator, followed by inputs from various channels as SerializedPage instances. // Inputs contain information about channels, so no need to distinguish and store them per-channel. private final List states = new ArrayList<>(); + // Consolidated time taken to serialize the state + private long serTime; private SnapshotState(MarkerPage marker) { this.snapshotId = marker.getSnapshotId(); this.resuming = marker.isResuming(); this.markedChannels = new HashSet<>(); + this.serTime = 0; + } + + public void addState(MultiInputRestorable restorable, PagesSerde pagesSerde) + { + Stopwatch timer = Stopwatch.createStarted(); + Object state = restorable.capture(pagesSerde); + timer.stop(); + serTime += timer.elapsed(TimeUnit.MILLISECONDS); + states.add(state); + } + + public void addInputState(Page page, PagesSerde pagesSerde) + { + Stopwatch timer = Stopwatch.createStarted(); + Object channelSnapshot = pagesSerde.serialize(page).capture(pagesSerde); + timer.stop(); + serTime += timer.elapsed(TimeUnit.MILLISECONDS); + states.add(channelSnapshot); } } } 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 99c33afd9..6a86531d0 100644 --- a/presto-main/src/main/java/io/prestosql/snapshot/QuerySnapshotManager.java +++ b/presto-main/src/main/java/io/prestosql/snapshot/QuerySnapshotManager.java @@ -38,7 +38,10 @@ import java.util.OptionalLong; import java.util.Set; import java.util.Timer; import java.util.TimerTask; +import java.util.concurrent.atomic.AtomicLong; +import java.util.function.BiConsumer; import java.util.function.Consumer; +import java.util.function.Function; import java.util.stream.Collectors; import static com.google.common.base.Preconditions.checkArgument; @@ -59,11 +62,12 @@ public class QuerySnapshotManager private final Set unfinishedTasks = Sets.newConcurrentHashSet(); // LinkedHashMap can be used to keep ordering private final Map> captureComponentCounters = Collections.synchronizedMap(new LinkedHashMap<>()); - private final Map captureResults = Collections.synchronizedMap(new LinkedHashMap<>()); + private final Map captureResults = Collections.synchronizedMap(new LinkedHashMap<>()); private final Set consolidatedFilePaths = Collections.synchronizedSet(new HashSet<>()); private final Map> restoreComponentCounters = Collections.synchronizedMap(new LinkedHashMap<>()); - private final RestoreResult restoreResult = new RestoreResult(); + private RestoreResult restoreResult = new RestoreResult(); private final List> restoreCompleteListeners = Collections.synchronizedList(new ArrayList<>()); + private final List restoreStats = Collections.synchronizedList(new ArrayList<>()); private final long maxRetry; private final long retryTimeout; @@ -118,7 +122,8 @@ public class QuerySnapshotManager public void snapshotInitiated(long snapshotId) { - captureResults.put(snapshotId, SnapshotResult.IN_PROGRESS); + updateSnapshotStatus(snapshotId, SnapshotResult.IN_PROGRESS); + setSnapshotBeginTime(snapshotId, System.currentTimeMillis()); initiatedSnapshotId.add(snapshotId); } @@ -178,21 +183,22 @@ public class QuerySnapshotManager } synchronized (captureResults) { - List> entryList = new ArrayList<>(captureResults.entrySet()); + List> entryList = new ArrayList<>(captureResults.entrySet()); // iterate in reverse order for (int i = entryList.size() - 1; i >= 0; i--) { long snapshotId = entryList.get(i).getKey(); - SnapshotResult snapshotResult = entryList.get(i).getValue(); + SnapshotInfo info = entryList.get(i).getValue(); + SnapshotResult snapshotResult = info.getSnapshotResult(); // Update the snapshot result to n/a where snapshotId > resumeSnapshotId && snapshotId <= beforeThis if (snapshotId == localBeforeThis.getAsLong()) { - captureResults.put(snapshotId, SnapshotResult.NA); + updateSnapshotStatus(snapshotId, SnapshotResult.NA); } else if (snapshotId < localBeforeThis.getAsLong()) { if (snapshotResult == SnapshotResult.SUCCESSFUL) { result = OptionalLong.of(snapshotId); break; } - captureResults.put(snapshotId, SnapshotResult.NA); + updateSnapshotStatus(snapshotId, SnapshotResult.NA); } } @@ -219,7 +225,7 @@ public class QuerySnapshotManager { synchronized (captureResults) { for (Long snapshotId : captureResults.keySet()) { - captureResults.put(snapshotId, SnapshotResult.NA); + updateSnapshotStatus(snapshotId, SnapshotResult.NA); } } } @@ -244,18 +250,25 @@ public class QuerySnapshotManager return true; } - private void queryRestoreComplete(RestoreResult restoreResult) + private void queryRestoreComplete() { if (!retryTimer.isPresent()) { return; } - if (restoreResult.getSnapshotResult() == SnapshotResult.SUCCESSFUL) { + if (restoreResult.getSnapshotInfo().getSnapshotResult() == SnapshotResult.SUCCESSFUL) { + synchronized (this.restoreResult) { + SnapshotInfo info = restoreResult.getSnapshotInfo(); + info.setEndTime(System.currentTimeMillis()); + restoreStats.add(restoreResult); + } cancelRestoreTimer(); + // reset retry count on successful restore + retryCount = 0; if (lastTriedId.isPresent()) { // Successfully resumed from this snapshot id. Avoid resuming from it again. // See HashBuilderOperator#finish(), which depends on this behavior. - captureResults.put(lastTriedId.getAsLong(), SnapshotResult.FAILED); + updateSnapshotStatus(lastTriedId.getAsLong(), SnapshotResult.FAILED); lastTriedId = OptionalLong.empty(); } } @@ -324,9 +337,7 @@ public class QuerySnapshotManager // clear all maps related to this query unfinishedTasks.clear(); captureComponentCounters.clear(); - captureResults.clear(); restoreComponentCounters.clear(); - restoreResult.setSnapshotResult(0, SnapshotResult.IN_PROGRESS); restoreCompleteListeners.clear(); cancelRestoreTimer(); } @@ -337,18 +348,20 @@ public class QuerySnapshotManager } // Update capture results based on TaskInfo - public void updateQueryCapture(TaskId taskId, Map captureResult) + public void updateQueryCapture(TaskId taskId, Map captureResult) { - for (Map.Entry entry : captureResult.entrySet()) { + for (Map.Entry entry : captureResult.entrySet()) { Long snapshotId = entry.getKey(); - SnapshotResult result = entry.getValue(); + SnapshotInfo info = entry.getValue(); + SnapshotResult result = info.getSnapshotResult(); + if (snapshotId < 0) { // Special case. Task will never receive any marker. Add it to the "finished" list checkArgument(result == SnapshotResult.SUCCESSFUL); updateCapturedComponents(ImmutableList.of(taskId), false); } else { - if (updateQueryCapture(taskId, entry.getKey(), entry.getValue())) { + if (updateQueryCapture(taskId, entry.getKey(), info)) { // if the capture works, then that means a consolidated file was created and we need to add it to the list addConsolidatedFileToList(TaskSnapshotManager.createConsolidatedId(snapshotId, taskId).toString()); } @@ -357,13 +370,14 @@ public class QuerySnapshotManager } // Update capture results based on TaskSnapshotManager running on coordinator - public boolean updateQueryCapture(TaskId taskId, long snapshotId, SnapshotResult result) + public boolean updateQueryCapture(TaskId taskId, long snapshotId, SnapshotInfo snapshotInfo) { + SnapshotResult result = snapshotInfo.getSnapshotResult(); if (result == SnapshotResult.FAILED) { - return updateQueryCapture(snapshotId, taskId, SnapshotComponentCounter.ComponentState.FAILED); + return updateQueryCapture(snapshotId, taskId, snapshotInfo, SnapshotComponentCounter.ComponentState.FAILED); } else if (result == SnapshotResult.SUCCESSFUL) { - return updateQueryCapture(snapshotId, taskId, SnapshotComponentCounter.ComponentState.SUCCESSFUL); + return updateQueryCapture(snapshotId, taskId, snapshotInfo, SnapshotComponentCounter.ComponentState.SUCCESSFUL); } return false; } @@ -372,28 +386,29 @@ public class QuerySnapshotManager public void updateQueryRestore(TaskId taskId, Optional restoreResult) { if (restoreResult.isPresent()) { - SnapshotResult result = restoreResult.get().getSnapshotResult(); + SnapshotInfo snapshotInfo = restoreResult.get().getSnapshotInfo(); + SnapshotResult result = snapshotInfo.getSnapshotResult(); long snapshotId = restoreResult.get().getSnapshotId(); if (snapshotId < 0) { synchronized (restoreComponentCounters) { // Special case. Task will never receive any marker. Treat as finished. checkArgument(result == SnapshotResult.SUCCESSFUL); for (Long sid : restoreComponentCounters.keySet()) { - updateQueryRestore(sid, taskId, SnapshotComponentCounter.ComponentState.SUCCESSFUL); + updateQueryRestore(sid, taskId, snapshotInfo, SnapshotComponentCounter.ComponentState.SUCCESSFUL); } } } else { if (result == SnapshotResult.FAILED) { LOG.debug("[FATAL] Failed to resume for: " + taskId + ", snapshot " + snapshotId); - updateQueryRestore(snapshotId, taskId, SnapshotComponentCounter.ComponentState.FAILED); + updateQueryRestore(snapshotId, taskId, snapshotInfo, SnapshotComponentCounter.ComponentState.FAILED); } else if (result == SnapshotResult.FAILED_FATAL) { LOG.debug("Failed to resume for: " + taskId + ", snapshot " + snapshotId); - updateQueryRestore(snapshotId, taskId, SnapshotComponentCounter.ComponentState.FAILED_FATAL); + updateQueryRestore(snapshotId, taskId, snapshotInfo, SnapshotComponentCounter.ComponentState.FAILED_FATAL); } else if (result == SnapshotResult.SUCCESSFUL) { - updateQueryRestore(snapshotId, taskId, SnapshotComponentCounter.ComponentState.SUCCESSFUL); + updateQueryRestore(snapshotId, taskId, snapshotInfo, SnapshotComponentCounter.ComponentState.SUCCESSFUL); } } } @@ -407,9 +422,9 @@ public class QuerySnapshotManager private void saveQuerySnapshotResult() { if (!captureResults.isEmpty()) { - Map doneResult = captureResults.entrySet() + Map doneResult = captureResults.entrySet() .stream() - .filter(e -> e.getValue().isDone()) + .filter(e -> e.getValue().getSnapshotResult().isDone()) .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue, (e1, e2) -> e1, LinkedHashMap::new)); try { @@ -422,19 +437,21 @@ public class QuerySnapshotManager } } - private boolean updateQueryCapture(long snapshotId, TaskId taskId, SnapshotComponentCounter.ComponentState componentState) + private boolean updateQueryCapture(long snapshotId, TaskId taskId, SnapshotInfo snapshotInfo, SnapshotComponentCounter.ComponentState componentState) { SnapshotComponentCounter counter = captureComponentCounters.computeIfAbsent(snapshotId, k -> // A snapshot is considered complete if tasks either finished their snapshots or have completed new SnapshotComponentCounter<>(ids -> ids.containsAll(unfinishedTasks))); - if (counter.updateComponent(taskId, componentState)) { SnapshotResult snapshotResult = counter.getSnapshotResult(); synchronized (captureResults) { - if (captureResults.get(snapshotId) != SnapshotResult.NA) { + SnapshotInfo info = captureResults.get(snapshotId); + if (info.getSnapshotResult() != SnapshotResult.NA) { LOG.debug("Finished capturing snapshot %d for task %s", snapshotId, taskId); - SnapshotResult oldResult = captureResults.put(snapshotId, snapshotResult); + updateTaskCaptureStats(snapshotId, snapshotInfo); + SnapshotResult oldResult = updateSnapshotStatus(snapshotId, snapshotResult); if (snapshotResult != oldResult && snapshotResult.isDone()) { + setSnapshotEndTime(snapshotId, System.currentTimeMillis(), snapshotResult); LOG.debug("Finished capturing snapshot %d for query %s. Result is %s.", snapshotId, queryId.getId(), snapshotResult); } return true; @@ -444,7 +461,15 @@ public class QuerySnapshotManager return false; } - private void updateQueryRestore(long snapshotId, TaskId taskId, SnapshotComponentCounter.ComponentState componentState) + private void updateTaskCaptureStats(long snapshotId, SnapshotInfo snapshotInfo) + { + synchronized (captureResults) { + SnapshotInfo info = captureResults.get(snapshotId); + info.updateStats(snapshotInfo); + } + } + + private void updateQueryRestore(long snapshotId, TaskId taskId, SnapshotInfo curSnapshotInfo, SnapshotComponentCounter.ComponentState componentState) { // update queryToRestoredSnapshotComponentCounterMap SnapshotComponentCounter counter = restoreComponentCounters.computeIfAbsent(snapshotId, k -> @@ -454,6 +479,8 @@ public class QuerySnapshotManager if (counter.updateComponent(taskId, componentState)) { LOG.debug("Finished restoring snapshot %d for task %s", snapshotId, taskId); + // Update stats + updateRestoreStats(curSnapshotInfo); // update queryToRestoreReportMap; SnapshotResult snapshotResult = counter.getSnapshotResult(); boolean changed; @@ -464,17 +491,25 @@ public class QuerySnapshotManager if (snapshotResult.isDone()) { LOG.debug("Finished restoring snapshot %d for query %s. Result is %s.", snapshotId, queryId.getId(), snapshotResult); // inform the listeners(ie schedulers) if query snapshot result is finished - queryRestoreComplete(restoreResult); + queryRestoreComplete(); } else if (snapshotResult == SnapshotResult.IN_PROGRESS_FAILED || snapshotResult == SnapshotResult.IN_PROGRESS_FAILED_FATAL) { LOG.debug("Failed to restore snapshot %d for query %s. Result is %s.", snapshotId, queryId.getId(), snapshotResult); // inform the listeners(ie schedulers) if query snapshot result is finished - queryRestoreComplete(restoreResult); + queryRestoreComplete(); } } } } + private void updateRestoreStats(SnapshotInfo curSnapshotInfo) + { + synchronized (restoreResult) { + SnapshotInfo curRestoreStats = restoreResult.getSnapshotInfo(); + curRestoreStats.updateStats(curSnapshotInfo); + } + } + public int computeSnapshotIndex(OptionalLong snapshotId) { if (!snapshotId.isPresent()) { @@ -507,13 +542,44 @@ public class QuerySnapshotManager // Update ongoing snapshots for (Long snapshotId : captureComponentCounters.keySet()) { for (TaskId taskId : capturedTasks) { - updateQueryCapture(taskId, ImmutableMap.of(snapshotId, SnapshotResult.SUCCESSFUL)); + updateQueryCapture(taskId, ImmutableMap.of(snapshotId, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL))); } } } } } + private SnapshotResult updateSnapshotStatus(long snapshotId, SnapshotResult newStatus) + { + synchronized (captureResults) { + SnapshotInfo snapshotInfo = captureResults.computeIfAbsent(snapshotId, k -> getNewSnapshotInfo(k)); + SnapshotResult oldStatus = snapshotInfo.getSnapshotResult(); + snapshotInfo.setSnapshotResult(newStatus); + return oldStatus; + } + } + + private SnapshotInfo getNewSnapshotInfo(long snapshotId) + { + return new SnapshotInfo(0, 0, 0, 0, SnapshotResult.IN_PROGRESS); + } + + private void setSnapshotBeginTime(long snapshotId, long currentTimeMillis) + { + SnapshotInfo snapshotInfo = captureResults.get(snapshotId); + snapshotInfo.setBeginTime(currentTimeMillis); + } + + private void setSnapshotEndTime(long snapshotId, long currentTimeMillis, SnapshotResult snapshotResult) + { + SnapshotInfo snapshotInfo = captureResults.get(snapshotId); + snapshotInfo.setEndTime(currentTimeMillis); + // Mark snapshot as complete to show in stats, Original result is altered during restore flow + if (snapshotResult == SnapshotResult.SUCCESSFUL) { + snapshotInfo.setCompleteSnapshot(true); + } + } + @VisibleForTesting RestoreResult getQuerySnapshotRestoreResult() { @@ -534,4 +600,44 @@ public class QuerySnapshotManager rescheduler = null; } } + + public long collectSnapshotCaptureStats(Function> eachUpdater, Function> lastUpdater) + { + AtomicLong lastSnapshotId = new AtomicLong(0L); + if (!captureResults.isEmpty()) { + captureResults.forEach( + (snapshotId, snapshotInfo) -> { + if (snapshotId > 0 && snapshotInfo.isCompleteSnapshot()) { + if (snapshotId.compareTo(lastSnapshotId.get()) > 0) { + // Get last successful snapshot id + lastSnapshotId.set(snapshotId); + } + long wallTime = snapshotInfo.getEndTime() - snapshotInfo.getBeginTime(); + eachUpdater.apply(snapshotInfo.getSizeBytes()).accept(wallTime, snapshotInfo.getCpuTime()); + } + }); + // Skip if there is no successful snapshot so far + if (lastSnapshotId.get() > 0) { + SnapshotInfo lastSnapshotInfo = captureResults.get(lastSnapshotId.get()); + long wallTime = lastSnapshotInfo.getEndTime() - lastSnapshotInfo.getBeginTime(); + lastUpdater.apply(lastSnapshotInfo.getSizeBytes()).accept(wallTime, lastSnapshotInfo.getCpuTime()); + } + } + return lastSnapshotId.longValue(); + } + + public void setRestoreStartTime(long curTime) + { + // Beginning restore process, reset restore result and init with Begin time + if (retryCount == 0) { + restoreResult = new RestoreResult(); + SnapshotInfo info = restoreResult.getSnapshotInfo(); + info.setBeginTime(curTime); + } + } + + public List getRestoreStats() + { + return restoreStats; + } } diff --git a/presto-main/src/main/java/io/prestosql/snapshot/RestoreResult.java b/presto-main/src/main/java/io/prestosql/snapshot/RestoreResult.java index fdb9d90a6..0eba8bba4 100644 --- a/presto-main/src/main/java/io/prestosql/snapshot/RestoreResult.java +++ b/presto-main/src/main/java/io/prestosql/snapshot/RestoreResult.java @@ -17,25 +17,27 @@ package io.prestosql.snapshot; import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; +import static com.google.common.base.MoreObjects.toStringHelper; + /** * RestoreResult contains information of restoring process from snapshot, and report to coordinator */ public class RestoreResult { private long snapshotId; - private SnapshotResult snapshotResult; + private SnapshotInfo snapshotInfo; public RestoreResult() { - this(0, SnapshotResult.IN_PROGRESS); + this(0, SnapshotInfo.withStatus(SnapshotResult.IN_PROGRESS)); } @JsonCreator public RestoreResult(@JsonProperty("snapshotId") long snapshotId, - @JsonProperty("snapshotResult") SnapshotResult snapshotResult) + @JsonProperty("snapshotInfo") SnapshotInfo snapshotInfo) { this.snapshotId = snapshotId; - this.snapshotResult = snapshotResult; + this.snapshotInfo = snapshotInfo; } @JsonProperty @@ -45,9 +47,9 @@ public class RestoreResult } @JsonProperty - public SnapshotResult getSnapshotResult() + public SnapshotInfo getSnapshotInfo() { - return snapshotResult; + return snapshotInfo; } boolean setSnapshotResult(long snapshotId, SnapshotResult snapshotResult) @@ -57,8 +59,8 @@ public class RestoreResult this.snapshotId = snapshotId; changed = true; } - if (this.snapshotResult != snapshotResult) { - this.snapshotResult = snapshotResult; + if (this.snapshotInfo.getSnapshotResult() != snapshotResult) { + this.snapshotInfo.setSnapshotResult(snapshotResult); changed = true; } return changed; @@ -75,7 +77,16 @@ public class RestoreResult } RestoreResult that = (RestoreResult) o; return snapshotId == that.snapshotId && - snapshotResult == that.snapshotResult; + snapshotInfo.getSnapshotResult() == that.snapshotInfo.getSnapshotResult(); + } + + @Override + public String toString() + { + return toStringHelper(this) + .add("snapshotId", snapshotId) + .add("snapshotInfo", snapshotInfo) + .toString(); } @Override diff --git a/presto-main/src/main/java/io/prestosql/snapshot/SingleInputSnapshotState.java b/presto-main/src/main/java/io/prestosql/snapshot/SingleInputSnapshotState.java index 2684291e1..db53adc07 100644 --- a/presto-main/src/main/java/io/prestosql/snapshot/SingleInputSnapshotState.java +++ b/presto-main/src/main/java/io/prestosql/snapshot/SingleInputSnapshotState.java @@ -14,6 +14,7 @@ */ package io.prestosql.snapshot; +import com.google.common.base.Stopwatch; import io.airlift.log.Logger; import io.hetu.core.transport.execution.buffer.PagesSerde; import io.prestosql.memory.context.LocalMemoryContext; @@ -28,6 +29,7 @@ import java.util.LinkedList; import java.util.List; import java.util.Optional; import java.util.Queue; +import java.util.concurrent.TimeUnit; import java.util.function.Function; import static java.util.Objects.requireNonNull; @@ -122,7 +124,9 @@ public class SingleInputSnapshotState LOG.error("BUG! State of component %s has never been stored successfully before snapshot %d", restorableId, snapshotId); } else { + Stopwatch timer = Stopwatch.createStarted(); restorable.restore(state.get(), pagesSerde); + timer.stop(); boolean successful = true; if (restorable instanceof Spillable && ((Spillable) restorable).isSpilled()) { Boolean result = loadSpilledFiles(snapshotId, (Spillable) restorable); @@ -139,7 +143,7 @@ public class SingleInputSnapshotState } if (successful) { LOG.debug("Successfully restored state to snapshot %d for %s", snapshotId, restorableId); - snapshotManager.succeededToRestore(componentId); + snapshotManager.succeededToRestore(componentId, timer.elapsed(TimeUnit.MILLISECONDS)); } } // Previous pending snapshots no longer need to be carried out @@ -173,12 +177,7 @@ public class SingleInputSnapshotState return; } try { - if (restorable.supportsConsolidatedWrites()) { - snapshotManager.storeConsolidatedState(componentId, restorable.capture(pagesSerde)); - } - else { - snapshotManager.storeState(componentId, restorable.capture(pagesSerde)); - } + storeState(componentId); if (restorable instanceof Spillable && ((Spillable) restorable).isSpilled()) { storeSpilledFiles(snapshotId, (Spillable) restorable); } @@ -199,6 +198,22 @@ public class SingleInputSnapshotState } } + private void storeState(SnapshotStateId componentId) + throws Exception + { + Stopwatch timer = Stopwatch.createStarted(); + Object state = restorable.capture(pagesSerde); + timer.stop(); + long serTime = timer.elapsed(TimeUnit.MILLISECONDS); + + if (restorable.supportsConsolidatedWrites()) { + snapshotManager.storeConsolidatedState(componentId, state, serTime); + } + else { + snapshotManager.storeState(componentId, state, serTime); + } + } + public boolean hasMarker() { return !markers.isEmpty(); diff --git a/presto-main/src/main/java/io/prestosql/snapshot/SnapshotDataCollector.java b/presto-main/src/main/java/io/prestosql/snapshot/SnapshotDataCollector.java new file mode 100644 index 000000000..ce0f8872b --- /dev/null +++ b/presto-main/src/main/java/io/prestosql/snapshot/SnapshotDataCollector.java @@ -0,0 +1,38 @@ +/* + * Copyright (C) 2018-2022. Huawei Technologies Co., Ltd. All rights reserved. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.prestosql.snapshot; + +public interface SnapshotDataCollector +{ + default void updateSnapshotCaptureSize(long snapshotId, long sizeBytes) + { + throw new UnsupportedOperationException(getClass().getName() + " does not support updateSnapshotSize()"); + } + + default void updateSnapshotCaptureCpuTime(long snapshotId, long time) + { + throw new UnsupportedOperationException(getClass().getName() + " does not support updateSnapshotTime()"); + } + + default void updateSnapshotRestoreSize(long sizeBytes) + { + throw new UnsupportedOperationException(getClass().getName() + " does not support updateSnapshotRestoreSize()"); + } + + default void updateSnapshotRestoreCpuTime(long time) + { + throw new UnsupportedOperationException(getClass().getName() + " does not support updateSnapshotRestoreCpuTime()"); + } +} diff --git a/presto-main/src/main/java/io/prestosql/snapshot/SnapshotFileBasedClient.java b/presto-main/src/main/java/io/prestosql/snapshot/SnapshotFileBasedClient.java index ee6dd1400..729727b5a 100644 --- a/presto-main/src/main/java/io/prestosql/snapshot/SnapshotFileBasedClient.java +++ b/presto-main/src/main/java/io/prestosql/snapshot/SnapshotFileBasedClient.java @@ -14,6 +14,7 @@ */ package io.prestosql.snapshot; +import com.google.common.base.Stopwatch; import com.google.common.io.ByteStreams; import io.airlift.log.Logger; import io.prestosql.spi.filesystem.HetuFileSystemClient; @@ -31,6 +32,7 @@ import java.util.List; import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.concurrent.TimeUnit; /** * SnapshotStoreFileBased is an implementation of SnapshotStoreClient. @@ -53,9 +55,10 @@ public class SnapshotFileBasedClient } @Override - public void storeState(SnapshotStateId snapshotStateId, Object state) + public void storeState(SnapshotStateId snapshotStateId, Object state, SnapshotDataCollector dataCollector) throws IOException { + Stopwatch timer = Stopwatch.createStarted(); Path file = SnapshotUtils.createStatePath(rootPath, snapshotStateId.getHierarchy()); fsClient.createDirectories(file.getParent()); @@ -63,26 +66,45 @@ public class SnapshotFileBasedClient try (OutputStream outputStream = fsClient.newOutputStream(file)) { SnapshotUtils.serializeState(state, outputStream, useKryo); } + timer.stop(); + if (dataCollector != null) { + long snapshotId = snapshotStateId.getSnapshotId(); + Long size = (Long) fsClient.getAttribute(file, "size"); + if (size != null) { + dataCollector.updateSnapshotCaptureSize(snapshotId, size.longValue()); + } + dataCollector.updateSnapshotCaptureCpuTime(snapshotId, timer.elapsed(TimeUnit.MILLISECONDS)); + } } @Override - public Optional loadState(SnapshotStateId snapshotStateId) + public Optional loadState(SnapshotStateId snapshotStateId, SnapshotDataCollector dataCollector) throws IOException, ClassNotFoundException { + Optional result; + Stopwatch timer = Stopwatch.createStarted(); Path file = SnapshotUtils.createStatePath(rootPath, snapshotStateId.getHierarchy()); if (!fsClient.exists(file)) { return Optional.empty(); } try (InputStream inputStream = fsClient.newInputStream(file)) { - return Optional.of(SnapshotUtils.deserializeState(inputStream, useKryo)); + result = Optional.of(SnapshotUtils.deserializeState(inputStream, useKryo)); } + timer.stop(); + if (dataCollector != null) { + Long size = (Long) fsClient.getAttribute(file, "size"); + dataCollector.updateSnapshotRestoreSize(size.longValue()); + dataCollector.updateSnapshotRestoreCpuTime(timer.elapsed(TimeUnit.MILLISECONDS)); + } + return result; } @Override - public void storeFile(SnapshotStateId snapshotStateId, Path sourceFile) + public void storeFile(SnapshotStateId snapshotStateId, Path sourceFile, SnapshotDataCollector dataCollector) throws IOException { + Stopwatch timer = Stopwatch.createStarted(); List hierarchy = new ArrayList<>(snapshotStateId.getHierarchy()); hierarchy.add(sourceFile.getFileName().toString()); Path file = SnapshotUtils.createStatePath(rootPath, hierarchy); @@ -93,12 +115,22 @@ public class SnapshotFileBasedClient InputStream inputStream = Files.newInputStream(sourceFile)) { ByteStreams.copy(inputStream, outputStream); } + timer.stop(); + if (dataCollector != null) { + long snapshotId = snapshotStateId.getSnapshotId(); + Long size = (Long) fsClient.getAttribute(file, "size"); + if (size != null) { + dataCollector.updateSnapshotCaptureSize(snapshotId, size.longValue()); + } + dataCollector.updateSnapshotCaptureCpuTime(snapshotId, timer.elapsed(TimeUnit.MILLISECONDS)); + } } @Override - public boolean loadFile(SnapshotStateId snapshotStateId, Path targetPath) + public boolean loadFile(SnapshotStateId snapshotStateId, Path targetPath, SnapshotDataCollector dataCollector) throws IOException { + Stopwatch timer = Stopwatch.createStarted(); List hierarchy = new ArrayList<>(snapshotStateId.getHierarchy()); String fileName = targetPath.getFileName().toString(); hierarchy.add(fileName); @@ -115,6 +147,12 @@ public class SnapshotFileBasedClient OutputStream outputStream = Files.newOutputStream(targetPath)) { ByteStreams.copy(inputStream, outputStream); } + timer.stop(); + if (dataCollector != null) { + Long size = (Long) fsClient.getAttribute(file, "size"); + dataCollector.updateSnapshotRestoreSize(size.longValue()); + dataCollector.updateSnapshotRestoreCpuTime(timer.elapsed(TimeUnit.MILLISECONDS)); + } return true; } @@ -127,7 +165,7 @@ public class SnapshotFileBasedClient } @Override - public void storeSnapshotResult(String queryId, Map result) + public void storeSnapshotResult(String queryId, Map result) throws IOException { Path file = SnapshotUtils.createStatePath(rootPath, queryId, "result"); @@ -140,7 +178,7 @@ public class SnapshotFileBasedClient } @Override - public Map loadSnapshotResult(String queryId) + public Map loadSnapshotResult(String queryId) throws IOException, ClassNotFoundException { Path file = SnapshotUtils.createStatePath(rootPath, queryId, "result"); @@ -150,7 +188,7 @@ public class SnapshotFileBasedClient } try (ObjectInputStream ois = new ObjectInputStream(fsClient.newInputStream(file))) { - return (LinkedHashMap) ois.readObject(); + return (LinkedHashMap) ois.readObject(); } } diff --git a/presto-main/src/main/java/io/prestosql/snapshot/SnapshotInfo.java b/presto-main/src/main/java/io/prestosql/snapshot/SnapshotInfo.java new file mode 100644 index 000000000..4d1ec82d1 --- /dev/null +++ b/presto-main/src/main/java/io/prestosql/snapshot/SnapshotInfo.java @@ -0,0 +1,154 @@ +/* + * Copyright (C) 2018-2022. Huawei Technologies Co., Ltd. All rights reserved. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.prestosql.snapshot; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; + +import java.io.Serializable; +import java.util.concurrent.atomic.AtomicLong; + +import static com.google.common.base.MoreObjects.toStringHelper; + +public class SnapshotInfo + implements Serializable +{ + // Snapshot status + private SnapshotResult snapshotResult; + // captured/restored size in bytes (consolidated to task level) + private AtomicLong sizeBytes; + // CPU time taken while capturing/restoring the state + private AtomicLong cpuTime; + // To track snapshot capture/restore begin and end time (Wall time) at query level + private long beginTime; + private long endTime; + // Used for capture result, to mark snapshot capture was succesful or not + private boolean completeSnapshot; + + @JsonCreator + public SnapshotInfo( + @JsonProperty("sizeBytes") long sizeBytes, + @JsonProperty("cpuTime") long cpuTime, + @JsonProperty("beginTime") long beginTime, + @JsonProperty("endTime") long endTime, + @JsonProperty("snapshotResult") SnapshotResult snapshotResult) + { + this.sizeBytes = new AtomicLong(sizeBytes); + this.cpuTime = new AtomicLong(cpuTime); + this.beginTime = beginTime; + this.endTime = endTime; + this.snapshotResult = snapshotResult; + this.completeSnapshot = false; + } + + @JsonProperty + public SnapshotResult getSnapshotResult() + { + return snapshotResult; + } + + @JsonProperty + public long getSizeBytes() + { + return sizeBytes.get(); + } + + @JsonProperty + public long getCpuTime() + { + return cpuTime.get(); + } + + @JsonProperty + public long getBeginTime() + { + return beginTime; + } + + @JsonProperty + public long getEndTime() + { + return endTime; + } + + @JsonProperty + public void setSnapshotResult(SnapshotResult snapshotResult) + { + this.snapshotResult = snapshotResult; + } + + @JsonProperty + public void setBeginTime(long beginTime) + { + this.beginTime = beginTime; + } + + @JsonProperty + public void setEndTime(long endTime) + { + this.endTime = endTime; + } + + @JsonProperty + public void updateSizeBytes(long sizeBytes) + { + this.sizeBytes.addAndGet(sizeBytes); + } + + @JsonProperty + public void updateCpuTime(long cpuTime) + { + this.cpuTime.addAndGet(cpuTime); + } + + @JsonProperty + public boolean isCompleteSnapshot() + { + return completeSnapshot; + } + + @JsonProperty + public void setCompleteSnapshot(boolean restoreCompleted) + { + this.completeSnapshot = restoreCompleted; + } + + public static SnapshotInfo withStatus(SnapshotResult result) + { + SnapshotInfo info = new SnapshotInfo(0, 0, 0, 0, result); + info.setSnapshotResult(result); + return info; + } + + public void updateStats(SnapshotInfo curSnapshotInfo) + { + // Update only Size and CpuTime, which to be accumulated from task level + sizeBytes.addAndGet(curSnapshotInfo.getSizeBytes()); + cpuTime.addAndGet(curSnapshotInfo.getCpuTime()); + } + + @Override + public String toString() + { + return toStringHelper(this) + .add("snapshotResult", snapshotResult) + .add("sizeBytes", sizeBytes) + .add("cpuTime", cpuTime) + .add("beginTime", beginTime) + .add("endTime", endTime) + .add("completeSnapshot", completeSnapshot) + .toString(); + } +} diff --git a/presto-main/src/main/java/io/prestosql/snapshot/SnapshotStoreClient.java b/presto-main/src/main/java/io/prestosql/snapshot/SnapshotStoreClient.java index 81c69d368..a5881807b 100644 --- a/presto-main/src/main/java/io/prestosql/snapshot/SnapshotStoreClient.java +++ b/presto-main/src/main/java/io/prestosql/snapshot/SnapshotStoreClient.java @@ -27,25 +27,25 @@ public interface SnapshotStoreClient /** * Store state in snapshot store */ - void storeState(SnapshotStateId snapshotStateId, Object state) + void storeState(SnapshotStateId snapshotStateId, Object state, SnapshotDataCollector dataCollector) throws Exception; /** * Load state from snapshot store. Optional.empty() is returned if state doesn't exist. */ - Optional loadState(SnapshotStateId snapshotStateId) + Optional loadState(SnapshotStateId snapshotStateId, SnapshotDataCollector dataCollector) throws Exception; /** * Store file from sourcePath to snapshotStateId of snapshot store */ - void storeFile(SnapshotStateId snapshotStateId, Path sourcePath) + void storeFile(SnapshotStateId snapshotStateId, Path sourcePath, SnapshotDataCollector dataCollector) throws Exception; /** * Load file from snapshotStateId of snapshot store to targetPath */ - boolean loadFile(SnapshotStateId snapshotStateId, Path targetPath) + boolean loadFile(SnapshotStateId snapshotStateId, Path targetPath, SnapshotDataCollector dataCollector) throws Exception; /** @@ -57,13 +57,13 @@ public interface SnapshotStoreClient /** * Store snapshot result of query */ - void storeSnapshotResult(String queryId, Map result) + void storeSnapshotResult(String queryId, Map result) throws Exception; /** * Load snapshot result of query */ - Map loadSnapshotResult(String queryId) + Map loadSnapshotResult(String queryId) throws Exception; /** 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 723e12c64..db4b1e928 100644 --- a/presto-main/src/main/java/io/prestosql/snapshot/SnapshotUtils.java +++ b/presto-main/src/main/java/io/prestosql/snapshot/SnapshotUtils.java @@ -146,13 +146,13 @@ public class SnapshotUtils /** * Store the state of snapshotStateId in snapshot store */ - public void storeState(SnapshotStateId snapshotStateId, Object state) + public void storeState(SnapshotStateId snapshotStateId, Object state, SnapshotDataCollector dataCollector) throws Exception { requireNonNull(snapshotStoreClient); requireNonNull(state); - snapshotStoreClient.storeState(snapshotStateId, state); + snapshotStoreClient.storeState(snapshotStateId, state, dataCollector); } /** @@ -161,38 +161,38 @@ public class SnapshotUtils * - NO_STATE: bug situation * - Other object: previously saved state */ - public Optional loadState(SnapshotStateId snapshotStateId) + public Optional loadState(SnapshotStateId snapshotStateId, SnapshotDataCollector dataCollector) throws Exception { requireNonNull(snapshotStoreClient); - return snapshotStoreClient.loadState(snapshotStateId); + return snapshotStoreClient.loadState(snapshotStateId, dataCollector); } - public void storeFile(SnapshotStateId snapshotStateId, Path sourceFile) + public void storeFile(SnapshotStateId snapshotStateId, Path sourceFile, SnapshotDataCollector dataCollector) throws Exception { requireNonNull(snapshotStoreClient); requireNonNull(sourceFile); - snapshotStoreClient.storeFile(snapshotStateId, sourceFile); + snapshotStoreClient.storeFile(snapshotStateId, sourceFile, dataCollector); } - public Boolean loadFile(SnapshotStateId snapshotStateId, Path targetFile) + public Boolean loadFile(SnapshotStateId snapshotStateId, Path targetFile, SnapshotDataCollector dataCollector) throws Exception { requireNonNull(snapshotStoreClient); requireNonNull(targetFile); - return snapshotStoreClient.loadFile(snapshotStateId, targetFile); + return snapshotStoreClient.loadFile(snapshotStateId, targetFile, dataCollector); } - public void storeSnapshotResult(String queryId, Map result) + public void storeSnapshotResult(String queryId, Map result) throws Exception { snapshotStoreClient.storeSnapshotResult(queryId, result); } - public Map loadSnapshotResult(String queryId) + public Map loadSnapshotResult(String queryId) throws Exception { return snapshotStoreClient.loadSnapshotResult(queryId); diff --git a/presto-main/src/main/java/io/prestosql/snapshot/TaskSnapshotManager.java b/presto-main/src/main/java/io/prestosql/snapshot/TaskSnapshotManager.java index 2fd95b290..8555d7fe2 100644 --- a/presto-main/src/main/java/io/prestosql/snapshot/TaskSnapshotManager.java +++ b/presto-main/src/main/java/io/prestosql/snapshot/TaskSnapshotManager.java @@ -40,6 +40,7 @@ import static java.util.Objects.requireNonNull; * TaskSnapshotManager keeps track of snapshot status of task components */ public class TaskSnapshotManager + implements SnapshotDataCollector { private static final Logger LOG = Logger.get(TaskSnapshotManager.class); public static final Object NO_STATE = new Object(); @@ -52,7 +53,7 @@ public class TaskSnapshotManager private int totalComponents = -1; // LinkedHashMap can be used to keep ordering private final Map> captureComponentCounters = Collections.synchronizedMap(new LinkedHashMap<>()); - private final Map captureResults = new LinkedHashMap<>(); + private final Map captureResults = new LinkedHashMap<>(); private final Map> restoreComponentCounters = Collections.synchronizedMap(new LinkedHashMap<>()); private final RestoreResult restoreResult = new RestoreResult(); @@ -86,22 +87,25 @@ public class TaskSnapshotManager return SnapshotStateId.forTaskComponent(snapshotId, taskId, CONSOLIDATED_STATE_COMPONENT); } - public void storeConsolidatedState(SnapshotStateId snapshotStateId, Object state) + public void storeConsolidatedState(SnapshotStateId snapshotStateId, Object state, long serCpuTime) { Map map = storeCache.computeIfAbsent(snapshotStateId.getSnapshotId(), (x) -> Collections.synchronizedMap(new HashMap<>())); map.put(snapshotStateId.toString(), state); + updateSnapshotCaptureCpuTime(snapshotStateId.getSnapshotId(), serCpuTime); } /** * Store the state of snapshotStateId in snapshot store */ - public void storeState(SnapshotStateId snapshotStateId, Object state) + public void storeState(SnapshotStateId snapshotStateId, Object state, long serCpuTime) throws Exception { - snapshotUtils.storeState(snapshotStateId, state); + snapshotUtils.storeState(snapshotStateId, state, this); + // store dummy value Map map = storeCache.computeIfAbsent(snapshotStateId.getSnapshotId(), (x) -> Collections.synchronizedMap(new HashMap<>())); map.put(snapshotStateId.toString(), snapshotStateId.toString()); + updateSnapshotCaptureCpuTime(snapshotStateId.getSnapshotId(), serCpuTime); } private void loadMapIfNecessary(long snapshotId, TaskId taskId) @@ -113,7 +117,7 @@ public class TaskSnapshotManager if (!loadCache.containsKey(snapshotId)) { String queryId = taskId.getQueryId().getId(); SnapshotStateId stateId = createConsolidatedId(snapshotId, taskId); - Optional loadedState = snapshotUtils.loadState(stateId); + Optional loadedState = snapshotUtils.loadState(stateId, this); if (createdConsolidatedFiles == null) { createdConsolidatedFiles = snapshotUtils.loadConsolidatedFiles(queryId); } @@ -143,7 +147,7 @@ public class TaskSnapshotManager Optional state; loadMapIfNecessary(snapshotId, snapshotStateIdTaskId); state = Optional.ofNullable(loadCache.get(snapshotId).get(newSnapshotStateId.toString())); - Map snapshotToSnapshotResultMap = null; + Map snapshotToSnapshotResultMap = null; while (!state.isPresent()) { // Snapshot is complete but no entry for this id, then the component must have finished // before the snapshot was taken. Look at previous complete snapshots for last saved state. @@ -186,7 +190,7 @@ public class TaskSnapshotManager { Optional loadedValue = loadWithBacktrack(snapshotStateId); if (loadedValue.isPresent() && loadedValue.get() != NO_STATE) { - return snapshotUtils.loadState(SnapshotStateId.fromString((String) loadedValue.get())); + return snapshotUtils.loadState(SnapshotStateId.fromString((String) loadedValue.get()), this); } return loadedValue; } @@ -194,7 +198,7 @@ public class TaskSnapshotManager public void storeFile(SnapshotStateId snapshotStateId, Path sourceFile) throws Exception { - snapshotUtils.storeFile(snapshotStateId, sourceFile); + snapshotUtils.storeFile(snapshotStateId, sourceFile, this); // store dummy value Map map = storeCache.computeIfAbsent(snapshotStateId.getSnapshotId(), (x) -> Collections.synchronizedMap(new HashMap<>())); map.put(snapshotStateId.toString(), snapshotStateId.toString()); @@ -212,16 +216,16 @@ public class TaskSnapshotManager if (loadedValue.get() == NO_STATE) { return null; } - return snapshotUtils.loadFile(SnapshotStateId.fromString((String) loadedValue.get()), targetFile); + return snapshotUtils.loadFile(SnapshotStateId.fromString((String) loadedValue.get()), targetFile, this); } - private OptionalLong getPreviousSnapshotIdIfComplete(Map snapshotToSnapshotResultMap, long snapshotId) + private OptionalLong getPreviousSnapshotIdIfComplete(Map snapshotToSnapshotResultMap, long snapshotId) { try { - List> entryList = new ArrayList<>(snapshotToSnapshotResultMap.entrySet()); + List> entryList = new ArrayList<>(snapshotToSnapshotResultMap.entrySet()); for (int i = entryList.size() - 1; i >= 0; i--) { long sId = entryList.get(i).getKey(); - SnapshotResult snapshotRestoreResult = entryList.get(i).getValue(); + SnapshotResult snapshotRestoreResult = entryList.get(i).getValue().getSnapshotResult(); if (sId < snapshotId) { if (snapshotRestoreResult == SnapshotResult.SUCCESSFUL) { return OptionalLong.of(sId); @@ -260,8 +264,9 @@ public class TaskSnapshotManager updateCapture(componentId, SnapshotComponentCounter.ComponentState.FAILED); } - public void succeededToRestore(SnapshotStateId componentId) + public void succeededToRestore(SnapshotStateId componentId, long deserCpuTime) { + updateSnapshotRestoreCpuTime(deserCpuTime); updateRestore(componentId, SnapshotComponentCounter.ComponentState.SUCCESSFUL); } @@ -276,13 +281,13 @@ public class TaskSnapshotManager } } - public Map getSnapshotCaptureResult() + public Map getSnapshotCaptureResult() { if (totalComponents == 0) { // Special case: don't expect any more markers from this task. // It's as if this task has finished. // Use -1 to indicate "all snapshots". - return ImmutableMap.of(-1L, SnapshotResult.SUCCESSFUL); + return ImmutableMap.of(-1L, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL)); } // Need to make a copy, otherwise there may be concurrent modification errors synchronized (captureResults) { @@ -308,7 +313,7 @@ public class TaskSnapshotManager // update capturedSnapshotResultMap SnapshotResult snapshotResult = counter.getSnapshotResult(); synchronized (captureResults) { - SnapshotResult oldResult = captureResults.put(snapshotId, snapshotResult); + SnapshotResult oldResult = updateSnapshotStatus(snapshotId, snapshotResult); if (snapshotResult != oldResult && snapshotResult.isDone()) { if (snapshotResult == SnapshotResult.SUCCESSFUL) { // All components for the task have captured their states successfully. @@ -323,12 +328,12 @@ public class TaskSnapshotManager else { map = Collections.emptyMap(); } - snapshotUtils.storeState(newId, map); + snapshotUtils.storeState(newId, map, this); } catch (Exception e) { LOG.error(e, "Failed to store state for " + newId); snapshotResult = SnapshotResult.FAILED; - captureResults.put(snapshotId, snapshotResult); + updateSnapshotStatus(snapshotId, snapshotResult); } } if (snapshotUtils.isCoordinator()) { @@ -338,7 +343,8 @@ public class TaskSnapshotManager if (snapshotResult == SnapshotResult.SUCCESSFUL) { querySnapshotManager.addConsolidatedFileToList(createConsolidatedId(snapshotId, componentIdTaskId).toString()); } - querySnapshotManager.updateQueryCapture(componentIdTaskId, snapshotId, snapshotResult); + SnapshotInfo snapshotInfo = captureResults.get(snapshotId); + querySnapshotManager.updateQueryCapture(componentIdTaskId, snapshotId, snapshotInfo); } } LOG.debug("Finished capturing snapshot %d for task %s. Result is %s.", snapshotId, componentIdTaskId, snapshotResult); @@ -400,6 +406,55 @@ public class TaskSnapshotManager checkState(totalComponents >= 0); } + @Override + public void updateSnapshotCaptureSize(long snapshotId, long sizeBytes) + { + synchronized (captureResults) { + SnapshotInfo snapshotInfo = captureResults.computeIfAbsent(snapshotId, k -> getNewSnapshotInfo(k)); + snapshotInfo.updateSizeBytes(sizeBytes); + } + } + + @Override + public void updateSnapshotCaptureCpuTime(long snapshotId, long time) + { + synchronized (captureResults) { + SnapshotInfo snapshotInfo = captureResults.computeIfAbsent(snapshotId, k -> getNewSnapshotInfo(k)); + snapshotInfo.updateCpuTime(time); + } + } + + @Override + public void updateSnapshotRestoreSize(long sizeBytes) + { + synchronized (restoreResult) { + SnapshotInfo snapshotInfo = restoreResult.getSnapshotInfo(); + snapshotInfo.updateSizeBytes(sizeBytes); + } + } + + @Override + public void updateSnapshotRestoreCpuTime(long time) + { + synchronized (restoreResult) { + SnapshotInfo snapshotInfo = restoreResult.getSnapshotInfo(); + snapshotInfo.updateCpuTime(time); + } + } + + private SnapshotResult updateSnapshotStatus(long snapshotId, SnapshotResult newStatus) + { + SnapshotInfo snapshotInfo = captureResults.computeIfAbsent(snapshotId, k -> getNewSnapshotInfo(k)); + SnapshotResult oldStatus = snapshotInfo.getSnapshotResult(); + snapshotInfo.setSnapshotResult(newStatus); + return oldStatus; + } + + private SnapshotInfo getNewSnapshotInfo(long snapshotId) + { + return new SnapshotInfo(0, 0, 0, 0, SnapshotResult.IN_PROGRESS); + } + @Override public String toString() { diff --git a/presto-main/src/test/java/io/prestosql/execution/buffer/TestPartitionedOutputBuffer.java b/presto-main/src/test/java/io/prestosql/execution/buffer/TestPartitionedOutputBuffer.java index 0d31d6e50..f66a90c69 100644 --- a/presto-main/src/test/java/io/prestosql/execution/buffer/TestPartitionedOutputBuffer.java +++ b/presto-main/src/test/java/io/prestosql/execution/buffer/TestPartitionedOutputBuffer.java @@ -24,6 +24,7 @@ import io.prestosql.operator.PageAssertions; import io.prestosql.operator.TaskContext; import io.prestosql.snapshot.SnapshotStateId; import io.prestosql.snapshot.SnapshotUtils; +import io.prestosql.snapshot.TaskSnapshotManager; import io.prestosql.spi.Page; import io.prestosql.spi.snapshot.MarkerPage; import io.prestosql.spi.snapshot.SnapshotTestUtil; @@ -331,13 +332,15 @@ public class TestPartitionedOutputBuffer ArgumentCaptor idArgument = ArgumentCaptor.forClass(SnapshotStateId.class); ArgumentCaptor stateArgument = ArgumentCaptor.forClass(Object.class); + ArgumentCaptor collectorArgument = ArgumentCaptor.forClass(TaskSnapshotManager.class); // storeState is called once for each partition - verify(snapshotUtils, times(3)).storeState(idArgument.capture(), stateArgument.capture()); + verify(snapshotUtils, times(3)).storeState(idArgument.capture(), stateArgument.capture(), collectorArgument.capture()); List ids = idArgument.getAllValues(); List states = stateArgument.getAllValues(); - when(snapshotUtils.loadState(ids.get(0))).thenReturn(Optional.of(states.get(0))); - when(snapshotUtils.loadState(ids.get(1))).thenReturn(Optional.of(states.get(1))); - when(snapshotUtils.loadState(ids.get(2))).thenReturn(Optional.of(states.get(2))); + List snapshotManagers = collectorArgument.getAllValues(); + when(snapshotUtils.loadState(ids.get(0), snapshotManagers.get(0))).thenReturn(Optional.of(states.get(0))); + when(snapshotUtils.loadState(ids.get(1), snapshotManagers.get(1))).thenReturn(Optional.of(states.get(1))); + when(snapshotUtils.loadState(ids.get(2), snapshotManagers.get(2))).thenReturn(Optional.of(states.get(2))); buffer = createPartitionedBuffer( createInitialEmptyOutputBuffers(PARTITIONED) @@ -352,7 +355,7 @@ public class TestPartitionedOutputBuffer // Resume both partitions buffer.enqueue(firstPartition, ImmutableList.of(PAGES_SERDE.serialize(resume)), channel1); - verify(snapshotUtils, times(3)).loadState(anyObject()); + verify(snapshotUtils, times(3)).loadState(anyObject(), anyObject()); // Newly added page (page2) should be received after the resume marker buffer.enqueue(firstPartition, ImmutableList.of(PAGES_SERDE.serialize(page2)), channel1); diff --git a/presto-main/src/test/java/io/prestosql/operator/TestPartitionedOutputOperator.java b/presto-main/src/test/java/io/prestosql/operator/TestPartitionedOutputOperator.java index 047b83b30..380558dd3 100644 --- a/presto-main/src/test/java/io/prestosql/operator/TestPartitionedOutputOperator.java +++ b/presto-main/src/test/java/io/prestosql/operator/TestPartitionedOutputOperator.java @@ -86,15 +86,14 @@ public class TestPartitionedOutputOperator operator.addInput(input.get(0)); operator.addInput(marker); ArgumentCaptor stateArgument = ArgumentCaptor.forClass(Object.class); - verify(snapshotUtils, times(1)).storeState(anyObject(), stateArgument.capture()); + verify(snapshotUtils, times(1)).storeState(anyObject(), stateArgument.capture(), anyObject()); Object snapshot = stateArgument.getValue(); - - when(snapshotUtils.loadState(anyObject())).thenReturn(Optional.of(snapshot)); + when(snapshotUtils.loadState(anyObject(), anyObject())).thenReturn(Optional.of(snapshot)); operator.addInput(input.get(1)); operator.addInput(resume); operator.addInput(marker2); - verify(snapshotUtils, times(2)).storeState(anyObject(), stateArgument.capture()); + verify(snapshotUtils, times(2)).storeState(anyObject(), stateArgument.capture(), anyObject()); snapshot = stateArgument.getValue(); Object snapshotEntry = ((Map) snapshot).get("query/2/1/1/0/0/0"); assertEquals(SnapshotTestUtil.toFullSnapshotMapping(snapshotEntry), createExpectedMappingBeforeFinish()); @@ -102,7 +101,7 @@ public class TestPartitionedOutputOperator operator.addInput(input.get(1)); operator.finish(); operator.addInput(marker3); - verify(snapshotUtils, times(3)).storeState(anyObject(), stateArgument.capture()); + verify(snapshotUtils, times(3)).storeState(anyObject(), stateArgument.capture(), anyObject()); snapshot = stateArgument.getValue(); snapshotEntry = ((Map) snapshot).get("query/3/1/1/0/0/0"); assertEquals(SnapshotTestUtil.toFullSnapshotMapping(snapshotEntry), createExpectedMappingAfterFinish()); diff --git a/presto-main/src/test/java/io/prestosql/snapshot/TestMultiInputSnapshotState.java b/presto-main/src/test/java/io/prestosql/snapshot/TestMultiInputSnapshotState.java index 62842c81c..5a0e1e7f3 100644 --- a/presto-main/src/test/java/io/prestosql/snapshot/TestMultiInputSnapshotState.java +++ b/presto-main/src/test/java/io/prestosql/snapshot/TestMultiInputSnapshotState.java @@ -79,6 +79,7 @@ public class TestMultiInputSnapshotState private MultiInputSnapshotState state; private final ArgumentCaptor argument = ArgumentCaptor.forClass(List.class); + private final ArgumentCaptor timeArgument = ArgumentCaptor.forClass(Long.class); @BeforeMethod public void setup() @@ -166,7 +167,7 @@ public class TestMultiInputSnapshotState ret = processPage(source2, marker1); assertFalse(ret.isPresent()); - verify(snapshotManager).storeState(eq(snapshotId1), argument.capture()); + verify(snapshotManager).storeState(eq(snapshotId1), argument.capture(), timeArgument.capture()); List savedState = argument.getValue(); assertEquals(savedState.size(), 2); assertEquals(savedState.get(0), saved); @@ -190,7 +191,7 @@ public class TestMultiInputSnapshotState ret = processSerializedPage(source2, serializedMarker); assertFalse(ret.isPresent()); - verify(snapshotManager).storeState(eq(snapshotId1), argument.capture()); + verify(snapshotManager).storeState(eq(snapshotId1), argument.capture(), timeArgument.capture()); List savedState = argument.getValue(); assertEquals(savedState.size(), 2); assertEquals(savedState.get(0), saved); @@ -213,7 +214,7 @@ public class TestMultiInputSnapshotState processPage(source2, regularPage); processPage(source2, marker1); - verify(snapshotManager).storeState(eq(snapshotId1), argument.capture()); + verify(snapshotManager).storeState(eq(snapshotId1), argument.capture(), timeArgument.capture()); List savedState = argument.getValue(); assertEquals(savedState.size(), 2); assertEquals(savedState.get(0), saved1); @@ -223,7 +224,7 @@ public class TestMultiInputSnapshotState ret = processPage(source2, marker2); assertFalse(ret.isPresent()); - verify(snapshotManager).storeState(eq(snapshotId2), argument.capture()); + verify(snapshotManager).storeState(eq(snapshotId2), argument.capture(), timeArgument.capture()); savedState = argument.getValue(); assertEquals(savedState.size(), 3); assertEquals(savedState.get(0), saved2); @@ -243,7 +244,7 @@ public class TestMultiInputSnapshotState processPage(source2, regularPage); processPage(source2, regularPage); processPage(source2, marker1); - verify(snapshotManager).storeState(eq(snapshotId1), argument.capture()); + verify(snapshotManager).storeState(eq(snapshotId1), argument.capture(), timeArgument.capture()); when(snapshotManager.loadState(snapshotId1)).thenReturn(Optional.of(argument.getValue())); Optional ret = processPageKeepState(source2, resume1); @@ -268,7 +269,7 @@ public class TestMultiInputSnapshotState processPage(inputSnapshotState, source1, marker1); processPage(inputSnapshotState, source2, marker1); - verify(snapshotManager, never()).storeState(anyObject(), anyObject()); + verify(snapshotManager, never()).storeState(anyObject(), anyObject(), timeArgument.capture()); } @Test(expectedExceptions = IllegalStateException.class) @@ -290,7 +291,7 @@ public class TestMultiInputSnapshotState { processPage(source1, marker1); processPage(source2, marker1); - verify(snapshotManager).storeState(eq(snapshotId1), argument.capture()); + verify(snapshotManager).storeState(eq(snapshotId1), argument.capture(), timeArgument.capture()); when(snapshotManager.loadState(snapshotId1)).thenReturn(Optional.of(argument.getValue())); processPage(source1, resume1); @@ -305,7 +306,7 @@ public class TestMultiInputSnapshotState { processPage(source1, marker1); processPage(source2, marker1); - verify(snapshotManager).storeState(eq(snapshotId1), argument.capture()); + verify(snapshotManager).storeState(eq(snapshotId1), argument.capture(), timeArgument.capture()); when(snapshotManager.loadState(snapshotId1)).thenReturn(Optional.of(argument.getValue())); processPage(source1, resume1); @@ -351,7 +352,7 @@ public class TestMultiInputSnapshotState assertEquals(pages.get(0), regularPage); assertEquals(pages.get(1), regularPage); - verify(snapshotManager).storeState(eq(snapshotId1), argument.capture()); + verify(snapshotManager).storeState(eq(snapshotId1), argument.capture(), timeArgument.capture()); when(snapshotManager.loadState(snapshotId1)).thenReturn(Optional.of(argument.getValue())); pages = processPages(source1, Arrays.asList(resume1)); @@ -445,7 +446,7 @@ public class TestMultiInputSnapshotState ret = processPage(null, resume1).get(); assertEquals(ret, resume1); - verify(snapshotManager, never()).storeState(anyObject(), anyObject()); + verify(snapshotManager, never()).storeState(anyObject(), anyObject(), timeArgument.capture()); verify(snapshotManager, never()).loadState(anyObject()); } @@ -458,8 +459,8 @@ public class TestMultiInputSnapshotState restorable.setSupportsConsolidatedWrites(true); processPage(state, source1, marker1); processPage(state, source2, marker1); - verify(snapshotManager, never()).storeState(anyObject(), anyObject()); - verify(snapshotManager, times(1)).storeConsolidatedState(anyObject(), argument.capture()); + verify(snapshotManager, never()).storeState(anyObject(), anyObject(), timeArgument.capture()); + verify(snapshotManager, times(1)).storeConsolidatedState(anyObject(), argument.capture(), timeArgument.capture()); when(snapshotManager.loadConsolidatedState(anyObject())).thenReturn(Optional.of(argument.getValue())); processPage(state, source1, resume1); @@ -476,8 +477,8 @@ public class TestMultiInputSnapshotState restorable.setSupportsConsolidatedWrites(false); processPage(state, source1, marker1); processPage(state, source2, marker1); - verify(snapshotManager, times(0)).storeConsolidatedState(anyObject(), anyObject()); - verify(snapshotManager, times(1)).storeState(anyObject(), argument.capture()); + verify(snapshotManager, times(0)).storeConsolidatedState(anyObject(), anyObject(), timeArgument.capture()); + verify(snapshotManager, times(1)).storeState(anyObject(), argument.capture(), timeArgument.capture()); when(snapshotManager.loadState(anyObject())).thenReturn(Optional.of(argument.getValue())); processPage(state, source1, resume1); diff --git a/presto-main/src/test/java/io/prestosql/snapshot/TestQuerySnapshotManager.java b/presto-main/src/test/java/io/prestosql/snapshot/TestQuerySnapshotManager.java index 8c120522d..47f7ac065 100644 --- a/presto-main/src/test/java/io/prestosql/snapshot/TestQuerySnapshotManager.java +++ b/presto-main/src/test/java/io/prestosql/snapshot/TestQuerySnapshotManager.java @@ -87,6 +87,7 @@ public class TestQuerySnapshotManager queryId = new QueryId("resumeid"); QuerySnapshotManager snapshotManager = new QuerySnapshotManager(queryId, snapshotUtils, TEST_SNAPSHOT_SESSION); TaskId taskId = new TaskId(queryId.getId(), 2, 3); + SnapshotInfo info = SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL); // Try1: no id is available yet snapshotManager.addNewTask(taskId); @@ -95,22 +96,24 @@ public class TestQuerySnapshotManager // Try2: setup some successful snapshots snapshotManager.addNewTask(taskId); - snapshotManager.updateQueryCapture(taskId, Collections.singletonMap(1L, SnapshotResult.SUCCESSFUL)); - snapshotManager.updateQueryCapture(taskId, Collections.singletonMap(2L, SnapshotResult.SUCCESSFUL)); + snapshotManager.snapshotInitiated(1L); + snapshotManager.updateQueryCapture(taskId, Collections.singletonMap(1L, info)); + snapshotManager.snapshotInitiated(2L); + snapshotManager.updateQueryCapture(taskId, Collections.singletonMap(2L, info)); sid = snapshotManager.getResumeSnapshotId(); assertEquals(sid.getAsLong(), 2); // Try3: get available snapshot before 2 snapshotManager.addNewTask(taskId); - snapshotManager.updateQueryCapture(taskId, Collections.singletonMap(1L, SnapshotResult.SUCCESSFUL)); - snapshotManager.updateQueryCapture(taskId, Collections.singletonMap(2L, SnapshotResult.SUCCESSFUL)); + snapshotManager.updateQueryCapture(taskId, Collections.singletonMap(1L, info)); + snapshotManager.updateQueryCapture(taskId, Collections.singletonMap(2L, info)); sid = snapshotManager.getResumeSnapshotId(); assertEquals(sid.getAsLong(), 1); // Try4: get available snapshot before 1 snapshotManager.addNewTask(taskId); - snapshotManager.updateQueryCapture(taskId, Collections.singletonMap(1L, SnapshotResult.SUCCESSFUL)); - snapshotManager.updateQueryCapture(taskId, Collections.singletonMap(2L, SnapshotResult.SUCCESSFUL)); + snapshotManager.updateQueryCapture(taskId, Collections.singletonMap(1L, info)); + snapshotManager.updateQueryCapture(taskId, Collections.singletonMap(2L, info)); sid = snapshotManager.getResumeSnapshotId(); assertFalse(sid.isPresent()); } @@ -125,9 +128,10 @@ public class TestQuerySnapshotManager snapshotManager.setRescheduler(rescheduler); snapshotManager.addNewTask(taskId); - snapshotManager.updateQueryCapture(taskId, Collections.singletonMap(1L, SnapshotResult.SUCCESSFUL)); + snapshotManager.snapshotInitiated(1L); + snapshotManager.updateQueryCapture(taskId, Collections.singletonMap(1L, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL))); OptionalLong sid = snapshotManager.getResumeSnapshotId(); - snapshotManager.updateQueryRestore(taskId, Optional.of(new RestoreResult(sid.getAsLong(), SnapshotResult.FAILED))); + snapshotManager.updateQueryRestore(taskId, Optional.of(new RestoreResult(sid.getAsLong(), SnapshotInfo.withStatus(SnapshotResult.FAILED)))); verify(rescheduler).run(); } @@ -146,7 +150,8 @@ public class TestQuerySnapshotManager snapshotManager.setRescheduler(() -> future.set(null)); snapshotManager.addNewTask(taskId); - snapshotManager.updateQueryCapture(taskId, Collections.singletonMap(1L, SnapshotResult.SUCCESSFUL)); + snapshotManager.snapshotInitiated(1L); + snapshotManager.updateQueryCapture(taskId, Collections.singletonMap(1L, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL))); snapshotManager.getResumeSnapshotId(); future.get(1, TimeUnit.SECONDS); @@ -166,8 +171,10 @@ public class TestQuerySnapshotManager snapshotManager.setRescheduler(rescheduler); snapshotManager.addNewTask(taskId); - snapshotManager.updateQueryCapture(taskId, Collections.singletonMap(1L, SnapshotResult.SUCCESSFUL)); - snapshotManager.updateQueryCapture(taskId, Collections.singletonMap(2L, SnapshotResult.SUCCESSFUL)); + snapshotManager.snapshotInitiated(1L); + snapshotManager.updateQueryCapture(taskId, Collections.singletonMap(1L, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL))); + snapshotManager.snapshotInitiated(2L); + snapshotManager.updateQueryCapture(taskId, Collections.singletonMap(2L, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL))); assertTrue(snapshotManager.getResumeSnapshotId().isPresent()); try { snapshotManager.getResumeSnapshotId(); @@ -185,15 +192,18 @@ public class TestQuerySnapshotManager QuerySnapshotManager snapshotManager = new QuerySnapshotManager(queryId, snapshotUtils, TEST_SNAPSHOT_SESSION); TaskId taskId = new TaskId(queryId.getId(), 2, 3); snapshotManager.addNewTask(taskId); + snapshotManager.snapshotInitiated(1L); - snapshotManager.updateQueryRestore(taskId, Optional.of(new RestoreResult(1, SnapshotResult.SUCCESSFUL))); - assertEquals(snapshotManager.getQuerySnapshotRestoreResult().getSnapshotResult(), SnapshotResult.SUCCESSFUL); + snapshotManager.updateQueryRestore(taskId, Optional.of(new RestoreResult(1, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL)))); + assertEquals(snapshotManager.getQuerySnapshotRestoreResult().getSnapshotInfo().getSnapshotResult(), SnapshotResult.SUCCESSFUL); - snapshotManager.updateQueryRestore(taskId, Optional.of(new RestoreResult(2, SnapshotResult.FAILED))); - assertEquals(snapshotManager.getQuerySnapshotRestoreResult().getSnapshotResult(), SnapshotResult.FAILED); + snapshotManager.snapshotInitiated(2L); + snapshotManager.updateQueryRestore(taskId, Optional.of(new RestoreResult(2, SnapshotInfo.withStatus(SnapshotResult.FAILED)))); + assertEquals(snapshotManager.getQuerySnapshotRestoreResult().getSnapshotInfo().getSnapshotResult(), SnapshotResult.FAILED); - snapshotManager.updateQueryRestore(taskId, Optional.of(new RestoreResult(3, SnapshotResult.FAILED_FATAL))); - assertEquals(snapshotManager.getQuerySnapshotRestoreResult().getSnapshotResult(), SnapshotResult.FAILED_FATAL); + snapshotManager.snapshotInitiated(3L); + snapshotManager.updateQueryRestore(taskId, Optional.of(new RestoreResult(3, SnapshotInfo.withStatus(SnapshotResult.FAILED_FATAL)))); + assertEquals(snapshotManager.getQuerySnapshotRestoreResult().getSnapshotInfo().getSnapshotResult(), SnapshotResult.FAILED_FATAL); } @Test @@ -206,9 +216,10 @@ public class TestQuerySnapshotManager TaskId taskId2 = new TaskId(queryId.getId(), 3, 4); snapshotManager.addNewTask(taskId1); snapshotManager.addNewTask(taskId2); + snapshotManager.snapshotInitiated(1L); snapshotManager.updateFinishedQueryComponents(ImmutableList.of(taskId2)); - snapshotManager.updateQueryCapture(taskId1, Collections.singletonMap(1L, SnapshotResult.SUCCESSFUL)); + snapshotManager.updateQueryCapture(taskId1, Collections.singletonMap(1L, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL))); assertEquals(snapshotManager.getResumeSnapshotId().getAsLong(), 1); } @@ -220,7 +231,8 @@ public class TestQuerySnapshotManager TaskId taskId1 = new TaskId(queryId.getId(), 2, 3); snapshotManager.addNewTask(taskId1); - snapshotManager.updateQueryCapture(taskId1, Collections.singletonMap(1L, SnapshotResult.SUCCESSFUL)); + snapshotManager.snapshotInitiated(1L); + snapshotManager.updateQueryCapture(taskId1, Collections.singletonMap(1L, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL))); assertTrue(snapshotManager.getResumeSnapshotId().isPresent()); snapshotManager.invalidateAllSnapshots(); diff --git a/presto-main/src/test/java/io/prestosql/snapshot/TestRestoreResult.java b/presto-main/src/test/java/io/prestosql/snapshot/TestRestoreResult.java index 5c66b3010..74f7e4bf6 100644 --- a/presto-main/src/test/java/io/prestosql/snapshot/TestRestoreResult.java +++ b/presto-main/src/test/java/io/prestosql/snapshot/TestRestoreResult.java @@ -30,13 +30,13 @@ public class TestRestoreResult RestoreResult restoreResult = new RestoreResult(); restoreResult.setSnapshotResult(snapshotId, SnapshotResult.IN_PROGRESS); Assert.assertEquals(restoreResult.getSnapshotId(), snapshotId); - Assert.assertEquals(restoreResult.getSnapshotResult(), SnapshotResult.IN_PROGRESS); + Assert.assertEquals(restoreResult.getSnapshotInfo().getSnapshotResult(), SnapshotResult.IN_PROGRESS); restoreResult.setSnapshotResult(snapshotId, SnapshotResult.SUCCESSFUL); - Assert.assertEquals(restoreResult.getSnapshotResult(), SnapshotResult.SUCCESSFUL); + Assert.assertEquals(restoreResult.getSnapshotInfo().getSnapshotResult(), SnapshotResult.SUCCESSFUL); // Test equals - RestoreResult restoreResult2 = new RestoreResult(snapshotId, SnapshotResult.SUCCESSFUL); - RestoreResult restoreResult3 = new RestoreResult(snapshotId, SnapshotResult.FAILED); + RestoreResult restoreResult2 = new RestoreResult(snapshotId, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL)); + RestoreResult restoreResult3 = new RestoreResult(snapshotId, SnapshotInfo.withStatus(SnapshotResult.FAILED)); Assert.assertTrue(restoreResult.equals(restoreResult2)); Assert.assertFalse(restoreResult.equals(restoreResult3)); diff --git a/presto-main/src/test/java/io/prestosql/snapshot/TestSingleInputSnapshotState.java b/presto-main/src/test/java/io/prestosql/snapshot/TestSingleInputSnapshotState.java index a64c00bee..544b738a2 100644 --- a/presto-main/src/test/java/io/prestosql/snapshot/TestSingleInputSnapshotState.java +++ b/presto-main/src/test/java/io/prestosql/snapshot/TestSingleInputSnapshotState.java @@ -42,6 +42,7 @@ import static java.util.concurrent.Executors.newScheduledThreadPool; import static org.mockito.Matchers.anyBoolean; import static org.mockito.Matchers.anyLong; import static org.mockito.Matchers.anyObject; +import static org.mockito.Matchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; @@ -241,7 +242,7 @@ public class TestSingleInputSnapshotState singleInputSnapshotState.processPage(resume1); verify(snapshotManager, times(2)).storeFile(anyObject(), anyObject()); verify(snapshotManager, times(4)).loadFile(anyObject(), anyObject()); - verify(snapshotManager, times(1)).succeededToRestore(anyObject()); + verify(snapshotManager, times(1)).succeededToRestore(anyObject(), eq(0L)); verify(snapshotManager, times(2)).failedToRestore(anyObject(), anyBoolean()); } @@ -261,8 +262,8 @@ public class TestSingleInputSnapshotState singleInputSnapshotState.processPage(marker1); when(snapshotManager.loadConsolidatedState(anyObject())).thenReturn(Optional.of(0)); singleInputSnapshotState.processPage(resume1); - verify(snapshotManager, times(0)).storeState(anyObject(), anyObject()); - verify(snapshotManager, times(1)).storeConsolidatedState(anyObject(), anyObject()); + verify(snapshotManager, times(0)).storeState(anyObject(), anyObject(), eq(0L)); + verify(snapshotManager, times(1)).storeConsolidatedState(anyObject(), anyObject(), eq(0L)); verify(snapshotManager, times(0)).loadState(anyObject()); verify(snapshotManager, times(1)).loadConsolidatedState(anyObject()); } @@ -283,8 +284,8 @@ public class TestSingleInputSnapshotState singleInputSnapshotState.processPage(marker1); when(snapshotManager.loadState(anyObject())).thenReturn(Optional.of(0)); singleInputSnapshotState.processPage(resume1); - verify(snapshotManager, times(0)).storeConsolidatedState(anyObject(), anyObject()); - verify(snapshotManager, times(1)).storeState(anyObject(), anyObject()); + verify(snapshotManager, times(0)).storeConsolidatedState(anyObject(), anyObject(), eq(0L)); + verify(snapshotManager, times(1)).storeState(anyObject(), anyObject(), eq(0L)); verify(snapshotManager, times(0)).loadConsolidatedState(anyObject()); verify(snapshotManager, times(1)).loadState(anyObject()); } diff --git a/presto-main/src/test/java/io/prestosql/snapshot/TestSnapshotFileBasedClient.java b/presto-main/src/test/java/io/prestosql/snapshot/TestSnapshotFileBasedClient.java index af7f694b5..5cfe195f1 100644 --- a/presto-main/src/test/java/io/prestosql/snapshot/TestSnapshotFileBasedClient.java +++ b/presto-main/src/test/java/io/prestosql/snapshot/TestSnapshotFileBasedClient.java @@ -38,16 +38,19 @@ public class TestSnapshotFileBasedClient { SnapshotFileBasedClient client = new SnapshotFileBasedClient(new HetuLocalFileSystemClient(new LocalConfig(new Properties()), Paths.get(ROOT_PATH_STR)), Paths.get(ROOT_PATH_STR), false); String queryId = "query1"; - LinkedHashMap map = new LinkedHashMap<>(); - map.put(3L, SnapshotResult.SUCCESSFUL); - map.put(1L, SnapshotResult.FAILED); - map.put(5L, SnapshotResult.FAILED_FATAL); - map.put(8L, SnapshotResult.SUCCESSFUL); + LinkedHashMap map = new LinkedHashMap<>(); + map.put(3L, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL)); + map.put(1L, SnapshotInfo.withStatus(SnapshotResult.FAILED)); + map.put(5L, SnapshotInfo.withStatus(SnapshotResult.FAILED_FATAL)); + map.put(8L, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL)); // Test store and Load client.storeSnapshotResult(queryId, map); - LinkedHashMap resultMap = (LinkedHashMap) client.loadSnapshotResult(queryId); - Assert.assertEquals(map, resultMap); + LinkedHashMap resultMap = (LinkedHashMap) client.loadSnapshotResult(queryId); + Assert.assertEquals(map.get(3L).getSnapshotResult(), resultMap.get(3L).getSnapshotResult()); + Assert.assertEquals(map.get(1L).getSnapshotResult(), resultMap.get(1L).getSnapshotResult()); + Assert.assertEquals(map.get(5L).getSnapshotResult(), resultMap.get(5L).getSnapshotResult()); + Assert.assertEquals(map.get(8L).getSnapshotResult(), resultMap.get(8L).getSnapshotResult()); } /** @@ -69,9 +72,9 @@ public class TestSnapshotFileBasedClient map.put(8L, SnapshotResult.SUCCESSFUL); // Test store and Load - client.storeState(snapshotStateId, map); - client.loadState(snapshotStateId); - Assert.assertEquals(map, client.loadState(snapshotStateId).get()); + client.storeState(snapshotStateId, map, null); + client.loadState(snapshotStateId, null); + Assert.assertEquals(map, client.loadState(snapshotStateId, null).get()); } /** @@ -93,8 +96,8 @@ public class TestSnapshotFileBasedClient map.put(8L, SnapshotResult.SUCCESSFUL); // Test store and Load - client.storeState(snapshotStateId, map); - client.loadState(snapshotStateId); - Assert.assertEquals(map, client.loadState(snapshotStateId).get()); + client.storeState(snapshotStateId, map, null); + client.loadState(snapshotStateId, null); + Assert.assertEquals(map, client.loadState(snapshotStateId, null).get()); } } diff --git a/presto-main/src/test/java/io/prestosql/snapshot/TestTaskSnapshotManager.java b/presto-main/src/test/java/io/prestosql/snapshot/TestTaskSnapshotManager.java index d1fa4f1b4..057d61719 100644 --- a/presto-main/src/test/java/io/prestosql/snapshot/TestTaskSnapshotManager.java +++ b/presto-main/src/test/java/io/prestosql/snapshot/TestTaskSnapshotManager.java @@ -97,7 +97,7 @@ public class TestTaskSnapshotManager // Test operator state MockState operatorState = new MockState("operator-state"); SnapshotStateId operatorStateId = SnapshotStateId.forOperator(1L, taskId1, 3, 4, 5); - snapshotManager.storeState(operatorStateId, operatorState); + snapshotManager.storeState(operatorStateId, operatorState, 0); snapshotManager.setTotalComponents(1); snapshotManager.succeededToCapture(operatorStateId); MockState newOperatorState = (MockState) snapshotManager.loadState(operatorStateId).get(); @@ -107,7 +107,7 @@ public class TestTaskSnapshotManager MockState taskState = new MockState("task-state"); TaskId taskId2 = new TaskId(queryId.getId(), 3, 4); SnapshotStateId taskStateId = SnapshotStateId.forOperator(3L, taskId2, 5, 6, 7); - snapshotManager.storeState(taskStateId, taskState); + snapshotManager.storeState(taskStateId, taskState, 0); snapshotManager.succeededToCapture(taskStateId); MockState newTaskState = (MockState) snapshotManager.loadState(taskStateId).get(); Assert.assertEquals(taskState.getState(), newTaskState.getState()); @@ -127,14 +127,15 @@ public class TestTaskSnapshotManager // Save operator state MockState state = new MockState("state"); SnapshotStateId stateId = SnapshotStateId.forOperator(1L, taskId, 3, 4, 5); - snapshotManager.storeState(stateId, state); + snapshotManager.storeState(stateId, state, 0); snapshotManager.setTotalComponents(1); snapshotManager.succeededToCapture(stateId); // Make snapshot manager think that snapshot #1 is complete QuerySnapshotManager querySnapshotManager = new QuerySnapshotManager(queryId, snapshotUtils, TEST_SNAPSHOT_SESSION); querySnapshotManager.addNewTask(taskId); - querySnapshotManager.updateQueryCapture(taskId, Collections.singletonMap(1L, SnapshotResult.SUCCESSFUL)); + querySnapshotManager.snapshotInitiated(1L); + querySnapshotManager.updateQueryCapture(taskId, Collections.singletonMap(1L, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL))); querySnapshotManager.getResumeSnapshotId(); SnapshotStateId newStateId = stateId.withSnapshotId(2); @@ -205,7 +206,8 @@ public class TestTaskSnapshotManager assertNull(snapshotManager.loadFile(id4load, targetPath)); // Try2: Previous snapshots are setup, so load should be successful - querySnapshotManager.updateQueryCapture(taskId, Collections.singletonMap(2L, SnapshotResult.SUCCESSFUL)); + querySnapshotManager.snapshotInitiated(2L); + querySnapshotManager.updateQueryCapture(taskId, Collections.singletonMap(2L, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL))); querySnapshotManager.getResumeSnapshotId(); assertTrue(snapshotManager.loadFile(id4load, targetPath)); @@ -213,9 +215,10 @@ public class TestTaskSnapshotManager Assert.assertEquals(output, fileContent); // Try3: Previous snapshot failed - querySnapshotManager.updateQueryRestore(taskId, Optional.of(new RestoreResult(2, SnapshotResult.SUCCESSFUL))); - querySnapshotManager.updateQueryCapture(taskId, Collections.singletonMap(2L, SnapshotResult.FAILED)); - querySnapshotManager.updateQueryCapture(taskId, Collections.singletonMap(3L, SnapshotResult.SUCCESSFUL)); + querySnapshotManager.updateQueryRestore(taskId, Optional.of(new RestoreResult(2, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL)))); + querySnapshotManager.updateQueryCapture(taskId, Collections.singletonMap(2L, SnapshotInfo.withStatus(SnapshotResult.FAILED))); + querySnapshotManager.snapshotInitiated(3L); + querySnapshotManager.updateQueryCapture(taskId, Collections.singletonMap(3L, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL))); querySnapshotManager.getResumeSnapshotId(); assertFalse(snapshotManager.loadFile(id4load, targetPath)); } @@ -236,15 +239,15 @@ public class TestTaskSnapshotManager // Test capture successfully snapshotManager1.succeededToCapture(new SnapshotStateId(1, taskId1, "component1")); - Assert.assertEquals(snapshotManager1.getSnapshotCaptureResult().get(1L), SnapshotResult.IN_PROGRESS); + Assert.assertEquals(snapshotManager1.getSnapshotCaptureResult().get(1L).getSnapshotResult(), SnapshotResult.IN_PROGRESS); snapshotManager1.succeededToCapture(new SnapshotStateId(1, taskId1, "component2")); - Assert.assertEquals(snapshotManager1.getSnapshotCaptureResult().get(1L), SnapshotResult.SUCCESSFUL); + Assert.assertEquals(snapshotManager1.getSnapshotCaptureResult().get(1L).getSnapshotResult(), SnapshotResult.SUCCESSFUL); // Test capture failed snapshotManager2.failedToCapture(new SnapshotStateId(1, taskId2, "component1")); - Assert.assertEquals(snapshotManager2.getSnapshotCaptureResult().get(1L), SnapshotResult.IN_PROGRESS_FAILED); + Assert.assertEquals(snapshotManager2.getSnapshotCaptureResult().get(1L).getSnapshotResult(), SnapshotResult.IN_PROGRESS_FAILED); snapshotManager2.failedToCapture(new SnapshotStateId(1, taskId2, "component2")); - Assert.assertEquals(snapshotManager2.getSnapshotCaptureResult().get(1L), SnapshotResult.FAILED); + Assert.assertEquals(snapshotManager2.getSnapshotCaptureResult().get(1L).getSnapshotResult(), SnapshotResult.FAILED); } @Test @@ -259,12 +262,12 @@ public class TestTaskSnapshotManager snapshotManager2.setTotalComponents(3); // Test restore successfully - snapshotManager1.succeededToRestore(new SnapshotStateId(1, taskId1, "component1")); + snapshotManager1.succeededToRestore(new SnapshotStateId(1, taskId1, "component1"), 0); RestoreResult result = snapshotManager1.getSnapshotRestoreResult(); Assert.assertEquals(result.getSnapshotId(), 1L); - Assert.assertEquals(result.getSnapshotResult(), SnapshotResult.IN_PROGRESS); - snapshotManager1.succeededToRestore(new SnapshotStateId(1, taskId1, "component2")); - Assert.assertEquals(snapshotManager1.getSnapshotRestoreResult().getSnapshotResult(), SnapshotResult.SUCCESSFUL); + Assert.assertEquals(result.getSnapshotInfo().getSnapshotResult(), SnapshotResult.IN_PROGRESS); + snapshotManager1.succeededToRestore(new SnapshotStateId(1, taskId1, "component2"), 0); + Assert.assertEquals(snapshotManager1.getSnapshotRestoreResult().getSnapshotInfo().getSnapshotResult(), SnapshotResult.SUCCESSFUL); // Test restore failed try { @@ -273,16 +276,16 @@ public class TestTaskSnapshotManager catch (Exception e) { // Ignore } - Assert.assertEquals(snapshotManager2.getSnapshotRestoreResult().getSnapshotResult(), SnapshotResult.IN_PROGRESS_FAILED); + Assert.assertEquals(snapshotManager2.getSnapshotRestoreResult().getSnapshotInfo().getSnapshotResult(), SnapshotResult.IN_PROGRESS_FAILED); try { snapshotManager2.failedToRestore(new SnapshotStateId(1, taskId2, "component2"), true); } catch (Exception e) { // Ignore } - Assert.assertEquals(snapshotManager2.getSnapshotRestoreResult().getSnapshotResult(), SnapshotResult.IN_PROGRESS_FAILED_FATAL); - snapshotManager2.succeededToRestore(new SnapshotStateId(1, taskId2, "component3")); - Assert.assertEquals(snapshotManager2.getSnapshotRestoreResult().getSnapshotResult(), SnapshotResult.FAILED_FATAL); + Assert.assertEquals(snapshotManager2.getSnapshotRestoreResult().getSnapshotInfo().getSnapshotResult(), SnapshotResult.IN_PROGRESS_FAILED_FATAL); + snapshotManager2.succeededToRestore(new SnapshotStateId(1, taskId2, "component3"), 0); + Assert.assertEquals(snapshotManager2.getSnapshotRestoreResult().getSnapshotInfo().getSnapshotResult(), SnapshotResult.FAILED_FATAL); } @Test @@ -297,7 +300,7 @@ public class TestTaskSnapshotManager sm.updateFinishedComponents(ImmutableList.of(mock(Operator.class))); sm.succeededToCapture(new SnapshotStateId(1, taskId)); - assertEquals(sm.getSnapshotCaptureResult().get(1L), SnapshotResult.SUCCESSFUL); + assertEquals(sm.getSnapshotCaptureResult().get(1L).getSnapshotResult(), SnapshotResult.SUCCESSFUL); } @Test @@ -313,7 +316,7 @@ public class TestTaskSnapshotManager MockState state = new MockState("mockstate"); SnapshotStateId stateId = SnapshotStateId.forOperator(1L, taskId1, 3, 4, 5); - snapshotManager.storeConsolidatedState(stateId, state); + snapshotManager.storeConsolidatedState(stateId, state, 0); snapshotManager.succeededToCapture(stateId); MockState newState = (MockState) snapshotManager.loadConsolidatedState(stateId).get(); Assert.assertEquals(state.getState(), newState.getState()); @@ -337,16 +340,16 @@ public class TestTaskSnapshotManager long firstSnapshotId = 1L; MockState state = new MockState("mockstate"); SnapshotStateId stateId = SnapshotStateId.forOperator(firstSnapshotId, taskId1, 3, 4, 5); - snapshotManager.storeConsolidatedState(stateId, state); + snapshotManager.storeConsolidatedState(stateId, state, 0); snapshotManager.succeededToCapture(stateId); // second store (null) SnapshotStateId secondId = SnapshotStateId.forOperator(2L, taskId1, 3, 4, 5); - snapshotManager.storeConsolidatedState(secondId, null); + snapshotManager.storeConsolidatedState(secondId, null, 0); snapshotManager.succeededToCapture(secondId); - Map queryCaptureResult = new LinkedHashMap<>(); - queryCaptureResult.put(firstSnapshotId, SnapshotResult.SUCCESSFUL); + Map queryCaptureResult = new LinkedHashMap<>(); + queryCaptureResult.put(firstSnapshotId, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL)); snapshotUtils.storeSnapshotResult(queryId.getId(), queryCaptureResult); MockState newState = (MockState) snapshotManager.loadConsolidatedState(secondId).get(); @@ -380,12 +383,12 @@ public class TestTaskSnapshotManager SnapshotStateId stateId22 = SnapshotStateId.forOperator(2L, taskId2, 3, 4, 5); // save 11 (part of snapshot 1), then 12 and 22 (snapshot 2), then finish snapshot 1 (21) - snapshotManager1.storeConsolidatedState(stateId11, state11); - snapshotManager2.storeConsolidatedState(stateId22, state22); - snapshotManager1.storeConsolidatedState(stateId12, state12); + snapshotManager1.storeConsolidatedState(stateId11, state11, 0); + snapshotManager2.storeConsolidatedState(stateId22, state22, 0); + snapshotManager1.storeConsolidatedState(stateId12, state12, 0); snapshotManager1.succeededToCapture(stateId12); snapshotManager2.succeededToCapture(stateId22); - snapshotManager2.storeConsolidatedState(stateId21, state21); + snapshotManager2.storeConsolidatedState(stateId21, state21, 0); snapshotManager1.succeededToCapture(stateId11); snapshotManager2.succeededToCapture(stateId21); @@ -419,16 +422,16 @@ public class TestTaskSnapshotManager long firstSnapshotId = 1L; MockState state = new MockState("mockstate"); SnapshotStateId stateId = SnapshotStateId.forOperator(firstSnapshotId, taskId1, 3, 4, 5); - snapshotManager.storeState(stateId, state); + snapshotManager.storeState(stateId, state, 0); snapshotManager.succeededToCapture(stateId); // second store, then deleted MockState secondState = new MockState("secondState"); SnapshotStateId secondId = SnapshotStateId.forOperator(2L, taskId1, 3, 4, 5); - snapshotManager.storeState(secondId, secondState); + snapshotManager.storeState(secondId, secondState, 0); snapshotManager.succeededToCapture(secondId); - - querySnapshotManager.updateQueryCapture(taskId1, Collections.singletonMap(firstSnapshotId, SnapshotResult.SUCCESSFUL)); + querySnapshotManager.snapshotInitiated(firstSnapshotId); + querySnapshotManager.updateQueryCapture(taskId1, Collections.singletonMap(firstSnapshotId, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL))); assertTrue(snapshotManager.loadState(secondId).isPresent()); File second = new File("/tmp/test_snapshot_manager/" + queryId + "/2/1/0/3/4/5"); @@ -450,7 +453,7 @@ public class TestTaskSnapshotManager queryId = new QueryId("failedstoreconsolidatedquery"); SnapshotUtils faultySnapshotUtils = mock(SnapshotUtils.class); - doThrow(new NullPointerException()).when(faultySnapshotUtils).storeState(any(), any()); + doThrow(new NullPointerException()).when(faultySnapshotUtils).storeState(any(), any(), any()); TaskId taskId1 = new TaskId(queryId.getId(), 1, 0); TaskSnapshotManager snapshotManager = new TaskSnapshotManager(taskId1, 0, faultySnapshotUtils); @@ -459,11 +462,11 @@ public class TestTaskSnapshotManager MockState state = new MockState("mockstate"); SnapshotStateId stateId = SnapshotStateId.forOperator(1L, taskId1, 3, 4, 5); - snapshotManager.storeConsolidatedState(stateId, state); + snapshotManager.storeConsolidatedState(stateId, state, 0); // Error messages will print. This is normal because we are failing the store on purpose snapshotManager.succeededToCapture(stateId); - Assert.assertEquals(snapshotManager.getSnapshotCaptureResult().get(1L), SnapshotResult.FAILED); + Assert.assertEquals(snapshotManager.getSnapshotCaptureResult().get(1L).getSnapshotResult(), SnapshotResult.FAILED); } @Test @@ -503,7 +506,8 @@ public class TestTaskSnapshotManager SnapshotStateId secondId = SnapshotStateId.forOperator(2L, taskId1, 3, 4, 5); snapshotManager.storeFile(secondId, secondFile.toPath()); snapshotManager.succeededToCapture(secondId); - querySnapshotManager.updateQueryCapture(taskId1, Collections.singletonMap(firstSnapshotId, SnapshotResult.SUCCESSFUL)); + querySnapshotManager.snapshotInitiated(firstSnapshotId); + querySnapshotManager.updateQueryCapture(taskId1, Collections.singletonMap(firstSnapshotId, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL))); File secondFileOperator = new File("/tmp/test_snapshot_manager/" + queryId + "/2/1/0/3/4/5/secondFile"); assertTrue(secondFileOperator.exists());