handle exclusion of incomplete spill files for snapshots

This commit is contained in:
Nitin Kashyap 2022-02-23 20:34:58 +05:30
parent 9f1c40ca2b
commit cfe310724c
No known key found for this signature in database
GPG Key ID: 7C4222EE7DDEFC35
7 changed files with 288 additions and 7 deletions

View File

@ -31,6 +31,7 @@ import io.prestosql.spi.type.Type;
import io.prestosql.spiller.Spiller;
import io.prestosql.spiller.SpillerFactory;
import io.prestosql.sql.gen.OrderingCompiler;
import org.apache.commons.lang3.tuple.Pair;
import java.io.Serializable;
import java.nio.file.Path;
@ -413,12 +414,14 @@ public class OrderByOperator
}
primarySpillRunning = true;
pageIndex.sort(sortChannels, sortOrder);
spillInProgress = spiller.get().spill(pageIndex.getSortedPages());
Pair<ListenableFuture<?>, Runnable> spillState = spiller.get().spillUnCommit(pageIndex.getSortedPages());
spillInProgress = spillState.getLeft();
LOG.debug("spilling to disk initiated by Order by operator using primary spiller");
finishMemoryRevoke = () -> {
pageIndex.clear();
updateMemoryUsage(true);
primarySpillRunning = false;
spillState.getRight().run();
};
return spillInProgress;
}
@ -523,12 +526,14 @@ public class OrderByOperator
verify(spiller.isPresent(), "spiller not present");
secondarySpillRunning = true;
secondaryPageIndex.sort(sortChannels, sortOrder);
spill2InProgress = spiller.get().spill(secondaryPageIndex.getSortedPages());
Pair<ListenableFuture<?>, Runnable> spillState = spiller.get().spillUnCommit(secondaryPageIndex.getSortedPages());
spill2InProgress = spillState.getLeft();
LOG.debug("spilling to disk initiated by Order by operator using secondary spiller");
finishMemoryRevoke2 = () -> {
secondaryPageIndex.clear();
updateMemoryUsage(false);
secondarySpillRunning = false;
spillState.getRight().run();
};
}

View File

@ -22,6 +22,8 @@ import io.prestosql.spi.Page;
import io.prestosql.spi.snapshot.BlockEncodingSerdeProvider;
import io.prestosql.spi.snapshot.RestorableConfig;
import io.prestosql.spi.type.Type;
import org.apache.commons.lang3.tuple.ImmutablePair;
import org.apache.commons.lang3.tuple.Pair;
import javax.annotation.concurrent.NotThreadSafe;
@ -31,14 +33,15 @@ import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.stream.Collectors;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.stream.IntStream;
import static com.google.common.base.Preconditions.checkState;
import static java.util.Objects.requireNonNull;
import static java.util.stream.Collectors.toList;
@NotThreadSafe
@RestorableConfig(uncapturedFields = {"types", "spillContext", "aggregatedMemoryContext", "singleStreamSpillerFactory", "closer", "previousSpill"})
@RestorableConfig(uncapturedFields = {"types", "spillContext", "aggregatedMemoryContext", "singleStreamSpillerFactory", "closer", "previousSpill", "spillCommitted"})
public class GenericSpiller
implements Spiller
{
@ -49,6 +52,7 @@ public class GenericSpiller
private final Closer closer = Closer.create();
private ListenableFuture<?> previousSpill = Futures.immediateFuture(null);
private final List<SingleStreamSpiller> singleStreamSpillers = new ArrayList<>();
private final List<AtomicBoolean> spillCommitted = new ArrayList<>();
public GenericSpiller(
List<Type> types,
@ -68,10 +72,28 @@ public class GenericSpiller
SingleStreamSpiller singleStreamSpiller = singleStreamSpillerFactory.create(types, spillContext, aggregatedMemoryContext.newLocalMemoryContext(GenericSpiller.class.getSimpleName()));
closer.register(singleStreamSpiller);
singleStreamSpillers.add(singleStreamSpiller);
spillCommitted.add(new AtomicBoolean(true));
previousSpill = singleStreamSpiller.spill(pageIterator);
return previousSpill;
}
/**
* Initiate spilling of pages stream. Returns completed future once spilling has finished with commit function.
*
* @param pageIterator
*/
@Override
public Pair<ListenableFuture<?>, Runnable> spillUnCommit(Iterator<Page> pageIterator)
{
SingleStreamSpiller singleStreamSpiller = singleStreamSpillerFactory.create(types, spillContext, aggregatedMemoryContext.newLocalMemoryContext(GenericSpiller.class.getSimpleName()));
closer.register(singleStreamSpiller);
singleStreamSpillers.add(singleStreamSpiller);
AtomicBoolean isCommitted = new AtomicBoolean(false);
spillCommitted.add(isCommitted);
previousSpill = singleStreamSpiller.spill(pageIterator);
return ImmutablePair.of(previousSpill, () -> isCommitted.set(true));
}
@Override
public List<Iterator<Page>> getSpills()
{
@ -110,15 +132,18 @@ public class GenericSpiller
@Override
public List<Path> getSpilledFilePaths()
{
return singleStreamSpillers.stream().map(s -> s.getFile()).collect(Collectors.toList());
return IntStream.range(0, spillCommitted.size()).filter(i -> spillCommitted.get(i).get()).mapToObj(o -> singleStreamSpillers.get(o).getFile()).collect(toList());
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{
GenericSpillerState myState = new GenericSpillerState();
for (SingleStreamSpiller s : singleStreamSpillers) {
myState.singleStreamSpillers.add(s.capture(serdeProvider));
for (int i = 0; i < singleStreamSpillers.size(); i++) {
if (spillCommitted.get(i).get()) {
SingleStreamSpiller s = singleStreamSpillers.get(i);
myState.singleStreamSpillers.add(s.capture(serdeProvider));
}
}
return myState;
}
@ -131,6 +156,7 @@ public class GenericSpiller
SingleStreamSpiller singleStreamSpiller = singleStreamSpillerFactory.create(types, spillContext, aggregatedMemoryContext.newLocalMemoryContext(GenericSpiller.class.getSimpleName()));
singleStreamSpiller.restore(s, serdeProvider);
this.singleStreamSpillers.add(singleStreamSpiller);
this.spillCommitted.add(new AtomicBoolean(true));
this.closer.register(singleStreamSpiller);
}
}

View File

@ -69,4 +69,14 @@ public interface SingleStreamSpiller
void deleteFile();
Path getFile();
default void commit()
{
/* do nothing */
}
default boolean isCommitted()
{
return true;
}
}

View File

@ -17,6 +17,8 @@ import com.google.common.collect.ImmutableList;
import com.google.common.util.concurrent.ListenableFuture;
import io.prestosql.spi.Page;
import io.prestosql.spi.snapshot.Restorable;
import org.apache.commons.lang3.tuple.ImmutablePair;
import org.apache.commons.lang3.tuple.Pair;
import java.io.Closeable;
import java.nio.file.Path;
@ -31,6 +33,14 @@ public interface Spiller
*/
ListenableFuture<?> spill(Iterator<Page> pageIterator);
/**
* Initiate spilling of pages stream. Returns completed future once spilling has finished with commit function.
*/
default Pair<ListenableFuture<?>, Runnable> spillUnCommit(Iterator<Page> pageIterator)
{
return ImmutablePair.of(spill(pageIterator), () -> {});
}
/**
* Returns list of previously spilled Pages streams.
*/

View File

@ -17,14 +17,19 @@ import com.google.common.collect.ImmutableList;
import com.google.common.util.concurrent.ListenableFuture;
import io.prestosql.memory.context.AggregatedMemoryContext;
import io.prestosql.spi.Page;
import io.prestosql.spi.snapshot.BlockEncodingSerdeProvider;
import io.prestosql.spi.snapshot.RestorableConfig;
import io.prestosql.spi.type.Type;
import io.prestosql.spiller.Spiller;
import io.prestosql.spiller.SpillerFactory;
import org.apache.commons.lang3.tuple.ImmutablePair;
import org.apache.commons.lang3.tuple.Pair;
import java.io.Serializable;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.util.concurrent.Futures.immediateFuture;
@ -43,15 +48,27 @@ public class DummySpillerFactory
private final RestorableConfig restorableConfig = null;
private final List<Iterable<Page>> spills = new ArrayList<>();
private final List<AtomicBoolean> spillCommitted = new ArrayList<>();
@Override
public ListenableFuture<?> spill(Iterator<Page> pageIterator)
{
spillsCount++;
spills.add(ImmutableList.copyOf(pageIterator));
spillCommitted.add(new AtomicBoolean(true));
return immediateFuture(null);
}
@Override
public Pair<ListenableFuture<?>, Runnable> spillUnCommit(Iterator<Page> pageIterator)
{
spillsCount++;
spills.add(ImmutableList.copyOf(pageIterator));
AtomicBoolean isCommitted = new AtomicBoolean(false);
spillCommitted.add(isCommitted);
return ImmutablePair.of(immediateFuture(null), () -> isCommitted.set(true));
}
@Override
public List<Iterator<Page>> getSpills()
{
@ -65,6 +82,49 @@ public class DummySpillerFactory
{
spills.clear();
}
/**
* Capture this object's internal state, so it can be used later to restore to the same state.
*
* @param serdeProvider
* @return An object representing internal state of the current object
*/
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{
DummySpillerState myState = new DummySpillerState();
for (int i = 0; i < spills.size(); i++) {
if (spillCommitted.get(i).get()) {
List<Page> pages = new ArrayList<>();
spills.get(i).forEach(pg -> pages.add(pg));
myState.spills.add(pages);
}
}
return myState;
}
/**
* Restore this object's internal state according to the snapshot
*
* @param state an object that represents this object's snapshot state
* @param serdeProvider
*/
@Override
public void restore(Object state, BlockEncodingSerdeProvider serdeProvider)
{
DummySpillerState myState = (DummySpillerState) state;
this.spills.clear();
for (List<Page> s : myState.spills) {
this.spills.add(s);
this.spillCommitted.add(new AtomicBoolean(true));
}
}
class DummySpillerState
implements Serializable
{
List<List<Page>> spills = new ArrayList<>();
}
};
}

View File

@ -309,6 +309,54 @@ public class TestOrderByOperator
assertOperatorEquals(operatorFactory, driverContext, input, expected, revokeMemoryWhenAddingPages);
}
@Test
public void testReverseOrderWithSnapshot()
{
List<Page> input = rowPagesBuilder(BIGINT, DOUBLE)
.row(1L, 0.1)
.row(2L, 0.2)
.pageBreak()
.row(-1L, -0.1)
.row(4L, 0.4)
.build();
OrderByOperatorFactory operatorFactory = new OrderByOperatorFactory(
0,
new PlanNodeId("test"),
ImmutableList.of(BIGINT, DOUBLE),
ImmutableList.of(0),
10,
ImmutableList.of(0),
ImmutableList.of(DESC_NULLS_LAST),
new PagesIndex.TestingFactory(false),
true,
Optional.of(spillerFactory),
new OrderingCompiler(),
true);
DriverContext driverContext = createDriverContext(8, TEST_SESSION);
MaterializedResult expected = resultBuilder(driverContext.getSession(), BIGINT)
.row(4L)
.row(2L)
.row(1L)
.row(-1L)
.build();
assertOperatorEqualsWithSimpleSelfStateComparison(operatorFactory, driverContext, input, expected, true, createExpectedMappingRevoke());
}
private Map<String, Object> createExpectedMappingRevoke()
{
Map<String, Object> expectedMapping = new HashMap<>();
expectedMapping.put("operatorContext", 0);
expectedMapping.put("revocableMemoryContext", 1288L);
expectedMapping.put("localUserMemoryContext", 0L);
expectedMapping.put("secondaryMemoryContext", 0L);
expectedMapping.put("secondarySpillRunning", false);
expectedMapping.put("primarySpillRunning", false);
return expectedMapping;
}
@Test(expectedExceptions = ExceededMemoryLimitException.class, expectedExceptionsMessageRegExp = "Query exceeded per-node user memory limit of 10B.*")
public void testMemoryLimit()
{

View File

@ -14,6 +14,7 @@
package io.prestosql.spiller;
import com.google.common.collect.ImmutableList;
import com.google.common.util.concurrent.ListenableFuture;
import io.hetu.core.transport.execution.buffer.PagesSerde;
import io.hetu.core.transport.execution.buffer.PagesSerdeFactory;
import io.prestosql.RowPagesBuilder;
@ -23,12 +24,18 @@ import io.prestosql.spi.Page;
import io.prestosql.spi.block.BlockBuilder;
import io.prestosql.spi.type.Type;
import io.prestosql.sql.analyzer.FeaturesConfig;
import io.prestosql.testing.TestingPagesSerdeFactory;
import org.apache.commons.lang3.tuple.Pair;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;
import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.ExecutionException;
@ -44,6 +51,7 @@ import static io.prestosql.spi.type.VarbinaryType.VARBINARY;
import static io.prestosql.spi.type.VarcharType.VARCHAR;
import static java.lang.Double.doubleToLongBits;
import static java.nio.file.Files.createTempDirectory;
import static java.nio.file.StandardCopyOption.REPLACE_EXISTING;
import static org.testng.Assert.assertEquals;
@Test(singleThreaded = true)
@ -52,6 +60,7 @@ public class TestBinaryFileSpiller
private static final List<Type> TYPES = ImmutableList.of(BIGINT, VARCHAR, DOUBLE, BIGINT);
private final File spillPath = createTempDirectory(getClass().getName()).toFile();
private final File spillUploadPath = createTempDirectory(getClass().getName()).toFile();
private SpillerStats spillerStats;
private FileSingleStreamSpillerFactory singleStreamSpillerFactory;
private SpillerFactory factory;
@ -89,6 +98,9 @@ public class TestBinaryFileSpiller
{
singleStreamSpillerFactory.destroy();
deleteRecursively(spillPath.toPath(), ALLOW_INSECURE);
if (spillUploadPath.exists()) {
deleteRecursively(spillUploadPath.toPath(), ALLOW_INSECURE);
}
}
@Test
@ -172,4 +184,114 @@ public class TestBinaryFileSpiller
spiller.close();
assertEquals(memoryContext.getBytes(), 0);
}
@Test
public void testFileVarbinaryComittableSpiller()
throws Exception
{
List<Type> types = ImmutableList.of(BIGINT, DOUBLE, VARBINARY);
BlockBuilder col1 = BIGINT.createBlockBuilder(null, 1);
BlockBuilder col2 = DOUBLE.createBlockBuilder(null, 1);
BlockBuilder col3 = VARBINARY.createBlockBuilder(null, 1);
col1.writeLong(42).closeEntry();
col2.writeLong(doubleToLongBits(43.0)).closeEntry();
col3.writeLong(doubleToLongBits(43.0)).writeLong(1).closeEntry();
Page page = new Page(col1.build(), col2.build(), col3.build());
testSpillerUnCommit(types,
ImmutableList.of(page),
ImmutableList.of(page, page),
ImmutableList.of(page, page, page));
}
private void testSpillerUnCommit(List<Type> types, List<Page>... spills)
throws ExecutionException, InterruptedException, IOException
{
long spilledBytesBefore = spillerStats.getTotalSpilledBytes();
long spilledBytes = 0;
assertEquals(memoryContext.getBytes(), 0);
List<Runnable> runners = new ArrayList<>();
PagesSerde serde = TestingPagesSerdeFactory.testingPagesSerde();
Spiller spiller = factory.create(TYPES, bytes -> {}, memoryContext);
spilledBytes = doSpill(spiller, spilledBytes, runners, spills, 0);
spillUploadPath.mkdirs();
int counter = 1;
int runCount = runners.size();
List<Path> uploadedFile = new ArrayList<>();
for (counter = 1; counter <= runCount; counter++) {
runners.remove(0).run();
assertEquals(spiller.getSpilledFilePaths().size(), counter);
assertEquals(runners.size(), runCount - counter);
Object snapshot = spiller.capture(serde);
spiller.getSpilledFilePaths().stream().forEach(path -> {
try {
Files.copy(path, Paths.get(spillUploadPath.getPath(), path.getFileName().toString()), REPLACE_EXISTING);
uploadedFile.add(Paths.get(spillUploadPath.getPath(), path.getFileName().toString()));
}
catch (IOException e) {
e.printStackTrace();
}
});
spiller.close();
spiller = factory.create(TYPES, bytes -> {}, memoryContext);
spiller.restore(snapshot, serde);
uploadedFile.stream().forEach(path -> {
try {
Files.move(path, Paths.get(spillPath.getPath(), path.getFileName().toString()), REPLACE_EXISTING);
}
catch (IOException e) {
e.printStackTrace();
}
});
uploadedFile.clear();
doSpill(spiller, spilledBytes, runners, spills, counter);
}
//assertEquals(spillerStats.getTotalSpilledBytes() - spilledBytesBefore, spilledBytes);
// At this point, the buffers should still be accounted for in the memory context, because
// the spiller (FileSingleStreamSpiller) doesn't release its memory reservation until it's closed.
//assertEquals(memoryContext.getBytes(), spills.length * FileSingleStreamSpiller.BUFFER_SIZE);
List<Iterator<Page>> actualSpills = spiller.getSpills();
assertEquals(actualSpills.size(), spills.length);
for (int i = 0; i < actualSpills.size(); i++) {
List<Page> actualSpill = ImmutableList.copyOf(actualSpills.get(i));
List<Page> expectedSpill = spills[i];
assertEquals(actualSpill.size(), expectedSpill.size());
for (int j = 0; j < actualSpill.size(); j++) {
assertPageEquals(types, actualSpill.get(j), expectedSpill.get(j));
}
}
spiller.close();
assertEquals(memoryContext.getBytes(), 0);
}
private long doSpill(Spiller spiller, long spilledBytes, List<Runnable> runners, List<Page>[] spills, int count) throws InterruptedException, ExecutionException
{
runners.clear();
for (List<Page> spill : spills) {
if (count > 0) {
count--;
continue;
}
spilledBytes += spill.stream()
.mapToLong(page -> pagesSerde.serialize(page).getSizeInBytes())
.sum();
Pair<ListenableFuture<?>, Runnable> spillState = spiller.spillUnCommit(spill.iterator());
runners.add(spillState.getRight());
spillState.getLeft().get();
}
return spilledBytes;
}
}