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<?> reserveSystemMemory(long bytes) {
checkArgument(bytes >= 0, "bytes is negative");
ListenableFuture<?> future = systemMemoryPool.reserve(queryId, bytes);
systemReserved += bytes;
return future;
}
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 boolean tryReserveMemory(long bytes) {
checkArgument(bytes >= 0, "bytes is negative");
if (reserved + bytes > maxMemory) {
return false;
}
if (memoryPool.tryReserve(queryId, bytes)) {
reserved += bytes;
return true;
}
return false;
}
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;
}
// TODO: This method should be removed, and the correct limit set in the constructor. However, due
// to the way QueryContext is constructed the memory limit is not known in advance
public synchronized void setResourceOvercommit() {
// Allow the query to use the entire pool. This way the worker will kill the query, if it uses
// the entire local general pool.
// The coordinator will kill the query if the cluster runs out of memory.
maxMemory = memoryPool.getMaxBytes();
}
public synchronized void freeSystemMemory(long bytes) {
checkArgument(bytes >= 0, "bytes is negative");
checkArgument(systemReserved - bytes >= 0, "tried to free more system memory than is reserved");
systemReserved -= bytes;
systemMemoryPool.free(queryId, bytes);
}
public synchronized void freeMemory(long bytes) {
checkArgument(reserved - bytes >= 0, "tried to free more memory than is reserved");
reserved -= bytes;
memoryPool.free(queryId, bytes);
}