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

Track memory usage in SingleInputSnapshotState

This commit is contained in:
Kevin Wan 2021-11-18 14:48:06 -05:00
parent efa2bc3ef5
commit 0576c49961
46 changed files with 315 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
presto-benchmark/src/main/java/io/prestosql/benchmark/HandTpchQuery1.java
View File

@ -204,6 +204,14 @@ public class HandTpchQuery1
finishing = true;
}
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
}
@Override
public boolean isFinished()
{

3
presto-main/src/main/java/io/prestosql/operator/AggregationOperator.java
View File

@ -135,6 +135,9 @@ public class AggregationOperator
{
userMemoryContext.setBytes(0);
systemMemoryContext.close();
if (snapshotState != null) {
snapshotState.close();
}
}
@Override

8
presto-main/src/main/java/io/prestosql/operator/AssignUniqueIdOperator.java
View File

@ -196,6 +196,14 @@ public class AssignUniqueIdOperator
return block.build();
}
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{

8
presto-main/src/main/java/io/prestosql/operator/DistinctLimitOperator.java
View File

@ -289,6 +289,14 @@ public class DistinctLimitOperator
return groupByHash.getCapacity();
}
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{

8
presto-main/src/main/java/io/prestosql/operator/DynamicFilterSourceOperator.java
View File

@ -231,6 +231,14 @@ public class DynamicFilterSourceOperator
}
}
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{

8
presto-main/src/main/java/io/prestosql/operator/EnforceSingleRowOperator.java
View File

@ -161,6 +161,14 @@ public class EnforceSingleRowOperator
return snapshotState.nextMarker();
}
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{

8
presto-main/src/main/java/io/prestosql/operator/ExplainAnalyzeOperator.java
View File

@ -226,6 +226,14 @@ public class ExplainAnalyzeOperator
}
}
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{

8
presto-main/src/main/java/io/prestosql/operator/FilterAndProjectOperator.java
View File

@ -134,6 +134,14 @@ public class FilterAndProjectOperator
return snapshotState.nextMarker();
}
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{

8
presto-main/src/main/java/io/prestosql/operator/GroupIdOperator.java
View File

@ -233,6 +233,14 @@ public class GroupIdOperator
return outputPage;
}
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{

3
presto-main/src/main/java/io/prestosql/operator/HashAggregationOperator.java
View File

@ -397,6 +397,9 @@ public class HashAggregationOperator
public void close()
{
closeAggregationBuilder();
if (snapshotState != null) {
snapshotState.close();
}
}
protected void closeAggregationBuilder()

5
presto-main/src/main/java/io/prestosql/operator/HashBuilderOperator.java
View File

@ -697,6 +697,11 @@ public class HashBuilderOperator
spiller.ifPresent(closer::register);
closer.register(() -> localUserMemoryContext.setBytes(0));
closer.register(() -> localRevocableMemoryContext.setBytes(0));
closer.register(() -> {
if (snapshotState != null) {
snapshotState.close();
}
});
}
catch (IOException e) {
throw new RuntimeException(e);

8
presto-main/src/main/java/io/prestosql/operator/HashSemiJoinOperator.java
View File

@ -237,6 +237,14 @@ public class HashSemiJoinOperator
return snapshotState.nextMarker();
}
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{

8
presto-main/src/main/java/io/prestosql/operator/LimitOperator.java
View File

@ -157,6 +157,14 @@ public class LimitOperator
return snapshotState.nextMarker();
}
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{

5
presto-main/src/main/java/io/prestosql/operator/LookupJoinOperator.java
View File

@ -589,6 +589,11 @@ public class LookupJoinOperator
closer.register(pageBuilder::reset);
closer.register(() -> Optional.ofNullable(lookupSourceProvider).ifPresent(LookupSourceProvider::close));
closer.register(() -> {
if (snapshotState != null) {
snapshotState.close();
}
});
spiller.ifPresent(closer::register);
}
catch (IOException e) {

3
presto-main/src/main/java/io/prestosql/operator/LookupOuterOperator.java
View File

@ -290,6 +290,9 @@ public class LookupOuterOperator
if (closed) {
return;
}
if (snapshotState != null) {
snapshotState.close();
}
closed = true;
pageBuilder.reset();
onClose.run();

8
presto-main/src/main/java/io/prestosql/operator/MarkDistinctOperator.java
View File

@ -229,6 +229,14 @@ public class MarkDistinctOperator
return markDistinctHash.getCapacity();
}
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{

8
presto-main/src/main/java/io/prestosql/operator/NestedLoopBuildOperator.java
View File

@ -153,6 +153,14 @@ public class NestedLoopBuildOperator
operatorContext.recordOutput(page.getSizeInBytes(), page.getPositionCount());
}
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{

3
presto-main/src/main/java/io/prestosql/operator/NestedLoopJoinOperator.java
View File

@ -259,6 +259,9 @@ public class NestedLoopJoinOperator
if (closed) {
return;
}
if (snapshotState != null) {
snapshotState.close();
}
closed = true;
// `afterClose` must be run last.
afterClose.run();

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

@ -128,4 +128,13 @@ public interface Operator
{
close();
}
/**
* Estimate the total memory used in an operator state by copying the total memory used by this operator
*/
@Override
default long getUsedMemory()
{
return getOperatorContext().getTotalMemoryBytes();
}
}

11
presto-main/src/main/java/io/prestosql/operator/OperatorContext.java
View File

@ -267,6 +267,12 @@ public class OperatorContext
return new InternalLocalMemoryContext(operatorMemoryContext.newSystemMemoryContext(allocationTag), memoryFuture, this::updatePeakMemoryReservations, true);
}
// caller should close this context as it's a new context
public LocalMemoryContext newLocalUserMemoryContext(String allocationTag)
{
return new InternalLocalMemoryContext(operatorMemoryContext.newUserMemoryContext(allocationTag), memoryFuture, this::updatePeakMemoryReservations, true);
}
// caller shouldn't close this context as it's managed by the OperatorContext
public LocalMemoryContext localUserMemoryContext()
{
@ -339,6 +345,11 @@ public class OperatorContext
return operatorMemoryContext.getRevocableMemory();
}
public long getTotalMemoryBytes()
{
return operatorMemoryContext.getUserMemory() + operatorMemoryContext.getSystemMemory();
}
private static void updateMemoryFuture(ListenableFuture<?> memoryPoolFuture, AtomicReference<SettableFuture<?>> targetFutureReference)
{
if (!memoryPoolFuture.isDone()) {

3
presto-main/src/main/java/io/prestosql/operator/OrderByOperator.java
View File

@ -418,6 +418,9 @@ public class OrderByOperator
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
pageIndex.clear();
sortedPages = null;
spiller.ifPresent(Spiller::close);

8
presto-main/src/main/java/io/prestosql/operator/PartitionedOutputOperator.java
View File

@ -612,6 +612,14 @@ public class PartitionedOutputOperator
}
}
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{

8
presto-main/src/main/java/io/prestosql/operator/RowNumberOperator.java
View File

@ -393,6 +393,14 @@ public class RowNumberOperator
return groupByHash.map(GroupByHash::getCapacity).orElse(0);
}
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{

8
presto-main/src/main/java/io/prestosql/operator/SetBuilderOperator.java
View File

@ -236,6 +236,14 @@ public class SetBuilderOperator
return channelSetBuilder.getCapacity();
}
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{

3
presto-main/src/main/java/io/prestosql/operator/SortAggregationOperator.java
View File

@ -367,6 +367,9 @@ public class SortAggregationOperator
public void close()
{
closeAggregationBuilder();
if (snapshotState != null) {
snapshotState.close();
}
}
protected void closeAggregationBuilder()

3
presto-main/src/main/java/io/prestosql/operator/SpatialIndexBuilderOperator.java
View File

@ -263,6 +263,9 @@ public class SpatialIndexBuilderOperator
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
index.clear();
localUserMemoryContext.setBytes(index.getEstimatedSize().toBytes());
}

4
presto-main/src/main/java/io/prestosql/operator/SpatialJoinOperator.java
View File

@ -368,7 +368,9 @@ public class SpatialJoinOperator
return;
}
closed = true;
if (snapshotState != null) {
snapshotState.close();
}
pagesSpatialIndexFuture = null;
onClose.run();
}

8
presto-main/src/main/java/io/prestosql/operator/StatisticsWriterOperator.java
View File

@ -228,6 +228,14 @@ public class StatisticsWriterOperator
void writeStatistics(Collection<ComputedStatistics> computedStatistics);
}
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{

8
presto-main/src/main/java/io/prestosql/operator/StreamingAggregationOperator.java
View File

@ -329,6 +329,14 @@ public class StreamingAggregationOperator
return finishing && outputPages.isEmpty() && currentGroup == null && pageBuilder.isEmpty();
}
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{

3
presto-main/src/main/java/io/prestosql/operator/TableFinishOperator.java
View File

@ -369,6 +369,9 @@ public class TableFinishOperator
public void close()
throws Exception
{
if (snapshotState != null) {
snapshotState.close();
}
statisticsAggregationOperator.close();
}

3
presto-main/src/main/java/io/prestosql/operator/TableWriterOperator.java
View File

@ -383,6 +383,9 @@ public class TableWriterOperator
public void close()
throws Exception
{
if (snapshotState != null) {
snapshotState.close();
}
closeImpl(pageSink::abort);
}

8
presto-main/src/main/java/io/prestosql/operator/TaskOutputOperator.java
View File

@ -211,6 +211,14 @@ public class TaskOutputOperator
operatorContext.recordOutput(page.getSizeInBytes(), page.getPositionCount());
}
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{

8
presto-main/src/main/java/io/prestosql/operator/TopNRankingNumberOperator.java
View File

@ -334,6 +334,14 @@ public class TopNRankingNumberOperator
return types.build();
}
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{

8
presto-main/src/main/java/io/prestosql/operator/ValuesOperator.java
View File

@ -155,6 +155,14 @@ public class ValuesOperator
return marker;
}
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{

3
presto-main/src/main/java/io/prestosql/operator/WindowOperator.java
View File

@ -1272,6 +1272,9 @@ public class WindowOperator
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
driverWindowInfo.set(Optional.of(windowInfo.build()));
spillablePagesToPagesIndexes.ifPresent(SpillablePagesToPagesIndexes::closeSpiller);
}

3
presto-main/src/main/java/io/prestosql/operator/WorkProcessorOperatorAdapter.java
View File

@ -151,6 +151,9 @@ public class WorkProcessorOperatorAdapter
public void close()
throws Exception
{
if (snapshotState != null) {
snapshotState.close();
}
workProcessorOperator.close();
}

8
presto-main/src/main/java/io/prestosql/operator/dynamicfilter/CrossRegionDynamicFilterOperator.java
View File

@ -195,6 +195,14 @@ public class CrossRegionDynamicFilterOperator
}
}
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{

3
presto-main/src/main/java/io/prestosql/operator/exchange/LocalExchangeSinkOperator.java
View File

@ -177,6 +177,9 @@ public class LocalExchangeSinkOperator
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
finish();
}

8
presto-main/src/main/java/io/prestosql/operator/unnest/UnnestOperator.java
View File

@ -336,6 +336,14 @@ public class UnnestOperator
}
}
@Override
public void close()
{
if (snapshotState != null) {
snapshotState.close();
}
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{

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

@ -16,6 +16,7 @@ package io.prestosql.snapshot;
import io.airlift.log.Logger;
import io.hetu.core.transport.execution.buffer.PagesSerde;
import io.prestosql.memory.context.LocalMemoryContext;
import io.prestosql.operator.Operator;
import io.prestosql.operator.OperatorContext;
import io.prestosql.spi.Page;
@ -35,6 +36,7 @@ import static java.util.Objects.requireNonNull;
* This is a utility class used by non-source operators, which only receive inputs from a single source.
* When an input is received from addInput(Page), the operator first calls the processPage() method to perform snapshot related processing.
* When getOutput() is called, the operator calls the pollMarker() method to determine if a marker page needs to be returned.
* Any instance of this class created through forOperator needs to be closed
*/
public class SingleInputSnapshotState
{
@ -52,6 +54,8 @@ public class SingleInputSnapshotState
private final Function<Long, SnapshotStateId> spillStateIdGenerator;
// Markers to be returned to the restorable object. The "nextMarker" method polls this list.
private final Queue<MarkerPage> markers = new LinkedList<>();
// For recording snapshot memory usage
private final LocalMemoryContext snapshotMemoryContext;
public static SingleInputSnapshotState forOperator(Operator operator, OperatorContext operatorContext)
{
@ -60,14 +64,16 @@ public class SingleInputSnapshotState
operatorContext.getDriverContext().getPipelineContext().getTaskContext().getSnapshotManager(),
operatorContext.getDriverContext().getSerde(),
snapshotId -> SnapshotStateId.forOperator(snapshotId, operatorContext),
snapshotId -> SnapshotStateId.forDriverComponent(snapshotId, operatorContext, operatorContext.getOperatorId() + "-spill"));
snapshotId -> SnapshotStateId.forDriverComponent(snapshotId, operatorContext, operatorContext.getOperatorId() + "-spill"),
operatorContext.newLocalUserMemoryContext(SingleInputSnapshotState.class.getSimpleName()));
}
SingleInputSnapshotState(Restorable restorable,
TaskSnapshotManager snapshotManager,
PagesSerde pagesSerde,
Function<Long, SnapshotStateId> snapshotStateIdGenerator,
Function<Long, SnapshotStateId> spillStateIdGenerator)
TaskSnapshotManager snapshotManager,
PagesSerde pagesSerde,
Function<Long, SnapshotStateId> snapshotStateIdGenerator,
Function<Long, SnapshotStateId> spillStateIdGenerator,
LocalMemoryContext snapshotMemoryContext)
{
this.restorable = requireNonNull(restorable, "restorable is null");
this.restorableId = String.format("%s (%s)", restorable.getClass().getSimpleName(), snapshotStateIdGenerator.apply(0L).getId());
@ -75,6 +81,12 @@ public class SingleInputSnapshotState
this.snapshotStateIdGenerator = requireNonNull(snapshotStateIdGenerator, "snapshotStateIdGenerator is null");
this.spillStateIdGenerator = requireNonNull(spillStateIdGenerator, "spillStateIdGenerator is null");
this.pagesSerde = pagesSerde;
this.snapshotMemoryContext = snapshotMemoryContext;
}
public void close()
{
snapshotMemoryContext.close();
}
/**
@ -154,6 +166,12 @@ public class SingleInputSnapshotState
private void captureState(long snapshotId, boolean record)
{
SnapshotStateId componentId = snapshotStateIdGenerator.apply(snapshotId);
long stateMemory = restorable.getUsedMemory();
if (!snapshotMemoryContext.trySetBytes(stateMemory)) {
LOG.warn("Insufficient memory on worker node to take snapshot");
snapshotManager.failedToCapture(componentId);
return;
}
try {
if (restorable.supportsConsolidatedWrites()) {
snapshotManager.storeConsolidatedState(componentId, restorable.capture(pagesSerde));
@ -176,6 +194,9 @@ public class SingleInputSnapshotState
LOG.warn(e, "Failed to capture and store snapshot state");
snapshotManager.failedToCapture(componentId);
}
finally {
snapshotMemoryContext.setBytes(0);
}
}
public boolean hasMarker()

1
presto-main/src/test/java/io/prestosql/execution/TestSqlStageExecution.java
View File

@ -68,6 +68,7 @@ import static org.testng.Assert.assertFalse;
import static org.testng.Assert.assertSame;
import static org.testng.Assert.assertTrue;
@Test(singleThreaded = true)
public class TestSqlStageExecution
{
private ExecutorService executor;

6
presto-main/src/test/java/io/prestosql/operator/TestOrderByOperator.java
View File

@ -84,6 +84,8 @@ import static org.testng.Assert.assertTrue;
@Test(singleThreaded = true)
public class TestOrderByOperator
{
private static final long defaultMemoryLimit = 1L << 28;
private ExecutorService executor;
private ScheduledExecutorService scheduledExecutor;
private DummySpillerFactory spillerFactory;
@ -404,7 +406,7 @@ public class TestOrderByOperator
Optional.of(spillerFactory),
new OrderingCompiler());
DriverContext driverContext = createDriverContext(8, TEST_SNAPSHOT_SESSION);
DriverContext driverContext = createDriverContext(defaultMemoryLimit, TEST_SNAPSHOT_SESSION);
driverContext.getPipelineContext().getTaskContext().getSnapshotManager().setTotalComponents(1);
OrderByOperator orderByOperator = (OrderByOperator) operatorFactory.createOperator(driverContext);
@ -426,7 +428,7 @@ public class TestOrderByOperator
}
// Step5: assume the task is rescheduled due to failure and everything is re-constructed
driverContext = createDriverContext(8, TEST_SNAPSHOT_SESSION);
driverContext = createDriverContext(defaultMemoryLimit, TEST_SNAPSHOT_SESSION);
operatorFactory = new OrderByOperatorFactory(
0,
new PlanNodeId("test"),

6
presto-main/src/test/java/io/prestosql/operator/TestWindowOperator.java
View File

@ -120,6 +120,8 @@ public class TestWindowOperator
private static final List<WindowFunctionDefinition> LEAD = ImmutableList.of(
window(new ReflectionWindowFunctionSupplier<>("lead", VARCHAR, ImmutableList.of(VARCHAR, BIGINT, VARCHAR), LeadFunction.class), VARCHAR, UNBOUNDED_FRAME, 1, 3, 4));
private static final long defaultMemoryLimit = 1L << 28;
private ExecutorService executor;
private ScheduledExecutorService scheduledExecutor;
private DummySpillerFactory spillerFactory;
@ -1223,7 +1225,7 @@ public class TestWindowOperator
spillerFactory,
new OrderingCompiler());
DriverContext driverContext = createDriverContext(0, TEST_SNAPSHOT_SESSION);
DriverContext driverContext = createDriverContext(defaultMemoryLimit, TEST_SNAPSHOT_SESSION);
WindowOperator windowOperator = (WindowOperator) operatorFactory.createOperator(driverContext);
// Step1: add the first 2 pages
@ -1340,7 +1342,7 @@ public class TestWindowOperator
ImmutableList.copyOf(new SortOrder[] {SortOrder.ASC_NULLS_LAST}),
false);
DriverContext driverContext = createDriverContext(0, TEST_SNAPSHOT_SESSION);
DriverContext driverContext = createDriverContext(defaultMemoryLimit, TEST_SNAPSHOT_SESSION);
WindowOperator windowOperator = (WindowOperator) operatorFactory.createOperator(driverContext);
// Step1: add the first 2 pages

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

@ -516,6 +516,12 @@ public class TestMultiInputSnapshotState
return this.supportsConsolidatedWrites;
}
@Override
public long getUsedMemory()
{
return 0;
}
public void setSupportsConsolidatedWrites(boolean supportsConsolidatedWrites)
{
this.supportsConsolidatedWrites = supportsConsolidatedWrites;

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

@ -16,6 +16,7 @@ package io.prestosql.snapshot;
import com.google.common.collect.ImmutableList;
import io.prestosql.execution.TaskId;
import io.prestosql.memory.context.LocalMemoryContext;
import io.prestosql.operator.DriverContext;
import io.prestosql.operator.Operator;
import io.prestosql.operator.OperatorContext;
@ -39,6 +40,7 @@ import static io.prestosql.SessionTestUtils.TEST_SNAPSHOT_SESSION;
import static io.prestosql.testing.TestingTaskContext.createTaskContext;
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.Mockito.mock;
import static org.mockito.Mockito.times;
@ -64,15 +66,18 @@ public class TestSingleInputSnapshotState
private TaskSnapshotManager snapshotManager;
private TestingRestorable restorable;
private SingleInputSnapshotState state;
private LocalMemoryContext snapshotMemoryContext;
@BeforeMethod
public void setup()
throws Exception
{
snapshotManager = mock(TaskSnapshotManager.class);
snapshotMemoryContext = mock(LocalMemoryContext.class);
when(snapshotMemoryContext.trySetBytes(anyLong())).thenReturn(true);
restorable = new TestingRestorable();
restorable.state = 100;
state = new SingleInputSnapshotState(restorable, snapshotManager, null, TestSingleInputSnapshotState::createSnapshotStateId, TestSingleInputSnapshotState::createSnapshotStateId);
state = new SingleInputSnapshotState(restorable, snapshotManager, null, TestSingleInputSnapshotState::createSnapshotStateId, TestSingleInputSnapshotState::createSnapshotStateId, snapshotMemoryContext);
}
private boolean processPage(Page page)
@ -199,7 +204,7 @@ public class TestSingleInputSnapshotState
public void testResumeBacktrack()
throws Exception
{
SingleInputSnapshotState state = new SingleInputSnapshotState(restorable, snapshotManager, null, TestSingleInputSnapshotState::createSnapshotStateId, TestSingleInputSnapshotState::createSnapshotStateId);
SingleInputSnapshotState state = new SingleInputSnapshotState(restorable, snapshotManager, null, TestSingleInputSnapshotState::createSnapshotStateId, TestSingleInputSnapshotState::createSnapshotStateId, snapshotMemoryContext);
state.processPage(regularPage);
restorable.state++;
int saved1 = restorable.state;
@ -222,7 +227,8 @@ public class TestSingleInputSnapshotState
snapshotManager,
null,
TestSingleInputSnapshotState::createSnapshotStateId,
TestSingleInputSnapshotState::createSnapshotStateId);
TestSingleInputSnapshotState::createSnapshotStateId,
snapshotMemoryContext);
state.processPage(marker1);
when(snapshotManager.loadState(anyObject())).thenReturn(Optional.of(1));
when(snapshotManager.loadFile(anyObject(), anyObject()))
@ -250,7 +256,8 @@ public class TestSingleInputSnapshotState
snapshotManager,
null,
TestSingleInputSnapshotState::createSnapshotStateId,
TestSingleInputSnapshotState::createSnapshotStateId);
TestSingleInputSnapshotState::createSnapshotStateId,
snapshotMemoryContext);
state.processPage(marker1);
when(snapshotManager.loadConsolidatedState(anyObject())).thenReturn(Optional.of(0));
state.processPage(resume1);
@ -271,7 +278,8 @@ public class TestSingleInputSnapshotState
snapshotManager,
null,
TestSingleInputSnapshotState::createSnapshotStateId,
TestSingleInputSnapshotState::createSnapshotStateId);
TestSingleInputSnapshotState::createSnapshotStateId,
snapshotMemoryContext);
state.processPage(marker1);
when(snapshotManager.loadState(anyObject())).thenReturn(Optional.of(0));
state.processPage(resume1);
@ -306,6 +314,12 @@ public class TestSingleInputSnapshotState
return this.supportsConsolidatedWrites;
}
@Override
public long getUsedMemory()
{
return 0;
}
public void setSupportsConsolidatedWrites(boolean supportsConsolidatedWrites)
{
this.supportsConsolidatedWrites = supportsConsolidatedWrites;

10
presto-spi/src/main/java/io/prestosql/spi/snapshot/Restorable.java
View File

@ -50,4 +50,14 @@ public interface Restorable
{
return true;
}
/**
* Finds the memory used to capture this object in a snapshot
*
* @return The size of the object created to be created in capture in bytes
*/
default long getUsedMemory()
{
throw new UnsupportedOperationException();
}
}