@PostConstruct public void start() { registerPeriodicCheck(); registerPoolListeners(); }
private void registerPeriodicCheck() { this.scheduledFuture = taskManagementExecutor.scheduleWithFixedDelay(() -> { try { requestMemoryRevokingIfNeeded(); } catch (Throwable e) { log.error(e, "Error requesting system memory revoking"); } }, 1, 1, SECONDS); }
@VisibleForTesting MemoryRevokingScheduler( List<MemoryPool> memoryPools, Supplier<? extends Collection<SqlTask>> currentTasksSupplier, ScheduledExecutorService taskManagementExecutor, double memoryRevokingThreshold, double memoryRevokingTarget) { this.memoryPools = ImmutableList.copyOf(requireNonNull(memoryPools, "memoryPools is null")); this.currentTasksSupplier = requireNonNull(currentTasksSupplier, "currentTasksSupplier is null"); this.taskManagementExecutor = requireNonNull(taskManagementExecutor, "taskManagementExecutor is null"); this.memoryRevokingThreshold = checkFraction(memoryRevokingThreshold, "memoryRevokingThreshold"); this.memoryRevokingTarget = checkFraction(memoryRevokingTarget, "memoryRevokingTarget"); checkArgument( memoryRevokingTarget <= memoryRevokingThreshold, "memoryRevokingTarget should be less than or equal memoryRevokingThreshold, but got %s and %s respectively", memoryRevokingTarget, memoryRevokingThreshold); }
private synchronized void runMemoryRevoking() { if (checkPending.getAndSet(false)) { Collection<SqlTask> sqlTasks = null; for (MemoryPool memoryPool : memoryPools) { if (!memoryRevokingNeeded(memoryPool)) { continue; } if (sqlTasks == null) { sqlTasks = requireNonNull(currentTasksSupplier.get()); } requestMemoryRevoking(memoryPool, sqlTasks); } } }
private void requestMemoryRevoking(MemoryPool memoryPool, Collection<SqlTask> sqlTasks) { long remainingBytesToRevoke = (long) (-memoryPool.getFreeBytes() + (memoryPool.getMaxBytes() * (1.0 - memoryRevokingTarget))); remainingBytesToRevoke -= getMemoryAlreadyBeingRevoked(sqlTasks, memoryPool); requestRevoking(memoryPool, sqlTasks, remainingBytesToRevoke); }
/** * Test that when a {@link MemoryPool} is over-allocated, revocable memory is revoked without delay (although asynchronously). */ @Test public void testImmediateMemoryRevoking() throws Exception { // Given SqlTask sqlTask = newSqlTask(); OperatorContext operatorContext = createContexts(sqlTask); allOperatorContexts = ImmutableSet.of(operatorContext); List<SqlTask> tasks = ImmutableList.of(sqlTask); MemoryRevokingScheduler scheduler = new MemoryRevokingScheduler(singletonList(memoryPool), () -> tasks, executor, 1.0, 1.0); scheduler.registerPoolListeners(); // no periodic check initiated // When operatorContext.localRevocableMemoryContext().setBytes(12); awaitAsynchronousCallbacksRun(); // Then assertMemoryRevokingRequestedFor(operatorContext); }
private void onMemoryReserved(MemoryPool memoryPool) { try { if (!memoryRevokingNeeded(memoryPool)) { return; } if (checkPending.compareAndSet(false, true)) { log.debug("Scheduling check for %s", memoryPool); scheduleRevoking(); } } catch (Throwable e) { log.error(e, "Error when acting on memory pool reservation"); } }
@Test public void testCountAlreadyRevokedMemoryWithinAPool() throws Exception { // Given SqlTask sqlTask1 = newSqlTask(); MemoryPool anotherMemoryPool = new MemoryPool(new MemoryPoolId("test"), new DataSize(10, BYTE)); sqlTask1.getQueryContext().setMemoryPool(anotherMemoryPool); OperatorContext operatorContext1 = createContexts(sqlTask1); SqlTask sqlTask2 = newSqlTask(); OperatorContext operatorContext2 = createContexts(sqlTask2); List<SqlTask> tasks = ImmutableList.of(sqlTask1, sqlTask2); MemoryRevokingScheduler scheduler = new MemoryRevokingScheduler(asList(memoryPool, anotherMemoryPool), () -> tasks, executor, 1.0, 1.0); allOperatorContexts = ImmutableSet.of(operatorContext1, operatorContext2); /* * sqlTask1 fills its pool */ operatorContext1.localRevocableMemoryContext().setBytes(12); requestMemoryRevoking(scheduler); assertMemoryRevokingRequestedFor(operatorContext1); /* * When sqlTask2 fills its pool */ operatorContext2.localRevocableMemoryContext().setBytes(12); requestMemoryRevoking(scheduler); /* * Then sqlTask2 should be asked to revoke its memory too */ assertMemoryRevokingRequestedFor(operatorContext1, operatorContext2); }
@Inject public MemoryRevokingScheduler( LocalMemoryManager localMemoryManager, SqlTaskManager sqlTaskManager, TaskManagementExecutor taskManagementExecutor, FeaturesConfig config) { this( ImmutableList.copyOf(getMemoryPools(localMemoryManager)), requireNonNull(sqlTaskManager, "sqlTaskManager cannot be null")::getAllTasks, requireNonNull(taskManagementExecutor, "taskManagementExecutor cannot be null").getExecutor(), config.getMemoryRevokingThreshold(), config.getMemoryRevokingTarget()); }
/** * Test that when a {@link MemoryPool} is over-allocated, revocable memory is revoked without delay (although asynchronously). */ @Test public void testImmediateMemoryRevoking() throws Exception { // Given SqlTask sqlTask = newSqlTask(); OperatorContext operatorContext = createContexts(sqlTask); allOperatorContexts = ImmutableSet.of(operatorContext); List<SqlTask> tasks = ImmutableList.of(sqlTask); MemoryRevokingScheduler scheduler = new MemoryRevokingScheduler(singletonList(memoryPool), () -> tasks, executor, 1.0, 1.0); scheduler.registerPoolListeners(); // no periodic check initiated // When operatorContext.localRevocableMemoryContext().setBytes(12); awaitAsynchronousCallbacksRun(); // Then assertMemoryRevokingRequestedFor(operatorContext); }
private void requestMemoryRevoking(MemoryPool memoryPool, Collection<SqlTask> sqlTasks) { long remainingBytesToRevoke = (long) (-memoryPool.getFreeBytes() + (memoryPool.getMaxBytes() * (1.0 - memoryRevokingTarget))); remainingBytesToRevoke -= getMemoryAlreadyBeingRevoked(sqlTasks, memoryPool); requestRevoking(memoryPool, sqlTasks, remainingBytesToRevoke); }
private void onMemoryReserved(MemoryPool memoryPool) { try { if (!memoryRevokingNeeded(memoryPool)) { return; } if (checkPending.compareAndSet(false, true)) { log.debug("Scheduling check for %s", memoryPool); scheduleRevoking(); } } catch (Throwable e) { log.error(e, "Error when acting on memory pool reservation"); } }
private synchronized void runMemoryRevoking() { if (checkPending.getAndSet(false)) { Collection<SqlTask> sqlTasks = null; for (MemoryPool memoryPool : memoryPools) { if (!memoryRevokingNeeded(memoryPool)) { continue; } if (sqlTasks == null) { sqlTasks = requireNonNull(currentTasksSupplier.get()); } requestMemoryRevoking(memoryPool, sqlTasks); } } }
@Test public void testCountAlreadyRevokedMemoryWithinAPool() throws Exception { // Given SqlTask sqlTask1 = newSqlTask(); MemoryPool anotherMemoryPool = new MemoryPool(new MemoryPoolId("test"), new DataSize(10, BYTE)); sqlTask1.getQueryContext().setMemoryPool(anotherMemoryPool); OperatorContext operatorContext1 = createContexts(sqlTask1); SqlTask sqlTask2 = newSqlTask(); OperatorContext operatorContext2 = createContexts(sqlTask2); List<SqlTask> tasks = ImmutableList.of(sqlTask1, sqlTask2); MemoryRevokingScheduler scheduler = new MemoryRevokingScheduler(asList(memoryPool, anotherMemoryPool), () -> tasks, executor, 1.0, 1.0); allOperatorContexts = ImmutableSet.of(operatorContext1, operatorContext2); /* * sqlTask1 fills its pool */ operatorContext1.localRevocableMemoryContext().setBytes(12); requestMemoryRevoking(scheduler); assertMemoryRevokingRequestedFor(operatorContext1); /* * When sqlTask2 fills its pool */ operatorContext2.localRevocableMemoryContext().setBytes(12); requestMemoryRevoking(scheduler); /* * Then sqlTask2 should be asked to revoke its memory too */ assertMemoryRevokingRequestedFor(operatorContext1, operatorContext2); }
@Inject public MemoryRevokingScheduler( LocalMemoryManager localMemoryManager, SqlTaskManager sqlTaskManager, TaskManagementExecutor taskManagementExecutor, FeaturesConfig config) { this( ImmutableList.copyOf(getMemoryPools(localMemoryManager)), requireNonNull(sqlTaskManager, "sqlTaskManager cannot be null")::getAllTasks, requireNonNull(taskManagementExecutor, "taskManagementExecutor cannot be null").getExecutor(), config.getMemoryRevokingThreshold(), config.getMemoryRevokingTarget()); }
@PostConstruct public void start() { registerPeriodicCheck(); registerPoolListeners(); }
private void registerPeriodicCheck() { this.scheduledFuture = taskManagementExecutor.scheduleWithFixedDelay(() -> { try { requestMemoryRevokingIfNeeded(); } catch (Throwable e) { log.error(e, "Error requesting system memory revoking"); } }, 1, 1, SECONDS); }
MemoryRevokingScheduler scheduler = new MemoryRevokingScheduler(singletonList(memoryPool), () -> tasks, executor, 1.0, 1.0);
@VisibleForTesting MemoryRevokingScheduler( List<MemoryPool> memoryPools, Supplier<? extends Collection<SqlTask>> currentTasksSupplier, ScheduledExecutorService taskManagementExecutor, double memoryRevokingThreshold, double memoryRevokingTarget) { this.memoryPools = ImmutableList.copyOf(requireNonNull(memoryPools, "memoryPools is null")); this.currentTasksSupplier = requireNonNull(currentTasksSupplier, "currentTasksSupplier is null"); this.taskManagementExecutor = requireNonNull(taskManagementExecutor, "taskManagementExecutor is null"); this.memoryRevokingThreshold = checkFraction(memoryRevokingThreshold, "memoryRevokingThreshold"); this.memoryRevokingTarget = checkFraction(memoryRevokingTarget, "memoryRevokingTarget"); checkArgument( memoryRevokingTarget <= memoryRevokingThreshold, "memoryRevokingTarget should be less than or equal memoryRevokingThreshold, but got %s and %s respectively", memoryRevokingTarget, memoryRevokingThreshold); }
private void requestMemoryRevoking(MemoryRevokingScheduler scheduler) throws Exception { scheduler.requestMemoryRevokingIfNeeded(); awaitAsynchronousCallbacksRun(); }