Examples with VoidTraversingQueryContextVisitor io.prestosql.memory.VoidTraversingQueryContextVisitor used on opensource projects

Example 1 with VoidTraversingQueryContextVisitor

use of io.prestosql.memory.VoidTraversingQueryContextVisitor in project hetu-core by openlookeng.

the class MemoryRevokingScheduler method requestRevoking.

private void requestRevoking(MemoryPool memoryPool, Collection<SqlTask> sqlTasks, long remainingBytesToRevoke) {
    AtomicLong remainingBytesToRevokeAtomic = new AtomicLong(remainingBytesToRevoke);
    Map<Long, List<SqlTask>> revocableMemoryMapping = new HashMap<>();
    // This approach disregards create time
    if (prioritizeLargerSpiltsMemoryRevoke) {
        sqlTasks.stream().filter(task -> task.getTaskStatus().getState() == TaskState.RUNNING).filter(task -> task.getQueryContext().getMemoryPool() == memoryPool).forEach(task -> task.getQueryContext().accept(new VoidTraversingQueryContextVisitor<AtomicLong>() {

            @Override
            public Void visitQueryContext(QueryContext queryContext, AtomicLong remainingBytesToRevoke) {
                if (remainingBytesToRevoke.get() < 0) {
                    // exit immediately if no work needs to be done
                    return null;
                }
                return super.visitQueryContext(queryContext, remainingBytesToRevoke);
            }

            @Override
            public Void visitOperatorContext(OperatorContext operatorContext, AtomicLong remainingBytesToRevoke) {
                if (remainingBytesToRevoke.get() <= 0) {
                    return null;
                }
                long revocableBytes = operatorContext.getRevocableMemory();
                if (revocableBytes > 0) {
                    if (revocableBytes >= revocableMemorySelectionThreshold) {
                        remainingBytesToRevoke.addAndGet(-revocableBytes);
                    }
                    revocableMemoryMapping.computeIfAbsent(revocableBytes, list -> new ArrayList()).add(task);
                }
                return null;
            }
        }, remainingBytesToRevokeAtomic));
        if (revocableMemoryMapping.size() > 0) {
            revocableMemoryMapping.entrySet().stream().sorted(Map.Entry.comparingByKey(Comparator.reverseOrder())).flatMap(list -> list.getValue().stream()).filter(task -> task.getTaskStatus().getState() == TaskState.RUNNING).forEach(task -> task.getQueryContext().accept(new VoidTraversingQueryContextVisitor<AtomicLong>() {

                @Override
                public Void visitQueryContext(QueryContext queryContext, AtomicLong remainingBytesToRevoke) {
                    if (remainingBytesToRevoke.get() < 0) {
                        // exit immediately if no work needs to be done
                        return null;
                    }
                    return super.visitQueryContext(queryContext, remainingBytesToRevoke);
                }

                @Override
                public Void visitOperatorContext(OperatorContext operatorContext, AtomicLong remainingBytesToRevoke) {
                    if (remainingBytesToRevoke.get() > 0) {
                        long revokedBytes = operatorContext.requestMemoryRevoking();
                        if (revokedBytes > 0) {
                            remainingBytesToRevoke.addAndGet(-revokedBytes);
                            log.debug("memoryPool=%s: requested revoking %s; remaining %s", memoryPool.getId(), revokedBytes, remainingBytesToRevoke.get());
                        }
                    }
                    return null;
                }
            }, remainingBytesToRevokeAtomic));
        }
    } else {
        sqlTasks.stream().filter(task -> task.getTaskStatus().getState() == TaskState.RUNNING).filter(task -> task.getQueryContext().getMemoryPool() == memoryPool).sorted(ORDER_BY_CREATE_TIME).forEach(task -> task.getQueryContext().accept(new VoidTraversingQueryContextVisitor<AtomicLong>() {

            @Override
            public Void visitQueryContext(QueryContext queryContext, AtomicLong remainingBytesToRevoke) {
                if (remainingBytesToRevoke.get() < 0) {
                    // exit immediately if no work needs to be done
                    return null;
                }
                return super.visitQueryContext(queryContext, remainingBytesToRevoke);
            }

            @Override
            public Void visitOperatorContext(OperatorContext operatorContext, AtomicLong remainingBytesToRevoke) {
                if (remainingBytesToRevoke.get() > 0) {
                    long revokedBytes = operatorContext.requestMemoryRevoking();
                    if (revokedBytes > 0) {
                        remainingBytesToRevoke.addAndGet(-revokedBytes);
                        log.debug("memoryPool=%s: requested revoking %s; remaining %s", memoryPool.getId(), revokedBytes, remainingBytesToRevoke.get());
                    }
                }
                return null;
            }
        }, remainingBytesToRevokeAtomic));
    }
}