Examples with RestorableConfig io.prestosql.spi.snapshot.RestorableConfig used on opensource projects

Example 1 with RestorableConfig

use of io.prestosql.spi.snapshot.RestorableConfig in project hetu-core by openlookeng.

the class TestSnapshotCompleteness method getAllFields.

private void getAllFields(List<String> allFields, Class<?> clazz) throws Exception {
    // Find fields in non-Restorable base class.
    if (clazz.getSuperclass() != null && !Restorable.class.isAssignableFrom(clazz.getSuperclass())) {
        getAllFields(allFields, clazz.getSuperclass());
    }
    // My own fields
    List<String> currentClassFields = Arrays.stream(clazz.getDeclaredFields()).filter(field -> !Modifier.isStatic(field.getModifiers()) && !(Modifier.isFinal(field.getModifiers()) && isImmutable(field.getGenericType()))).map(field -> clazz.getName() + "." + field.getName()).collect(Collectors.toList());
    allFields.addAll(currentClassFields);
    // Get the uncapturedFields and baseState field name
    RestorableConfig config = getConfigAnnotation(clazz);
    String[] uncapturedFields = {};
    String baseClassStateName = "baseState";
    if (config != null) {
        uncapturedFields = config.uncapturedFields();
        baseClassStateName = config.baseClassStateName();
    }
    // If base class is Restorable and has any captured fields, the current class needs to have a field for base class state.
    if (clazz.getSuperclass() != null && Restorable.class.isAssignableFrom(clazz.getSuperclass())) {
        int baseFieldCount = test(clazz.getSuperclass());
        if (baseFieldCount > 0) {
            allFields.add(baseClassStateName);
        }
    }
    outer: for (String uncaptured : uncapturedFields) {
        int fragmentCount = uncaptured.split("\\.").length;
        if (fragmentCount > 2) {
            // Assumed to be "full-package-name.class.field"
            if (!allFields.remove(uncaptured)) {
                unusedUncapturedFields.computeIfAbsent(clazz, ignored -> new HashSet<>()).add(uncaptured);
            }
        } else {
            // Try to look for a match in all base class fields from back to front.
            for (int i = allFields.size() - 1; i >= 0; i--) {
                // $ before inner class names, . for others
                if (allFields.get(i).endsWith("$" + uncaptured) || allFields.get(i).endsWith("." + uncaptured)) {
                    allFields.remove(i);
                    continue outer;
                }
            }
            unusedUncapturedFields.computeIfAbsent(clazz, ignored -> new HashSet<>()).add(uncaptured);
        }
    }
}

Example 2 with RestorableConfig

use of io.prestosql.spi.snapshot.RestorableConfig in project hetu-core by openlookeng.

the class MergeSortedPages method buildPage.

private static WorkProcessor<Page> buildPage(WorkProcessor<PageWithPosition> pageWithPositions, List<Integer> outputChannels, List<Type> outputTypes, BiPredicate<PageBuilder, PageWithPosition> pageBreakPredicate, boolean updateMemoryAfterEveryPosition, AggregatedMemoryContext aggregatedMemoryContext, DriverYieldSignal yieldSignal) {
    LocalMemoryContext memoryContext = aggregatedMemoryContext.newLocalMemoryContext(MergeSortedPages.class.getSimpleName());
    PageBuilder pageBuilder = new PageBuilder(outputTypes);
    return pageWithPositions.yielding(yieldSignal::isSet).transform(new WorkProcessor.Transformation<PageWithPosition, Page>() {

        @RestorableConfig(stateClassName = "BuildPageState", uncapturedFields = { "val$outputTypes", "val$pageBreakPredicate", "val$outputChannels" })
        private final RestorableConfig restorableConfig = null;

        @Override
        public TransformationState<Page> process(@Nullable PageWithPosition pageWithPosition) {
            boolean finished = pageWithPosition == null;
            if (finished && pageBuilder.isEmpty()) {
                memoryContext.close();
                return TransformationState.finished();
            }
            if (finished || pageBreakPredicate.test(pageBuilder, pageWithPosition)) {
                if (!updateMemoryAfterEveryPosition) {
                    // update memory usage just before producing page to cap from top
                    memoryContext.setBytes(pageBuilder.getRetainedSizeInBytes());
                }
                Page page = pageBuilder.build();
                pageBuilder.reset();
                if (!finished) {
                    pageWithPosition.appendTo(pageBuilder, outputChannels, outputTypes);
                }
                if (updateMemoryAfterEveryPosition) {
                    memoryContext.setBytes(pageBuilder.getRetainedSizeInBytes());
                }
                return TransformationState.ofResult(page, !finished);
            }
            pageWithPosition.appendTo(pageBuilder, outputChannels, outputTypes);
            if (updateMemoryAfterEveryPosition) {
                memoryContext.setBytes(pageBuilder.getRetainedSizeInBytes());
            }
            return TransformationState.needsMoreData();
        }

        @Override
        public Object captureResult(Page result, BlockEncodingSerdeProvider serdeProvider) {
            if (result != null) {
                SerializedPage serializedPage = ((PagesSerde) serdeProvider).serialize(result);
                return serializedPage.capture(serdeProvider);
            }
            return null;
        }

        @Override
        public Page restoreResult(Object resultState, BlockEncodingSerdeProvider serdeProvider) {
            if (resultState != null) {
                return ((PagesSerde) serdeProvider).deserialize(SerializedPage.restoreSerializedPage(resultState));
            }
            return null;
        }

        @Override
        public Object captureInput(PageWithPosition input, BlockEncodingSerdeProvider serdeProvider) {
            if (input != null) {
                return input.capture(serdeProvider);
            }
            return null;
        }

        @Override
        public PageWithPosition restoreInput(Object inputState, PageWithPosition input, BlockEncodingSerdeProvider serdeProvider) {
            if (inputState != null) {
                return PageWithPosition.restorePageWithPosition(inputState, serdeProvider);
            }
            return null;
        }

        @Override
        public Object capture(BlockEncodingSerdeProvider serdeProvider) {
            BuildPageState myState = new BuildPageState();
            myState.pageBuilder = pageBuilder.capture(serdeProvider);
            myState.memoryContext = memoryContext.getBytes();
            return myState;
        }

        @Override
        public void restore(Object state, BlockEncodingSerdeProvider serdeProvider) {
            BuildPageState myState = (BuildPageState) state;
            pageBuilder.restore(myState.pageBuilder, serdeProvider);
            memoryContext.setBytes(myState.memoryContext);
        }
    });
}

Example 3 with RestorableConfig

use of io.prestosql.spi.snapshot.RestorableConfig in project hetu-core by openlookeng.

the class DummySpillerFactory method create.

@Override
public Spiller create(List<Type> types, SpillContext spillContext, AggregatedMemoryContext memoryContext) {
    return new Spiller() {

        @RestorableConfig(unsupported = true)
        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() {
            return spills.stream().map(Iterable::iterator).collect(toImmutableList());
        }

        @Override
        public void close() {
            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<>();
        }
    };
}

Example 4 with RestorableConfig

use of io.prestosql.spi.snapshot.RestorableConfig in project hetu-core by openlookeng.

the class WorkProcessorUtils method flatMap.

static <T, R> WorkProcessor<R> flatMap(WorkProcessor<T> processor, WorkProcessor.RestorableFunction<T, WorkProcessor<R>> mapper) {
    requireNonNull(processor, "processor is null");
    requireNonNull(mapper, "mapper is null");
    return processor.flatTransform(new Transformation<T, WorkProcessor<R>>() {

        @RestorableConfig(uncapturedFields = "mapper")
        private final RestorableConfig restorableConfig = null;

        T savedElement;

        @Override
        public TransformationState<WorkProcessor<R>> process(@org.jetbrains.annotations.Nullable T element) {
            if (element == null) {
                return TransformationState.finished();
            }
            savedElement = element;
            return TransformationState.ofResult(mapper.apply(element));
        }

        @Override
        public Object captureResult(WorkProcessor<R> result, BlockEncodingSerdeProvider serdeProvider) {
            return result.capture(serdeProvider);
        }

        @Override
        public WorkProcessor<R> restoreResult(Object resultState, BlockEncodingSerdeProvider serdeProvider) {
            if (savedElement != null) {
                WorkProcessor<R> restoredProcessor = mapper.apply(savedElement);
                restoredProcessor.restore(resultState, serdeProvider);
                return restoredProcessor;
            }
            return null;
        }

        @Override
        public Object captureInput(T input, BlockEncodingSerdeProvider serdeProvider) {
            return mapper.captureInput(input, serdeProvider);
        }

        @Override
        public T restoreInput(Object inputState, T input, BlockEncodingSerdeProvider serdeProvider) {
            return mapper.restoreInput(inputState, serdeProvider);
        }

        @Override
        public Object capture(BlockEncodingSerdeProvider serdeProvider) {
            return mapper.captureInput(savedElement, serdeProvider);
        }

        @Override
        public void restore(Object state, BlockEncodingSerdeProvider serdeProvider) {
            this.savedElement = mapper.restoreInput(state, serdeProvider);
        }
    });
}

Example 5 with RestorableConfig

use of io.prestosql.spi.snapshot.RestorableConfig in project hetu-core by openlookeng.

the class WorkProcessorUtils method transform.

static <T, R> WorkProcessor<R> transform(WorkProcessor<T> processor, Transformation<T, R> transformation) {
    requireNonNull(processor, "processor is null");
    requireNonNull(transformation, "transformation is null");
    return create(new WorkProcessor.Process<R>() {

        @RestorableConfig(stateClassName = "TransformState", uncapturedFields = { "element" })
        private final RestorableConfig restorableConfig = null;

        T element;

        @Override
        public ProcessState<R> process() {
            while (true) {
                if (element == null && !processor.isFinished()) {
                    if (processor.process()) {
                        if (!processor.isFinished()) {
                            element = requireNonNull(processor.getResult(), "result is null");
                        }
                    } else if (processor.isBlocked()) {
                        return ProcessState.blocked(processor.getBlockedFuture());
                    } else {
                        return ProcessState.yield();
                    }
                }
                TransformationState<R> state = requireNonNull(transformation.process(element), "state is null");
                if (state.isNeedsMoreData()) {
                    checkState(!processor.isFinished(), "Cannot request more data when base processor is finished");
                    // set element to empty() in order to fetch a new one
                    element = null;
                }
                // pass-through transformation state if it doesn't require new data
                switch(state.getType()) {
                    case NEEDS_MORE_DATA:
                        break;
                    case BLOCKED:
                        return ProcessState.blocked(state.getBlocked());
                    case YIELD:
                        return ProcessState.yield();
                    case RESULT:
                        return ProcessState.ofResult(state.getResult());
                    case FINISHED:
                        return ProcessState.finished();
                }
            }
        }

        @Override
        public Object capture(BlockEncodingSerdeProvider serdeProvider) {
            TransformState myState = new TransformState();
            myState.transformation = transformation.capture(serdeProvider);
            myState.validElement = element != null;
            myState.processor = processor.capture(serdeProvider);
            return myState;
        }

        @Override
        public void restore(Object state, BlockEncodingSerdeProvider serdeProvider) {
            TransformState myState = (TransformState) state;
            transformation.restore(myState.transformation, serdeProvider);
            processor.restore(myState.processor, serdeProvider);
            if (myState.validElement) {
                checkArgument(processor.getResult() != null);
                element = processor.getResult();
            }
        }

        @Override
        public Object captureResult(R result, BlockEncodingSerdeProvider serdeProvider) {
            return transformation.captureResult(result, serdeProvider);
        }

        @Override
        public R restoreResult(Object resultState, BlockEncodingSerdeProvider serdeProvider) {
            return transformation.restoreResult(resultState, serdeProvider);
        }
    });
}