public synchronized ListenableFuture<?> reserveSystemMemory(long bytes) {
checkArgument(bytes >= 0, "bytes is negative");
ListenableFuture<?> future = systemMemoryPool.reserve(queryId, bytes);
systemReserved += bytes;
return future;
}
public synchronized void setMemoryPool(MemoryPool pool) {
requireNonNull(pool, "pool is null");
if (pool.getId().equals(memoryPool.getId())) {
// Don't unblock our tasks and thrash the pools, if this is a no-op
return;
}
MemoryPool originalPool = memoryPool;
long originalReserved = reserved;
memoryPool = pool;
ListenableFuture<?> future = pool.reserve(queryId, reserved);
Futures.addCallback(
future,
new FutureCallback<Object>() {
@Override
public void onSuccess(Object result) {
originalPool.free(queryId, originalReserved);
// Unblock all the tasks, if they were waiting for memory, since we're in a new pool.
taskContexts.stream().forEach(TaskContext::moreMemoryAvailable);
}
@Override
public void onFailure(Throwable t) {
originalPool.free(queryId, originalReserved);
// Unblock all the tasks, if they were waiting for memory, since we're in a new pool.
taskContexts.stream().forEach(TaskContext::moreMemoryAvailable);
}
});
}
public synchronized ListenableFuture<?> reserveMemory(long bytes) {
checkArgument(bytes >= 0, "bytes is negative");
if (reserved + bytes > maxMemory) {
throw exceededLocalLimit(
new DataSize(maxMemory, DataSize.Unit.BYTE).convertToMostSuccinctDataSize());
}
ListenableFuture<?> future = memoryPool.reserve(queryId, bytes);
reserved += bytes;
return future;
}
public synchronized ListenableFuture<?> reserveMemory(long bytes) {
checkArgument(bytes >= 0, "bytes is negative");
if (reserved + bytes > maxMemory) {
throw exceededLocalLimit(succinctBytes(maxMemory));
}
ListenableFuture<?> future = memoryPool.reserve(queryId, bytes);
reserved += bytes;
// Never block queries using a trivial amount of memory
if (reserved < GUARANTEED_MEMORY) {
return NOT_BLOCKED;
}
return future;
}