jffwang
/
hetu-core
0
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!1393 Changes to record snapshot capture metrics 

Merge pull request !1393 from i-robot/pull291
This commit is contained in:
i-robot 2022-03-02 08:24:31 +00:00 committed by Gitee
parent 2e513c05c2 2ff0fb2f30
commit 94e52ebd1f
No known key found for this signature in database
GPG Key ID: 173E9B9CA92EEF8F
33 changed files with 1113 additions and 222 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
presto-cli/src/main/java/io/prestosql/cli/StatusPrinter.java
View File

@ -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]

283
presto-client/src/main/java/io/prestosql/client/SnapshotStats.java Normal file
View File

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

22
presto-client/src/main/java/io/prestosql/client/StatementStats.java
View File

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

2
presto-jdbc/src/test/java/io/prestosql/jdbc/TestProgressMonitor.java
View File

@ -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,

6
presto-main/src/main/java/io/prestosql/execution/QueryExecution.java
View File

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

6
presto-main/src/main/java/io/prestosql/execution/QueryManager.java
View File

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

8
presto-main/src/main/java/io/prestosql/execution/SqlQueryExecution.java
View File

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

7
presto-main/src/main/java/io/prestosql/execution/SqlQueryManager.java
View File

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

4
presto-main/src/main/java/io/prestosql/execution/SqlTask.java
View File

@ -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<Lifespan> completedDriverGroups = ImmutableSet.of();
long fullGcCount = 0;
Duration fullGcTime = new Duration(0, MILLISECONDS);
Map<Long, SnapshotResult> snapshotCaptureResult = ImmutableMap.of();
Map<Long, SnapshotInfo> snapshotCaptureResult = ImmutableMap.of();
Optional<RestoreResult> snapshotRestoreResult = Optional.empty();
TaskInfo finalTaskInfo = taskHolder.getFinalTaskInfo();
if (finalTaskInfo != null) {

8
presto-main/src/main/java/io/prestosql/execution/TaskStatus.java
View File

@ -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<Long, SnapshotResult> snapshotCaptureResult;
private final Map<Long, SnapshotInfo> snapshotCaptureResult;
private final Optional<RestoreResult> snapshotRestoreResult;
@JsonCreator
@ -102,7 +102,7 @@ public class TaskStatus
@JsonProperty("revocableMemoryReservation") DataSize revocableMemoryReservation,
@JsonProperty("fullGcCount") long fullGcCount,
@JsonProperty("fullGcTime") Duration fullGcTime,
@JsonProperty("snapshotCaptureResult") Map<Long, SnapshotResult> snapshotCaptureResult,
@JsonProperty("snapshotCaptureResult") Map<Long, SnapshotInfo> snapshotCaptureResult,
@JsonProperty("snapshotRestoreResult") Optional<RestoreResult> snapshotRestoreResult)
{
this.taskId = requireNonNull(taskId, "taskId is null");
@ -241,7 +241,7 @@ public class TaskStatus
}
@JsonProperty
public Map<Long, SnapshotResult> getSnapshotCaptureResult()
public Map<Long, SnapshotInfo> getSnapshotCaptureResult()
{
return snapshotCaptureResult;
}

1
presto-main/src/main/java/io/prestosql/execution/scheduler/SqlQueryScheduler.java
View File

@ -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(() -> {

2
presto-main/src/main/java/io/prestosql/server/protocol/PagePublisherQueryManager.java
View File

@ -62,7 +62,7 @@ public class PagePublisherQueryManager
private final Set<String> queries = Sets.newConcurrentHashSet();
private final Map<String, PagePublisherQueryRunner> 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;

6
presto-main/src/main/java/io/prestosql/server/protocol/PagePublisherQueryRunner.java
View File

@ -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<Comparable<Duration>> WAIT_ORDERING = Ordering.natural().nullsLast();
private static final Duration MAX_WAIT_TIME = new Duration(1, SECONDS);

76
presto-main/src/main/java/io/prestosql/server/protocol/Query.java
View File

@ -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<RestoreResult> 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) {

4
presto-main/src/main/java/io/prestosql/server/remotetask/ContinuousTaskStatusFetcher.java
View File

@ -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<Long, SnapshotResult> captureResult, Optional<RestoreResult> restoreResult)
private void updateSnapshots(Map<Long, SnapshotInfo> captureResult, Optional<RestoreResult> restoreResult)
{
snapshotManager.updateQueryCapture(taskId, captureResult);
snapshotManager.updateQueryRestore(taskId, restoreResult);

38
presto-main/src/main/java/io/prestosql/snapshot/MultiInputSnapshotState.java
View File

@ -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<Object> 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);
}
}
}

176
presto-main/src/main/java/io/prestosql/snapshot/QuerySnapshotManager.java
View File

@ -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<TaskId> unfinishedTasks = Sets.newConcurrentHashSet();
// LinkedHashMap can be used to keep ordering
private final Map<Long, SnapshotComponentCounter<TaskId>> captureComponentCounters = Collections.synchronizedMap(new LinkedHashMap<>());
private final Map<Long, SnapshotResult> captureResults = Collections.synchronizedMap(new LinkedHashMap<>());
private final Map<Long, SnapshotInfo> captureResults = Collections.synchronizedMap(new LinkedHashMap<>());
private final Set<String> consolidatedFilePaths = Collections.synchronizedSet(new HashSet<>());
private final Map<Long, SnapshotComponentCounter<TaskId>> restoreComponentCounters = Collections.synchronizedMap(new LinkedHashMap<>());
private final RestoreResult restoreResult = new RestoreResult();
private RestoreResult restoreResult = new RestoreResult();
private final List<Consumer<RestoreResult>> restoreCompleteListeners = Collections.synchronizedList(new ArrayList<>());
private final List<RestoreResult> 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<Map.Entry<Long, SnapshotResult>> entryList = new ArrayList<>(captureResults.entrySet());
List<Map.Entry<Long, SnapshotInfo>> 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<Long, SnapshotResult> captureResult)
public void updateQueryCapture(TaskId taskId, Map<Long, SnapshotInfo> captureResult)
{
for (Map.Entry<Long, SnapshotResult> entry : captureResult.entrySet()) {
for (Map.Entry<Long, SnapshotInfo> 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> 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<Long, SnapshotResult> doneResult = captureResults.entrySet()
Map<Long, SnapshotInfo> 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<TaskId> 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<TaskId> 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<Long, BiConsumer<Long, Long>> eachUpdater, Function<Long, BiConsumer<Long, Long>> 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<RestoreResult> getRestoreStats()
{
return restoreStats;
}
}

29
presto-main/src/main/java/io/prestosql/snapshot/RestoreResult.java
View File

@ -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

29
presto-main/src/main/java/io/prestosql/snapshot/SingleInputSnapshotState.java
View File

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

38
presto-main/src/main/java/io/prestosql/snapshot/SnapshotDataCollector.java Normal file
View File

@ -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()");
}
}

54
presto-main/src/main/java/io/prestosql/snapshot/SnapshotFileBasedClient.java
View File

@ -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<Object> loadState(SnapshotStateId snapshotStateId)
public Optional<Object> loadState(SnapshotStateId snapshotStateId, SnapshotDataCollector dataCollector)
throws IOException, ClassNotFoundException
{
Optional<Object> 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<String> 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<String> 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<Long, SnapshotResult> result)
public void storeSnapshotResult(String queryId, Map<Long, SnapshotInfo> result)
throws IOException
{
Path file = SnapshotUtils.createStatePath(rootPath, queryId, "result");
@ -140,7 +178,7 @@ public class SnapshotFileBasedClient
}
@Override
public Map<Long, SnapshotResult> loadSnapshotResult(String queryId)
public Map<Long, SnapshotInfo> 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<Long, SnapshotResult>) ois.readObject();
return (LinkedHashMap<Long, SnapshotInfo>) ois.readObject();
}
}

154
presto-main/src/main/java/io/prestosql/snapshot/SnapshotInfo.java Normal file
View File

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

12
presto-main/src/main/java/io/prestosql/snapshot/SnapshotStoreClient.java
View File

@ -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<Object> loadState(SnapshotStateId snapshotStateId)
Optional<Object> 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<Long, SnapshotResult> result)
void storeSnapshotResult(String queryId, Map<Long, SnapshotInfo> result)
throws Exception;
/**
* Load snapshot result of query
*/
Map<Long, SnapshotResult> loadSnapshotResult(String queryId)
Map<Long, SnapshotInfo> loadSnapshotResult(String queryId)
throws Exception;
/**

20
presto-main/src/main/java/io/prestosql/snapshot/SnapshotUtils.java
View File

@ -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<Object> loadState(SnapshotStateId snapshotStateId)
public Optional<Object> 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<Long, SnapshotResult> result)
public void storeSnapshotResult(String queryId, Map<Long, SnapshotInfo> result)
throws Exception
{
snapshotStoreClient.storeSnapshotResult(queryId, result);
}
public Map<Long, SnapshotResult> loadSnapshotResult(String queryId)
public Map<Long, SnapshotInfo> loadSnapshotResult(String queryId)
throws Exception
{
return snapshotStoreClient.loadSnapshotResult(queryId);

93
presto-main/src/main/java/io/prestosql/snapshot/TaskSnapshotManager.java
View File

@ -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<Long, SnapshotComponentCounter<SnapshotStateId>> captureComponentCounters = Collections.synchronizedMap(new LinkedHashMap<>());
private final Map<Long, SnapshotResult> captureResults = new LinkedHashMap<>();
private final Map<Long, SnapshotInfo> captureResults = new LinkedHashMap<>();
private final Map<Long, SnapshotComponentCounter<SnapshotStateId>> 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<String, Object> 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<String, Object> 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<Object> loadedState = snapshotUtils.loadState(stateId);
Optional<Object> loadedState = snapshotUtils.loadState(stateId, this);
if (createdConsolidatedFiles == null) {
createdConsolidatedFiles = snapshotUtils.loadConsolidatedFiles(queryId);
}
@ -143,7 +147,7 @@ public class TaskSnapshotManager
Optional<Object> state;
loadMapIfNecessary(snapshotId, snapshotStateIdTaskId);
state = Optional.ofNullable(loadCache.get(snapshotId).get(newSnapshotStateId.toString()));
Map<Long, SnapshotResult> snapshotToSnapshotResultMap = null;
Map<Long, SnapshotInfo> 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<Object> 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<String, Object> 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<Long, SnapshotResult> snapshotToSnapshotResultMap, long snapshotId)
private OptionalLong getPreviousSnapshotIdIfComplete(Map<Long, SnapshotInfo> snapshotToSnapshotResultMap, long snapshotId)
{
try {
List<Map.Entry<Long, SnapshotResult>> entryList = new ArrayList<>(snapshotToSnapshotResultMap.entrySet());
List<Map.Entry<Long, SnapshotInfo>> 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<Long, SnapshotResult> getSnapshotCaptureResult()
public Map<Long, SnapshotInfo> 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()
{

13
presto-main/src/test/java/io/prestosql/execution/buffer/TestPartitionedOutputBuffer.java
View File

@ -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<SnapshotStateId> idArgument = ArgumentCaptor.forClass(SnapshotStateId.class);
ArgumentCaptor<Object> stateArgument = ArgumentCaptor.forClass(Object.class);
ArgumentCaptor<TaskSnapshotManager> 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<SnapshotStateId> ids = idArgument.getAllValues();
List<Object> 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<TaskSnapshotManager> 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);

9
presto-main/src/test/java/io/prestosql/operator/TestPartitionedOutputOperator.java
View File

@ -86,15 +86,14 @@ public class TestPartitionedOutputOperator
operator.addInput(input.get(0));
operator.addInput(marker);
ArgumentCaptor<Object> 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<String, Object>) 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<String, Object>) snapshot).get("query/3/1/1/0/0/0");
assertEquals(SnapshotTestUtil.toFullSnapshotMapping(snapshotEntry), createExpectedMappingAfterFinish());

29
presto-main/src/test/java/io/prestosql/snapshot/TestMultiInputSnapshotState.java
View File

@ -79,6 +79,7 @@ public class TestMultiInputSnapshotState
private MultiInputSnapshotState state;
private final ArgumentCaptor<List> argument = ArgumentCaptor.forClass(List.class);
private final ArgumentCaptor<Long> 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<Object> 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<Object> 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<Object> 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<Page> 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);

50
presto-main/src/test/java/io/prestosql/snapshot/TestQuerySnapshotManager.java
View File

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

8
presto-main/src/test/java/io/prestosql/snapshot/TestRestoreResult.java
View File

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

11
presto-main/src/test/java/io/prestosql/snapshot/TestSingleInputSnapshotState.java
View File

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

29
presto-main/src/test/java/io/prestosql/snapshot/TestSnapshotFileBasedClient.java
View File

@ -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<Long, SnapshotResult> 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<Long, SnapshotInfo> 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<Long, SnapshotResult> resultMap = (LinkedHashMap<Long, SnapshotResult>) client.loadSnapshotResult(queryId);
Assert.assertEquals(map, resultMap);
LinkedHashMap<Long, SnapshotInfo> resultMap = (LinkedHashMap<Long, SnapshotInfo>) 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());
}
}

80
presto-main/src/test/java/io/prestosql/snapshot/TestTaskSnapshotManager.java
View File

@ -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<Long, SnapshotResult> queryCaptureResult = new LinkedHashMap<>();
queryCaptureResult.put(firstSnapshotId, SnapshotResult.SUCCESSFUL);
Map<Long, SnapshotInfo> 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());