Examples with LocalMemoryContext io.trino.memory.context.LocalMemoryContext used on opensource projects

Example 1 with LocalMemoryContext

use of io.trino.memory.context.LocalMemoryContext in project trino by trinodb.

the class SpoolingExchangeOutputBuffer method updateMemoryUsage.

private void updateMemoryUsage(long bytes) {
    LocalMemoryContext context = getSystemMemoryContextOrNull();
    if (context != null) {
        context.setBytes(bytes);
    }
    updatePeakMemoryUsage(bytes);
}

Example 2 with LocalMemoryContext

use of io.trino.memory.context.LocalMemoryContext in project trino by trinodb.

the class OutputBufferMemoryManager method close.

public synchronized void close() {
    updateMemoryUsage(-bufferedBytes.get());
    LocalMemoryContext memoryContext = getMemoryContextOrNull();
    if (memoryContext != null) {
        memoryContext.close();
    }
    closed = true;
}

Example 3 with LocalMemoryContext

use of io.trino.memory.context.LocalMemoryContext in project trino by trinodb.

the class OutputBufferMemoryManager method updateMemoryUsage.

public void updateMemoryUsage(long bytesAdded) {
    // If the memoryContext doesn't exist, the task is probably already
    // aborted, so we can just return (see the comment in getMemoryContextOrNull()).
    LocalMemoryContext memoryContext = getMemoryContextOrNull();
    if (memoryContext == null) {
        return;
    }
    ListenableFuture<Void> waitForMemory = null;
    SettableFuture<Void> notifyUnblocked = null;
    long currentBufferedBytes;
    synchronized (this) {
        // we can also safely ignore any calls after OutputBufferMemoryManager is closed.
        if (closed) {
            return;
        }
        currentBufferedBytes = bufferedBytes.updateAndGet(bytes -> {
            long result = bytes + bytesAdded;
            checkArgument(result >= 0, "bufferedBytes (%s) plus delta (%s) would be negative", bytes, bytesAdded);
            return result;
        });
        ListenableFuture<Void> blockedOnMemory = memoryContext.setBytes(currentBufferedBytes);
        if (!blockedOnMemory.isDone()) {
            if (this.blockedOnMemory != blockedOnMemory) {
                this.blockedOnMemory = blockedOnMemory;
                // only register a callback when blocked and the future is different
                waitForMemory = blockedOnMemory;
            }
        } else {
            this.blockedOnMemory = NOT_BLOCKED;
            if (currentBufferedBytes <= maxBufferedBytes || !blockOnFull.get()) {
                // Complete future in a new thread to avoid making a callback on the caller thread.
                // This make is easier for callers to use this class since they can update the memory
                // usage while holding locks.
                notifyUnblocked = this.bufferBlockedFuture;
                this.bufferBlockedFuture = null;
            }
        }
    }
    peakMemoryUsage.accumulateAndGet(currentBufferedBytes, Math::max);
    // Notify listeners outside of the critical section
    notifyListener(notifyUnblocked);
    if (waitForMemory != null) {
        waitForMemory.addListener(this::onMemoryAvailable, notificationExecutor);
    }
}

Example 4 with LocalMemoryContext

use of io.trino.memory.context.LocalMemoryContext in project trino by trinodb.

the class MergeSortedPages method pageWithPositions.

private static WorkProcessor<PageWithPosition> pageWithPositions(WorkProcessor<Page> pages, AggregatedMemoryContext aggregatedMemoryContext) {
    return pages.flatMap(page -> {
        LocalMemoryContext memoryContext = aggregatedMemoryContext.newLocalMemoryContext(MergeSortedPages.class.getSimpleName());
        memoryContext.setBytes(page.getRetainedSizeInBytes());
        return WorkProcessor.create(new WorkProcessor.Process<PageWithPosition>() {

            int position;

            @Override
            public ProcessState<PageWithPosition> process() {
                if (position >= page.getPositionCount()) {
                    memoryContext.close();
                    return ProcessState.finished();
                }
                return ProcessState.ofResult(new PageWithPosition(page, position++));
            }
        });
    });
}

Example 5 with LocalMemoryContext

use of io.trino.memory.context.LocalMemoryContext in project trino by trinodb.

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