@GuardedBy("exclusiveLock") private void handleMemoryRevoke() { for (int i = 0; i < activeOperators.size() && !driverContext.isDone(); i++) { Operator operator = activeOperators.get(i); if (revokingOperators.containsKey(operator)) { checkOperatorFinishedRevoking(operator); } else if (operator.getOperatorContext().isMemoryRevokingRequested()) { ListenableFuture<?> future = operator.startMemoryRevoke(); revokingOperators.put(operator, future); checkOperatorFinishedRevoking(operator); } } }
/** * Returns how much revocable memory will be revoked by the operator */ public long requestMemoryRevoking() { long revokedMemory = 0L; Runnable listener = null; synchronized (this) { if (!isMemoryRevokingRequested() && operatorMemoryContext.getRevocableMemory() > 0) { memoryRevokingRequested = true; revokedMemory = operatorMemoryContext.getRevocableMemory(); listener = memoryRevocationRequestListener; } } if (listener != null) { runListener(listener); } return revokedMemory; }
private void assertMemoryRevokingRequestedFor(OperatorContext... operatorContexts) { ImmutableSet<OperatorContext> operatorContextsSet = ImmutableSet.copyOf(operatorContexts); operatorContextsSet.forEach( operatorContext -> assertTrue(operatorContext.isMemoryRevokingRequested(), "expected memory requested for operator " + operatorContext.getOperatorId())); Sets.difference(allOperatorContexts, operatorContextsSet).forEach( operatorContext -> assertFalse(operatorContext.isMemoryRevokingRequested(), "expected memory not requested for operator " + operatorContext.getOperatorId())); }
private Predicate<OperatorContext> waitingForRevocableSystemMemory() { return (OperatorContext operatorContext) -> !operatorContext.isWaitingForRevocableMemory().isDone() && !operatorContext.isMemoryRevokingRequested(); }
@Test public void testOperatorMemoryRevocation() { AtomicInteger counter = new AtomicInteger(); OperatorContext operatorContext = TestingOperatorContext.create(scheduledExecutor); LocalMemoryContext revocableMemoryContext = operatorContext.localRevocableMemoryContext(); revocableMemoryContext.setBytes(1000); operatorContext.setMemoryRevocationRequestListener(() -> counter.incrementAndGet()); operatorContext.requestMemoryRevoking(); assertTrue(operatorContext.isMemoryRevokingRequested()); assertEquals(counter.get(), 1); // calling resetMemoryRevokingRequested() should clear the memory revoking requested flag operatorContext.resetMemoryRevokingRequested(); assertFalse(operatorContext.isMemoryRevokingRequested()); operatorContext.requestMemoryRevoking(); assertEquals(counter.get(), 2); assertTrue(operatorContext.isMemoryRevokingRequested()); }
@Test public void testRevocationAlreadyRequested() { AtomicInteger counter = new AtomicInteger(); OperatorContext operatorContext = TestingOperatorContext.create(scheduledExecutor); LocalMemoryContext revocableMemoryContext = operatorContext.localRevocableMemoryContext(); revocableMemoryContext.setBytes(1000); // when memory revocation is already requested setting a listener should immediately execute it operatorContext.requestMemoryRevoking(); operatorContext.setMemoryRevocationRequestListener(() -> counter.incrementAndGet()); assertTrue(operatorContext.isMemoryRevokingRequested()); assertEquals(counter.get(), 1); }
private static void triggerMemoryRevokingAndWait(HashBuilderOperator operator, TaskStateMachine taskStateMachine) throws Exception { // When there is background thread running Driver, we must delegate memory revoking to that thread operator.getOperatorContext().requestMemoryRevoking(); while (operator.getOperatorContext().isMemoryRevokingRequested()) { checkErrors(taskStateMachine); Thread.sleep(10); } checkErrors(taskStateMachine); checkState(operator.getState() == HashBuilderOperator.State.SPILLING_INPUT || operator.getState() == HashBuilderOperator.State.INPUT_SPILLED); }