private MemoryPoolAssignmentsRequest updateAssignments(Iterable<QueryExecution> queries) {
ClusterMemoryPool reservedPool = pools.get(RESERVED_POOL);
ClusterMemoryPool generalPool = pools.get(GENERAL_POOL);
long version = memoryPoolAssignmentsVersion.incrementAndGet();
// Check that all previous assignments have propagated to the visible nodes. This doesn't
// account for temporary network issues,
// and is more of a safety check than a guarantee
if (reservedPool != null && generalPool != null && allAssignmentsHavePropagated(queries)) {
if (reservedPool.getQueries() == 0 && generalPool.getBlockedNodes() > 0) {
QueryExecution biggestQuery = null;
long maxMemory = -1;
for (QueryExecution queryExecution : queries) {
long bytesUsed = queryExecution.getTotalMemoryReservation();
if (bytesUsed > maxMemory) {
biggestQuery = queryExecution;
maxMemory = bytesUsed;
}
}
for (QueryExecution queryExecution : queries) {
if (queryExecution.getQueryId().equals(biggestQuery.getQueryId())) {
queryExecution.setMemoryPool(new VersionedMemoryPoolId(RESERVED_POOL, version));
}
}
}
}
ImmutableList.Builder<MemoryPoolAssignment> assignments = ImmutableList.builder();
for (QueryExecution queryExecution : queries) {
assignments.add(
new MemoryPoolAssignment(
queryExecution.getQueryId(), queryExecution.getMemoryPool().getId()));
}
return new MemoryPoolAssignmentsRequest(coordinatorId, version, assignments.build());
}
public synchronized void process(Iterable<QueryExecution> queries) {
if (!enabled) {
return;
}
boolean outOfMemory = isClusterOutOfMemory();
if (!outOfMemory) {
lastTimeNotOutOfMemory = System.nanoTime();
}
boolean queryKilled = false;
long totalBytes = 0;
for (QueryExecution query : queries) {
long bytes = query.getTotalMemoryReservation();
DataSize sessionMaxQueryMemory = getQueryMaxMemory(query.getSession());
long queryMemoryLimit = Math.min(maxQueryMemory.toBytes(), sessionMaxQueryMemory.toBytes());
totalBytes += bytes;
if (resourceOvercommit(query.getSession()) && outOfMemory) {
// If a query has requested resource overcommit, only kill it if the cluster has run out of
// memory
DataSize memory = succinctBytes(bytes);
query.fail(
new PrestoException(
CLUSTER_OUT_OF_MEMORY,
format(
"The cluster is out of memory and %s=true, so this query was killed. It was using %s of memory",
RESOURCE_OVERCOMMIT, memory)));
queryKilled = true;
}
if (!resourceOvercommit(query.getSession()) && bytes > queryMemoryLimit) {
DataSize maxMemory = succinctBytes(queryMemoryLimit);
query.fail(exceededGlobalLimit(maxMemory));
queryKilled = true;
}
}
clusterMemoryUsageBytes.set(totalBytes);
if (killOnOutOfMemory) {
boolean shouldKillQuery =
nanosSince(lastTimeNotOutOfMemory).compareTo(killOnOutOfMemoryDelay) > 0 && outOfMemory;
boolean lastKilledQueryIsGone = (lastKilledQuery == null);
if (!lastKilledQueryIsGone) {
ClusterMemoryPool generalPool = pools.get(GENERAL_POOL);
if (generalPool != null) {
lastKilledQueryIsGone =
generalPool.getQueryMemoryReservations().containsKey(lastKilledQuery);
}
}
if (shouldKillQuery && lastKilledQueryIsGone && !queryKilled) {
// Kill the biggest query in the general pool
QueryExecution biggestQuery = null;
long maxMemory = -1;
for (QueryExecution query : queries) {
long bytesUsed = query.getTotalMemoryReservation();
if (bytesUsed > maxMemory && query.getMemoryPool().getId().equals(GENERAL_POOL)) {
biggestQuery = query;
maxMemory = bytesUsed;
}
}
if (biggestQuery != null) {
biggestQuery.fail(
new PrestoException(
CLUSTER_OUT_OF_MEMORY,
"The cluster is out of memory, and your query was killed. Please try again in a few minutes."));
queriesKilledDueToOutOfMemory.incrementAndGet();
lastKilledQuery = biggestQuery.getQueryId();
}
}
}
Map<MemoryPoolId, Integer> countByPool = new HashMap<>();
for (QueryExecution query : queries) {
MemoryPoolId id = query.getMemoryPool().getId();
countByPool.put(id, countByPool.getOrDefault(id, 0) + 1);
}
updatePools(countByPool);
updateNodes(updateAssignments(queries));
// check if CPU usage is over limit
for (QueryExecution query : queries) {
Duration cpuTime = query.getTotalCpuTime();
Duration sessionLimit = getQueryMaxCpuTime(query.getSession());
Duration limit = maxQueryCpuTime.compareTo(sessionLimit) < 0 ? maxQueryCpuTime : sessionLimit;
if (cpuTime.compareTo(limit) > 0) {
query.fail(new ExceededCpuLimitException(limit));
}
}
}